[Mlir-commits] [mlir] [mlir][gpu] Allow insert/extract of signed ints in gpu.mma_matrix (PR #221800)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 7 11:49:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: fabrizio-indirli
<details>
<summary>Changes</summary>
`gpu.mma_matrix` can contain elements of both signed, unsigned and signless integer types, in addition to floating point types. However, the current implementation of the
`gpu.subgroup_mma_constant_matrix`, `.subgroup_mma_extract_thread_local` and `.subgroup_mma_insert_thread_local` operations only supported signless integers (and Fp types).
This patch extends such operations to support also signed and unsigned integer types, by defining a `GPU_MMAMatrixElementType` type constraint in TableGen.
Assisted-by: codex
---
Full diff: https://github.com/llvm/llvm-project/pull/221800.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt (+4)
- (modified) mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h (+4)
- (modified) mlir/include/mlir/Dialect/GPU/IR/GPUOps.td (+15-7)
- (modified) mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp (+13)
- (modified) mlir/lib/Dialect/GPU/CMakeLists.txt (+1)
- (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+19-3)
- (modified) mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir (+10-2)
``````````diff
diff --git a/mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt
index 32323ab62332a..0effb291ed2c4 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/GPU/IR/CMakeLists.txt
@@ -12,6 +12,10 @@ mlir_tablegen(ParallelLoopMapperEnums.cpp.inc -gen-enum-defs)
add_mlir_dialect_tablegen_target(MLIRParallelLoopMapperEnumsGen)
set(LLVM_TARGET_DEFINITIONS GPUOps.td)
+mlir_tablegen(GPUTypeConstraints.h.inc -gen-type-constraint-decls)
+mlir_tablegen(GPUTypeConstraints.cpp.inc -gen-type-constraint-defs)
+add_mlir_dialect_tablegen_target(MLIRGPUTypeConstraintsIncGen)
+
mlir_tablegen(GPUOpsEnums.h.inc -gen-enum-decls)
mlir_tablegen(GPUOpsEnums.cpp.inc -gen-enum-defs)
add_mlir_dialect_tablegen_target(MLIRGPUOpsEnumsGen)
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
index 5880c78136566..19bd42937f02e 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
@@ -219,6 +219,10 @@ class SparseSpGEMMOpHandleType
} // namespace gpu
} // namespace mlir
+namespace mlir::gpu {
+#include "mlir/Dialect/GPU/IR/GPUTypeConstraints.h.inc"
+} // namespace mlir::gpu
+
#include "mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc"
#include "mlir/Dialect/GPU/IR/GPUOpsDialect.h.inc"
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 1066dcb3b7308..4697e16e64dfc 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -38,6 +38,13 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
class GPU_Op<string mnemonic, list<Trait> traits = []> :
Op<GPU_Dialect, mnemonic, traits>;
+// Scalar element types supported by gpu.mma_matrix. The generated predicate
+// is also used by the hand-written MMAMatrixType verifier.
+def GPU_MMAMatrixElementType :
+ AnyTypeOf<[SI8, UI8, AnyI32, F16, F32, F64]> {
+ let cppFunctionName = "isValidMMAMatrixElementType";
+}
+
class GPU_IndexOp<string mnemonic, list<Trait> traits = []> :
GPU_Op<mnemonic, !listconcat(traits, [
Pure,
@@ -2073,7 +2080,7 @@ def GPU_SubgroupMmaConstantMatrixOp : GPU_Op<"subgroup_mma_constant_matrix",
```
}];
- let arguments = (ins AnyTypeOf<[SI8, UI8, I32, F16, F32]>:$value);
+ let arguments = (ins GPU_MMAMatrixElementType:$value);
let results = (outs GPU_MMAMatrix:$res);
@@ -2124,7 +2131,7 @@ def GPU_SubgroupMmaExtractThreadLocalOp : GPU_Op<"subgroup_mma_extract_thread_lo
let arguments = (ins GPU_MMAMatrix:$matrix, Variadic<Index>:$indices);
- let results = (outs AnyIntegerOrFloat:$res);
+ let results = (outs GPU_MMAMatrixElementType:$res);
let assemblyFormat = [{
$matrix`[`$indices`]` attr-dict `:` type($matrix) `->` type($res)
@@ -2132,10 +2139,7 @@ def GPU_SubgroupMmaExtractThreadLocalOp : GPU_Op<"subgroup_mma_extract_thread_lo
}
def GPU_SubgroupMmaInsertThreadLocalOp : GPU_Op<"subgroup_mma_insert_thread_local",
- [Pure,
- TypesMatchWith<"value type matches element type of mma_matrix",
- "matrix", "value",
- "::llvm::cast<gpu::MMAMatrixType>($_self).getElementType()"> ]>{
+ [Pure, AllTypesMatch<["matrix", "res"]>]> {
let summary = "Insert a value into GPU warp by invocation and indices";
@@ -2145,6 +2149,8 @@ def GPU_SubgroupMmaInsertThreadLocalOp : GPU_Op<"subgroup_mma_insert_thread_loca
This operation takes scalar value as its first operand and `!gpu.mma_matrix`
as its second operand. The op inserts the scalar value to the matrix.
+ A signless `i8` value may be inserted into an `si8` or `ui8` matrix; this
+ preserves the payload bits while the matrix type carries the signedness.
Since `matrix` is packed into the the threads within a subgroup, `indices` are
the indices into the values stored by each thread. That is, an index of 0 (or [0, 0])
@@ -2167,7 +2173,8 @@ def GPU_SubgroupMmaInsertThreadLocalOp : GPU_Op<"subgroup_mma_insert_thread_loca
```
}];
- let arguments = (ins AnyIntegerOrFloat:$value, GPU_MMAMatrix:$matrix,
+ let arguments = (ins AnyTypeOf<[I8, GPU_MMAMatrixElementType]>:$value,
+ GPU_MMAMatrix:$matrix,
Variadic<Index>:$indices);
let results = (outs GPU_MMAMatrix:$res);
@@ -2175,6 +2182,7 @@ def GPU_SubgroupMmaInsertThreadLocalOp : GPU_Op<"subgroup_mma_insert_thread_loca
let assemblyFormat = [{
$value`,` $matrix`[`$indices`]` attr-dict `:` type($value)`,` type($matrix) `->` type($res)
}];
+ let hasVerifier = 1;
}
def GPU_ElementwiseOpAddF : I32EnumAttrCase<"ADDF", 0, "addf">;
diff --git a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
index f3d57cf317c2a..d4aee40c79c37 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
@@ -219,6 +219,19 @@ struct WmmaInsertOpToSPIRVLowering final
if (!coopType)
return rewriter.notifyMatchFailure(op, "type conversion failed");
+ Type elementType =
+ cast<spirv::CooperativeMatrixType>(coopType).getElementType();
+ if (value.getType() != elementType) {
+ if (value.getType().getIntOrFloatBitWidth() !=
+ elementType.getIntOrFloatBitWidth())
+ return rewriter.notifyMatchFailure(op, "MMA element bitwidth mismatch");
+
+ // Memory provides signless bytes; the MMA matrix type carries the
+ // signedness required by the cooperative matrix operation.
+ value =
+ spirv::BitcastOp::create(rewriter, op.getLoc(), elementType, value);
+ }
+
SmallVector<int32_t> intValues;
for (Value val : op.getIndices()) {
if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>()) {
diff --git a/mlir/lib/Dialect/GPU/CMakeLists.txt b/mlir/lib/Dialect/GPU/CMakeLists.txt
index 547812da0ab97..a25ac25fcce3c 100644
--- a/mlir/lib/Dialect/GPU/CMakeLists.txt
+++ b/mlir/lib/Dialect/GPU/CMakeLists.txt
@@ -9,6 +9,7 @@ add_mlir_dialect_library(MLIRGPUDialect
DEPENDS
MLIRGPUOpsIncGen
MLIRGPUOpsAttributesIncGen
+ MLIRGPUTypeConstraintsIncGen
MLIRGPUOpsEnumsGen
MLIRGPUOpInterfacesIncGen
MLIRGPUCompilationAttrInterfacesIncGen
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 22c5e3c9b86ad..1452076e99edc 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -49,6 +49,10 @@ using namespace mlir::gpu;
#include "mlir/Dialect/GPU/IR/GPUOpsDialect.cpp.inc"
+namespace mlir::gpu {
+#include "mlir/Dialect/GPU/IR/GPUTypeConstraints.cpp.inc"
+} // namespace mlir::gpu
+
//===----------------------------------------------------------------------===//
// GPU Device Mapping Attributes
//===----------------------------------------------------------------------===//
@@ -210,9 +214,7 @@ Type MMAMatrixType::getElementType() const { return getImpl()->elementType; }
StringRef MMAMatrixType::getOperand() const { return getImpl()->getOperand(); }
bool MMAMatrixType::isValidElementType(Type elementType) {
- return elementType.isF16() || elementType.isF32() || elementType.isF64() ||
- elementType.isUnsignedInteger(8) || elementType.isSignedInteger(8) ||
- elementType.isInteger(32);
+ return isValidMMAMatrixElementType(elementType);
}
LogicalResult
@@ -2210,6 +2212,20 @@ LogicalResult SubgroupMmaComputeOp::verify() {
return success();
}
+LogicalResult SubgroupMmaInsertThreadLocalOp::verify() {
+ Type valueType = getValue().getType();
+ Type elementType =
+ cast<gpu::MMAMatrixType>(getMatrix().getType()).getElementType();
+ if (valueType == elementType)
+ return success();
+ if ((elementType.isSignedInteger() || elementType.isUnsignedInteger()) &&
+ valueType.isSignlessInteger(elementType.getIntOrFloatBitWidth()))
+ return success();
+ return emitOpError("expected value type to match matrix element type, or "
+ "be a matching-width signless integer for a signed or "
+ "unsigned integer matrix");
+}
+
LogicalResult MemcpyOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<::mlir::OpFoldResult> &results) {
return memref::foldMemRefCast(*this);
diff --git a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
index ad7f0546ba279..5d41e0b7b333d 100644
--- a/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/wmma-ops-to-spirv-khr-coop-matrix.mlir
@@ -81,19 +81,27 @@ module attributes {
// CHECK-SAME: !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixB>
// CHECK-SAME: !spirv.coopmatrix<16x16xi32, Subgroup, MatrixAcc>
gpu.func @gpu_wmma_signed_i8_mma_op(
+ %value: i8,
%A: !gpu.mma_matrix<16x16xsi8, "AOp">,
%B: !gpu.mma_matrix<16x16xsi8, "BOp">,
%C: !gpu.mma_matrix<16x16xi32, "COp">,
%ptr: memref<16x16xi32, #spirv.storage_class<StorageBuffer>>) kernel
attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [32, 4, 1]>} {
+ %i = arith.constant 0 : index
+ // CHECK: %[[SIGNED_VALUE:.*]] = spirv.Bitcast %{{.*}} : i8 to si8
+ // CHECK: spirv.CompositeInsert %[[SIGNED_VALUE]], %{{.*}}[0 : i32] : si8 into !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixA>
+ %insertedA = gpu.subgroup_mma_insert_thread_local %value, %A[%i] : i8, !gpu.mma_matrix<16x16xsi8, "AOp"> -> !gpu.mma_matrix<16x16xsi8, "AOp">
+ // CHECK: spirv.CompositeExtract %{{.*}}[0 : i32] : !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixA>
+ %extractedA = gpu.subgroup_mma_extract_thread_local %insertedA[%i] : !gpu.mma_matrix<16x16xsi8, "AOp"> -> si8
+ // CHECK: %[[REINSERTED_A:.*]] = spirv.CompositeInsert %{{.*}}, %{{.*}}[0 : i32] : si8 into !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixA>
+ %reinsertedA = gpu.subgroup_mma_insert_thread_local %extractedA, %insertedA[%i] : si8, !gpu.mma_matrix<16x16xsi8, "AOp"> -> !gpu.mma_matrix<16x16xsi8, "AOp">
// CHECK: spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}}, <ASigned|BSigned> :
// CHECK-SAME: !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixA>,
// CHECK-SAME: !spirv.coopmatrix<16x16xsi8, Subgroup, MatrixB>
// CHECK-SAME: -> !spirv.coopmatrix<16x16xi32, Subgroup, MatrixAcc>
- %D = gpu.subgroup_mma_compute %A, %B, %C : !gpu.mma_matrix<16x16xsi8, "AOp">,
+ %D = gpu.subgroup_mma_compute %reinsertedA, %B, %C : !gpu.mma_matrix<16x16xsi8, "AOp">,
!gpu.mma_matrix<16x16xsi8, "BOp">
-> !gpu.mma_matrix<16x16xi32, "COp">
- %i = arith.constant 0 : index
// CHECK: spirv.KHR.CooperativeMatrixStore %{{.+}}, %{{.+}}, %{{.+}}, <RowMajor>
gpu.subgroup_mma_store_matrix %D, %ptr[%i, %i] leadDimension 32 :
!gpu.mma_matrix<16x16xi32, "COp">, memref<16x16xi32, #spirv.storage_class<StorageBuffer>>
``````````
</details>
https://github.com/llvm/llvm-project/pull/221800
More information about the Mlir-commits
mailing list