583 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s2// Verify the printed output can be parsed.3// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s4// 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 %s6 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.terminator15 gpu.terminator16 }17 return18 }19 20 // CHECK-LABEL:func @launch_with_module_func_attr(%{{.*}}: index)21 func.func @launch_with_module_func_attr(%sz : index) {22 // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) module(@test_module) function(@test_kernel_func)23 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)24 threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz)25 module(@test_module) function(@test_kernel_func) {26 // CHECK: gpu.terminator27 gpu.terminator28 }29 return30 }31 32 // CHECK-LABEL:func @args(%{{.*}}: index, %{{.*}}: index, %{{.*}}: f32, %{{.*}}: memref<?xf32, 1>) {33 func.func @args(%blk : index, %thrd : index, %float : f32, %data : memref<?xf32,1>) {34 // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})35 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)36 threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {37 "use"(%float) : (f32) -> ()38 "use"(%data) : (memref<?xf32,1>) -> ()39 // CHECK: gpu.terminator40 gpu.terminator41 }42 return43 }44 45 // CHECK-LABEL:func @launch_async(%{{.*}}: index, %{{.*}}: index) {46 func.func @launch_async(%blk : index, %thrd : index) {47 // CHECK: gpu.launch async [%{{.+}}] blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})48 %t = gpu.wait async49 %name = gpu.launch async [%t] blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)50 threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {51 gpu.terminator52 }53 return54 }55 56 // CHECK-LABEL:func @launch_async_no_deps(%{{.*}}: index, %{{.*}}: index) {57 func.func @launch_async_no_deps(%blk : index, %thrd : index) {58 // CHECK: %{{.*}} = gpu.launch async blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})59 %t0 = gpu.launch async blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)60 threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {61 gpu.terminator62 }63 // CHECK: gpu.launch async blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}, %{{.*}} = %{{.*}})64 %t1 = gpu.launch async [] blocks(%arg0, %arg1, %arg2) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)65 threads(%arg3, %arg4, %arg5) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd) {66 gpu.terminator67 }68 return69 }70 71 // CHECK-LABEL: func @launch_with_attributions(72 func.func @launch_with_attributions(%blk : index, %thrd : index, %float : f32, %data : memref<?xf32,1>) {73 // CHECK: gpu.launch74 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %blk, %grid_y = %blk, %grid_z = %blk)75 threads(%tx, %ty, %tz) in (%block_x = %thrd, %block_y = %thrd, %block_z = %thrd)76 // CHECK-SAME: workgroup(%[[WGROUP1:.*]] : memref<42xf32, 3>, %[[WGROUP2:.*]] : memref<2xf32, 3>)77 workgroup(%arg1: memref<42xf32, 3>, %arg2: memref<2xf32, 3>)78 // CHECK-SAME: private(%[[PRIVATE1:.*]] : memref<2xf32, 5>, %[[PRIVATE2:.*]] : memref<1xf32, 5>)79 private(%arg3: memref<2xf32, 5>, %arg4: memref<1xf32, 5>)80 {81 "use"(%float) : (f32) -> ()82 "use"(%data) : (memref<?xf32,1>) -> ()83 // CHECK: "use"(%[[WGROUP1]], %[[WGROUP2]])84 "use"(%arg1, %arg2) : (memref<42xf32, 3>, memref<2xf32, 3>) -> ()85 // CHECK: "use"(%[[PRIVATE1]])86 "use"(%arg3) : (memref<2xf32, 5>) -> ()87 // CHECK: "use"(%[[PRIVATE2]])88 "use"(%arg4) : (memref<1xf32, 5>) -> ()89 // CHECK: gpu.terminator90 gpu.terminator91 }92 return93 }94 95 96 gpu.module @kernels {97 gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {98 %tIdX = gpu.thread_id x99 // CHECK: thread_id_x100 %tIdY = gpu.thread_id y101 // CHECK-NEXT: thread_id_y102 %tIdZ = gpu.thread_id z103 // CHECK-NEXT: thread_id_z104 105 %bDimX = gpu.block_dim x106 // CHECK-NEXT: block_dim_x107 %bDimY = gpu.block_dim y108 // CHECK-NEXT: block_dim_y109 %bDimZ = gpu.block_dim z110 // CHECK-NEXT: block_dim_z111 112 %bIdX = gpu.block_id x113 // CHECK-NEXT: block_id_x114 %bIdY = gpu.block_id y115 // CHECK-NEXT: block_id_y116 %bIdZ = gpu.block_id z117 // CHECK-NEXT: block_id_z118 119 %gDimX = gpu.grid_dim x120 // CHECK-NEXT: grid_dim_x121 %gDimY = gpu.grid_dim y122 // CHECK-NEXT: grid_dim_y123 %gDimZ = gpu.grid_dim z124 // CHECK-NEXT: grid_dim_z125 126 %gIdX = gpu.global_id x127 // CHECK-NEXT: global_id_x128 %gIdY = gpu.global_id y129 // CHECK-NEXT: global_id_y130 %gIdZ = gpu.global_id z131 // CHECK-NEXT: global_id_z132 133 %sgId = gpu.subgroup_id : index134 %numSg = gpu.num_subgroups : index135 %SgSi = gpu.subgroup_size : index136 137 %one = arith.constant 1.0 : f32138 139 %vec = vector.broadcast %arg0 : f32 to vector<4xf32>140 141 // CHECK: %{{.*}} = gpu.all_reduce add %{{.*}} {142 // CHECK-NEXT: } : (f32) -> f32143 %sum = gpu.all_reduce add %one {} : (f32) -> (f32)144 145 // CHECK: %{{.*}} = gpu.all_reduce add %{{.*}} uniform {146 // CHECK-NEXT: } : (f32) -> f32147 %sum1 = gpu.all_reduce add %one uniform {} : (f32) -> f32148 149 // CHECK: %{{.*}} = gpu.all_reduce %{{.*}} {150 // CHECK-NEXT: ^{{.*}}(%{{.*}}: f32, %{{.*}}: f32):151 // CHECK-NEXT: %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32152 // CHECK-NEXT: gpu.yield %{{.*}} : f32153 // CHECK-NEXT: } : (f32) -> f32154 %sum2 = gpu.all_reduce %one {155 ^bb(%lhs : f32, %rhs : f32):156 %tmp = arith.addf %lhs, %rhs : f32157 gpu.yield %tmp : f32158 } : (f32) -> (f32)159 160 // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (f32) -> f32161 %sum_subgroup = gpu.subgroup_reduce add %one : (f32) -> f32162 163 // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} uniform : (f32) -> f32164 %sum_subgroup1 = gpu.subgroup_reduce add %one uniform : (f32) -> f32165 166 // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (vector<4xf32>) -> vector<4xf32>167 %sum_subgroup2 = gpu.subgroup_reduce add %vec : (vector<4xf32>) -> vector<4xf32>168 169 %width = arith.constant 7 : i32170 %offset = arith.constant 3 : i32171 // CHECK: gpu.shuffle xor %{{.*}}, %{{.*}}, %{{.*}} : f32172 %shfl, %pred = gpu.shuffle xor %arg0, %offset, %width : f32173 // CHECK: gpu.shuffle up %{{.*}}, %{{.*}}, %{{.*}} : f32174 %shfl1, %pred1 = gpu.shuffle up %arg0, %offset, %width : f32175 // CHECK: gpu.shuffle down %{{.*}}, %{{.*}}, %{{.*}} : f32176 %shfl2, %pred2 = gpu.shuffle down %arg0, %offset, %width : f32177 // CHECK: gpu.shuffle idx %{{.*}}, %{{.*}}, %{{.*}} : f32178 %shfl3, %pred3 = gpu.shuffle idx %arg0, %offset, %width : f32179 180 // CHECK: gpu.rotate %{{.*}}, 3, 16 : f32181 %rotate, %pred4 = gpu.rotate %arg0, 3, 16 : f32182 183 "gpu.barrier"() : () -> ()184 185 "some_op"(%bIdX, %tIdX) : (index, index) -> ()186 %42 = memref.load %arg1[%bIdX] : memref<?xf32, 1>187 gpu.return188 }189 190 gpu.func @kernel_2() kernel {191 gpu.return192 }193 }194 195 gpu.binary @binary_1 [#gpu.object<#nvvm.target, "">]196 197 gpu.binary @binary_2 <#gpu.select_object<#nvvm.target<chip = "sm_90">>> [#gpu.object<#nvvm.target, "">, #gpu.object<#nvvm.target<chip = "sm_90">, "">]198 199 gpu.binary @binary_3 <#gpu.select_object<1>> [#gpu.object<#nvvm.target, "">, #gpu.object<#nvvm.target<chip = "sm_90">, "">]200 201 gpu.binary @binary_4 [#gpu.object<#nvvm.target, bin = "">,202 #gpu.object<#nvvm.target, assembly = "">,203 #gpu.object<#nvvm.target, offload = "">,204 #gpu.object<#nvvm.target, properties = { O = 3 : i32 }, offload = "">205 ]206 207 // Check that fatbin gets ellided as it's the default format.208 // CHECK: gpu.binary @binary_5 [#gpu.object<#nvvm.target, properties = {O = 3 : i32}, "">]209 gpu.binary @binary_5 [#gpu.object<#nvvm.target, properties = {O = 3 : i32}, fatbin = "">]210 211 func.func private @two_value_generator() -> (f32, memref<?xf32, 1>)212 213 func.func @foo() {214 %0 = "op"() : () -> (f32)215 %1 = "op"() : () -> (memref<?xf32, 1>)216 // CHECK: %{{.*}} = arith.constant 8217 %cst = arith.constant 8 : index218 %cstI64 = arith.constant 8 : i64219 %c0 = arith.constant 0 : i32220 %t0 = gpu.wait async221 %lowStream = llvm.mlir.zero : !llvm.ptr222 223 // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)224 gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%0 : f32, %1 : memref<?xf32, 1>)225 226 // CHECK: gpu.launch_func @kernels::@kernel_1 clusters in (%{{.*}}, %{{.*}}, %{{.*}}) blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)227 gpu.launch_func @kernels::@kernel_1 clusters in (%cst, %cst, %cst) blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%0 : f32, %1 : memref<?xf32, 1>)228 229 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>)230 231 // CHECK: gpu.launch_func @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})232 gpu.launch_func @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)233 234 // CHECK: %{{.*}} = gpu.launch_func async [%{{.*}}] @kernels::@kernel_2 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}})235 %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)236 237 // CHECK: gpu.launch_func <%{{.*}} : !llvm.ptr> @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i64 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)238 gpu.launch_func <%lowStream : !llvm.ptr> @kernels::@kernel_1 blocks in (%cstI64, %cstI64, %cstI64) threads in (%cstI64, %cstI64, %cstI64) : i64 args(%0 : f32, %1 : memref<?xf32, 1>)239 240 // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i32 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)241 gpu.launch_func @kernels::@kernel_1 blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) : i32 args(%0 : f32, %1 : memref<?xf32, 1>)242 243 // CHECK: gpu.launch_func @binary_1::@kernel blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i32 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)244 gpu.launch_func @binary_1::@kernel blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) : i32 args(%0 : f32, %1 : memref<?xf32, 1>)245 246 // CHECK: %[[VALUES:.*]]:2 = call247 %values:2 = func.call @two_value_generator() : () -> (f32, memref<?xf32, 1>)248 // CHECK: gpu.launch_func @kernels::@kernel_1 {{.*}} args(%[[VALUES]]#0 : f32, %[[VALUES]]#1 : memref<?xf32, 1>)249 gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%values#0 : f32, %values#1 : memref<?xf32, 1>)250 251 return252 }253 254 gpu.module @gpu_funcs {255 // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32)256 gpu.func @kernel_1(%arg0: f32)257 // CHECK: workgroup(%[[WGROUP1:.*]] : memref<42xf32, 3>, %[[WGROUP2:.*]] : memref<2xf32, 3>)258 workgroup(%arg1: memref<42xf32, 3>, %arg2: memref<2xf32, 3>)259 // CHECK: private(%[[PRIVATE1:.*]] : memref<2xf32, 5>, %[[PRIVATE2:.*]] : memref<1xf32, 5>)260 private(%arg3: memref<2xf32, 5>, %arg4: memref<1xf32, 5>)261 kernel262 // CHECK: attributes {foo = "bar"}263 attributes {foo = "bar"} {264 // CHECK: "use"(%[[WGROUP1]], %[[WGROUP2]])265 "use"(%arg1, %arg2) : (memref<42xf32, 3>, memref<2xf32, 3>) -> ()266 // CHECK: "use"(%[[PRIVATE1]])267 "use"(%arg3) : (memref<2xf32, 5>) -> ()268 // CHECK: "use"(%[[PRIVATE2]])269 "use"(%arg4) : (memref<1xf32, 5>) -> ()270 gpu.return271 }272 273 // CHECK-LABEL: gpu.func @printf_test274 // CHECK: (%[[ARG0:.*]]: i32)275 // CHECK: gpu.printf "Value: %d", %[[ARG0]] : i32276 gpu.func @printf_test(%arg0 : i32) {277 gpu.printf "Value: %d", %arg0 : i32278 gpu.return279 }280 281 // CHECK-LABEL: gpu.func @printf_empty282 // CHECK: gpu.printf "]"283 // CHECK: scf.if284 // CHECK: gpu.printf ", "285 gpu.func @printf_empty(%arg0 : i32) {286 gpu.printf "]"287 %1 = arith.cmpi slt, %arg0, %arg0 : i32288 scf.if %1 {289 gpu.printf ", "290 }291 gpu.return292 }293 294 // CHECK-LABEL: gpu.func @no_attribution295 // CHECK: {296 gpu.func @no_attribution(%arg0: f32) {297 gpu.return298 }299 300 // CHECK-LABEL: @no_attribution_attrs301 // CHECK: attributes302 // CHECK: {303 gpu.func @no_attribution_attrs(%arg0: f32) attributes {foo="bar"} {304 gpu.return305 }306 307 // CHECK-LABEL: @workgroup_only308 // CHECK: workgroup({{.*}}: {{.*}})309 // CHECK: {310 gpu.func @workgroup_only() workgroup(%arg0: memref<42xf32, 3>) {311 gpu.return312 }313 // CHECK-LABEL: @private_only314 // CHECK: private({{.*}}: {{.*}})315 // CHECK: {316 gpu.func @private_only() private(%arg0: memref<2xf32, 5>) {317 gpu.return318 }319 320 // CHECK-LABEL: @empty_attribution321 // CHECK: {322 gpu.func @empty_attribution(%arg0: f32) workgroup() private() {323 gpu.return324 }325 }326 327 gpu.module @explicit_attributions {328 // CHECK-LABEL: gpu.func @kernel_1({{.*}}: f32, {{.*}}: memref<?xf32>) workgroup({{.*}}: memref<5xf32, 3>) private({{.*}}: memref<5xf32, 5>)329 "gpu.func"() ({330 ^bb0(%arg0: f32, %arg1: memref<?xf32>, %arg2: memref<5xf32, 3>, %arg3: memref<5xf32, 5>):331 "gpu.return"() : () -> ()332 } ) {function_type = (f32, memref<?xf32>) -> (), gpu.kernel, sym_name = "kernel_1", workgroup_attributions = 1: i64} : () -> ()333 }334 335 func.func @alloc() {336 // CHECK-LABEL: func @alloc()337 338 // CHECK: %[[m0:.*]] = gpu.alloc () : memref<13xf32, 1>339 %m0 = gpu.alloc () : memref<13xf32, 1>340 // CHECK: gpu.dealloc %[[m0]] : memref<13xf32, 1>341 gpu.dealloc %m0 : memref<13xf32, 1>342 343 %t0 = gpu.wait async344 // CHECK: %[[m1:.*]], %[[t1:.*]] = gpu.alloc async [{{.*}}] () : memref<13xf32, 1>345 %m1, %t1 = gpu.alloc async [%t0] () : memref<13xf32, 1>346 // CHECK: gpu.dealloc async [%[[t1]]] %[[m1]] : memref<13xf32, 1>347 %t2 = gpu.dealloc async [%t1] %m1 : memref<13xf32, 1>348 349 // CHECK: %[[m2:.*]] = gpu.alloc host_shared () : memref<13xf32, 1>350 %m2 = gpu.alloc host_shared () : memref<13xf32, 1>351 // CHECK: gpu.dealloc %[[m2]] : memref<13xf32, 1>352 gpu.dealloc %m2 : memref<13xf32, 1>353 354 return355 }356 357 func.func @async_token(%arg0 : !gpu.async.token) -> !gpu.async.token {358 // CHECK-LABEL: func @async_token({{.*}}: !gpu.async.token)359 // CHECK: return {{.*}} : !gpu.async.token360 return %arg0 : !gpu.async.token361 }362 363 func.func @async_wait() {364 // CHECK-LABEL: func @async_wait365 // CHECK: %[[t0:.*]] = gpu.wait async366 %0 = gpu.wait async367 // CHECK: %[[t1:.*]] = gpu.wait async [%[[t0]]]368 %1 = gpu.wait async [%0]369 // CHECK: %{{.*}} = gpu.wait async [%[[t0]], %[[t1]]]370 %2 = gpu.wait async [%0, %1]371 // CHECK: gpu.wait [%[[t0]], %[[t1]]]372 // CHECK-NOT: async373 gpu.wait [%0, %1]374 // CHECK: gpu.wait375 // CHECK-NOT: async376 gpu.wait // Valid, but a no-op.377 return378 }379 380 func.func @memcpy(%dst : memref<3x7xf32>, %src : memref<3x7xf32, 1>) {381 // CHECK-LABEL: func @memcpy382 // CHECK: gpu.memcpy {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>383 gpu.memcpy %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>384 // CHECK: %[[t0:.*]] = gpu.wait async385 %0 = gpu.wait async386 // CHECK: {{.*}} = gpu.memcpy async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, memref<3x7xf32, 1>387 %1 = gpu.memcpy async [%0] %dst, %src : memref<3x7xf32>, memref<3x7xf32, 1>388 return389 }390 391 func.func @memset(%dst : memref<3x7xf32>, %value : f32) {392 // CHECK-LABEL: func @memset393 // CHECK: gpu.memset {{.*}}, {{.*}} : memref<3x7xf32>, f32394 gpu.memset %dst, %value : memref<3x7xf32>, f32395 // CHECK: %[[t0:.*]] = gpu.wait async396 %0 = gpu.wait async397 // CHECK: {{.*}} = gpu.memset async [%[[t0]]] {{.*}}, {{.*}} : memref<3x7xf32>, f32398 %1 = gpu.memset async [%0] %dst, %value : memref<3x7xf32>, f32399 return400 }401 402 func.func @mmamatrix_valid_scalar_element_type(%src : memref<32x32xf16, affine_map<(d0, d1) -> (d0 * 64 + d1)>>){403 // CHECK-LABEL: func @mmamatrix_valid_scalar_element_type404 %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>405 // CHECK: %[[wg:.*]] = memref.alloca()406 %i = arith.constant 16 : index407 // CHECK: %[[i:.*]] = arith.constant 16 : index408 %cst = arith.constant 1.000000e+00 : f32409 // CHECK: %[[cst:.*]] = arith.constant 1.000000e+00 : f32410 %0 = gpu.subgroup_mma_load_matrix %wg[%i, %i] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">411 // CHECK: gpu.subgroup_mma_load_matrix %[[wg]][%[[i]], %[[i]]] {leadDimension = 32 : index} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">412 %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">413 // CHECK: gpu.subgroup_mma_load_matrix %{{.*}}[%[[i]], %[[i]]] {leadDimension = 64 : index} : memref<32x32xf16, #{{.*}}> -> !gpu.mma_matrix<16x16xf16, "AOp">414 %1 = gpu.subgroup_mma_constant_matrix %cst : !gpu.mma_matrix<16x16xf32, "COp">415 // CHECK: gpu.subgroup_mma_elementwise addf %{{.*}}, %{{.*}} : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">416 %2 = gpu.subgroup_mma_elementwise addf %1, %1 : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">417 // CHECK: gpu.subgroup_mma_elementwise maxf %{{.*}}, %{{.*}} : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">418 %3 = gpu.subgroup_mma_elementwise maxf %2, %1 : (!gpu.mma_matrix<16x16xf32, "COp">, !gpu.mma_matrix<16x16xf32, "COp">) -> !gpu.mma_matrix<16x16xf32, "COp">419 return420 }421 422 // CHECK-LABEL: func @mmamatrix_valid_vector_element_type423 func.func @mmamatrix_valid_vector_element_type(%src : memref<32x4xvector<4xf32>>, %i : index) {424 // CHECK: gpu.subgroup_mma_load_matrix425 %s = gpu.subgroup_mma_load_matrix %src[%i, %i] {leadDimension = 4 : index} : memref<32x4xvector<4xf32>> -> !gpu.mma_matrix<16x16xf16, "COp">426 // CHECK: gpu.subgroup_mma_store_matrix427 gpu.subgroup_mma_store_matrix %s, %src[%i, %i] {leadDimension = 4 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<32x4xvector<4xf32>>428 return429 }430 431 // CHECK-LABEL: func @set_default_device432 func.func @set_default_device(%arg0: i32) {433 // CHECK: gpu.set_default_device434 gpu.set_default_device %arg0435 return436 }437 438 // CHECK-LABEL: func @sparse_ops439 func.func @sparse_ops(%arg0: index) {440 // CHECK: gpu.wait async441 %token0 = gpu.wait async442 // CHECK: gpu.alloc async443 %mem1, %token1 = gpu.alloc async [%token0] (%arg0) : memref<?xindex>444 // CHECK: gpu.alloc async445 %mem2, %token2 = gpu.alloc async [%token1] (%arg0) : memref<?xf64>446 // CHECK: gpu.create_coo async447 %spmat, %token4 = gpu.create_coo async [%token2] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>448 // CHECK: gpu.create_csr async449 %spmat2, %token5 = gpu.create_csr async [%token4] %arg0, %arg0, %arg0, %mem1, %mem1, %mem2 : memref<?xindex>, memref<?xindex>, memref<?xf64>450 // CHECK: gpu.create_dn_tensor async451 %dnvec, %token6 = gpu.create_dn_tensor async [%token5] %mem2, %arg0 : index into memref<?xf64>452 // CHECK: gpu.spmv_buffer_size async453 %bufferSz, %token7 = gpu.spmv_buffer_size async [%token6] %spmat, %dnvec, %dnvec into f64454 // CHECK: gpu.spmv async455 %token8 = gpu.spmv async [%token7] %spmat, %dnvec, %dnvec, %mem2 : memref<?xf64> into f64456 // CHECK: gpu.create_dn_tensor async457 %dnmat, %token9 = gpu.create_dn_tensor async [%token8] %mem2, %arg0, %arg0 : index, index into memref<?xf64>458 // CHECK: gpu.spmm_buffer_size async459 %bufferSz2, %token10 = gpu.spmm_buffer_size async [%token9] %spmat, %dnmat, %dnmat : index into f64460 // CHECK: gpu.spmm async461 %token11 = gpu.spmm async [%token10] %spmat, %dnmat, %dnmat, %mem2 : memref<?xf64> into f64462 // CHECK: gpu.sddmm_buffer_size async463 %bufferSz3, %token12 = gpu.sddmm_buffer_size async [%token11] %dnmat, %dnmat, %spmat into f64464 // CHECK: gpu.sddmm async465 %token13 = gpu.sddmm async [%token12] %dnmat, %dnmat, %spmat, %mem2 : memref<?xf64> into f64466 // CHECK: gpu.destroy_dn_tensor async467 %token14 = gpu.destroy_dn_tensor async [%token13] %dnmat468 // CHECK: gpu.destroy_sp_mat async469 %token15 = gpu.destroy_sp_mat async [%token14] %spmat470 // CHECK: gpu.destroy_dn_tensor async471 %token16 = gpu.destroy_dn_tensor async [%token15] %dnvec472 // CHECK: gpu.wait473 gpu.wait [%token16]474 return475 }476 477 // CHECK-LABEL: func @extract_insert_mma478 func.func @extract_insert_mma(%src : !gpu.mma_matrix<16x16xf32, "COp">,479 %ptr: memref<16x16xf32>) {480 %zero = arith.constant 0.0 : f32481 %c0 = arith.constant 0 : index482 // CHECK: gpu.subgroup_mma_extract_thread_local483 %val = gpu.subgroup_mma_extract_thread_local %src[%c0] : !gpu.mma_matrix<16x16xf32, "COp"> -> f32484 %m = gpu.subgroup_mma_constant_matrix %zero : !gpu.mma_matrix<16x16xf32, "COp">485 // CHECK: gpu.subgroup_mma_insert_thread_local486 %s0 = gpu.subgroup_mma_insert_thread_local %val, %m[%c0] : f32, !gpu.mma_matrix<16x16xf32, "COp"> -> !gpu.mma_matrix<16x16xf32, "COp">487 gpu.subgroup_mma_store_matrix %s0, %ptr[%c0, %c0] {leadDimension = 16 : index} : !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>488 return489 }490}491 492// Just check that this doesn't crash.493gpu.module @module {494 "gpu.func"() ({495 gpu.return496 }) {function_type = () -> (), sym_name = "func"} : () -> ()497}498 499// Check that this doesn't crash.500gpu.module @module_with_one_target [#nvvm.target] {501 gpu.func @kernel(%arg0 : f32) kernel {502 gpu.return503 }504}505 506gpu.module @module_with_two_target [#nvvm.target, #rocdl.target<chip = "gfx90a">] {507 gpu.func @kernel(%arg0 : f32) kernel {508 gpu.return509 }510}511 512gpu.module @module_with_offload_handler <#gpu.select_object<0>> [#nvvm.target] {513}514 515// Test kernel attributes516gpu.binary @kernel_attrs_1 [517 #gpu.object<#rocdl.target<chip = "gfx900">,518 kernels = #gpu.kernel_table<[519 #gpu.kernel_metadata<"kernel0", (i32, f32) -> (), metadata = {sgpr_count = 255}>,520 #gpu.kernel_metadata<"kernel1", (i32) -> (), arg_attrs = [{llvm.read_only}]>521 ]>,522 bin = "BLOB">523 ]524 525// Verify the kernels are sorted526// CHECK-LABEL: gpu.binary @kernel_attrs_2527gpu.binary @kernel_attrs_2 [528 // CHECK: [#gpu.kernel_metadata<"a_kernel", () -> ()>, #gpu.kernel_metadata<"m_kernel", () -> ()>, #gpu.kernel_metadata<"z_kernel", () -> ()>]529 #gpu.object<#rocdl.target<chip = "gfx900">,530 kernels = #gpu.kernel_table<[531 #gpu.kernel_metadata<"z_kernel", () -> ()>,532 #gpu.kernel_metadata<"m_kernel", () -> ()>,533 #gpu.kernel_metadata<"a_kernel", () -> ()>534 ]>,535 bin = "BLOB">536 ]537 538// CHECK-LABEL: func @warp_execute_on_lane_0(539func.func @warp_execute_on_lane_0(%laneid: index) {540// CHECK-NEXT: gpu.warp_execute_on_lane_0(%{{.*}})[32] {541 gpu.warp_execute_on_lane_0(%laneid)[32] {542// CHECK-NEXT: }543 }544// CHECK-NEXT: return545 return546}547 548// CHECK-LABEL: func.func @warp_execute_on_lane_0_2d549func.func @warp_execute_on_lane_0_2d(%laneid: index) {550 // CHECK: gpu.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1x4xi32>)551 %2 = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x4xi32>) {552 %0 = arith.constant dense<2>: vector<4x32xi32>553 // CHECK: gpu.yield %{{.+}} : vector<4x32xi32>554 gpu.yield %0 : vector<4x32xi32>555 }556 return557}558 559// CHECK-LABEL: func @warp_operand_result(560func.func @warp_operand_result(%laneid: index, %v0 : vector<4xi32>) -> (vector<4xi32>) {561// CHECK-NEXT: %{{.*}} = gpu.warp_execute_on_lane_0(%{{.*}})[32] args(%{{.*}} : vector<4xi32>) -> (vector<4xi32>) {562 %2 = gpu.warp_execute_on_lane_0(%laneid)[32]563 args(%v0 : vector<4xi32>) -> (vector<4xi32>) {564 ^bb0(%arg0 : vector<128xi32>) :565 %0 = arith.constant dense<2>: vector<128xi32>566 %1 = arith.addi %arg0, %0 : vector<128xi32>567// CHECK: gpu.yield %{{.*}} : vector<128xi32>568 gpu.yield %1 : vector<128xi32>569// CHECK-NEXT: }570 }571 return %2 : vector<4xi32>572}573 574// CHECK-LABEL: func @subgroup_broadcast575// CHECK-SAME: (%[[ARG:.*]]: f32, %[[IDX:.*]]: i32)576func.func @subgroup_broadcast(%arg0 : f32, %arg1 : i32) -> (f32, f32) {577 // CHECK: gpu.subgroup_broadcast %[[ARG]], first_active_lane : f32578 %0 = gpu.subgroup_broadcast %arg0, first_active_lane : f32579 // CHECK: gpu.subgroup_broadcast %[[ARG]], specific_lane %[[IDX]] : f32580 %1 = gpu.subgroup_broadcast %arg0, specific_lane %arg1 : f32581 func.return %0, %1 : f32, f32582}583