[Mlir-commits] [mlir] [mlir][SPIR-V] Make Stride optional in KHR CooperativeMatrixLoad/Store per spec (PR #214194)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 5 03:55:54 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Arseniy Obolenskiy (aobolensk)
<details>
<summary>Changes</summary>
The SPIR-V spec marks Stride as an optional operand for OpCooperativeMatrixLoadKHR/StoreKHR, since only RowMajor and ColumnMajor layouts require it
Add a verifier check for that requirement, and hand-write (de)serialization since the auto-generated logic cannot place an optional operand ahead of a mandatory attribute
---
Full diff: https://github.com/llvm/llvm-project/pull/214194.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td (+16-12)
- (modified) mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp (+17-7)
- (modified) mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp (+160)
- (modified) mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp (+78)
- (modified) mlir/test/Dialect/SPIRV/IR/khr-cooperative-matrix-ops.mlir (+38)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
index e8124b8b0bed9..8afbfc4a3e339 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
@@ -97,9 +97,6 @@ def SPIRV_KHRCooperativeMatrixLoadOp : SPIRV_KhrVendorOp<"CooperativeMatrixLoad"
All invocations in a given scope instance must be active or all must be
inactive.
- TODO: In the SPIR-V spec, `stride` is an optional argument. We should also
- support this optionality in the SPIR-V dialect.
-
#### Example:
```
@@ -114,8 +111,8 @@ def SPIRV_KHRCooperativeMatrixLoadOp : SPIRV_KhrVendorOp<"CooperativeMatrixLoad"
}];
let assemblyFormat = [{
- $pointer `,` $stride `,` $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
- type(operands) `->` type($result)
+ $pointer `,` ( $stride^ `,` )? $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
+ type($pointer) (`,` type($stride)^)? `->` type($result)
}];
let availability = [
@@ -125,11 +122,16 @@ def SPIRV_KHRCooperativeMatrixLoadOp : SPIRV_KhrVendorOp<"CooperativeMatrixLoad"
Capability<[SPIRV_C_CooperativeMatrixKHR]>
];
+ // The auto-generated (de)serialization cannot handle an optional operand
+ // (stride) followed by a mandatory attribute (matrix_layout), so it is
+ // hand-written. See Serializer/Deserializer::processOp<KHRCooperativeMatrixLoadOp>.
+ let autogenSerialization = 0;
+
// TODO: Add scope operand for MakePointer*. See #145485.
let arguments = (ins
SPIRV_AnyPtr:$pointer,
SPIRV_KHR_CooperativeMatrixLayoutAttr:$matrix_layout,
- SPIRV_Integer:$stride,
+ Optional<SPIRV_Integer>:$stride,
OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand,
OptionalAttr<IntValidAlignment<I32Attr>>:$alignment
);
@@ -183,9 +185,6 @@ def SPIRV_KHRCooperativeMatrixStoreOp : SPIRV_KhrVendorOp<"CooperativeMatrixStor
All invocations in a given scope instance must be active or all must be
inactive.
- TODO: In the SPIR-V spec, `stride` is an optional argument. We should also
- support this optionality in the SPIR-V dialect.
-
#### Example:
```
@@ -198,8 +197,8 @@ def SPIRV_KHRCooperativeMatrixStoreOp : SPIRV_KhrVendorOp<"CooperativeMatrixStor
}];
let assemblyFormat = [{
- $pointer `,` $object `,` $stride `,` $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
- type(operands)
+ $pointer `,` $object `,` ( $stride^ `,` )? $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
+ type($pointer) `,` type($object) (`,` type($stride)^)?
}];
let availability = [
@@ -209,12 +208,17 @@ def SPIRV_KHRCooperativeMatrixStoreOp : SPIRV_KhrVendorOp<"CooperativeMatrixStor
Capability<[SPIRV_C_CooperativeMatrixKHR]>
];
+ // The auto-generated (de)serialization cannot handle an optional operand
+ // (stride) followed by a mandatory attribute (matrix_layout), so it is
+ // hand-written. See Serializer/Deserializer::processOp<KHRCooperativeMatrixStoreOp>.
+ let autogenSerialization = 0;
+
// TODO: Add scope operand for MakePointer*. See #145485.
let arguments = (ins
SPIRV_AnyPtr:$pointer,
SPIRV_AnyCooperativeMatrix:$object,
SPIRV_KHR_CooperativeMatrixLayoutAttr:$matrix_layout,
- SPIRV_Integer:$stride,
+ Optional<SPIRV_Integer>:$stride,
OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand,
OptionalAttr<IntValidAlignment<I32Attr>>:$alignment
);
diff --git a/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp b/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
index eccc01137576a..ca41e1361214c 100644
--- a/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
@@ -22,7 +22,8 @@ namespace mlir::spirv {
static LogicalResult
verifyCoopMatrixAccess(Operation *op, Type pointer, Type coopMatrix,
- spirv::MemoryAccessAttr memoryOperand,
+ spirv::CooperativeMatrixLayoutKHR matrixLayout,
+ Value stride, spirv::MemoryAccessAttr memoryOperand,
IntegerAttr alignment) {
auto pointerType = cast<PointerType>(pointer);
Type pointeeType = pointerType.getPointeeType();
@@ -32,6 +33,15 @@ verifyCoopMatrixAccess(Operation *op, Type pointer, Type coopMatrix,
<< pointeeType;
}
+ // Per the SPIR-V spec, Stride is required for the RowMajor and ColumnMajor
+ // layouts.
+ if (!stride &&
+ (matrixLayout == spirv::CooperativeMatrixLayoutKHR::RowMajor ||
+ matrixLayout == spirv::CooperativeMatrixLayoutKHR::ColumnMajor)) {
+ return op->emitOpError("Stride is required for '")
+ << stringifyCooperativeMatrixLayoutKHR(matrixLayout) << "'";
+ }
+
if (memoryOperand) {
spirv::MemoryAccess operandSet = memoryOperand.getValue();
@@ -76,9 +86,9 @@ verifyCoopMatrixAccess(Operation *op, Type pointer, Type coopMatrix,
//===----------------------------------------------------------------------===//
LogicalResult KHRCooperativeMatrixLoadOp::verify() {
- return verifyCoopMatrixAccess(*this, getPointer().getType(),
- getResult().getType(), getMemoryOperandAttr(),
- getAlignmentAttr());
+ return verifyCoopMatrixAccess(
+ *this, getPointer().getType(), getResult().getType(), getMatrixLayout(),
+ getStride(), getMemoryOperandAttr(), getAlignmentAttr());
}
//===----------------------------------------------------------------------===//
@@ -86,9 +96,9 @@ LogicalResult KHRCooperativeMatrixLoadOp::verify() {
//===----------------------------------------------------------------------===//
LogicalResult KHRCooperativeMatrixStoreOp::verify() {
- return verifyCoopMatrixAccess(*this, getPointer().getType(),
- getObject().getType(), getMemoryOperandAttr(),
- getAlignmentAttr());
+ return verifyCoopMatrixAccess(
+ *this, getPointer().getType(), getObject().getType(), getMatrixLayout(),
+ getStride(), getMemoryOperandAttr(), getAlignmentAttr());
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp b/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
index 0e6d199eeabf0..1ad4072016671 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
@@ -753,6 +753,166 @@ Deserializer::processOp<spirv::CopyMemoryOp>(ArrayRef<uint32_t> words) {
return success();
}
+// The auto-generated deserialization only supports a Variadic<>/Optional<>
+// operand as the last ODS argument, but `stride` sits between `pointer`/
+// `object` and the mandatory `matrix_layout` attribute. Hand-write the
+// deserialization since whether Stride is present cannot be inferred purely
+// from word count without also decoding MemoryLayout first.
+template <>
+LogicalResult Deserializer::processOp<spirv::KHRCooperativeMatrixLoadOp>(
+ ArrayRef<uint32_t> words) {
+ SmallVector<Type, 1> resultTypes;
+ size_t wordIndex = 0;
+
+ if (wordIndex >= words.size())
+ return emitError(unknownLoc,
+ "expected result type <id> while deserializing "
+ "spirv::KHRCooperativeMatrixLoadOp");
+ Type resultType = getType(words[wordIndex]);
+ if (!resultType)
+ return emitError(unknownLoc, "unknown type result <id> : ")
+ << words[wordIndex];
+ resultTypes.push_back(resultType);
+ wordIndex++;
+
+ if (wordIndex >= words.size())
+ return emitError(unknownLoc, "expected result <id> while deserializing "
+ "spirv::KHRCooperativeMatrixLoadOp");
+ uint32_t valueID = words[wordIndex++];
+
+ SmallVector<Value, 4> operands;
+ SmallVector<NamedAttribute, 4> attributes;
+
+ // Consumes the next word as a value <id>, appending it to `operands`.
+ // No-op once `words` is exhausted, since `stride` is optional.
+ auto readOperand = [&]() -> LogicalResult {
+ if (wordIndex >= words.size())
+ return success();
+ Value arg = getValue(words[wordIndex]);
+ if (!arg)
+ return emitError(unknownLoc, "unknown result <id> : ")
+ << words[wordIndex];
+ operands.push_back(arg);
+ wordIndex++;
+ return success();
+ };
+
+ // Consumes the next word into `word`. Returns false once `words` is
+ // exhausted, since `memory_operand` and `alignment` are optional.
+ auto tryConsumeWord = [&](uint32_t &word) {
+ if (wordIndex >= words.size())
+ return false;
+ word = words[wordIndex++];
+ return true;
+ };
+
+ if (failed(readOperand())) // pointer
+ return failure();
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(opBuilder.getNamedAttr(
+ "matrix_layout",
+ opBuilder.getAttr<spirv::CooperativeMatrixLayoutKHRAttr>(
+ static_cast<spirv::CooperativeMatrixLayoutKHR>(
+ getConstantInt(word).getValue().getZExtValue()))));
+
+ if (failed(readOperand())) // stride
+ return failure();
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(opBuilder.getNamedAttr(
+ "memory_operand", opBuilder.getAttr<spirv::MemoryAccessAttr>(
+ static_cast<spirv::MemoryAccess>(word))));
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(
+ opBuilder.getNamedAttr("alignment", opBuilder.getI32IntegerAttr(word)));
+
+ if (wordIndex != words.size())
+ return emitError(unknownLoc,
+ "found more operands than expected when deserializing "
+ "spirv::KHRCooperativeMatrixLoadOp, only ")
+ << wordIndex << " of " << words.size() << " processed";
+
+ if (decorations.count(valueID)) {
+ auto attrs = decorations[valueID].getAttrs();
+ attributes.append(attrs.begin(), attrs.end());
+ }
+ Location loc = createFileLineColLoc(opBuilder);
+ auto op = spirv::KHRCooperativeMatrixLoadOp::create(
+ opBuilder, loc, resultTypes, operands, attributes);
+ valueMap[valueID] = op.getResult();
+
+ return success();
+}
+
+template <>
+LogicalResult Deserializer::processOp<spirv::KHRCooperativeMatrixStoreOp>(
+ ArrayRef<uint32_t> words) {
+ size_t wordIndex = 0;
+ SmallVector<Value, 4> operands;
+ SmallVector<NamedAttribute, 4> attributes;
+
+ // Consumes the next word as a value <id>, appending it to `operands`.
+ // No-op once `words` is exhausted, since `stride` is optional.
+ auto readOperand = [&]() -> LogicalResult {
+ if (wordIndex >= words.size())
+ return success();
+ Value arg = getValue(words[wordIndex]);
+ if (!arg)
+ return emitError(unknownLoc, "unknown result <id> : ")
+ << words[wordIndex];
+ operands.push_back(arg);
+ wordIndex++;
+ return success();
+ };
+
+ // Consumes the next word into `word`. Returns false once `words` is
+ // exhausted, since `memory_operand` and `alignment` are optional.
+ auto tryConsumeWord = [&](uint32_t &word) {
+ if (wordIndex >= words.size())
+ return false;
+ word = words[wordIndex++];
+ return true;
+ };
+
+ if (failed(readOperand())) // pointer
+ return failure();
+ if (failed(readOperand())) // object
+ return failure();
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(opBuilder.getNamedAttr(
+ "matrix_layout",
+ opBuilder.getAttr<spirv::CooperativeMatrixLayoutKHRAttr>(
+ static_cast<spirv::CooperativeMatrixLayoutKHR>(
+ getConstantInt(word).getValue().getZExtValue()))));
+
+ if (failed(readOperand())) // stride
+ return failure();
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(opBuilder.getNamedAttr(
+ "memory_operand", opBuilder.getAttr<spirv::MemoryAccessAttr>(
+ static_cast<spirv::MemoryAccess>(word))));
+
+ if (uint32_t word; tryConsumeWord(word))
+ attributes.push_back(
+ opBuilder.getNamedAttr("alignment", opBuilder.getI32IntegerAttr(word)));
+
+ if (wordIndex != words.size())
+ return emitError(unknownLoc,
+ "found more operands than expected when deserializing "
+ "spirv::KHRCooperativeMatrixStoreOp, only ")
+ << wordIndex << " of " << words.size() << " processed";
+
+ Location loc = createFileLineColLoc(opBuilder);
+ spirv::KHRCooperativeMatrixStoreOp::create(opBuilder, loc, TypeRange(),
+ operands, attributes);
+
+ return success();
+}
+
template <>
LogicalResult Deserializer::processOp<spirv::GenericCastToPtrExplicitOp>(
ArrayRef<uint32_t> words) {
diff --git a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
index 32fa603e6c23d..b81af1778f76c 100644
--- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
@@ -1174,6 +1174,84 @@ LogicalResult Serializer::processOp<spirv::GenericCastToPtrExplicitOp>(
return success();
}
+// The auto-generated serialization only supports a Variadic<>/Optional<>
+// operand as the last ODS argument, but `stride` sits between `pointer`/
+// `object` and the mandatory `matrix_layout` attribute. Hand-write the
+// serialization to emit Stride only when present.
+template <>
+LogicalResult Serializer::processOp<spirv::KHRCooperativeMatrixLoadOp>(
+ spirv::KHRCooperativeMatrixLoadOp op) {
+ SmallVector<uint32_t, 6> operands;
+ SmallVector<StringRef, 3> elidedAttrs;
+ uint32_t resultTypeID = 0;
+ if (failed(processType(op.getLoc(), op.getType(), resultTypeID)))
+ return failure();
+ operands.push_back(resultTypeID);
+
+ uint32_t resultID = getNextID();
+ valueIDMap[op.getResult()] = resultID;
+ operands.push_back(resultID);
+
+ operands.push_back(getValueID(op.getPointer()));
+
+ operands.push_back(prepareConstantInt(
+ op.getLoc(), Builder(op).getI32IntegerAttr(
+ static_cast<uint32_t>(op.getMatrixLayout()))));
+ elidedAttrs.push_back("matrix_layout");
+
+ if (Value stride = op.getStride())
+ operands.push_back(getValueID(stride));
+
+ if (auto attr = op.getMemoryOperandAttr())
+ operands.push_back(static_cast<uint32_t>(attr.getValue()));
+ elidedAttrs.push_back("memory_operand");
+
+ if (auto attr = op.getAlignmentAttr())
+ operands.push_back(static_cast<uint32_t>(attr.getValue().getZExtValue()));
+ elidedAttrs.push_back("alignment");
+
+ if (failed(emitDebugLine(functionBody, op.getLoc())))
+ return failure();
+ encodeInstructionInto(functionBody, spirv::Opcode::OpCooperativeMatrixLoadKHR,
+ operands);
+
+ for (auto attr : op->getAttrs()) {
+ if (llvm::is_contained(elidedAttrs, attr.getName()))
+ continue;
+ if (failed(processDecoration(op.getLoc(), resultID, attr)))
+ return failure();
+ }
+ return success();
+}
+
+template <>
+LogicalResult Serializer::processOp<spirv::KHRCooperativeMatrixStoreOp>(
+ spirv::KHRCooperativeMatrixStoreOp op) {
+ SmallVector<uint32_t, 6> operands;
+
+ operands.push_back(getValueID(op.getPointer()));
+ operands.push_back(getValueID(op.getObject()));
+
+ operands.push_back(prepareConstantInt(
+ op.getLoc(), Builder(op).getI32IntegerAttr(
+ static_cast<uint32_t>(op.getMatrixLayout()))));
+
+ if (Value stride = op.getStride())
+ operands.push_back(getValueID(stride));
+
+ if (auto attr = op.getMemoryOperandAttr())
+ operands.push_back(static_cast<uint32_t>(attr.getValue()));
+
+ if (auto attr = op.getAlignmentAttr())
+ operands.push_back(static_cast<uint32_t>(attr.getValue().getZExtValue()));
+
+ if (failed(emitDebugLine(functionBody, op.getLoc())))
+ return failure();
+ encodeInstructionInto(functionBody,
+ spirv::Opcode::OpCooperativeMatrixStoreKHR, operands);
+ return success();
+}
+
// Pull in auto-generated Serializer::dispatchToAutogenSerialization() and
// various Serializer::processOp<...>() specializations.
#define GET_SERIALIZATION_FNS
diff --git a/mlir/test/Dialect/SPIRV/IR/khr-cooperative-matrix-ops.mlir b/mlir/test/Dialect/SPIRV/IR/khr-cooperative-matrix-ops.mlir
index 69235eab8d0dc..2566ca557b724 100644
--- a/mlir/test/Dialect/SPIRV/IR/khr-cooperative-matrix-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/khr-cooperative-matrix-ops.mlir
@@ -129,6 +129,24 @@ spirv.func @cooperative_matrix_load_missing_attr(%ptr : !spirv.ptr<i32, StorageB
// -----
+spirv.func @cooperative_matrix_load_missing_stride_row_major(%ptr : !spirv.ptr<i32, StorageBuffer>) "None" {
+ // expected-error @+1 {{op Stride is required for 'RowMajor'}}
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, <RowMajor> :
+ !spirv.ptr<i32, StorageBuffer> -> !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+ spirv.Return
+}
+
+// -----
+
+spirv.func @cooperative_matrix_load_missing_stride_column_major(%ptr : !spirv.ptr<i32, StorageBuffer>) "None" {
+ // expected-error @+1 {{op Stride is required for 'ColumnMajor'}}
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, <ColumnMajor> :
+ !spirv.ptr<i32, StorageBuffer> -> !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+ spirv.Return
+}
+
+// -----
+
spirv.func @cooperative_matrix_load_bad_operad(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
// expected-error @+1 {{op not compatible with memory operand 'MakePointerAvailable'}}
%0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor>, <MakePointerAvailable> :
@@ -166,6 +184,26 @@ spirv.func @cooperative_matrix_store_missing_attr(%ptr : !spirv.ptr<i32, Storage
// -----
+spirv.func @cooperative_matrix_store_missing_stride_row_major(%ptr : !spirv.ptr<i32, StorageBuffer>,
+ %m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
+ // expected-error @+1 {{op Stride is required for 'RowMajor'}}
+ spirv.KHR.CooperativeMatrixStore %ptr, %m, <RowMajor> :
+ !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+ spirv.Return
+}
+
+// -----
+
+spirv.func @cooperative_matrix_store_missing_stride_column_major(%ptr : !spirv.ptr<i32, StorageBuffer>,
+ %m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
+ // expected-error @+1 {{op Stride is required for 'ColumnMajor'}}
+ spirv.KHR.CooperativeMatrixStore %ptr, %m, <ColumnMajor> :
+ !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+ spirv.Return
+}
+
+// -----
+
spirv.func @cooperative_matrix_store_missing_attr(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32,
%m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
// expected-error @+1 {{expected '<'}}
``````````
</details>
https://github.com/llvm/llvm-project/pull/214194
More information about the Mlir-commits
mailing list