[mlir][GPUToSPIRV] Modify the lowering of gpu.block_dim to be consistent with Vulkan SPEC

[antiagainst]

LGTM.

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

71

s/cause/because/

72

This explanation is a bit confusing to me. SPIR-V spec allows variables to have initializers; it's the Vulkan spec requires WorkgroupSize to be decorated onto constants. What about using the following:

... 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 infromation ...

257

Nit: getLaunchConfigIndex ? This actually does not check what the op is; it's just querying the launch configuration index.

438

Sort these alphabetically

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

83

Let's add some notes here to explain where this magic number comes from.

[antiagainst]

Please also fix the test.

[mravishankar]

Addressing comments

[mravishankar]

Addressing missed comment.

[mravishankar]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

257

Are you saying it should. It just returns {} if it doesnt find the right attribute and attribute value.

[mehdi_amini]

Committed in https://reviews.llvm.org/rGaaddca1efd1d5a66d014023b9649cb273a84a7ae

[herhut]

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

262

Maybe llvm::StringSwitch?

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

83

nit: from

95

I think we should either require that the gpu.launch has the same constant argument than what the SPIR-V lowering uses or update the description of the operation to state that the provided dimensions are just guidelines.

I do not understand where the need to pass the workgroup size comes from. Wouldn't it be sufficient to require that the arguments to the launch are constants? The pass that constructs the launch and earlier passes that do the tiling need to know this value anyway. So they could manifest them in IR as constant arguments to launch.

115

Maybe use a constant other than 1 here to differentiate from the default.

mlir/lib/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.cpp

 public:
   using SPIRVOpLowering<SourceOp>::SPIRVOpLowering;
 
   PatternMatchResult
   matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Pattern lowering gpu::BlockDimOp. This is separate cause the workgroup size
