1// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
2// Verify the printed output can be parsed.
3// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
4// Verify the generic form can be parsed.
5// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
6
7module attributes {gpu.container_module} {
8
9  // CHECK-LABEL:func @no_args(%{{.*}}: index)
10  func.func @no_args(%sz : index) {
11    // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
12    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)
13               threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {
14      // CHECK: gpu.terminator
15      gpu.terminator
16    }
17    return
18  }
19
20  // CHECK-LABEL:func @args(%{{.*}}: index, %{{.*}}: index, %{{.*}}: f32, %{{.*}}: memref<?xf32, 1>) {
21  func.func @args(%blk : index, %thrd : index, %float : f32, %data : memref<?xf32,1>) {
22    // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
23    gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
24               threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
25      "use"(%float) : (f32) -> ()
26      "use"(%data) : (memref<?xf32,1>) -> ()
27      // CHECK: gpu.terminator
28      gpu.terminator
29    }
30    return
31  }
32
33  // CHECK-LABEL:func @launch_async(%{{.*}}: index, %{{.*}}: index) {
34  func.func @launch_async(%blk : index, %thrd : index) {
35    // CHECK: gpu.launch async [%{{.+}}] blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
36    %t = gpu.wait async
37    %name = gpu.launch async [%t] blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
38               threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
39      gpu.terminator
40    }
41    return
42  }
43
44  // CHECK-LABEL:func @launch_async_no_deps(%{{.*}}: index, %{{.*}}: index) {
45  func.func @launch_async_no_deps(%blk : index, %thrd : index) {
46    // CHECK: %{{.*}} = gpu.launch async blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
47    %t0 = gpu.launch async blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
48               threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
49      gpu.terminator
50    }
51    // CHECK: gpu.launch async blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})
52    %t1 = gpu.launch async [] blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)
53               threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {
54      gpu.terminator
55    }
56    return
57  }
58
59  gpu.module @kernels {
60    gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {
61      %tIdX = gpu.thread_id x
62      %tIdY = gpu.thread_id y
63      %tIdZ = gpu.thread_id z
64
65      %bDimX = gpu.block_dim x
66      %bDimY = gpu.block_dim y
67      %bDimZ = gpu.block_dim z
68
69      %bIdX = gpu.block_id x
70      %bIdY = gpu.block_id y
71      %bIdZ = gpu.block_id z
72
73      %gDimX = gpu.grid_dim x
74      %gDimY = gpu.grid_dim y
75      %gDimZ = gpu.grid_dim z
76
77      %gIdX = gpu.global_id x
78      %gIdY = gpu.global_id y
79      %gIdZ = gpu.global_id z
80
81      %sgId = gpu.subgroup_id : index
82      %numSg = gpu.num_subgroups : index
83      %SgSi = gpu.subgroup_size : index
84
85      %one = arith.constant 1.0 : f32
86      %sum = gpu.all_reduce add %one {} : (f32) -> (f32)
87
88      %width = arith.constant 7 : i32
89      %offset = arith.constant 3 : i32
90      // CHECK: gpu.shuffle xor %{{.*}}, %{{.*}}, %{{.*}} : f32
91      %shfl, %pred = gpu.shuffle xor %arg0, %offset, %width : f32
92      // CHECK: gpu.shuffle up %{{.*}}, %{{.*}}, %{{.*}} : f32
93      %shfl1, %pred1 = gpu.shuffle up %arg0, %offset, %width : f32
94      // CHECK: gpu.shuffle down %{{.*}}, %{{.*}}, %{{.*}} : f32
95      %shfl2, %pred2 = gpu.shuffle down %arg0, %offset, %width : f32
96      // CHECK: gpu.shuffle idx %{{.*}}, %{{.*}}, %{{.*}} : f32
97      %shfl3, %pred3 = gpu.shuffle idx %arg0, %offset, %width : f32
98
99      "gpu.barrier"() : () -> ()
100
101      "some_op"(%bIdX, %tIdX) : (index, index) -> ()
102      %42 = memref.load %arg1[%bIdX] : memref<?xf32, 1>
103      gpu.return
104    }
105
106    gpu.func @kernel_2() kernel {
107      gpu.return
108    }
109  }
110
111  func.func @foo() {
112    %0 = "op"() : () -> (f32)
113    %1 = "op"() : () -> (memref<?xf32, 1>)
114    // CHECK: %{{.*}} = arith.constant 8
115    %cst = arith.constant 8 : index
116    %c0 = arith.constant 0 : i32
117    %t0 = gpu.wait async
118
119    // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
120    gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%0 : f32, %1 : memref<?xf32, 1>)
121
122    gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) dynamic_shared_memory_size %c0 args(%0 : f32, %1 : memref<?xf32, 1>)
123
124    // CHECK: gpu.launch_func @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
125    gpu.launch_func @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
126
127    // CHECK: %{{.*}} = gpu.launch_func async [%{{.*}}] @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})
128    %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2  blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
129
130    return
131  }
132
133  gpu.module @gpu_funcs {
134    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)
135    // CHECK:       workgroup
136    // CHECK:       private
137    // CHECK:       attributes
138    gpu.func @kernel_1(%arg0: f32)
139        workgroup(%arg1: memref<42xf32, 3>)
140        private(%arg2: memref<2xf32, 5>, %arg3: memref<1xf32, 5>)
141        kernel
142        attributes {foo="bar"} {
143      "use"(%arg1) : (memref<42xf32, 3>) -> ()
144      "use"(%arg2) : (memref<2xf32, 5>) -> ()
145      "use"(%arg3) : (memref<1xf32, 5>) -> ()
146      gpu.return
147    }
148
149    // CHECK-LABEL gpu.func @printf_test
150    // CHECK: (%[[ARG0:.*]]: i32)
151    // CHECK: gpu.printf "Value: %d" %[[ARG0]] : i32
152    gpu.func @printf_test(%arg0 : i32) {
153      gpu.printf "Value: %d" %arg0 : i32
154      gpu.return
155    }
156
157    // CHECK-LABEL: gpu.func @no_attribution
158    // CHECK: {
159    gpu.func @no_attribution(%arg0: f32) {
160      gpu.return
161    }
162
163    // CHECK-LABEL: @no_attribution_attrs
164    // CHECK:       attributes
165    // CHECK:       {
166    gpu.func @no_attribution_attrs(%arg0: f32) attributes {foo="bar"} {
167      gpu.return
168    }
169
170    // CHECK-LABEL: @workgroup_only
171    // CHECK:       workgroup({{.*}}: {{.*}})
172    // CHECK:       {
173    gpu.func @workgroup_only() workgroup(%arg0: memref<42xf32, 3>) {
174      gpu.return
175    }
176    // CHECK-LABEL: @private_only
177    // CHECK:       private({{.*}}: {{.*}})
178    // CHECK:       {
179    gpu.func @private_only() private(%arg0: memref<2xf32, 5>) {
180      gpu.return
181    }
182
183    // CHECK-LABEL: @empty_attribution
184    // CHECK:       {
185    gpu.func @empty_attribution(%arg0: f32) workgroup() private() {
186      gpu.return
187    }
188  }
189
190  gpu.module @explicit_attributions {
191    // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32, {{.*}}: memref<?xf32>) workgroup({{.*}}: memref<5xf32, 3>) private({{.*}}: memref<5xf32, 5>)
192    "gpu.func"() ({
193    ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):
194      "gpu.return"() : () -> ()
195    } ) {function_type = (f32, memref<?xf32>) -> (), gpu.kernel, sym_name = "kernel_1", workgroup_attributions = 1: i64} : () -> ()
196  }
197
198  func.func @alloc() {
199    // CHECK-LABEL: func @alloc()
200
201    // CHECK: %[[m0:.*]] = gpu.alloc () : memref<13xf32, 1>
202    %m0 = gpu.alloc () : memref<13xf32, 1>
203    // CHECK: gpu.dealloc %[[m0]] : memref<13xf32, 1>
204    gpu.dealloc %m0 : memref<13xf32, 1>
205
206    %t0 = gpu.wait async
207    // CHECK: %[[m1:.*]], %[[t1:.*]] = gpu.alloc async [{{.*}}] () : memref<13xf32, 1>
208    %m1, %t1 = gpu.alloc async [%t0] () : memref<13xf32, 1>
209    // CHECK: gpu.dealloc async [%[[t1]]] %[[m1]] : memref<13xf32, 1>
210    %t2 = gpu.dealloc async [%t1] %m1 : memref<13xf32, 1>
211
212    return
213  }
214
215  func.func @async_token(%arg0 : !gpu.async.token) -> !gpu.async.token {
216    // CHECK-LABEL: func @async_token({{.*}}: !gpu.async.token)
217    // CHECK: return {{.*}} : !gpu.async.token
218    return %arg0 : !gpu.async.token
219  }
220
221  func.func @async_wait() {
222    // CHECK-LABEL: func @async_wait
223    // CHECK: %[[t0:.*]] = gpu.wait async
224    %0 = gpu.wait async
225    // CHECK: %[[t1:.*]] = gpu.wait async [%[[t0]]]
226    %1 = gpu.wait async [%0]
227    // CHECK: %{{.*}} = gpu.wait async [%[[t0]], %[[t1]]]
228    %2 = gpu.wait async [%0, %1]
229    // CHECK: gpu.wait [%[[t0]], %[[t1]]]
230    // CHECK-NOT: async
231    gpu.wait [%0, %1]
232    // CHECK: gpu.wait
233    // CHECK-NOT: async
234    gpu.wait // Valid, but a no-op.
235    return
236  }
237
238  func.func @memcpy(%dst : memref<3x7xf32>, %src : memref<3x7xf32, 1>) {
239    // CHECK-LABEL: func @memcpy
240    // CHECK: gpu.memcpy {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>
241    gpu.memcpy %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
242    // CHECK: %[[t0:.*]] = gpu.wait async
243    %0 = gpu.wait async
244    // CHECK: {{.*}} = gpu.memcpy async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>
245    %1 = gpu.memcpy async [%0] %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>
246    return
247  }
248
249  func.func @memset(%dst : memref<3x7xf32>, %value : f32) {
250    // CHECK-LABEL: func @memset
251    // CHECK: gpu.memset {{.*}}, {{.*}} : memref<3x7xf32>, f32
252    gpu.memset %dst, %value : memref<3x7xf32>, f32
253    // CHECK: %[[t0:.*]] = gpu.wait async
254    %0 = gpu.wait async
255    // CHECK: {{.*}} = gpu.memset async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, f32
256    %1 = gpu.memset async [%0] %dst, %value : memref<3x7xf32>, f32
257    return
258  }
259
260  func.func @mmamatrix_valid_element_type(%src : memref<32x32xf16, affine_map<(d0, d1) -> (d0 * 64 + d1)>>){
261    // CHECK-LABEL: func @mmamatrix_valid_element_type
262    %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
263    // CHECK: %[[wg:.*]] = memref.alloca()
264    %i = arith.constant 16 : index
265    // CHECK: %[[i:.*]] = arith.constant 16 : index
266     %cst = arith.constant 1.000000e+00 : f32
267    // CHECK: %[[cst:.*]] = arith.constant 1.000000e+00 : f32
268    %0 = gpu.subgroup_mma_load_matrix %wg[%i, %i] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
269    // CHECK: gpu.subgroup_mma_load_matrix %[[wg]][%[[i]], %[[i]]] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
270    %s = gpu.subgroup_mma_load_matrix %src[%i, %i] {leadDimension = 64 : index} : memref<32x32xf16, affine_map<(d0, d1) -> (d0 * 64 + d1)>> -> !gpu.mma_matrix<16x16xf16, "AOp">
271    // CHECK: gpu.subgroup_mma_load_matrix %{{.*}}[%[[i]], %[[i]]] {leadDimension = 64 : index} : memref<32x32xf16, #{{.*}}> -> !gpu.mma_matrix<16x16xf16, "AOp">
272    %1 = gpu.subgroup_mma_constant_matrix %cst : !gpu.mma_matrix<16x16xf32, "COp">
273    // CHECK: gpu.subgroup_mma_elementwise addf %{{.*}}, %{{.*}} : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">
274    %2 = gpu.subgroup_mma_elementwise addf %1, %1 : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">
275    // CHECK: gpu.subgroup_mma_elementwise maxf %{{.*}}, %{{.*}} : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">
276    %3 = gpu.subgroup_mma_elementwise maxf %2, %1 : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">
277    return
278  }
279
280  func.func @async_cp(%dst : memref<2x7x5xf32, 3>, %src : memref<4x5xf32>){
281    // CHECK-LABEL: func @async_cp
282    %c0 = arith.constant 0 : index
283    // CHECK: gpu.device_async_copy %{{.*}}[{{.*}}, {{.*}}], %{{.*}}[{{.*}}, {{.*}}, {{.*}}], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
284    %0 = gpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 : memref<4x5xf32> to memref<2x7x5xf32, 3>
285    // CHECK: %{{.*}} = gpu.device_async_create_group
286    %token = gpu.device_async_create_group %0
287    // CHECK: gpu.device_async_wait %{{.*}} {numGroups = 1 : i32}
288    gpu.device_async_wait %token {numGroups = 1 : i32}
289    return
290  }
291
292  // CHECK-LABEL: func @set_default_device
293  func.func @set_default_device(%arg0: i32) {
294    // CHECK: gpu.set_default_device
295    gpu.set_default_device %arg0
296    return
297  }
298}
299