[mlir][gpu] Introduce gpu.global_id op

[Hardcode84]

Summary

Introduce OpenCL-style global_id op and corresponding spirv lowering.

[Hardcode84]

fix op description

[bondhugula]

mlir/include/mlir/Dialect/GPU/GPUOps.td

117

There definitely has to be more description and explanation here -- at least on par with what you have above thread_id and block_id ops.

[Hardcode84]

expand description

[antiagainst]

mlir/include/mlir/Dialect/GPU/GPUOps.td

117

I'd suggest something like:

Returns the unique global workitem/thread id, i.e., the unique index of the current workitem/thread within all workgroups / grid along the x, y, or z `dimension`.

[Hardcode84]

rebase, update op description

mlir/include/mlir/Dialect/GPU/GPUOps.td

let description = [{
     ```mlir
     %sgId = gpu.subgroup_id : index
     ```
   }];
 
   let assemblyFormat = "attr-dict `:` type($result)";
 }
 
+def GPU_GlobalIdOp : GPU_IndexOp<"global_id"> {
+  let description = [{
+    Returns the unique global work-item ID.
+
+    Example:
+
+    ```mlir
+    %gidX = gpu.global_id x
+    ```
+  }];
+}
+
+
 def GPU_NumSubgroupsOp : GPU_Op<"num_subgroups", [NoSideEffect]>,
     Arguments<(ins)>, Results<(outs Index:$result)> {
   let description = [{
     Returns the number of subgroups within a workgroup.
 
     Example:
 
     ```mlir

mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp

void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
   patterns.add<
       GPUBarrierConversion, GPUFuncOpConversion, GPUModuleConversion,
       GPUModuleEndConversion, GPUReturnOpConversion,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>,
+      LaunchConfigConversion<gpu::GlobalIdOp,
+                             spirv::BuiltIn::GlobalInvocationId>,
       SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
                                       spirv::BuiltIn::SubgroupId>,
       SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
                                       spirv::BuiltIn::NumSubgroups>,
       SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
                                       spirv::BuiltIn::SubgroupSize>,
       WorkGroupSizeConversion>(typeConverter, patterns.getContext());
 }

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

gpu.func @builtin_workgroup_size_z() kernel
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
+  func @builtin() {
+    %c0 = arith.constant 1 : index
+    gpu.launch_func @kernels::@builtin_global_id_x
+        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_x() 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]]{{\[}}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} {
   // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
   // CHECK: spv.GlobalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize")
   gpu.module @kernels {
     gpu.func @builtin_subgroup_size() kernel
       attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
       // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[SUBGROUPSIZE]]
       // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
       %0 = gpu.subgroup_size : index
       gpu.return
     }
   }
 }

mlir/test/Dialect/GPU/ops.mlir

gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {
       %bIdX = gpu.block_id x
       %bIdY = gpu.block_id y
       %bIdZ = gpu.block_id z
 
       %gDimX = gpu.grid_dim x
       %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
 
       %one = arith.constant 1.0 : f32
       %sum = gpu.all_reduce add %one {} : (f32) -> (f32)
 
       %width = arith.constant 7 : i32