+/// is supposed to be constant for Vulkan. Currently the buitlin variable that
+/// would be generated would not have a constant initializer (unsupported by
+/// SPIR-V dialect). 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;
+
+  PatternMatchResult
+  matchAndRewrite(gpu::BlockDimOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Pattern to convert a kernel function in GPU dialect within a spv.module.
 class KernelFnConversion final : public SPIRVOpLowering<gpu::GPUFuncOp> {
 public:
   KernelFnConversion(MLIRContext *context, SPIRVTypeConverter &converter,
                      ArrayRef<int64_t> workGroupSize,
                      PatternBenefit benefit = 1)
       : SPIRVOpLowering<gpu::GPUFuncOp>(context, converter, benefit) {
     auto config = workGroupSize.take_front(3);
IfOpConversion::matchAndRewrite(loop::IfOp ifOp, ArrayRef<Value> operands,
   rewriter.eraseOp(ifOp);
   return matchSuccess();
 }
 
 //===----------------------------------------------------------------------===//
 // Builtins.
 //===----------------------------------------------------------------------===//
 
-template <typename SourceOp, spirv::BuiltIn builtin>
-PatternMatchResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
-    SourceOp op, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  auto dimAttr =
-      op.getOperation()->template getAttrOfType<StringAttr>("dimension");
+static Optional<int32_t> getIndexForLaunchConfigOp(Operation *op) {
+  auto dimAttr = op->getAttrOfType<StringAttr>("dimension");
   if (!dimAttr) {
-    return this->matchFailure();
+    return {};
   }
-  int32_t index = 0;
   if (dimAttr.getValue() == "x") {
-    index = 0;
+    return 0;
   } else if (dimAttr.getValue() == "y") {
-    index = 1;
+    return 1;
   } else if (dimAttr.getValue() == "z") {
-    index = 2;
-  } else {
-    return this->matchFailure();
-  }
+    return 2;
+  }
+  return {};
 
+template <typename SourceOp, spirv::BuiltIn builtin>
+PatternMatchResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
+    SourceOp op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto index = getIndexForLaunchConfigOp(op);
+  if (!index)
+    return this->matchFailure();
+
   // 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}));
+      rewriter.getI32ArrayAttr({index.getValue()}));
   return this->matchSuccess();
 }
 
+PatternMatchResult WorkGroupSizeConversion::matchAndRewrite(
+    gpu::BlockDimOp op, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto index = getIndexForLaunchConfigOp(op);
+  if (!index)
+    return matchFailure();
+
+  auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
+  auto val = workGroupSizeAttr.getValue<int32_t>(index.getValue());
+  auto convertedType = typeConverter.convertType(op.getResult().getType());
+  if (!convertedType)
+    return matchFailure();
+  rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
+      op, convertedType, IntegerAttr::get(convertedType, val));
+  return matchSuccess();
+}
+
 //===----------------------------------------------------------------------===//
 // GPUFuncOp
 //===----------------------------------------------------------------------===//
 
 // Legalizes a GPU function as an entry SPIR-V function.
 static FuncOp
 lowerAsEntryFunction(gpu::GPUFuncOp funcOp, SPIRVTypeConverter &typeConverter,
                      ConversionPatternRewriter &rewriter,

 
 void mlir::populateGPUToSPIRVPatterns(MLIRContext *context,
                                       SPIRVTypeConverter &typeConverter,
                                       OwningRewritePatternList &patterns,
                                       ArrayRef<int64_t> workGroupSize) {
   populateWithGenerated(context, &patterns);
   patterns.insert<KernelFnConversion>(context, typeConverter, workGroupSize);
   patterns.insert<
       ForOpConversion, GPUReturnOpConversion, IfOpConversion,
+      GPUModuleConversion, GPUReturnOpConversion, ForOpConversion,
       GPUModuleConversion,
-      GPUReturnOpConversion, ForOpConversion, GPUModuleConversion,
-      LaunchConfigConversion<gpu::BlockDimOp, spirv::BuiltIn::WorkgroupSize>,
       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
       LaunchConfigConversion<gpu::ThreadIdOp,
                              spirv::BuiltIn::LocalInvocationId>,
-      TerminatorOpConversion>(context, typeConverter);
+      TerminatorOpConversion, WorkGroupSizeConversion>(context, typeConverter);
 }

mlir/test/Conversion/GPUToSPIRV/builtins.mlir

-// RUN: mlir-opt -split-input-file -convert-gpu-to-spirv %s -o - | FileCheck %s
+// RUN: mlir-opt -split-input-file -pass-pipeline='convert-gpu-to-spirv{workgroup-size=32,4}' %s -o - | FileCheck %s
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_id_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 

 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_size_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
-  // CHECK: spv.globalVariable [[WORKGROUPSIZE:@.*]] built_in("WorkgroupSize")
   gpu.module @kernels {
     gpu.func @builtin_workgroup_size_x()
       attributes {gpu.kernel} {
-      // CHECK: [[ADDRESS:%.*]] = spv._address_of [[WORKGROUPSIZE]]
-      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
-      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
+      // CHECK: spv.constant 32 : i32
       %0 = "gpu.block_dim"() {dimension = "x"} : () -> index
       gpu.return
     }
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
   func @builtin() {
     %c0 = constant 1 : index
+    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_size_y", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
+    return
+  }
+
+  // CHECK-LABEL:  spv.module "Logical" "GLSL450"
+  gpu.module @kernels {
+    gpu.func @builtin_workgroup_size_y()
+      attributes {gpu.kernel} {
+      // CHECK: spv.constant 4 : i32
+      %0 = "gpu.block_dim"() {dimension = "y"} : () -> index
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  func @builtin() {
+    %c0 = constant 1 : index
+    "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_workgroup_size_z", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
+    return
+  }
+
+  // CHECK-LABEL:  spv.module "Logical" "GLSL450"
+  gpu.module @kernels {
+    gpu.func @builtin_workgroup_size_z()
+      attributes {gpu.kernel} {
+      // CHECK: spv.constant 1 : i32
+      %0 = "gpu.block_dim"() {dimension = "z"} : () -> index
+      gpu.return
+    }
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
+  func @builtin() {
+    %c0 = constant 1 : index
     "gpu.launch_func"(%c0, %c0, %c0, %c0, %c0, %c0) {kernel = "builtin_local_id_x", kernel_module = @kernels} : (index, index, index, index, index, index) -> ()
     return
   }
 
   // CHECK-LABEL:  spv.module "Logical" "GLSL450"
   // CHECK: spv.globalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
   gpu.module @kernels {
     gpu.func @builtin_local_id_x()