[Mlir-commits] [mlir] [mlir][SPIR-V] Make Stride optional in KHR CooperativeMatrixLoad/Store per spec (PR #214194)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Thu Oct 1 03:15:10 PDT 2026
https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/214194
>From 4b62c57daa0877134a785858bf0967b00f36cad9 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 5 Aug 2026 12:54:03 +0200
Subject: [PATCH 1/2] [mlir][SPIR-V] Make Stride optional in KHR
CooperativeMatrixLoad/Store per spec
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
---
.../SPIRV/IR/SPIRVCooperativeMatrixOps.td | 28 +--
.../Dialect/SPIRV/IR/CooperativeMatrixOps.cpp | 24 ++-
.../SPIRV/Deserialization/DeserializeOps.cpp | 160 ++++++++++++++++++
.../SPIRV/Serialization/SerializeOps.cpp | 78 +++++++++
.../SPIRV/IR/khr-cooperative-matrix-ops.mlir | 38 +++++
5 files changed, 309 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
index e8124b8b0bed964..8afbfc4a3e339c1 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 eccc01137576ac6..ca41e1361214c73 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 0e6d199eeabf0d6..1ad4072016671a8 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 32fa603e6c23d52..b81af1778f76c6f 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 69235eab8d0dc6d..2566ca557b7240f 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 '<'}}
>From 9c9b0654507e23979a4971450edc060c367a7066 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 1 Oct 2026 12:14:53 +0200
Subject: [PATCH 2/2] autogen
---
.../SPIRV/IR/SPIRVCooperativeMatrixOps.td | 10 --
.../SPIRV/Deserialization/DeserializeOps.cpp | 160 ------------------
.../SPIRV/Serialization/SerializeOps.cpp | 78 ---------
mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp | 9 +-
4 files changed, 4 insertions(+), 253 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
index 8afbfc4a3e339c1..feabefbe0bf573b 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
@@ -122,11 +122,6 @@ 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,
@@ -208,11 +203,6 @@ 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,
diff --git a/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp b/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
index 1ad4072016671a8..0e6d199eeabf0d6 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/DeserializeOps.cpp
@@ -753,166 +753,6 @@ 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 b81af1778f76c6f..32fa603e6c23d52 100644
--- a/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
+++ b/mlir/lib/Target/SPIRV/Serialization/SerializeOps.cpp
@@ -1174,84 +1174,6 @@ 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/tools/mlir-tblgen/SPIRVUtilsGen.cpp b/mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp
index 4de8c60dfb55aa1..784d29628fb7486 100644
--- a/mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp
+++ b/mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp
@@ -953,11 +953,10 @@ static void emitOperandDeserialization(const Operator &op, ArrayRef<SMLoc> loc,
auto argument = op.getArg(i);
if (auto *valueArg =
llvm::dyn_cast_if_present<NamedTypeConstraint *>(argument)) {
- if (valueArg->isVariableLength()) {
+ if (valueArg->isVariadic()) {
if (i != e - 1) {
- PrintFatalError(
- loc, "SPIR-V ops can have Variadic<..> or "
- "Optional<...> arguments only if it's the last argument");
+ PrintFatalError(loc, "SPIR-V ops can have Variadic<..> arguments "
+ "only if it's the last argument");
}
os << tabs
<< formatv("for (; {0} < {1}.size(); ++{0})", wordIndex, words);
@@ -975,7 +974,7 @@ static void emitOperandDeserialization(const Operator &op, ArrayRef<SMLoc> loc,
words, wordIndex);
os << tabs << " }\n";
os << tabs << formatv(" {0}.push_back(arg);\n", operands);
- if (!valueArg->isVariableLength()) {
+ if (!valueArg->isVariadic()) {
os << tabs << formatv(" {0}++;\n", wordIndex);
}
os << tabs << "}\n";
More information about the Mlir-commits
mailing list