[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
Fri Sep 11 03:37:30 PDT 2026
https://github.com/fabrizio-indirli updated https://github.com/llvm/llvm-project/pull/221800
>From 9b536260b795e61083c40994c45eaa5807dee191 Mon Sep 17 00:00:00 2001
From: Fabrizio Indirli <fabrizio.indirli at arm.com>
Date: Thu, 3 Sep 2026 16:23:16 +0100
Subject: [PATCH] [mlir][gpu] Allow insert/extract of signed ints in
gpu.mma_matrix
`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
Signed-off-by: Fabrizio Indirli <fabrizio.indirli at arm.com>
---
.../mlir/Dialect/GPU/IR/CMakeLists.txt | 4 +++
mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h | 4 +++
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 26 ++++++++++++-----
.../Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp | 15 ++++++++++
mlir/lib/Dialect/GPU/CMakeLists.txt | 1 +
mlir/lib/Dialect/GPU/IR/GPUDialect.cpp | 29 +++++++++++++++++--
.../wmma-ops-to-spirv-khr-coop-matrix.mlir | 12 ++++++--
7 files changed, 79 insertions(+), 12 deletions(-)
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 038f7e8763ee0..816e8df67427e 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/GPUOpInterfaces.h.inc"
#include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 2aa5219e90404..20ea1098c2bf9 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -38,6 +38,12 @@ 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
+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,
@@ -2082,7 +2088,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);
@@ -2133,7 +2139,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)
@@ -2141,10 +2147,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";
@@ -2154,6 +2157,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])
@@ -2176,14 +2181,21 @@ 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);
+ let extraClassDeclaration = [{
+ static bool canBitcastToMatrixElement(
+ ::mlir::Type valueType, ::mlir::Type elementType);
+ }];
+
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..142e2fb60b161 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/WmmaOpsToSPIRV.cpp
@@ -219,6 +219,21 @@ struct WmmaInsertOpToSPIRVLowering final
if (!coopType)
return rewriter.notifyMatchFailure(op, "type conversion failed");
+ Type elementType =
+ cast<spirv::CooperativeMatrixType>(coopType).getElementType();
+ if (value.getType() != elementType) {
+ Type valueType = value.getType();
+ if (!gpu::SubgroupMmaInsertThreadLocalOp::canBitcastToMatrixElement(
+ valueType, elementType))
+ return rewriter.notifyMatchFailure(
+ op, "expected a signless integer value and a matching-width "
+ "signed or unsigned matrix element");
+
+ // Allow signless integers into signed/unsigned matrix by bitcasting.
+ 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 87e09f1dbf47f..ab22a4fb34197 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 a39b43c9803eb..ea9ba818a2c39 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
@@ -2198,6 +2200,27 @@ LogicalResult SubgroupMmaComputeOp::verify() {
return success();
}
+bool SubgroupMmaInsertThreadLocalOp::canBitcastToMatrixElement(
+ Type valueType, Type elementType) {
+ return valueType.isSignlessInteger() &&
+ (elementType.isSignedInteger() || elementType.isUnsignedInteger()) &&
+ valueType.getIntOrFloatBitWidth() ==
+ elementType.getIntOrFloatBitWidth();
+}
+
+LogicalResult SubgroupMmaInsertThreadLocalOp::verify() {
+ Type valueType = getValue().getType();
+ Type elementType =
+ cast<gpu::MMAMatrixType>(getMatrix().getType()).getElementType();
+ if (valueType == elementType)
+ return success();
+ if (canBitcastToMatrixElement(valueType, elementType))
+ 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>>
More information about the Mlir-commits
mailing list