25 lines · plain
1// RUN: mlir-opt --test-gpu-rewrite -split-input-file %s | FileCheck %s2 3// CHECK-LABEL: func.func @subgroupId4// CHECK-SAME: (%[[SZ:.*]]: index, %[[MEM:.*]]: memref<index, 1>) {5func.func @subgroupId(%sz : index, %mem: memref<index, 1>) {6 gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %sz, %grid_y = %sz, %grid_z = %sz)7 threads(%tx, %ty, %tz) in (%block_x = %sz, %block_y = %sz, %block_z = %sz) {8 // CHECK: %[[DIMX:.*]] = gpu.block_dim x9 // CHECK-NEXT: %[[DIMY:.*]] = gpu.block_dim y10 // CHECK-NEXT: %[[TIDX:.*]] = gpu.thread_id x11 // CHECK-NEXT: %[[TIDY:.*]] = gpu.thread_id y12 // CHECK-NEXT: %[[TIDZ:.*]] = gpu.thread_id z13 // CHECK-NEXT: %[[T0:.*]] = arith.muli %[[DIMY]], %[[TIDZ]] : index14 // CHECK-NEXT: %[[T1:.*]] = arith.addi %[[T0]], %[[TIDY]] : index15 // CHECK-NEXT: %[[T2:.*]] = arith.muli %[[DIMX]], %[[T1]] : index16 // CHECK-NEXT: %[[T3:.*]] = arith.addi %[[TIDX]], %[[T2]] : index17 // CHECK-NEXT: %[[T4:.*]] = gpu.subgroup_size : index18 // CHECK-NEXT: %[[T5:.*]] = arith.divui %[[T3]], %[[T4]] : index19 %idz = gpu.subgroup_id : index20 memref.store %idz, %mem[] : memref<index, 1>21 gpu.terminator22 }23 return24}25