[Mlir-commits] [mlir] 9dbb6e1 - [mlir][spirv] Add target width to SPIR-V ABI (#88555)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Apr 15 09:59:21 PDT 2024
Author: Hsiangkai Wang
Date: 2024-04-15T17:59:17+01:00
New Revision: 9dbb6e1978d3d2d61ef65c2dac1fd8add5a4c7a2
URL: https://github.com/llvm/llvm-project/commit/9dbb6e1978d3d2d61ef65c2dac1fd8add5a4c7a2
DIFF: https://github.com/llvm/llvm-project/commit/9dbb6e1978d3d2d61ef65c2dac1fd8add5a4c7a2.diff
LOG: [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.
Added:
Modified:
mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp
mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
index 74d36445e31138..3a11284da05122 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAttributes.td
+++ b/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) `>`";
}
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
index c35a8c26c2bc9b..24574bfaf6199a 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/TargetAndABI.h
+++ b/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.
diff --git a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
index 5b7c0a59ba4200..bbc318e17300ad 100644
--- a/mlir/lib/Dialect/SPIRV/IR/TargetAndABI.cpp
+++ b/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) {
diff --git a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp b/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp
index 6150b5ee17851d..2024a2e5279ffc 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/LowerABIAttributesPass.cpp
+++ b/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);
diff --git a/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir b/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
index 99369d11a4ba39..8b14d6569f4864 100644
--- a/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/entry-point.mlir
@@ -1,5 +1,8 @@
// 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]>
@@ -7,6 +10,15 @@
// 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
diff --git a/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp b/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
index 47b06e4531152d..0e8dfb8e3ee9a0 100644
--- a/mlir/test/lib/Dialect/SPIRV/TestEntryPointAbi.cpp
+++ b/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)));
}
}
More information about the Mlir-commits
mailing list