[mlir][spirv] Add target width to SPIR-V ABI (#88555)

There are execution modes need target width as their extra operands.
SignedZeroInfNanPreserve is one of them. This patch adds `target width`
as one of SPIR-V ABI attributes.

Author: Hsiangkai

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td

@@ -32,10 +32,12 @@ class SPIRV_Attr<string attrName, string attrMnemonic>
 // points in the generated SPIR-V module:
 // 1) [optional] Requested workgroup size.
 // 2) [optional] Requested subgroup size.
+// 3) [optional] Requested target width.
 def SPIRV_EntryPointABIAttr : SPIRV_Attr<"EntryPointABI", "entry_point_abi"> {
   let parameters = (ins
     OptionalParameter<"DenseI32ArrayAttr">:$workgroup_size,
-    OptionalParameter<"std::optional<int>">:$subgroup_size
+    OptionalParameter<"std::optional<int>">:$subgroup_size,
+    OptionalParameter<"std::optional<int>">:$target_width
   );
   let assemblyFormat = "`<` struct(params) `>`";
 }

mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h

@@ -87,9 +87,14 @@ bool needsInterfaceVarABIAttrs(TargetEnvAttr targetAttr);
 StringRef getEntryPointABIAttrName();
 
 /// Gets the EntryPointABIAttr given its fields.
+/// targetWidth is used by several execution modes. It is the element width
+/// of floating-point operations.
+/// Refer to Execution Mode in SPIR-V specification.
+/// https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_execution_mode
 EntryPointABIAttr getEntryPointABIAttr(MLIRContext *context,
                                        ArrayRef<int32_t> workgroupSize = {},
-                                       std::optional<int> subgroupSize = {});
+                                       std::optional<int> subgroupSize = {},
+                                       std::optional<int> targetWidth = {});
 
 /// Queries the entry point ABI on the nearest function-like op containing the
 /// given `op`. Returns null attribute if not found.

mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp

@@ -120,17 +120,16 @@ bool spirv::needsInterfaceVarABIAttrs(spirv::TargetEnvAttr targetAttr) {
 
 StringRef spirv::getEntryPointABIAttrName() { return "spirv.entry_point_abi"; }
 
-spirv::EntryPointABIAttr
-spirv::getEntryPointABIAttr(MLIRContext *context,
-                            ArrayRef<int32_t> workgroupSize,
-                            std::optional<int> subgroupSize) {
+spirv::EntryPointABIAttr spirv::getEntryPointABIAttr(
+    MLIRContext *context, ArrayRef<int32_t> workgroupSize,
+    std::optional<int> subgroupSize, std::optional<int> targetWidth) {
   DenseI32ArrayAttr workgroupSizeAttr;
   if (!workgroupSize.empty()) {
     assert(workgroupSize.size() == 3);
     workgroupSizeAttr = DenseI32ArrayAttr::get(context, workgroupSize);
   }
-  return spirv::EntryPointABIAttr::get(context, workgroupSizeAttr,
-                                       subgroupSize);
+  return spirv::EntryPointABIAttr::get(context, workgroupSizeAttr, subgroupSize,
+                                       targetWidth);
 }
 
 spirv::EntryPointABIAttr spirv::lookupEntryPointABI(Operation *op) {

mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp

@@ -157,7 +157,7 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
       // Erase workgroup size.
       entryPointAttr = spirv::EntryPointABIAttr::get(
           entryPointAttr.getContext(), DenseI32ArrayAttr(),
-          entryPointAttr.getSubgroupSize());
+          entryPointAttr.getSubgroupSize(), entryPointAttr.getTargetWidth());
     }
   }
   if (std::optional<int> subgroupSize = entryPointAttr.getSubgroupSize()) {
@@ -170,10 +170,24 @@ static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
       // Erase subgroup size.
       entryPointAttr = spirv::EntryPointABIAttr::get(
           entryPointAttr.getContext(), entryPointAttr.getWorkgroupSize(),
-          std::nullopt);
+          std::nullopt, entryPointAttr.getTargetWidth());
     }
   }
