[mlir][gpu] Add subgroup Id/Size/Num to GPU dialect

[ThomasRaoux]

Summary

Add SubgroupId, SubgroupSize and NumSubgroups to GPU dialect ops and add the lowering of those ops to SPIRV.

[ftynse]

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

93

I think we use "workgroup" in the documentation. We should actually be consistent in the op names as well

101

let assemblyFormat = "attr-dict : type($result)" and changing the header to have Results<(outs Index:$result)> should give you custom syntax for almost free

[ThomasRaoux]

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

93

I can change it but all the Id Op above use block in comments and name (i.e. BlockDimOp, BlockIdOp) so it wold look inconsistent from this point of view. What do you think?

[ThomasRaoux]

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

101

Update the syntax.

[ftynse]

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

93

Other ops with inconsistent names are the reason why I suggested to change it in the documentation only. "Subgroup within the block" just sounds weird.

Arguably, the op names should also be changed, but that would require an RFC and make sure @herhut and @mravishankar sign off on it.

[antiagainst]

Cool, thanks Thomas! SPIR-V side looks good to me. :)

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

114

Nit: It seems the sgId needs to be properly named ? :)

130

Here too.

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

73

SingleDimLaunchConfigConversion?

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

185

I don't think we need this function for this test? Removing it can simplify the test itself. Less code is typically better. :)

Similarly for the rest. :)

[ThomasRaoux]

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

93

Thanks for explaining, changed it.

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

let description = [{
     Example:
 
     ```mlir
     %tIdX = "gpu.thread_id"() {dimension = "x"} : () -> (index)
     ```
   }];
 }
 
