[Mlir-commits] [mlir] [mlir][spirv] Fix coop matrix load (PR #65712)
Jakub Kuderski
llvmlistbot at llvm.org
Thu Sep 7 21:43:23 PDT 2023
https://github.com/kuhar created https://github.com/llvm/llvm-project/pull/65712:
- Allow for stride to be any integer type
- Use ODS for parsing/printing
- Update examples and tests
>From 7d64a7d9dbc2543b0001af705a5383260f41ff2f Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 8 Sep 2023 00:41:50 -0400
Subject: [PATCH] [mlir][spirv] Fix coop matrix load
- Allow for stride to be any integer type
- Use ODS for parsing/printing
- Update examples and tests
---
.../SPIRV/IR/SPIRVCooperativeMatrixOps.td | 19 +++++--
.../Dialect/SPIRV/IR/CooperativeMatrixOps.cpp | 43 ---------------
.../SPIRV/IR/cooperative-matrix-ops.mlir | 53 +++++++++++--------
3 files changed, 45 insertions(+), 70 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
index 7060aa80dc113ed..5d6c5f057ae6acc 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td
@@ -101,20 +101,29 @@ def SPIRV_KHRCooperativeMatrixLoadOp : SPIRV_KhrVendorOp<"CooperativeMatrixLoad"
``` {.ebnf}
cooperative-matrix-load-op ::= ssa-id `=` `spirv.KHR.CooperativeMatrixLoad`
ssa-use `,` ssa-use `,`
- cooperative-matrix-layout `,`
- (`[` memory-operand `]`)? ` : `
- pointer-type `as` cooperative-matrix-type
+ `<` cooperative-matrix-layout `>`
+ (`,` `<` memory-operand `>`)? `:`
+ pointer-type `as` cooperative-matrix-type `,` stride-type
```
#### Example:
```
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, RowMajor
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor>
: !spirv.ptr<i32, StorageBuffer>
- as !spirv.KHR.coopmatrix<16x8xi32, Workgroup, MatrixA>
+ as !spirv.KHR.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
+
+ %1 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor>, <Volatile>
+ : !spirv.ptr<f32, StorageBuffer>
+ as !spirv.KHR.coopmatrix<8x8xf32, Subgroup, MatrixAcc>, i64
```
}];
+ let assemblyFormat = [{
+ $pointer `,` $stride `,` $matrix_layout ( `,` $memory_operand^ )? attr-dict `:`
+ type($pointer) `as` type($result) `,` type($stride)
+ }];
+
let availability = [
MinVersion<SPIRV_V_1_6>,
MaxVersion<SPIRV_V_1_6>,
diff --git a/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp b/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
index bc1d30f55518300..4f986065d8d9cd3 100644
--- a/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/CooperativeMatrixOps.cpp
@@ -37,49 +37,6 @@ LogicalResult KHRCooperativeMatrixLengthOp::verify() {
// spirv.KHR.CooperativeMatrixLoad
//===----------------------------------------------------------------------===//
-ParseResult KHRCooperativeMatrixLoadOp::parse(OpAsmParser &parser,
- OperationState &result) {
- std::array<OpAsmParser::UnresolvedOperand, 2> operandInfo = {};
- if (parser.parseOperand(operandInfo[0]) || parser.parseComma())
- return failure();
- if (parser.parseOperand(operandInfo[1]) || parser.parseComma())
- return failure();
-
- CooperativeMatrixLayoutKHR layout;
- if (parseEnumKeywordAttr<CooperativeMatrixLayoutKHRAttr>(
- layout, parser, result, kKhrCooperativeMatrixLayoutAttrName)) {
- return failure();
- }
-
- if (parseMemoryAccessAttributes(parser, result, kMemoryOperandAttrName))
- return failure();
-
- Type ptrType;
- Type elementType;
- if (parser.parseColon() || parser.parseType(ptrType) ||
- parser.parseKeywordType("as", elementType)) {
- return failure();
- }
- result.addTypes(elementType);
-
- Type strideType = parser.getBuilder().getIntegerType(32);
- if (parser.resolveOperands(operandInfo, {ptrType, strideType},
- parser.getNameLoc(), result.operands)) {
- return failure();
- }
-
- return success();
-}
-
-void KHRCooperativeMatrixLoadOp::print(OpAsmPrinter &printer) {
- printer << " " << getPointer() << ", " << getStride() << ", "
- << getMatrixLayout();
- // Print optional memory operand attribute.
- if (auto memOperand = getMemoryOperand())
- printer << " [\"" << memOperand << "\"]";
- printer << " : " << getPointer().getType() << " as " << getType();
-}
-
static LogicalResult verifyPointerAndCoopMatrixType(Operation *op, Type pointer,
Type coopMatrix) {
auto pointerType = cast<PointerType>(pointer);
diff --git a/mlir/test/Dialect/SPIRV/IR/cooperative-matrix-ops.mlir b/mlir/test/Dialect/SPIRV/IR/cooperative-matrix-ops.mlir
index aa6e072b03c5d3d..4a34674b810309d 100644
--- a/mlir/test/Dialect/SPIRV/IR/cooperative-matrix-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/cooperative-matrix-ops.mlir
@@ -23,37 +23,46 @@ spirv.func @cooperative_matrix_length_wrong_matrix() -> i32 "None" {
// CHECK-LABEL: @cooperative_matrix_load
spirv.func @cooperative_matrix_load(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
- // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, RowMajor :
- // CHECK-SAME: !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, RowMajor :
- !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>
+ // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor> :
+ // CHECK-SAME: !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor> :
+ !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
spirv.Return
}
// CHECK-LABEL: @cooperative_matrix_load_memoperand
spirv.func @cooperative_matrix_load_memoperand(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
- // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, ColumnMajor ["Volatile"] :
- // CHECK-SAME: !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, ColumnMajor ["Volatile"] :
- !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>
+ // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <ColumnMajor>, <Volatile> :
+ // CHECK-SAME: !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor>, <Volatile> :
+ !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
spirv.Return
}
// CHECK-LABEL: @cooperative_matrix_load_vector_ptr_type
spirv.func @cooperative_matrix_load_vector_ptr_type(%ptr : !spirv.ptr<vector<4xi32>, StorageBuffer>, %stride : i32) "None" {
- // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, RowMajor ["Volatile"] :
- // CHECK-SAME: !spirv.ptr<vector<4xi32>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, RowMajor ["Volatile"] :
- !spirv.ptr<vector<4xi32>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+ // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor>, <Volatile> :
+ // CHECK-SAME: !spirv.ptr<vector<4xi32>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>, i32
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor>, <Volatile> :
+ !spirv.ptr<vector<4xi32>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>, i32
spirv.Return
}
// CHECK-LABEL: @cooperative_matrix_load_function
spirv.func @cooperative_matrix_load_function(%ptr : !spirv.ptr<i32, Function>, %stride : i32) "None" {
- // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, RowMajor :
+ // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor> :
// CHECK-SAME: !spirv.ptr<i32, Function> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixAcc>
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, RowMajor :
- !spirv.ptr<i32, Function> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixAcc>
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor> :
+ !spirv.ptr<i32, Function> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixAcc>, i32
+ spirv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_load_stride_i16
+spirv.func @cooperative_matrix_load_stride_i16(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i16) "None" {
+ // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor> :
+ // CHECK-SAME: !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i16
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor> :
+ !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i16
spirv.Return
}
@@ -82,8 +91,8 @@ spirv.func @cooperative_matrix_store_memoperand(%ptr : !spirv.ptr<i32, StorageBu
spirv.func @cooperative_matrix_load_bad_ptr(%ptr : !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer>, %stride : i32) "None" {
// expected-error @+1 {{Pointer must point to a scalar or vector type}}
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, ColumnMajor :
- !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor> :
+ !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>, i32
spirv.Return
}
@@ -92,16 +101,16 @@ spirv.func @cooperative_matrix_load_bad_ptr(%ptr : !spirv.ptr<!spirv.struct<(f32
spirv.func @cooperative_matrix_load_missing_attr(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
// expected-error @+1 {{expected ','}}
%0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride :
- !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+ !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>, i32
spirv.Return
}
// -----
spirv.func @cooperative_matrix_load_missing_attr(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
- // expected-error @+1 {{expected valid keyword}}
+ // expected-error @+1 {{expected '<'}}
%0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, :
- !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup, MatrixA>
+ !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup, MatrixA>, i32
spirv.Return
}
@@ -109,8 +118,8 @@ spirv.func @cooperative_matrix_load_missing_attr(%ptr : !spirv.ptr<i32, StorageB
spirv.func @cooperative_matrix_load_bad_result(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
// expected-error @+1 {{op result #0 must be any SPIR-V cooperative matrix type}}
- %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, ColumnMajor :
- !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>
+ %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor> :
+ !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>, i32
spirv.Return
}
More information about the Mlir-commits
mailing list