-  if (entryPointAttr.getWorkgroupSize() || entryPointAttr.getSubgroupSize())
+  if (std::optional<int> targetWidth = entryPointAttr.getTargetWidth()) {
+    std::optional<ArrayRef<spirv::Capability>> caps =
+        spirv::getCapabilities(spirv::ExecutionMode::SignedZeroInfNanPreserve);
+    if (!caps || targetEnv.allows(*caps)) {
+      builder.create<spirv::ExecutionModeOp>(
+          funcOp.getLoc(), funcOp,
+          spirv::ExecutionMode::SignedZeroInfNanPreserve, *targetWidth);
+      // Erase target width.
+      entryPointAttr = spirv::EntryPointABIAttr::get(
+          entryPointAttr.getContext(), entryPointAttr.getWorkgroupSize(),
+          entryPointAttr.getSubgroupSize(), std::nullopt);
+    }
+  }
+  if (entryPointAttr.getWorkgroupSize() || entryPointAttr.getSubgroupSize() ||
+      entryPointAttr.getTargetWidth())
     funcOp->setAttr(entryPointAttrName, entryPointAttr);
   else
     funcOp->removeAttr(entryPointAttrName);

mlir/test/Conversion/GPUToSPIRV/entry-point.mlir

@@ -1,12 +1,24 @@
 // RUN: mlir-opt -test-spirv-entry-point-abi %s | FileCheck %s -check-prefix=DEFAULT
 // RUN: mlir-opt -test-spirv-entry-point-abi="workgroup-size=32" %s | FileCheck %s -check-prefix=WG32
+// RUN: mlir-opt -test-spirv-entry-point-abi="subgroup-size=4" %s | FileCheck %s -check-prefix=SG4
+// RUN: mlir-opt -test-spirv-entry-point-abi="target-width=32" %s | FileCheck %s -check-prefix=TW32
+// RUN: mlir-opt -test-spirv-entry-point-abi="workgroup-size=32,8 subgroup-size=4 target-width=32" %s | FileCheck %s -check-prefix=WG32_8-SG4-TW32
 
 //      DEFAULT: gpu.func @foo()
 // DEFAULT-SAME: spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>
 
 //      WG32: gpu.func @foo()
 // WG32-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 1, 1]>
 
+//      SG4: gpu.func @foo()
+// SG4-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1], subgroup_size = 4>
+
+//      TW32: gpu.func @foo()
+// TW32-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1], target_width = 32>
+
+//      WG32_8-SG4-TW32: gpu.func @foo()
+// WG32_8-SG4-TW32-SAME:  spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 8, 1], subgroup_size = 4, target_width = 32>
+
 gpu.module @kernels {
   gpu.func @foo() kernel {
     gpu.return

mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp

@@ -45,6 +45,16 @@ struct TestSpirvEntryPointABIPass
           "Workgroup size to use for all gpu.func kernels in the module, "
           "specified with x-dimension first, y-dimension next and z-dimension "
           "last. Unspecified dimensions will be set to 1")};
+  Pass::Option<int> subgroupSize{
+      *this, "subgroup-size",
+      llvm::cl::desc(
+          "Subgroup size to use for all gpu.func kernels in the module"),
+      llvm::cl::init(0)};
+  Pass::Option<int> targetWidth{
+      *this, "target-width",
+      llvm::cl::desc(
+          "Specify the component width of floating-point instructions"),
+      llvm::cl::init(0)};
 };
 } // namespace
 
@@ -60,7 +70,12 @@ void TestSpirvEntryPointABIPass::runOnOperation() {
                                              workgroupSize.end());
     workgroupSizeVec.resize(3, 1);
     gpuFunc->setAttr(attrName,
-                     spirv::getEntryPointABIAttr(context, workgroupSizeVec));
+                     spirv::getEntryPointABIAttr(
+                         context, workgroupSizeVec,
+                         (subgroupSize == 0) ? std::nullopt
+                                             : std::optional<int>(subgroupSize),
+                         (targetWidth == 0) ? std::nullopt
+                                            : std::optional<int>(targetWidth)));
   }
 }