+def GPU_SubgroupIdOp : GPU_Op<"subgroup_id", [NoSideEffect]>,
+    Arguments<(ins)>, Results<(outs Index:$result)> {
+  let description = [{
+    Returns the subgroup id, i.e. the index of the current subgroup within the
+    workgroup.
+
+    Example:
+
+    ```mlir
+    %sgId = gpu.subgroup_id : index
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict `:` type($result)";
+  let verifier = [{ return success(); }];
+}
+
+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
+    %numSg = gpu.num_subgroups : index
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict `:` type($result)";
+  let verifier = [{ return success(); }];
+}
+
+def GPU_SubgroupSizeOp : GPU_Op<"subgroup_size", [NoSideEffect]>,
+    Arguments<(ins)>, Results<(outs Index:$result)> {
+  let description = [{
+    Returns the number of threads within a subgroup.
+
+    Example:
+
+    ```mlir
+    %sgSz = gpu.subgroup_size : index
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict `:` type($result)";
+  let verifier = [{ return success(); }];
+}
+
 def GPU_GPUFuncOp : GPU_Op<"func", [HasParent<"GPUModuleOp">,
                                     AutomaticAllocationScope, FunctionLike,
                                     IsolatedFromAbove, Symbol]> {
   let summary = "Function executable on a GPU";
 
   let description = [{
     Defines a function that can be executed on a GPU. This supports memory
     attribution and its body has a particular execution model.

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

 public:
   using SPIRVOpLowering<SourceOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Pattern lowering subgoup size/id to loading SPIR-V invocation
+/// builtin variables.
+template <typename SourceOp, spirv::BuiltIn builtin>
+class SingleDimLaunchConfigConversion : public SPIRVOpLowering<SourceOp> {
+public:
+  using SPIRVOpLowering<SourceOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// This is separate because in Vulkan workgroup size is exposed to shaders via
 /// a constant with WorkgroupSize decoration. So here we cannot generate a
 /// builtin variable; instead the information in the `spv.entry_point_abi`
 /// attribute on the surrounding FuncOp is used to replace the gpu::BlockDimOp.
 class WorkGroupSizeConversion : public SPIRVOpLowering<gpu::BlockDimOp> {
 public:
   using SPIRVOpLowering<gpu::BlockDimOp>::SPIRVOpLowering;
 
LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
   // SPIR-V invocation builtin variables are a vector of type <3xi32>
   auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
   rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
       op, rewriter.getIntegerType(32), spirvBuiltin,
       rewriter.getI32ArrayAttr({index.getValue()}));
   return success();
 }
 
+template <typename SourceOp, spirv::BuiltIn builtin>
+LogicalResult
+SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
+    SourceOp op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
+  rewriter.replaceOp(op, spirvBuiltin);
+  return success();
+}
+
 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
     gpu::BlockDimOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   auto index = getLaunchConfigIndex(op);
   if (!index)
     return failure();
 
   auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
void mlir::populateGPUToSPIRVPatterns(MLIRContext *context,
   populateWithGenerated(context, &patterns);
   patterns.insert<
       ForOpConversion, GPUFuncOpConversion, GPUModuleConversion,
       GPUReturnOpConversion, IfOpConversion,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>,
+      SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
+                                      spirv::BuiltIn::SubgroupId>,
+      SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
+                                      spirv::BuiltIn::NumSubgroups>,
+      SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
+                                      spirv::BuiltIn::SubgroupSize>,
       TerminatorOpConversion, WorkGroupSizeConversion>(context, typeConverter);
 }

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

case spirv::BuiltIn::GlobalInvocationId: {
     auto ptrType = spirv::PointerType::get(
         VectorType::get({3}, builder.getIntegerType(32)),
         spirv::StorageClass::Input);
     std::string name = getBuiltinVarName(builtin);
     newVarOp =
         builder.create<spirv::GlobalVariableOp>(loc, ptrType, name, builtin);
     break;
   }
+  case spirv::BuiltIn::SubgroupId:
+  case spirv::BuiltIn::NumSubgroups:
+  case spirv::BuiltIn::SubgroupSize: {
+    auto ptrType = spirv::PointerType::get(builder.getIntegerType(32),
+                                           spirv::StorageClass::Input);
+    std::string name = getBuiltinVarName(builtin);
+    newVarOp =
+        builder.create<spirv::GlobalVariableOp>(loc, ptrType, name, builtin);
+    break;
+  }
   default:
     emitError(loc, "unimplemented builtin variable generation for ")
         << stringifyBuiltIn(builtin);
   }
   return newVarOp;
 }
 
 Value mlir::spirv::getBuiltinVariableValue(Operation *op,

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

gpu.func @builtin_num_workgroups_x() kernel
       // CHECK: [[ADDRESS:%.*]] = spv._address_of [[NUMWORKGROUPS]]
       // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
       // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
       %0 = "gpu.grid_dim"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
+
+// -----
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL:  spv.module Logical GLSL450
+  // CHECK: spv.globalVariable [[SUBGROUPID:@.*]] built_in("SubgroupId")
+  gpu.module @kernels {
+    gpu.func @builtin_subgroup_id() kernel
+      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
+      // CHECK: [[ADDRESS:%.*]] = spv._address_of [[SUBGROUPID]]
+      // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
+      %0 = gpu.subgroup_id : index
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  // CHECK-LABEL:  spv.module Logical GLSL450
+  // CHECK: spv.globalVariable [[NUMSUBGROUPS:@.*]] built_in("NumSubgroups")
+  gpu.module @kernels {
+    gpu.func @builtin_num_subgroups() kernel
+      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
+      // CHECK: [[ADDRESS:%.*]] = spv._address_of [[NUMSUBGROUPS]]
+      // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
+      %0 = gpu.num_subgroups : index
+      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._address_of [[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"() {dimension = "x"} : () -> (index)
       %bIdY = "gpu.block_id"() {dimension = "y"} : () -> (index)
       %bIdZ = "gpu.block_id"() {dimension = "z"} : () -> (index)
 
       %gDimX = "gpu.grid_dim"() {dimension = "x"} : () -> (index)
       %gDimY = "gpu.grid_dim"() {dimension = "y"} : () -> (index)
       %gDimZ = "gpu.grid_dim"() {dimension = "z"} : () -> (index)
 
+      %sgId = gpu.subgroup_id : index
+      %numSg = gpu.num_subgroups : index
+      %SgSi = gpu.subgroup_size : index
+
       %one = constant 1.0 : f32
       %sum = "gpu.all_reduce"(%one) ({}) {op = "add"} : (f32) -> (f32)
 
       %width = constant 7 : i32
       %offset = constant 3 : i32
       // CHECK: gpu.shuffle %{{.*}}, %{{.*}}, %{{.*}} xor : f32
       %shfl, %pred = gpu.shuffle %arg0, %offset, %width xor : f32