825 lines · plain
1// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx950' -split-input-file | FileCheck %s2// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx950 allowed-dialects=func,arith,math' -split-input-file | FileCheck %s3// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx950 index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s4 5// CHECK-LABEL: @test_module6// CHECK-SAME: llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"7 8gpu.module @test_module {9 // CHECK-LABEL: func @gpu_index_ops()10 // CHECK32-LABEL: func @gpu_index_ops()11 func.func @gpu_index_ops()12 -> (index, index, index, index, index, index,13 index, index, index, index, index, index,14 index, index, index) {15 // CHECK32-NOT: = llvm.sext %{{.*}} : i32 to i6416 17 // CHECK: rocdl.workitem.id.x : i3218 // CHECK: = llvm.sext %{{.*}} : i32 to i6419 %tIdX = gpu.thread_id x20 // CHECK: rocdl.workitem.id.y : i3221 // CHECK: = llvm.sext %{{.*}} : i32 to i6422 %tIdY = gpu.thread_id y23 // CHECK: rocdl.workitem.id.z : i3224 // CHECK: = llvm.sext %{{.*}} : i32 to i6425 %tIdZ = gpu.thread_id z26 27 // CHECK: rocdl.workgroup.dim.x : i3228 // CHECK: = llvm.sext %{{.*}} : i32 to i6429 %bDimX = gpu.block_dim x30 // CHECK: rocdl.workgroup.dim.y : i3231 // CHECK: = llvm.sext %{{.*}} : i32 to i6432 %bDimY = gpu.block_dim y33 // CHECK: rocdl.workgroup.dim.z : i3234 // CHECK: = llvm.sext %{{.*}} : i32 to i6435 %bDimZ = gpu.block_dim z36 37 // CHECK: rocdl.workgroup.id.x : i3238 // CHECK: = llvm.sext %{{.*}} : i32 to i6439 %bIdX = gpu.block_id x40 // CHECK: rocdl.workgroup.id.y : i3241 // CHECK: = llvm.sext %{{.*}} : i32 to i6442 %bIdY = gpu.block_id y43 // CHECK: rocdl.workgroup.id.z : i3244 // CHECK: = llvm.sext %{{.*}} : i32 to i6445 %bIdZ = gpu.block_id z46 47 // CHECK: rocdl.grid.dim.x : i3248 // CHECK: = llvm.sext %{{.*}} : i32 to i6449 %gDimX = gpu.grid_dim x50 // CHECK: rocdl.grid.dim.y : i3251 // CHECK: = llvm.sext %{{.*}} : i32 to i6452 %gDimY = gpu.grid_dim y53 // CHECK: rocdl.grid.dim.z : i3254 // CHECK: = llvm.sext %{{.*}} : i32 to i6455 %gDimZ = gpu.grid_dim z56 57 // CHECK: = rocdl.mbcnt.lo %{{.*}}, %{{.*}} {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 32>}]} : (i32, i32) -> i3258 // CHECK: = rocdl.mbcnt.hi %{{.*}}, %{{.*}} {res_attrs = [{llvm.noundef, llvm.range = #llvm.constant_range<i32, 0, 64>}]} : (i32, i32) -> i3259 // CHECK: = llvm.sext %{{.*}} : i32 to i6460 %laneId = gpu.lane_id61 62 // CHECK: = rocdl.wavefrontsize : i3263 // CHECK: = llvm.sext %{{.*}} : i32 to i6464 %subgroupSize = gpu.subgroup_size : index65 66 // CHECK: = rocdl.wavefrontsize range <i32, 64, 65> : i3267 // CHECK: = llvm.sext %{{.*}} : i32 to i6468 %subgroupSize2 = gpu.subgroup_size upper_bound 64 : index69 70 func.return %tIdX, %tIdY, %tIdZ, %bDimX, %bDimY, %bDimZ,71 %bIdX, %bIdY, %bIdZ, %gDimX, %gDimY, %gDimZ,72 %laneId, %subgroupSize, %subgroupSize273 : index, index, index, index, index, index,74 index, index, index, index, index, index,75 index, index, index76 }77}78 79// -----80 81gpu.module @test_module {82 // CHECK-LABEL: func @gpu_index_ops_range83 // CHECK-SAME: rocdl.flat_work_group_size = "1536,1536"84 // CHECK-SAME: rocdl.reqd_work_group_size = array<i32: 8, 12, 16>85 gpu.func @gpu_index_ops_range(%place: memref<i32>) kernel attributes86 {known_block_size = array<i32: 8, 12, 16>,87 known_grid_size = array<i32: 20, 24, 28>} {88 89 // CHECK: rocdl.workitem.id.x range <i32, 0, 8> : i3290 %tIdX = gpu.thread_id x91 // CHECK: rocdl.workitem.id.y range <i32, 0, 12> : i3292 %tIdY = gpu.thread_id y93 // CHECK: rocdl.workitem.id.z range <i32, 0, 16> : i3294 %tIdZ = gpu.thread_id z95 96 // CHECK: rocdl.workgroup.id.x range <i32, 0, 20> : i3297 %bIdX = gpu.block_id x98 // CHECK: rocdl.workgroup.id.y range <i32, 0, 24> : i3299 %bIdY = gpu.block_id y100 // CHECK: rocdl.workgroup.id.z range <i32, 0, 28> : i32101 %bIdZ = gpu.block_id z102 103 // "Usage" to make the ID calls not die104 %0 = arith.addi %tIdX, %tIdY : index105 %1 = arith.addi %0, %tIdZ : index106 %2 = arith.addi %1, %bIdX : index107 %3 = arith.addi %2, %bIdY : index108 %4 = arith.addi %3, %bIdZ : index109 %5 = arith.index_cast %4 : index to i32110 memref.store %5, %place[] : memref<i32>111 gpu.return112 }113}114 115// -----116 117gpu.module @test_module {118 // CHECK-LABEL: func @gpu_index_comp119 // CHECK32-LABEL: func @gpu_index_comp120 func.func @gpu_index_comp(%idx : index) -> index {121 // CHECK: = llvm.add %{{.*}}, %{{.*}} : i64122 // CHECK32: = llvm.add %{{.*}}, %{{.*}} : i32123 %0 = arith.addi %idx, %idx : index124 // CHECK: llvm.return %{{.*}} : i64125 // CHECK32: llvm.return %{{.*}} : i32126 func.return %0 : index127 }128}129 130// -----131 132gpu.module @test_module {133 // CHECK-LABEL: func @gpu_sync()134 func.func @gpu_sync() {135 // CHECK: rocdl.barrier136 gpu.barrier137 func.return138 }139}140 141// -----142 143gpu.module @test_module {144 // CHECK-LABEL: func @gpu_sqrt145 func.func @gpu_sqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {146 %result16 = math.sqrt %arg_f16 : f16147 // CHECK: llvm.intr.sqrt(%{{.*}}) : (f16) -> f16148 %result32 = math.sqrt %arg_f32 : f32149 // CHECK: llvm.intr.sqrt(%{{.*}}) : (f32) -> f32150 %result64 = math.sqrt %arg_f64 : f64151 // CHECK: llvm.intr.sqrt(%{{.*}}) : (f64) -> f64152 func.return %result16, %result32, %result64 : f16, f32, f64153 }154}155 156// -----157 158gpu.module @test_module {159 // CHECK-LABEL: func @gpu_fabs160 func.func @gpu_fabs(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {161 %result16 = math.absf %arg_f16 : f16162 // CHECK: llvm.intr.fabs(%{{.*}}) : (f16) -> f16163 %result32 = math.absf %arg_f32 : f32164 // CHECK: llvm.intr.fabs(%{{.*}}) : (f32) -> f32165 %result64 = math.absf %arg_f64 : f64166 // CHECK: llvm.intr.fabs(%{{.*}}) : (f64) -> f64167 func.return %result16, %result32, %result64 : f16, f32, f64168 }169}170 171// -----172 173gpu.module @test_module {174 // CHECK: llvm.func @__ocml_exp_f16(f16) -> f16175 // CHECK: llvm.func @__ocml_exp_f64(f64) -> f64176 // CHECK-LABEL: func @gpu_exp177 func.func @gpu_exp(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {178 %result16 = math.exp %arg_f16 : f16179 // CHECK: llvm.call @__ocml_exp_f16(%{{.*}}) : (f16) -> f16180 %result32 = math.exp %arg_f32 : f32181 // CHECK: llvm.intr.exp(%{{.*}}) : (f32) -> f32182 %result64 = math.exp %arg_f64 : f64183 // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (f64) -> f64184 func.return %result16, %result32, %result64 : f16, f32, f64185 }186}187 188// -----189 190gpu.module @test_module {191 // CHECK: llvm.func @__ocml_log_f16(f16) -> f16192 // CHECK: llvm.func @__ocml_log_f64(f64) -> f64193 // CHECK-LABEL: func @gpu_log194 func.func @gpu_log(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {195 %result16 = math.log %arg_f16 : f16196 // CHECK: llvm.call @__ocml_log_f16(%{{.*}}) : (f16) -> f16197 %result32 = math.log %arg_f32 : f32198 // CHECK: llvm.intr.log(%{{.*}}) : (f32) -> f32199 %result64 = math.log %arg_f64 : f64200 // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64201 func.return %result16, %result32, %result64 : f16, f32, f64202 }203}204 205// -----206 207gpu.module @test_module {208 // CHECK: llvm.func @__ocml_cbrt_f16(f16) -> f16209 // CHECK: llvm.func @__ocml_cbrt_f32(f32) -> f32210 // CHECK: llvm.func @__ocml_cbrt_f64(f64) -> f64211 // CHECK-LABEL: func @gpu_cbrt212 func.func @gpu_cbrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {213 %result16 = math.cbrt %arg_f16 : f16214 // CHECK: llvm.call @__ocml_cbrt_f16(%{{.*}}) : (f16) -> f16215 %result32 = math.cbrt %arg_f32 : f32216 // CHECK: llvm.call @__ocml_cbrt_f32(%{{.*}}) : (f32) -> f32217 %result64 = math.cbrt %arg_f64 : f64218 // CHECK: llvm.call @__ocml_cbrt_f64(%{{.*}}) : (f64) -> f64219 func.return %result16, %result32, %result64 : f16, f32, f64220 }221}222 223// -----224 225gpu.module @test_module {226 // CHECK: llvm.func @__ocml_ceil_f16(f16) -> f16227 // CHECK: llvm.func @__ocml_ceil_f32(f32) -> f32228 // CHECK: llvm.func @__ocml_ceil_f64(f64) -> f64229 // CHECK-LABEL: func @gpu_ceil230 func.func @gpu_ceil(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {231 %result16 = math.ceil %arg_f16 : f16232 // CHECK: llvm.call @__ocml_ceil_f16(%{{.*}}) : (f16) -> f16233 %result32 = math.ceil %arg_f32 : f32234 // CHECK: llvm.call @__ocml_ceil_f32(%{{.*}}) : (f32) -> f32235 %result64 = math.ceil %arg_f64 : f64236 // CHECK: llvm.call @__ocml_ceil_f64(%{{.*}}) : (f64) -> f64237 func.return %result16, %result32, %result64 : f16, f32, f64238 }239}240 241// -----242 243gpu.module @test_module {244 // CHECK: llvm.func @__ocml_floor_f16(f16) -> f16245 // CHECK: llvm.func @__ocml_floor_f32(f32) -> f32246 // CHECK: llvm.func @__ocml_floor_f64(f64) -> f64247 // CHECK-LABEL: func @gpu_floor248 func.func @gpu_floor(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {249 %result16 = math.floor %arg_f16 : f16250 // CHECK: llvm.call @__ocml_floor_f16(%{{.*}}) : (f16) -> f16251 %result32 = math.floor %arg_f32 : f32252 // CHECK: llvm.call @__ocml_floor_f32(%{{.*}}) : (f32) -> f32253 %result64 = math.floor %arg_f64 : f64254 // CHECK: llvm.call @__ocml_floor_f64(%{{.*}}) : (f64) -> f64255 func.return %result16, %result32, %result64 : f16, f32, f64256 }257}258 259// -----260 261gpu.module @test_module {262 // CHECK: llvm.func @__ocml_cos_f16(f16) -> f16263 // CHECK: llvm.func @__ocml_cos_f32(f32) -> f32264 // CHECK: llvm.func @__ocml_cos_f64(f64) -> f64265 // CHECK-LABEL: func @gpu_cos266 func.func @gpu_cos(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {267 %result16 = math.cos %arg_f16 : f16268 // CHECK: llvm.call @__ocml_cos_f16(%{{.*}}) : (f16) -> f16269 %result32 = math.cos %arg_f32 : f32270 // CHECK: llvm.call @__ocml_cos_f32(%{{.*}}) : (f32) -> f32271 %result64 = math.cos %arg_f64 : f64272 // CHECK: llvm.call @__ocml_cos_f64(%{{.*}}) : (f64) -> f64273 func.return %result16, %result32, %result64 : f16, f32, f64274 }275}276 277// -----278 279gpu.module @test_module {280 // CHECK: llvm.func @__ocml_exp2_f16(f16) -> f16281 // CHECK: llvm.func @__ocml_exp2_f32(f32) -> f32282 // CHECK: llvm.func @__ocml_exp2_f64(f64) -> f64283 // CHECK-LABEL: func @gpu_exp2284 func.func @gpu_exp2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {285 %result16 = math.exp2 %arg_f16 : f16286 // CHECK: llvm.call @__ocml_exp2_f16(%{{.*}}) : (f16) -> f16287 %exp2_f32 = math.exp2 %arg_f32 : f32288 // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32289 %result32 = math.exp2 %exp2_f32 : f32290 // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32291 %result64 = math.exp2 %arg_f64 : f64292 // CHECK: llvm.call @__ocml_exp2_f64(%{{.*}}) : (f64) -> f64293 func.return %result16, %result32, %result64 : f16, f32, f64294 }295}296 297// -----298 299// Test that we handled properly operation with SymbolTable other than module op300gpu.module @test_module {301 "test.symbol_scope"() ({302 // CHECK: test.symbol_scope303 // CHECK: llvm.func @__ocml_sin_f16(f16) -> f16304 // CHECK: llvm.func @__ocml_sin_f32(f32) -> f32305 // CHECK: llvm.func @__ocml_sin_f64(f64) -> f64306 // CHECK-LABEL: func @gpu_sin307 func.func @gpu_sin(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {308 // CHECK: llvm.call @__ocml_sin_f16(%{{.*}}) : (f16) -> f16309 %result16 = math.sin %arg_f16 : f16310 // CHECK: llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32311 %result32 = math.sin %arg_f32 : f32312 // CHECK: llvm.call @__ocml_sin_f64(%{{.*}}) : (f64) -> f64313 %result64 = math.sin %arg_f64 : f64314 func.return %result16, %result32, %result64 : f16, f32, f64315 }316 "test.finish" () : () -> ()317 }) : () -> ()318}319 320// -----321 322gpu.module @test_module {323 // CHECK: llvm.func @__ocml_expm1_f16(f16) -> f16324 // CHECK: llvm.func @__ocml_expm1_f32(f32) -> f32325 // CHECK: llvm.func @__ocml_expm1_f64(f64) -> f64326 // CHECK-LABEL: func @gpu_expm1327 func.func @gpu_expm1(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {328 %result16 = math.expm1 %arg_f16 : f16329 // CHECK: llvm.call @__ocml_expm1_f16(%{{.*}}) : (f16) -> f16330 %expm1_f32 = math.expm1 %arg_f32 : f32331 // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32332 %result32 = math.expm1 %expm1_f32 : f32333 // CHECK: llvm.call @__ocml_expm1_f32(%{{.*}}) : (f32) -> f32334 %result64 = math.expm1 %arg_f64 : f64335 // CHECK: llvm.call @__ocml_expm1_f64(%{{.*}}) : (f64) -> f64336 func.return %result16, %result32, %result64 : f16, f32, f64337 }338}339 340// -----341 342gpu.module @test_module {343 // CHECK: llvm.func @__ocml_log_f16(f16) -> f16344 // CHECK: llvm.func @__ocml_log_f64(f64) -> f64345 // CHECK-LABEL: func @gpu_log346 func.func @gpu_log(%arg_f16 : f16, %arg_f64 : f64) -> (f16, f64) {347 %result16 = math.log %arg_f16 : f16348 // CHECK: llvm.call @__ocml_log_f16(%{{.*}}) : (f16) -> f16349 %result64 = math.log %arg_f64 : f64350 // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (f64) -> f64351 func.return %result16, %result64 : f16, f64352 }353}354 355// -----356 357gpu.module @test_module {358 // CHECK: llvm.func @__ocml_log1p_f16(f16) -> f16359 // CHECK: llvm.func @__ocml_log1p_f32(f32) -> f32360 // CHECK: llvm.func @__ocml_log1p_f64(f64) -> f64361 // CHECK-LABEL: func @gpu_log1p362 func.func @gpu_log1p(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {363 %result16 = math.log1p %arg_f16 : f16364 // CHECK: llvm.call @__ocml_log1p_f16(%{{.*}}) : (f16) -> f16365 %result32 = math.log1p %arg_f32 : f32366 // CHECK: llvm.call @__ocml_log1p_f32(%{{.*}}) : (f32) -> f32367 %result64 = math.log1p %arg_f64 : f64368 // CHECK: llvm.call @__ocml_log1p_f64(%{{.*}}) : (f64) -> f64369 func.return %result16, %result32, %result64 : f16, f32, f64370 }371}372 373// -----374 375gpu.module @test_module {376 // CHECK: llvm.func @__ocml_log10_f16(f16) -> f16377 // CHECK: llvm.func @__ocml_log10_f32(f32) -> f32378 // CHECK: llvm.func @__ocml_log10_f64(f64) -> f64379 // CHECK-LABEL: func @gpu_log10380 func.func @gpu_log10(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {381 %result16 = math.log10 %arg_f16 : f16382 // CHECK: llvm.call @__ocml_log10_f16(%{{.*}}) : (f16) -> f16383 %result32 = math.log10 %arg_f32 : f32384 // CHECK: llvm.call @__ocml_log10_f32(%{{.*}}) : (f32) -> f32385 %result64 = math.log10 %arg_f64 : f64386 // CHECK: llvm.call @__ocml_log10_f64(%{{.*}}) : (f64) -> f64387 func.return %result16, %result32, %result64 : f16, f32, f64388 }389}390 391// -----392 393gpu.module @test_module {394 // CHECK: llvm.func @__ocml_log2_f16(f16) -> f16395 // CHECK: llvm.func @__ocml_log2_f32(f32) -> f32396 // CHECK: llvm.func @__ocml_log2_f64(f64) -> f64397 // CHECK-LABEL: func @gpu_log2398 func.func @gpu_log2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {399 %result16 = math.log2 %arg_f16 : f16400 // CHECK: llvm.call @__ocml_log2_f16(%{{.*}}) : (f16) -> f16401 %result32 = math.log2 %arg_f32 : f32402 // CHECK: llvm.call @__ocml_log2_f32(%{{.*}}) : (f32) -> f32403 %result64 = math.log2 %arg_f64 : f64404 // CHECK: llvm.call @__ocml_log2_f64(%{{.*}}) : (f64) -> f64405 func.return %result16, %result32, %result64 : f16, f32, f64406 }407}408 409// -----410 411gpu.module @test_module {412 // CHECK: llvm.func @__ocml_rsqrt_f16(f16) -> f16413 // CHECK: llvm.func @__ocml_rsqrt_f32(f32) -> f32414 // CHECK: llvm.func @__ocml_rsqrt_f64(f64) -> f64415 // CHECK-LABEL: func @gpu_rsqrt416 func.func @gpu_rsqrt(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {417 %result16 = math.rsqrt %arg_f16 : f16418 // CHECK: llvm.call @__ocml_rsqrt_f16(%{{.*}}) : (f16) -> f16419 %result32 = math.rsqrt %arg_f32 : f32420 // CHECK: llvm.call @__ocml_rsqrt_f32(%{{.*}}) : (f32) -> f32421 %result64 = math.rsqrt %arg_f64 : f64422 // CHECK: llvm.call @__ocml_rsqrt_f64(%{{.*}}) : (f64) -> f64423 func.return %result16, %result32, %result64 : f16, f32, f64424 }425}426 427// -----428 429gpu.module @test_module {430 // CHECK: llvm.func @__ocml_tan_f16(f16) -> f16431 // CHECK: llvm.func @__ocml_tan_f32(f32) -> f32432 // CHECK: llvm.func @__ocml_tan_f64(f64) -> f64433 // CHECK-LABEL: func @gpu_tan434 func.func @gpu_tan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {435 %result16 = math.tan %arg_f16 : f16436 // CHECK: llvm.call @__ocml_tan_f16(%{{.*}}) : (f16) -> f16437 %result32 = math.tan %arg_f32 : f32438 // CHECK: llvm.call @__ocml_tan_f32(%{{.*}}) : (f32) -> f32439 %result64 = math.tan %arg_f64 : f64440 // CHECK: llvm.call @__ocml_tan_f64(%{{.*}}) : (f64) -> f64441 func.return %result16, %result32, %result64 : f16, f32, f64442 }443}444 445// -----446 447gpu.module @test_module {448 // CHECK: llvm.func @__ocml_tanh_f16(f16) -> f16449 // CHECK: llvm.func @__ocml_tanh_f32(f32) -> f32450 // CHECK: llvm.func @__ocml_tanh_f64(f64) -> f64451 // CHECK-LABEL: func @gpu_tanh452 func.func @gpu_tanh(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {453 %result16 = math.tanh %arg_f16 : f16454 // CHECK: llvm.call @__ocml_tanh_f16(%{{.*}}) : (f16) -> f16455 %result32 = math.tanh %arg_f32 : f32456 // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (f32) -> f32457 %result64 = math.tanh %arg_f64 : f64458 // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (f64) -> f64459 func.return %result16, %result32, %result64 : f16, f32, f64460 }461}462 463// -----464 465gpu.module @test_module {466 // CHECK: llvm.func @__ocml_atan_f16(f16) -> f16467 // CHECK: llvm.func @__ocml_atan_f32(f32) -> f32468 // CHECK: llvm.func @__ocml_atan_f64(f64) -> f64469 // CHECK-LABEL: func @gpu_atan470 func.func @gpu_atan(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {471 %result16 = math.atan %arg_f16 : f16472 // CHECK: llvm.call @__ocml_atan_f16(%{{.*}}) : (f16) -> f16473 %result32 = math.atan %arg_f32 : f32474 // CHECK: llvm.call @__ocml_atan_f32(%{{.*}}) : (f32) -> f32475 %result64 = math.atan %arg_f64 : f64476 // CHECK: llvm.call @__ocml_atan_f64(%{{.*}}) : (f64) -> f64477 func.return %result16, %result32, %result64 : f16, f32, f64478 }479}480 481// -----482 483gpu.module @test_module {484 // CHECK: llvm.func @__ocml_atan2_f16(f16, f16) -> f16485 // CHECK: llvm.func @__ocml_atan2_f32(f32, f32) -> f32486 // CHECK: llvm.func @__ocml_atan2_f64(f64, f64) -> f64487 // CHECK-LABEL: func @gpu_atan2488 func.func @gpu_atan2(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {489 %result16 = math.atan2 %arg_f16, %arg_f16 : f16490 // CHECK: llvm.call @__ocml_atan2_f16(%{{.*}}) : (f16, f16) -> f16491 %result32 = math.atan2 %arg_f32, %arg_f32 : f32492 // CHECK: llvm.call @__ocml_atan2_f32(%{{.*}}) : (f32, f32) -> f32493 %result64 = math.atan2 %arg_f64, %arg_f64 : f64494 // CHECK: llvm.call @__ocml_atan2_f64(%{{.*}}) : (f64, f64) -> f64495 func.return %result16, %result32, %result64 : f16, f32, f64496 }497}498 499// -----500 501gpu.module @test_module {502 // CHECK: llvm.func @__ocml_pow_f16(f16, f16) -> f16503 // CHECK: llvm.func @__ocml_pow_f32(f32, f32) -> f32504 // CHECK: llvm.func @__ocml_pow_f64(f64, f64) -> f64505 // CHECK-LABEL: func @gpu_pow506 func.func @gpu_pow(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {507 %result16 = math.powf %arg_f16, %arg_f16 : f16508 // CHECK: llvm.call @__ocml_pow_f16(%{{.*}}, %{{.*}}) : (f16, f16) -> f16509 %result32 = math.powf %arg_f32, %arg_f32 : f32510 // CHECK: llvm.call @__ocml_pow_f32(%{{.*}}, %{{.*}}) : (f32, f32) -> f32511 %result64 = math.powf %arg_f64, %arg_f64 : f64512 // CHECK: llvm.call @__ocml_pow_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64513 func.return %result16, %result32, %result64 : f16, f32, f64514 }515}516 517// -----518 519gpu.module @test_module {520 // CHECK: llvm.func @__ocml_erf_f16(f16) -> f16521 // CHECK: llvm.func @__ocml_erf_f32(f32) -> f32522 // CHECK: llvm.func @__ocml_erf_f64(f64) -> f64523 // CHECK-LABEL: func @gpu_erf524 func.func @gpu_erf(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {525 %result16 = math.erf %arg_f16 : f16526 // CHECK: llvm.call @__ocml_erf_f16(%{{.*}}) : (f16) -> f16527 %result32 = math.erf %arg_f32 : f32528 // CHECK: llvm.call @__ocml_erf_f32(%{{.*}}) : (f32) -> f32529 %result64 = math.erf %arg_f64 : f64530 // CHECK: llvm.call @__ocml_erf_f64(%{{.*}}) : (f64) -> f64531 func.return %result16, %result32, %result64 : f16, f32, f64532 }533}534 535// -----536 537gpu.module @test_module {538 // CHECK-LABEL: func @gpu_unroll539 func.func @gpu_unroll(%arg0 : vector<4xf32>) -> vector<4xf32> {540 %result = math.sin %arg0 : vector<4xf32>541 // CHECK: %[[V0:.+]] = llvm.mlir.poison : vector<4xf32>542 // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32543 // CHECK: %[[V1:.+]] = llvm.insertelement %[[CL]], %[[V0]]544 // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32545 // CHECK: %[[V2:.+]] = llvm.insertelement %[[CL]], %[[V1]]546 // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32547 // CHECK: %[[V3:.+]] = llvm.insertelement %[[CL]], %[[V2]]548 // CHECK: %[[CL:.+]] = llvm.call @__ocml_sin_f32(%{{.*}}) : (f32) -> f32549 // CHECK: %[[V4:.+]] = llvm.insertelement %[[CL]], %[[V3]]550 // CHECK: return %[[V4]]551 func.return %result : vector<4xf32>552 }553}554 555// -----556 557// Test that the bf16 type is passed through to LLVM.558 559gpu.module @test_module {560 // CHECK-LABEL: func @bf16_id561 func.func @bf16_id(%arg0 : bf16) -> bf16 {562 // CHECK-SAME: (%[[ARG0:.+]]: bf16)563 // CHECK-SAME: -> bf16564 // CHECK: return %[[ARG0]] : bf16565 func.return %arg0 : bf16566 }567 568 // CHECK-LABEL: func @bf16x4_id569 func.func @bf16x4_id(%arg0 : vector<4xbf16>) -> vector<4xbf16> {570 // CHECK-SAME: (%[[ARG0:.+]]: vector<4xbf16>)571 // CHECK-SAME: -> vector<4xbf16>572 // CHECK: return %[[ARG0]] : vector<4xbf16>573 func.return %arg0 : vector<4xbf16>574 }575 576}577 578// -----579 580gpu.module @test_module {581 // CHECK-LABEL: @kernel_func582 // CHECK: attributes583 // CHECK: gpu.kernel584 // CHECK: rocdl.kernel585 gpu.func @kernel_func() kernel {586 gpu.return587 }588}589 590// -----591 592gpu.module @module {593// CHECK-LABEL: @spirv_sin594// CHECK: llvm.call @__ocml_sin_f32595 spirv.func @spirv_sin(%arg0: vector<4xf32>) -> vector<4xf32> "None" {596 %0 = math.sin %arg0 : vector<4xf32>597 spirv.ReturnValue %0 : vector<4xf32>598 }599}600 601// -----602 603gpu.module @test_module {604 // CHECK-LABEL: func @gpu_all_reduce_op()605 gpu.func @gpu_all_reduce_op() {606 %arg0 = arith.constant 1.0 : f32607 // TODO: Check full IR expansion once lowering has settled.608 // CHECK: llvm.add609 // CHECK: llvm.and610 // CHECK: llvm.xor611 // CHECK: llvm.icmp "slt"612 // CHECK: llvm.select613 // CHECK: llvm.shl614 // CHECK: rocdl.ds_bpermute {{.*}}615 // CHECK: rocdl.barrier616 // CHECK: llvm.bitcast617 // CHECK: llvm.fadd618 %result = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)619 620 gpu.return621 }622}623 624 625// -----626 627gpu.module @test_module {628 // CHECK-LABEL: func @gpu_all_reduce_region()629 gpu.func @gpu_all_reduce_region() {630 %arg0 = arith.constant 1 : i32631 // TODO: Check full IR expansion once lowering has settled.632 // CHECK: llvm.add633 // CHECK: llvm.and634 // CHECK: llvm.xor635 // CHECK: llvm.icmp "slt"636 // CHECK: llvm.select637 // CHECK: llvm.shl638 // CHECK: rocdl.ds_bpermute {{.*}}639 // CHECK: rocdl.barrier640 %result = gpu.all_reduce %arg0 uniform {641 ^bb(%lhs : i32, %rhs : i32):642 %xor = arith.xori %lhs, %rhs : i32643 "gpu.yield"(%xor) : (i32) -> ()644 } : (i32) -> (i32)645 gpu.return646 }647}648 649// -----650 651gpu.module @test_module {652 // CHECK: llvm.func @__ocml_fmod_f16(f16, f16) -> f16653 // CHECK: llvm.func @__ocml_fmod_f32(f32, f32) -> f32654 // CHECK: llvm.func @__ocml_fmod_f64(f64, f64) -> f64655 // CHECK-LABEL: func @gpu_fmod656 func.func @gpu_fmod(%arg_f16 : f16, %arg_f32 : f32, %arg_f64 : f64) -> (f16, f32, f64) {657 %result16 = arith.remf %arg_f16, %arg_f16 : f16658 // CHECK: llvm.call @__ocml_fmod_f16(%{{.*}}, %{{.*}}) : (f16, f16) -> f16659 %result32 = arith.remf %arg_f32, %arg_f32 : f32660 // CHECK: llvm.call @__ocml_fmod_f32(%{{.*}}, %{{.*}}) : (f32, f32) -> f32661 %result64 = arith.remf %arg_f64, %arg_f64 : f64662 // CHECK: llvm.call @__ocml_fmod_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64663 func.return %result16, %result32, %result64 : f16, f32, f64664 }665}666 667// -----668 669gpu.module @test_module {670 // CHECK-LABEL: func @gpu_shuffle()671 func.func @gpu_shuffle() -> (f32, f32, f32, f32) {672 // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32673 %arg0 = arith.constant 1.0 : f32674 // CHECK: %[[#OFFSET:]] = llvm.mlir.constant(4 : i32) : i32675 %arg1 = arith.constant 4 : i32676 // CHECK: %[[#WIDTH:]] = llvm.mlir.constant(23 : i32) : i32677 %arg2 = arith.constant 23 : i32678 // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi679 // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32680 // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32681 // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32682 // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32683 // CHECK: %[[#XOR:]] = llvm.xor %[[#LANE_ID]], %{{.*}} : i32684 // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#XOR]], %[[#WARP_OR_ZERO]] : i32685 // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#XOR]], %{{.*}} : i1, i32686 // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32687 // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32688 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32689 // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32690 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32691 %shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : f32692 // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi693 // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32694 // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32695 // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32696 // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32697 // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#OFFSET]], %[[#WARP_OR_ZERO]] : i32698 // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#OFFSET]], %{{.*}} : i1, i32699 // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32700 // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32701 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32702 // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32703 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32704 %shfli, %predi = gpu.shuffle idx %arg0, %arg1, %arg2 : f32705 // *** UP mode shuffle ***706 // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi707 // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32708 // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32709 // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32710 // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32711 // CHECK: %[[#UP:]] = llvm.sub %[[#LANE_ID]], %{{.*}} : i32712 // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#UP]], %[[#WARP_OR_ZERO]] : i32713 // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#UP]], %{{.*}} : i1, i32714 // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32715 // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32716 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32717 // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32718 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32719 %shflu, %predu = gpu.shuffle up %arg0, %arg1, %arg2 : f32720 // CHECK: %[[#LANE_ID:]] = rocdl.mbcnt.hi721 // CHECK: %[[#ZERO:]] = llvm.mlir.constant(0 : i32) : i32722 // CHECK: %[[#NEG_WIDTH:]] = llvm.sub %[[#ZERO]], %[[#WIDTH]] : i32723 // CHECK: %[[#ADD:]] = llvm.add %[[#LANE_ID]], %[[#WIDTH]] : i32724 // CHECK: %[[#WARP_OR_ZERO:]] = llvm.and %[[#ADD]], %[[#NEG_WIDTH]] : i32725 // CHECK: %[[#DOWN:]] = llvm.add %[[#LANE_ID]], %{{.*}} : i32726 // CHECK: %[[#CMP:]] = llvm.icmp "slt" %[[#DOWN]], %[[#WARP_OR_ZERO]] : i32727 // CHECK: %[[#DST_LANE:]] = llvm.select %[[#CMP]], %[[#DOWN]], %{{.*}} : i1, i32728 // CHECK: %[[#TWO:]] = llvm.mlir.constant(2 : i32) : i32729 // CHECK: %[[#ALIGNED_DST_LANE:]] = llvm.shl %[[#DST_LANE]], %[[#TWO]] : i32730 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32731 // CHECK: %[[#PERMUTE:]] = rocdl.ds_bpermute %[[#ALIGNED_DST_LANE]], %[[#CAST_VALUE]] : (i32, i32) -> i32732 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32733 %shfld, %predd = gpu.shuffle down %arg0, %arg1, %arg2 : f32734 func.return %shfl, %shfli, %shflu, %shfld : f32, f32, f32, f32735 }736 737 // CHECK-LABEL: func @gpu_shuffle_promote()738 func.func @gpu_shuffle_promote() -> (f32, f32, f32) {739 // CHECK: %[[#VALUE:]] = llvm.mlir.constant(1.000000e+00 : f32) : f32740 %arg0 = arith.constant 1.0 : f32741 %arg1 = arith.constant 4 : i32742 %arg2 = arith.constant 16 : i32743 %arg3 = arith.constant 32 : i32744 %arg4 = arith.constant 64 : i32745 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32746 // CHECK: %[[#MASK:]] = llvm.mlir.constant(4127 : i32) : i32747 // CHECK: %[[#PERMUTE:]] = rocdl.ds_swizzle %[[#CAST_VALUE]], %[[#MASK]] : (i32, i32) -> i32748 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#PERMUTE]] : i32 to f32749 %shfl1, %pred1 = gpu.shuffle xor %arg0, %arg1, %arg4 : f32750 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32751 // CHECK: %[[#PERMUTE:]] = rocdl.permlane16.swap %[[#CAST_VALUE]], %[[#CAST_VALUE]], false, false : (i32, i32) -> <(i32, i32)>752 // CHECK: %[[#EXTRACT0:]] = llvm.extractvalue %[[#PERMUTE:]][0] : !llvm.struct<(i32, i32)>753 // CHECK: %[[#EXTRACT1:]] = llvm.extractvalue %[[#PERMUTE:]][1] : !llvm.struct<(i32, i32)>754 // CHECK: %[[#CMP:]] = llvm.icmp "eq" %[[#EXTRACT0]], %[[#CAST_VALUE]] : i32755 // CHECK: %[[#SEL:]] = llvm.select %[[#CMP]], %[[#EXTRACT1]], %[[#EXTRACT0]] : i1, i32756 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#SEL]] : i32 to f32757 %shfl2, %pred2 = gpu.shuffle xor %arg0, %arg2, %arg4 : f32758 // CHECK: %[[#CAST_VALUE:]] = llvm.bitcast %[[#VALUE]] : f32 to i32759 // CHECK: %[[#PERMUTE:]] = rocdl.permlane32.swap %[[#CAST_VALUE]], %[[#CAST_VALUE]], false, false : (i32, i32) -> <(i32, i32)>760 // CHECK: %[[#EXTRACT0:]] = llvm.extractvalue %[[#PERMUTE:]][0] : !llvm.struct<(i32, i32)>761 // CHECK: %[[#EXTRACT1:]] = llvm.extractvalue %[[#PERMUTE:]][1] : !llvm.struct<(i32, i32)>762 // CHECK: %[[#CMP:]] = llvm.icmp "eq" %[[#EXTRACT0]], %[[#CAST_VALUE]] : i32763 // CHECK: %[[#SEL:]] = llvm.select %[[#CMP]], %[[#EXTRACT1]], %[[#EXTRACT0]] : i1, i32764 // CHECK: %[[#CAST_SHFL_VALUE:]] = llvm.bitcast %[[#SEL]] : i32 to f32765 %shfl3, %pred3 = gpu.shuffle xor %arg0, %arg3, %arg4 : f32766 func.return %shfl1, %shfl2, %shfl3 : f32, f32, f32767 }768 769 // CHECK-LABEL: func @gpu_shuffle_vec770 // CHECK-SAME: (%[[ARG:.*]]: vector<4xf16>, %{{.*}}: i32, %{{.*}}: i32)771 func.func @gpu_shuffle_vec(%arg0: vector<4xf16>, %arg1: i32, %arg2: i32) -> vector<4xf16> {772 // CHECK: %[[CAST1:.*]] = llvm.bitcast %[[ARG]] : vector<4xf16> to vector<2xi32>773 // CHECK: %[[IDX0:.*]] = llvm.mlir.constant(0 : i32) : i32774 // CHECK: %[[ELEM0:.*]] = llvm.extractelement %[[CAST1]][%[[IDX0]] : i32] : vector<2xi32>775 // CHECK: %[[IDX1:.*]] = llvm.mlir.constant(1 : i32) : i32776 // CHECK: %[[ELEM1:.*]] = llvm.extractelement %[[CAST1]][%[[IDX1]] : i32] : vector<2xi32>777 // CHECK: %[[PERM0:.*]] = rocdl.ds_bpermute %{{.*}}, %[[ELEM0]] : (i32, i32) -> i32778 // CHECK: %[[PERM1:.*]] = rocdl.ds_bpermute %{{.*}}, %[[ELEM1]] : (i32, i32) -> i32779 // CHECK: %[[V0:.*]] = llvm.mlir.poison : vector<2xi32>780 // CHECK: %[[IDX0:.*]] = llvm.mlir.constant(0 : i32) : i32781 // CHECK: %[[V1:.*]] = llvm.insertelement %[[PERM0]], %[[V0]][%[[IDX0]] : i32] : vector<2xi32>782 // CHECK: %[[IDX1:.*]] = llvm.mlir.constant(1 : i32) : i32783 // CHECK: %[[V2:.*]] = llvm.insertelement %[[PERM1]], %[[V1]][%[[IDX1]] : i32] : vector<2xi32>784 // CHECK: %[[RES:.*]] = llvm.bitcast %[[V2]] : vector<2xi32> to vector<4xf16>785 // CHECK: llvm.return %[[RES]] : vector<4xf16>786 %shfl, %pred = gpu.shuffle xor %arg0, %arg1, %arg2 : vector<4xf16>787 func.return %shfl : vector<4xf16>788 }789}790 791// -----792 793// CHECK-LABEL: @test_custom_data_layout794// CHECK-SAME: llvm.data_layout = "e"795gpu.module @test_custom_data_layout attributes {llvm.data_layout = "e"} {796 797}798 799// -----800 801gpu.module @test_module {802 // CHECK32-LABEL: func @gpu_dim_int_max_upper_bound()803 func.func @gpu_dim_int_max_upper_bound()804 -> (index) {805 806 // CHECK32: rocdl.workgroup.dim.x range <i32, 1, 2147483647> : i32807 %bDimX = gpu.block_dim x upper_bound 2147483647808 func.return %bDimX : index809 }810}811 812// -----813 814gpu.module @test_module {815// CHECK-LABEL: func @broadcast816// CHECK-SAME: (%[[ARG:.*]]: i64, %[[IDX:.*]]: i32)817func.func @broadcast(%arg0 : index, %arg1 : i32) -> (index, index) {818// CHECK: %{{.*}} = rocdl.readfirstlane %[[ARG]] : i64819// CHECK: %{{.*}} = rocdl.readlane %[[ARG]], %[[IDX]] : (i64, i32) -> i64820 %0 = gpu.subgroup_broadcast %arg0, first_active_lane : index821 %1 = gpu.subgroup_broadcast %arg0, specific_lane %arg1 : index822 func.return %0, %1 : index, index823}824}825