1// RUN: mlir-opt -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s
2
3// CHECK: module attributes {gpu.container_module}
4
5// CHECK-LABEL: func @launch()
6func @launch() {
7  // CHECK: %[[ARG0:.*]] = "op"() : () -> f32
8  %0 = "op"() : () -> (f32)
9  // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>
10  %1 = "op"() : () -> (memref<?xf32, 1>)
11  // CHECK: %[[GDIMX:.*]] = constant 8
12  %gDimX = constant 8 : index
13  // CHECK: %[[GDIMY:.*]] = constant 12
14  %gDimY = constant 12 : index
15  // CHECK: %[[GDIMZ:.*]] = constant 16
16  %gDimZ = constant 16 : index
17  // CHECK: %[[BDIMX:.*]] = constant 20
18  %bDimX = constant 20 : index
19  // CHECK: %[[BDIMY:.*]] = constant 24
20  %bDimY = constant 24 : index
21  // CHECK: %[[BDIMZ:.*]] = constant 28
22  %bDimZ = constant 28 : index
23
24  // CHECK: "gpu.launch_func"(%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]], %[[BDIMX]], %[[BDIMY]], %[[BDIMZ]], %[[ARG0]], %[[ARG1]]) {kernel = "launch_kernel", kernel_module = @launch_kernel} : (index, index, index, index, index, index, f32, memref<?xf32, 1>) -> ()
25  // CHECK-NOT: gpu.launch blocks
26  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,
27                                       %grid_z = %gDimZ)
28             threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,
29                                        %block_z = %bDimZ)
30             args(%arg0 = %0, %arg1 = %1) : f32, memref<?xf32, 1> {
31    "use"(%arg0): (f32) -> ()
32    "some_op"(%bx, %block_x) : (index, index) -> ()
33    %42 = load %arg1[%tx] : memref<?xf32, 1>
34    gpu.return
35  }
36  return
37}
38
39
40// CHECK-LABEL: module @launch_kernel
41// CHECK-NEXT: gpu.func @launch_kernel
42// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)
43// CHECK-NEXT: %[[BID:.*]] = "gpu.block_id"() {dimension = "x"} : () -> index
44// CHECK-NEXT: = "gpu.block_id"() {dimension = "y"} : () -> index
45// CHECK-NEXT: = "gpu.block_id"() {dimension = "z"} : () -> index
46// CHECK-NEXT: %[[TID:.*]] = "gpu.thread_id"() {dimension = "x"} : () -> index
47// CHECK-NEXT: = "gpu.thread_id"() {dimension = "y"} : () -> index
48// CHECK-NEXT: = "gpu.thread_id"() {dimension = "z"} : () -> index
49// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "x"} : () -> index
50// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "y"} : () -> index
51// CHECK-NEXT: = "gpu.grid_dim"() {dimension = "z"} : () -> index
52// CHECK-NEXT: %[[BDIM:.*]] = "gpu.block_dim"() {dimension = "x"} : () -> index
53// CHECK-NEXT: = "gpu.block_dim"() {dimension = "y"} : () -> index
54// CHECK-NEXT: = "gpu.block_dim"() {dimension = "z"} : () -> index
55// CHECK-NEXT: "use"(%[[KERNEL_ARG0]]) : (f32) -> ()
56// CHECK-NEXT: "some_op"(%[[BID]], %[[BDIM]]) : (index, index) -> ()
57// CHECK-NEXT: = load %[[KERNEL_ARG1]][%[[TID]]] : memref<?xf32, 1>
58
59// -----
60
61// CHECK: module attributes {gpu.container_module}
62
63func @multiple_launches() {
64  // CHECK: %[[CST:.*]] = constant 8 : index
65  %cst = constant 8 : index
66  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel} : (index, index, index, index, index, index) -> ()
67  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
68                                       %grid_z = %cst)
69             threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
70                                        %block_z = %cst) {
71    gpu.return
72  }
73  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]]) {kernel = "multiple_launches_kernel", kernel_module = @multiple_launches_kernel_0} : (index, index, index, index, index, index) -> ()
74  gpu.launch blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,
75                                          %grid_z2 = %cst)
76             threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,
77                                           %block_z2 = %cst) {
78    gpu.return
79  }
80  return
81}
82
83// CHECK: module @multiple_launches_kernel
84// CHECK: func @multiple_launches_kernel
85// CHECK: module @multiple_launches_kernel_0
86// CHECK: func @multiple_launches_kernel
87
88// -----
89
90func @extra_constants(%arg0 : memref<?xf32>) {
91  // CHECK: %[[CST:.*]] = constant 8 : index
92  %cst = constant 8 : index
93  %cst2 = constant 2 : index
94  %cst3 = dim %arg0, 0 : memref<?xf32>
95  // CHECK: "gpu.launch_func"(%[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %[[CST]], %{{.*}}) {kernel = "extra_constants_kernel", kernel_module = @extra_constants_kernel} : (index, index, index, index, index, index, memref<?xf32>) -> ()
96  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
97                                       %grid_z = %cst)
98             threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
99                                        %block_z = %cst)
100             args(%kernel_arg0 = %cst2, %kernel_arg1 = %arg0, %kernel_arg2 = %cst3) : index, memref<?xf32>, index {
101    "use"(%kernel_arg0, %kernel_arg1, %kernel_arg2) : (index, memref<?xf32>, index) -> ()
102    gpu.return
103  }
104  return
105}
106
107// CHECK-LABEL: func @extra_constants_kernel(%{{.*}}: memref<?xf32>)
108// CHECK: constant
109// CHECK: constant
110
111// -----
112
113llvm.mlir.global internal @global(42 : i64) : !llvm.i64
114
115func @function_call(%arg0 : memref<?xf32>) {
116  %cst = constant 8 : index
117  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
118                                       %grid_z = %cst)
119             threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,
120                                        %block_z = %cst) {
121    call @device_function() : () -> ()
122    call @device_function() : () -> ()
123    %0 = llvm.mlir.addressof @global : !llvm<"i64*">
124    gpu.return
125  }
126  return
127}
128
129func @device_function() {
130  call @recursive_device_function() : () -> ()
131  gpu.return
132}
133
134func @recursive_device_function() {
135  call @recursive_device_function() : () -> ()
136  gpu.return
137}
138
139// CHECK: module @function_call_kernel attributes {gpu.kernel_module} {
140// CHECK:   gpu.func @function_call_kernel()
141// CHECK:     call @device_function() : () -> ()
142// CHECK:     call @device_function() : () -> ()
143// CHECK:     llvm.mlir.addressof @global : !llvm<"i64*">
144//
145// CHECK:   llvm.mlir.global internal @global(42 : i64) : !llvm.i64
146//
147// CHECK:   func @device_function()
148// CHECK:   func @recursive_device_function()
149// CHECK-NOT:   func @device_function
150