diff --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td @@ -112,6 +112,19 @@ let assemblyFormat = "attr-dict `:` type($result)"; } +def GPU_GlobalIdOp : GPU_IndexOp<"global_id"> { + let description = [{ + Returns the unique global work-item ID along the x, y, or z `dimension`. + + Example: + + ```mlir + %gidX = gpu.global_id x + ``` + }]; +} + + def GPU_NumSubgroupsOp : GPU_Op<"num_subgroups", [NoSideEffect]>, Arguments<(ins)>, Results<(outs Index:$result)> { let description = [{ diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -373,6 +373,8 @@ LaunchConfigConversion, LaunchConfigConversion, + LaunchConfigConversion, SingleDimLaunchConfigConversion, SingleDimLaunchConfigConversion: vector<3xi32>}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[GLOBALINVOCATIONID]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}} + %0 = gpu.global_id x + gpu.return + } + } +} + +// ----- + +module attributes {gpu.container_module} { + func @builtin() { + %c0 = arith.constant 1 : index + gpu.launch_func @kernels::@builtin_global_id_y + blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) + return + } + + // CHECK-LABEL: spv.module @{{.*}} Logical GLSL450 + // CHECK: spv.GlobalVariable [[GLOBALINVOCATIONID:@.*]] built_in("GlobalInvocationId") + gpu.module @kernels { + gpu.func @builtin_global_id_y() kernel + attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[GLOBALINVOCATIONID]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}} + %0 = gpu.global_id y + gpu.return + } + } +} + +// ----- + +module attributes {gpu.container_module} { + func @builtin() { + %c0 = arith.constant 1 : index + gpu.launch_func @kernels::@builtin_global_id_z + blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) + return + } + + // CHECK-LABEL: spv.module @{{.*}} Logical GLSL450 + // CHECK: spv.GlobalVariable [[GLOBALINVOCATIONID:@.*]] built_in("GlobalInvocationId") + gpu.module @kernels { + gpu.func @builtin_global_id_z() kernel + attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} { + // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[GLOBALINVOCATIONID]] + // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]] + // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}} + %0 = gpu.global_id z + gpu.return + } + } +} + + // ----- module attributes {gpu.container_module} { diff --git a/mlir/test/Dialect/GPU/ops.mlir b/mlir/test/Dialect/GPU/ops.mlir --- a/mlir/test/Dialect/GPU/ops.mlir +++ b/mlir/test/Dialect/GPU/ops.mlir @@ -44,6 +44,10 @@ %gDimY = gpu.grid_dim y %gDimZ = gpu.grid_dim z + %gIdX = gpu.global_id x + %gIdY = gpu.global_id y + %gIdZ = gpu.global_id z + %sgId = gpu.subgroup_id : index %numSg = gpu.num_subgroups : index %SgSi = gpu.subgroup_size : index