663 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -gpu-kernel-outlining -split-input-file -verify-diagnostics %s | FileCheck %s2// RUN: mlir-opt -allow-unregistered-dialect -gpu-launch-sink-index-computations -gpu-kernel-outlining=data-layout-str='#dlti.dl_spec<#dlti.dl_entry<index,32:i32>>' -split-input-file %s | FileCheck --check-prefix CHECK-DL %s3 4// CHECK: module attributes {gpu.container_module}5 6// CHECK-LABEL: func @launch()7func.func @launch() {8 // CHECK: %[[ARG0:.*]] = "op"() : () -> f329 %0 = "op"() : () -> (f32)10 // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>11 %1 = "op"() : () -> (memref<?xf32, 1>)12 // CHECK: %[[GDIMX:.*]] = arith.constant 813 %gDimX = arith.constant 8 : index14 // CHECK: %[[GDIMY:.*]] = arith.constant 1215 %gDimY = arith.constant 12 : index16 // CHECK: %[[GDIMZ:.*]] = arith.constant 1617 %gDimZ = arith.constant 16 : index18 // CHECK: %[[BDIMX:.*]] = arith.constant 2019 %bDimX = arith.constant 20 : index20 // CHECK: %[[BDIMY:.*]] = arith.constant 2421 %bDimY = arith.constant 24 : index22 // CHECK: %[[BDIMZ:.*]] = arith.constant 2823 %bDimZ = arith.constant 28 : index24 25 // CHECK: gpu.launch_func @launch_kernel::@launch_kernel blocks in (%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]]) threads in (%[[BDIMX]], %[[BDIMY]], %[[BDIMZ]]) args(%[[ARG0]] : f32, %[[ARG1]] : memref<?xf32, 1>)26 // CHECK-NOT: gpu.launch blocks27 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,28 %grid_z = %gDimZ)29 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,30 %block_z = %bDimZ) {31 "use"(%0): (f32) -> ()32 "some_op"(%bx, %block_x) : (index, index) -> ()33 %42 = memref.load %1[%tx] : memref<?xf32, 1>34 gpu.terminator35 }36 return37}38 39// CHECK-DL-LABEL: gpu.module @launch_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}40// CHECK-LABEL: gpu.module @launch_kernel41// CHECK-NEXT: gpu.func @launch_kernel42// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)43// CHECK-SAME: known_block_size = array<i32: 20, 24, 28>44// CHECK-SAME: known_grid_size = array<i32: 8, 12, 16>45// CHECK-NEXT: %[[BID:.*]] = gpu.block_id x46// CHECK-NEXT: = gpu.block_id y47// CHECK-NEXT: = gpu.block_id z48// CHECK-NEXT: %[[TID:.*]] = gpu.thread_id x49// CHECK-NEXT: = gpu.thread_id y50// CHECK-NEXT: = gpu.thread_id z51// CHECK-NEXT: = gpu.grid_dim x52// CHECK-NEXT: = gpu.grid_dim y53// CHECK-NEXT: = gpu.grid_dim z54// CHECK-NEXT: %[[BDIM:.*]] = gpu.block_dim x55// CHECK-NEXT: = gpu.block_dim y56// CHECK-NEXT: = gpu.block_dim z57// CHECK-NEXT: "use"(%[[KERNEL_ARG0]]) : (f32) -> ()58// CHECK-NEXT: "some_op"(%[[BID]], %[[BDIM]]) : (index, index) -> ()59// CHECK-NEXT: = memref.load %[[KERNEL_ARG1]][%[[TID]]] : memref<?xf32, 1>60 61// -----62 63// Verify that we can outline a CFG64// CHECK-LABEL: gpu.func @launchCFG_kernel(65// CHECK: cf.br66// CHECK: gpu.return67func.func @launchCFG() {68 %0 = "op"() : () -> (f32)69 %1 = "op"() : () -> (memref<?xf32, 1>)70 %gDimX = arith.constant 8 : index71 %gDimY = arith.constant 12 : index72 %gDimZ = arith.constant 16 : index73 %bDimX = arith.constant 20 : index74 %bDimY = arith.constant 24 : index75 %bDimZ = arith.constant 28 : index76 77 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,78 %grid_z = %gDimZ)79 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,80 %block_z = %bDimZ) {81 "use"(%0): (f32) -> ()82 cf.br ^bb183 ^bb1:84 "some_op"(%bx, %block_x) : (index, index) -> ()85 %42 = memref.load %1[%tx] : memref<?xf32, 1>86 gpu.terminator87 }88 return89}90 91 92// -----93 94// This test checks gpu-out-lining can handle gpu.launch kernel from an llvm.func95// CHECK-LABEL: @launch_from_llvm_func96llvm.func @launch_from_llvm_func() {97 // CHECK: %[[ARG0:.*]] = "op"() : () -> f3298 %0 = "op"() : () -> (f32)99 // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>100 %1 = "op"() : () -> (memref<?xf32, 1>)101 102 // CHECK: %[[DIM:.*]] = arith.constant 1103 %dim = arith.constant 1 : index104 105 // CHECK: gpu.launch_func @launch_from_llvm_func_kernel::@launch_from_llvm_func_kernel106 // CHECK-SAME: (%[[DIM]], %[[DIM]], %[[DIM]])107 // CHECK-SAME: (%[[DIM]], %[[DIM]], %[[DIM]]) args(%[[ARG0]] : f32, %[[ARG1]] : memref<?xf32, 1>)108 // CHECK-NEXT: llvm.return109 110 // CHECK: gpu.func {{.*}} kernel attributes111 // CHECK-SAME: known_block_size = array<i32: 1, 1, 1>112 // CHECK-SAME: known_grid_size = array<i32: 1, 1, 1>113 // CHECK: gpu.return114 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %dim, %grid_y = %dim,115 %grid_z = %dim)116 threads(%tx, %ty, %tz) in (%block_x = %dim, %block_y = %dim,117 %block_z = %dim) {118 "use"(%0): (f32) -> ()119 "some_op"(%bx, %block_x) : (index, index) -> ()120 %2 = memref.load %1[%tx] : memref<?xf32, 1>121 gpu.terminator122 }123 llvm.return124}125 126// CHECK-DL-LABEL: gpu.module @launch_from_llvm_func_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}127 128// -----129 130// CHECK: module attributes {gpu.container_module}131// CHECK-LABEL: @multiple_launches132func.func @multiple_launches() {133 // CHECK: %[[CST:.*]] = arith.constant 8 : index134 %cst = arith.constant 8 : index135 // CHECK: gpu.launch_func @multiple_launches_kernel::@multiple_launches_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]])136 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,137 %grid_z = %cst)138 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,139 %block_z = %cst) {140 gpu.terminator141 }142 // CHECK: gpu.launch_func @multiple_launches_kernel_0::@multiple_launches_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]])143 gpu.launch blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,144 %grid_z2 = %cst)145 threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,146 %block_z2 = %cst) {147 gpu.terminator148 }149 150 // With async and async deps.151 // CHECK: %[[TOKEN:.*]] = gpu.wait async152 // CHECK: gpu.launch_func async [%[[TOKEN]]] @multiple_launches_kernel_1::@multiple_launches_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]])153 %t = gpu.wait async154 %u = gpu.launch async [%t] blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,155 %grid_z2 = %cst)156 threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,157 %block_z2 = %cst) {158 gpu.terminator159 }160 161 // CHECK: gpu.launch_func async @multiple_launches_kernel_2::@multiple_launches_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]])162 %v = gpu.launch async blocks(%bx2, %by2, %bz2) in (%grid_x2 = %cst, %grid_y2 = %cst,163 %grid_z2 = %cst)164 threads(%tx2, %ty2, %tz2) in (%block_x2 = %cst, %block_y2 = %cst,165 %block_z2 = %cst) {166 gpu.terminator167 }168 169 return170}171 172// CHECK-DL-LABEL: gpu.module @multiple_launches_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}173// CHECK-DL-LABEL: gpu.module @multiple_launches_kernel_0 attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}174 175// CHECK: gpu.module @multiple_launches_kernel176// CHECK: func @multiple_launches_kernel177// CHECK: module @multiple_launches_kernel_0178// CHECK: func @multiple_launches_kernel179 180// -----181 182// CHECK-LABEL: @extra_constants_not_inlined183func.func @extra_constants_not_inlined(%arg0: memref<?xf32>) {184 // CHECK: %[[CST:.*]] = arith.constant 8 : index185 %cst = arith.constant 8 : index186 %cst2 = arith.constant 2 : index187 %c0 = arith.constant 0 : index188 %cst3 = "secret_constant"() : () -> index189 // CHECK: gpu.launch_func @extra_constants_not_inlined_kernel::@extra_constants_not_inlined_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]]) args({{.*}} : memref<?xf32>, {{.*}} : index)190 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,191 %grid_z = %cst)192 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,193 %block_z = %cst) {194 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()195 gpu.terminator196 }197 return198}199 200// CHECK-DL-LABEL: gpu.module @extra_constants_not_inlined_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}201 202// CHECK-LABEL: func @extra_constants_not_inlined_kernel(%{{.*}}: memref<?xf32>, %{{.*}}: index)203// CHECK: arith.constant 2204 205// -----206 207// CHECK-LABEL: @extra_constants208// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>209func.func @extra_constants(%arg0: memref<?xf32>) {210 // CHECK: %[[CST:.*]] = arith.constant 8 : index211 %cst = arith.constant 8 : index212 %cst2 = arith.constant 2 : index213 %c0 = arith.constant 0 : index214 %cst3 = memref.dim %arg0, %c0 : memref<?xf32>215 // CHECK: gpu.launch_func @extra_constants_kernel::@extra_constants_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]]) args(%[[ARG0]] : memref<?xf32>)216 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,217 %grid_z = %cst)218 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,219 %block_z = %cst) {220 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()221 gpu.terminator222 }223 return224}225 226// CHECK-DL-LABEL: gpu.module @extra_constants_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}227 228// CHECK-LABEL: func @extra_constants_kernel(229// CHECK-SAME: %[[KARG0:.*]]: memref<?xf32>230// CHECK: arith.constant 2231// CHECK: arith.constant 0232// CHECK: memref.dim %[[KARG0]]233 234// -----235 236// CHECK-LABEL: @extra_constants_noarg237// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>238func.func @extra_constants_noarg(%arg0: memref<?xf32>, %arg1: memref<?xf32>) {239 // CHECK: %[[CST:.*]] = arith.constant 8 : index240 %cst = arith.constant 8 : index241 %cst2 = arith.constant 2 : index242 %c0 = arith.constant 0 : index243 // CHECK: memref.dim %[[ARG1]]244 %cst3 = memref.dim %arg1, %c0 : memref<?xf32>245 // CHECK: gpu.launch_func @extra_constants_noarg_kernel::@extra_constants_noarg_kernel blocks in (%[[CST]], %[[CST]], %[[CST]]) threads in (%[[CST]], %[[CST]], %[[CST]]) args(%[[ARG0]] : memref<?xf32>, {{.*}} : index)246 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,247 %grid_z = %cst)248 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,249 %block_z = %cst) {250 "use"(%cst2, %arg0, %cst3) : (index, memref<?xf32>, index) -> ()251 gpu.terminator252 }253 return254}255 256// CHECK-DL-LABEL: gpu.module @extra_constants_noarg_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}257 258// CHECK-LABEL: func @extra_constants_noarg_kernel(259// CHECK-SAME: %[[KARG0:.*]]: memref<?xf32>, %[[KARG1:.*]]: index260// CHECK: %[[KCST:.*]] = arith.constant 2261// CHECK: "use"(%[[KCST]], %[[KARG0]], %[[KARG1]])262 263// -----264 265// CHECK-LABEL: @multiple_uses266func.func @multiple_uses(%arg0 : memref<?xf32>) {267 %c1 = arith.constant 1 : index268 %c2 = arith.constant 2 : index269 // CHECK: gpu.func {{.*}} {270 // CHECK: %[[C2:.*]] = arith.constant 2 : index271 // CHECK: "use1"(%[[C2]], %[[C2]])272 // CHECK: "use2"(%[[C2]])273 // CHECK: gpu.return274 // CHECK: }275 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,276 %grid_z = %c1)277 threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,278 %block_z = %c1) {279 "use1"(%c2, %c2) : (index, index) -> ()280 "use2"(%c2) : (index) -> ()281 gpu.terminator282 }283 return284}285 286// CHECK-DL-LABEL: gpu.module @multiple_uses_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}287 288// -----289 290// CHECK-LABEL: @multiple_uses2291func.func @multiple_uses2(%arg0 : memref<*xf32>) {292 %c1 = arith.constant 1 : index293 %c2 = arith.constant 2 : index294 %d = memref.dim %arg0, %c2 : memref<*xf32>295 // CHECK: gpu.func {{.*}} {296 // CHECK: %[[C2:.*]] = arith.constant 2 : index297 // CHECK: %[[D:.*]] = memref.dim %[[ARG:.*]], %[[C2]]298 // CHECK: "use1"(%[[D]])299 // CHECK: "use2"(%[[C2]], %[[C2]])300 // CHECK: "use3"(%[[ARG]])301 // CHECK: gpu.return302 // CHECK: }303 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,304 %grid_z = %c1)305 threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,306 %block_z = %c1) {307 "use1"(%d) : (index) -> ()308 "use2"(%c2, %c2) : (index, index) -> ()309 "use3"(%arg0) : (memref<*xf32>) -> ()310 gpu.terminator311 }312 return313}314 315// CHECK-DL-LABEL: gpu.module @multiple_uses2_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}316 317// -----318 319llvm.mlir.global internal @global(42 : i64) : i64320 321//CHECK-LABEL: @function_call322func.func @function_call(%arg0 : memref<?xf32>) {323 %cst = arith.constant 8 : index324 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,325 %grid_z = %cst)326 threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,327 %block_z = %cst) {328 func.call @device_function() : () -> ()329 func.call @device_function() : () -> ()330 %0 = llvm.mlir.addressof @global : !llvm.ptr331 gpu.terminator332 }333 return334}335 336func.func @device_function() {337 call @recursive_device_function() : () -> ()338 return339}340 341func.func @recursive_device_function() {342 call @recursive_device_function() : () -> ()343 return344}345 346// CHECK-DL-LABEL: gpu.module @function_call_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}347 348// CHECK: gpu.module @function_call_kernel {349// CHECK: gpu.func @function_call_kernel()350// CHECK: call @device_function() : () -> ()351// CHECK: call @device_function() : () -> ()352// CHECK: llvm.mlir.addressof @global : !llvm.ptr353// CHECK: gpu.return354//355// CHECK: llvm.mlir.global internal @global(42 : i64) {addr_space = 0 : i32} : i64356//357// CHECK: func @device_function()358// CHECK: func @recursive_device_function()359// CHECK-NOT: func @device_function360 361// -----362 363// CHECK-LABEL: @non_constant_launches364func.func @non_constant_launches(%arg0 : index) {365 // CHECK-NOT: known_block_size366 // CHECK-NOT: known_grid_size367 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %arg0, %grid_y = %arg0,368 %grid_z = %arg0)369 threads(%tx, %ty, %tz) in (%block_x = %arg0, %block_y = %arg0,370 %block_z = %arg0) {371 gpu.terminator372 }373 return374}375 376// CHECK-DL-LABEL: gpu.module @non_constant_launches_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}377 378// CHECK: module attributes {gpu.container_module}379 380// -----381 382// This test checks memory attributions for gpu.launch, using both workgroup and private attributions.383// CHECK-LABEL: func @launch_memory_attributions_0()384func.func @launch_memory_attributions_0() {385 %1 = "op"() : () -> (memref<?xf32, 1>)386 %128 = arith.constant 128 : index387 388 // CHECK: gpu.launch_func @launch_memory_attributions_0_kernel::@launch_memory_attributions_0_kernel389 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %128, %grid_y = %128,390 %grid_z = %128)391 threads(%tx, %ty, %tz) in (%block_x = %128, %block_y = %128,392 %block_z = %128)393 workgroup(%shared: memref<42xf32, 3>)394 private(%priv0: memref<2xf32, 5>, %priv1: memref<1xf32, 5>) {395 "some_op"(%bx, %block_x) : (index, index) -> ()396 %42 = memref.load %1[%tx] : memref<?xf32, 1>397 %43 = memref.load %shared[%tx] : memref<42xf32, 3>398 %44 = memref.load %priv1[%tx] : memref<1xf32, 5>399 gpu.terminator400 }401 return402}403 404// CHECK-DL-LABEL: gpu.module @launch_memory_attributions_0_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}405 406// CHECK-LABEL: gpu.module @launch_memory_attributions_0_kernel407// CHECK-NEXT: gpu.func @launch_memory_attributions_0_kernel408// CHECK-SAME: workgroup(%[[KERNEL_ARG1:.*]] : memref<42xf32, 3>)409// CHECK-SAME: private(%[[KERNEL_ARG2:.*]] : memref<2xf32, 5>, %[[KERNEL_ARG3:.*]] : memref<1xf32, 5>)410// CHECK: %[[TID:.*]] = gpu.thread_id x411// CHECK: = memref.load %[[KERNEL_ARG1]][%[[TID]]] : memref<42xf32, 3>412// CHECK-NEXT: = memref.load %[[KERNEL_ARG3]][%[[TID]]] : memref<1xf32, 5>413 414// -----415 416// This test checks correctness of private attributions in the absence of workgroup attributions.417// CHECK-LABEL: @launch_memory_attributions_1418func.func @launch_memory_attributions_1(%arg0 : memref<*xf32>) {419 %c1 = arith.constant 1 : index420 %c2 = arith.constant 2 : index421 %d = memref.dim %arg0, %c2 : memref<*xf32>422 // CHECK: gpu.func {{.*}} private(%[[KERNEL_ARG:.*]] : memref<3xf32, 5>) {{.*}} {423 // CHECK: %[[C2:.*]] = arith.constant 2 : index424 // CHECK: = memref.load %[[KERNEL_ARG]][%[[C2]]] : memref<3xf32, 5>425 // CHECK: gpu.return426 // CHECK: }427 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,428 %grid_z = %c1)429 threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,430 %block_z = %c1)431 private(%priv0: memref<3xf32, 5>) {432 %42 = memref.load %priv0[%c2] : memref<3xf32, 5>433 gpu.terminator434 }435 return436}437 438// CHECK-DL-LABEL: gpu.module @launch_memory_attributions_1_kernel attributes {dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>}439 440// -----441// CHECK: module attributes {gpu.container_module}442 443// CHECK-LABEL: func @launch_cluster()444func.func @launch_cluster() {445 // CHECK: %[[ARG0:.*]] = "op"() : () -> f32446 %0 = "op"() : () -> (f32)447 // CHECK: %[[ARG1:.*]] = "op"() : () -> memref<?xf32, 1>448 %1 = "op"() : () -> (memref<?xf32, 1>)449 // CHECK: %[[CDIMX:.*]] = arith.constant 1450 %cDimX = arith.constant 1 : index451 // CHECK: %[[CDIMY:.*]] = arith.constant 2452 %cDimY = arith.constant 2 : index453 // CHECK: %[[CDIMZ:.*]] = arith.constant 1454 %cDimZ = arith.constant 1 : index455 // CHECK: %[[GDIMX:.*]] = arith.constant 8456 %gDimX = arith.constant 8 : index457 // CHECK: %[[GDIMY:.*]] = arith.constant 12458 %gDimY = arith.constant 12 : index459 // CHECK: %[[GDIMZ:.*]] = arith.constant 16460 %gDimZ = arith.constant 16 : index461 // CHECK: %[[BDIMX:.*]] = arith.constant 20462 %bDimX = arith.constant 20 : index463 // CHECK: %[[BDIMY:.*]] = arith.constant 24464 %bDimY = arith.constant 24 : index465 // CHECK: %[[BDIMZ:.*]] = arith.constant 28466 %bDimZ = arith.constant 28 : index467 468 // CHECK: gpu.launch_func @launch_cluster_kernel::@launch_cluster_kernel clusters in (%[[CDIMX]], %[[CDIMY]], %[[CDIMZ]]) blocks in (%[[GDIMX]], %[[GDIMY]], %[[GDIMZ]]) threads in (%[[BDIMX]], %[[BDIMY]], %[[BDIMZ]]) args(%[[ARG0]] : f32, %[[ARG1]] : memref<?xf32, 1>)469 // CHECK-NOT: gpu.launch blocks470 gpu.launch clusters(%cx, %cy, %cz) in (%cluster_x = %cDimX, %cluster_y = %cDimY,471 %cluster_z = %cDimZ)472 blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY,473 %grid_z = %gDimZ)474 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY,475 %block_z = %bDimZ) {476 "use"(%0): (f32) -> ()477 "some_op"(%cx, %bx, %block_x) : (index, index, index) -> ()478 %42 = memref.load %1[%tx] : memref<?xf32, 1>479 gpu.terminator480 }481 return482}483 484// CHECK-LABEL: gpu.module @launch_cluster_kernel485// CHECK-NEXT: gpu.func @launch_cluster_kernel486// CHECK-SAME: (%[[KERNEL_ARG0:.*]]: f32, %[[KERNEL_ARG1:.*]]: memref<?xf32, 1>)487// CHECK-SAME: known_block_size = array<i32: 20, 24, 28>488// CHECK-SAME: known_grid_size = array<i32: 8, 12, 16>489// CHECK-NEXT: %[[BID:.*]] = gpu.block_id x490// CHECK-NEXT: = gpu.block_id y491// CHECK-NEXT: = gpu.block_id z492// CHECK-NEXT: %[[TID:.*]] = gpu.thread_id x493// CHECK-NEXT: = gpu.thread_id y494// CHECK-NEXT: = gpu.thread_id z495// CHECK-NEXT: = gpu.grid_dim x496// CHECK-NEXT: = gpu.grid_dim y497// CHECK-NEXT: = gpu.grid_dim z498// CHECK-NEXT: %[[BDIM:.*]] = gpu.block_dim x499// CHECK-NEXT: = gpu.block_dim y500// CHECK-NEXT: = gpu.block_dim z501// CHECK-NEXT: %[[CID:.*]] = gpu.cluster_id x502// CHECK-NEXT: = gpu.cluster_id y503// CHECK-NEXT: = gpu.cluster_id z504// CHECK-NEXT: %[[CDIM:.*]] = gpu.cluster_dim x505// CHECK-NEXT: = gpu.cluster_dim y506// CHECK-NEXT: = gpu.cluster_dim z507// CHECK-NEXT: "use"(%[[KERNEL_ARG0]]) : (f32) -> ()508// CHECK-NEXT: "some_op"(%[[CID]], %[[BID]], %[[BDIM]]) : (index, index, index) -> ()509// CHECK-NEXT: = memref.load %[[KERNEL_ARG1]][%[[TID]]] : memref<?xf32, 1>510 511// -----512// This test tests the two optional attributes `module` and `function` for gpu.launch513// CHECK-LABEL: func.func @testKernelAttributes()514// CHECK: gpu.launch_func @test_module::@test_kernel_func blocks in (%[[GRID_X:.*]], %[[GRID_Y:.*]], %[[GRID_Z:.*]]) threads in (%[[BLOCK_X:.*]], %[[BLOCK_Y:.*]], %[[BLOCK_Z:.*]])515// CHECK: gpu.module @test_module516// CHECK: gpu.func @test_kernel_func()517func.func @testKernelAttributes() {518 %gDimX = arith.constant 8 : index519 %gDimY = arith.constant 12 : index520 %gDimZ = arith.constant 16 : index521 %bDimX = arith.constant 32 : index522 %bDimY = arith.constant 16 : index523 %bDimZ = arith.constant 8 : index524 525 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, %grid_z = %gDimZ)526 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, %block_z = %bDimZ)527 module(@test_module) function(@test_kernel_func) {528 "some_op"(%bx, %tx) : (index, index) -> ()529 gpu.terminator530 }531 return532}533 534// -----535// This test tests the two optional attributes `module` and `function` for gpu.launch, when kernelModule already exists.536 537// CHECK-LABEL: gpu.module @existing_module538// CHECK: gpu.func @test_kernel_func()539// CHECK: gpu.func @test_kernel_func_0()540// CHECK-NOT: gpu.module @testExistingModule_kernel541// CHECK-NOT: gpu.func @testExistingModule_kernel()542// CHECK: func.func @testExistingModule()543// CHECK: gpu.launch_func @existing_module::@test_kernel_func_0 blocks in (%[[GRID_X:.*]], %[[GRID_Y:.*]], %[[GRID_Z:.*]]) threads in (%[[BLOCK_X:.*]], %[[BLOCK_Y:.*]], %[[BLOCK_Z:.*]])544 545gpu.module @existing_module {546 gpu.func @test_kernel_func() {547 gpu.return548 }549}550 551func.func @testExistingModule() {552 %gDimX = arith.constant 8 : index553 %gDimY = arith.constant 12 : index554 %gDimZ = arith.constant 16 : index555 %bDimX = arith.constant 32 : index556 %bDimY = arith.constant 16 : index557 %bDimZ = arith.constant 8 : index558 559 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, %grid_z = %gDimZ)560 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, %block_z = %bDimZ)561 module(@existing_module) function(@test_kernel_func) {562 "some_op"(%bx, %tx) : (index, index) -> ()563 gpu.terminator564 }565 return566}567 568// -----569// This test tests the optional attribute `module` for gpu.launch.570// CHECK-LABEL: func.func @testKernelModuleOnly()571// CHECK: gpu.launch_func @test_module::@testKernelModuleOnly_kernel blocks in (%[[GRID_X:.*]], %[[GRID_Y:.*]], %[[GRID_Z:.*]]) threads in (%[[BLOCK_X:.*]], %[[BLOCK_Y:.*]], %[[BLOCK_Z:.*]])572// CHECK: gpu.module @test_module573// CHECK: gpu.func @testKernelModuleOnly_kernel()574func.func @testKernelModuleOnly() {575 %gDimX = arith.constant 8 : index576 %gDimY = arith.constant 12 : index577 %gDimZ = arith.constant 16 : index578 %bDimX = arith.constant 32 : index579 %bDimY = arith.constant 16 : index580 %bDimZ = arith.constant 8 : index581 582 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, %grid_z = %gDimZ)583 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, %block_z = %bDimZ)584 module(@test_module) {585 "some_op"(%bx, %tx) : (index, index) -> ()586 gpu.terminator587 }588 return589}590 591// -----592// This test tests the optional attribute `function` for gpu.launch.593// CHECK-LABEL: func.func @testKernelFuncOnly()594// CHECK: gpu.launch_func @test_kernel_func::@test_kernel_func blocks in (%[[GRID_X:.*]], %[[GRID_Y:.*]], %[[GRID_Z:.*]]) threads in (%[[BLOCK_X:.*]], %[[BLOCK_Y:.*]], %[[BLOCK_Z:.*]])595 596// CHECK: gpu.module @test_kernel_func597// CHECK: gpu.func @test_kernel_func()598func.func @testKernelFuncOnly() {599 %gDimX = arith.constant 8 : index600 %gDimY = arith.constant 12 : index601 %gDimZ = arith.constant 16 : index602 %bDimX = arith.constant 32 : index603 %bDimY = arith.constant 16 : index604 %bDimZ = arith.constant 8 : index605 606 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, %grid_z = %gDimZ)607 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, %block_z = %bDimZ)608 function(@test_kernel_func) {609 "some_op"(%bx, %tx) : (index, index) -> ()610 gpu.terminator611 }612 return613}614 615// -----616// This test tests gpu.launch when optional attributes `module` and `function` are not specified.617// CHECK-LABEL: func.func @testNoAttributes()618// CHECK: gpu.launch_func @testNoAttributes_kernel::@testNoAttributes_kernel blocks in (%[[GRID_X:.*]], %[[GRID_Y:.*]], %[[GRID_Z:.*]]) threads in (%[[BLOCK_X:.*]], %[[BLOCK_Y:.*]], %[[BLOCK_Z:.*]])619 620// CHECK: gpu.module @testNoAttributes_kernel621// CHECK: gpu.func @testNoAttributes_kernel()622func.func @testNoAttributes() {623 %gDimX = arith.constant 8 : index624 %gDimY = arith.constant 12 : index625 %gDimZ = arith.constant 16 : index626 %bDimX = arith.constant 32 : index627 %bDimY = arith.constant 16 : index628 %bDimZ = arith.constant 8 : index629 630 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %gDimX, %grid_y = %gDimY, %grid_z = %gDimZ)631 threads(%tx, %ty, %tz) in (%block_x = %bDimX, %block_y = %bDimY, %block_z = %bDimZ) {632 "some_op"(%bx, %tx) : (index, index) -> ()633 gpu.terminator634 }635 return636}637 638// -----639 640// This test tests nested `gpu.launch`.641 642// CHECK-LABEL: func.func @nested_launch(643// CHECK-SAME: %[[ARG0:.*]]: index) {644// CHECK: gpu.launch_func @nested_launch_kernel_0::@nested_launch_kernel blocks in (%[[ARG0]], %[[ARG0]], %[[ARG0]]) threads in (%[[ARG0]], %[[ARG0]], %[[ARG0]]) args(%[[ARG0]] : index)645// CHECK: gpu.module @nested_launch_kernel646// CHECK: gpu.func @nested_launch_kernel() kernel647// CHECK: "some_op"648// CHECK: gpu.module @nested_launch_kernel_0649// CHECK: gpu.func @nested_launch_kernel(%[[VAL_0:.*]]: index) kernel650// CHECK: gpu.launch_func @nested_launch_kernel::@nested_launch_kernel blocks in (%[[VAL_0]], %[[VAL_0]], %[[VAL_0]]) threads in (%[[VAL_0]], %[[VAL_0]], %[[VAL_0]])651func.func @nested_launch(%sz : index) {652 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)653 threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {654 gpu.launch blocks(%bx1, %by1, %bz1) in (%grid_x1 = %sz, %grid_y1 = %sz, %grid_z1 = %sz)655 threads(%tx1, %ty1, %tz1) in (%block_x1 = %sz, %block_y1 = %sz, %block_z1 = %sz) {656 "some_op"(%bx1, %tx1) : (index, index) -> ()657 gpu.terminator658 }659 gpu.terminator660 }661 return662}663