MLIR  18.0.0git
CooperativeMatrixOps.cpp
Go to the documentation of this file.
1 //===- CooperativeMatrixOps.cpp - MLIR SPIR-V Cooperative Matrix Ops -----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Defines the Cooperative Matrix operations in the SPIR-V dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "SPIRVParsingUtils.h"
17 #include "llvm/ADT/STLExtras.h"
18 #include <cstdint>
19 
20 using namespace mlir::spirv::AttrNames;
21 
22 namespace mlir::spirv {
23 
24 static LogicalResult
25 verifyCoopMatrixAccess(Operation *op, Type pointer, Type coopMatrix,
26  spirv::MemoryAccessAttr memoryOperand) {
27  auto pointerType = cast<PointerType>(pointer);
28  Type pointeeType = pointerType.getPointeeType();
29  if (!isa<ScalarType, VectorType>(pointeeType)) {
30  return op->emitOpError(
31  "Pointer must point to a scalar or vector type but provided ")
32  << pointeeType;
33  }
34 
35  if (memoryOperand) {
36  spirv::MemoryAccess operandSet = memoryOperand.getValue();
37 
38  if (isa<spirv::KHRCooperativeMatrixLoadOp>(op) &&
39  spirv::bitEnumContainsAll(operandSet,
40  spirv::MemoryAccess::MakePointerAvailable)) {
41  return op->emitOpError(
42  "not compatible with memory operand 'MakePointerAvailable'");
43  }
44 
45  if (isa<spirv::KHRCooperativeMatrixStoreOp>(op) &&
46  spirv::bitEnumContainsAll(operandSet,
47  spirv::MemoryAccess::MakePointerVisible)) {
48  return op->emitOpError(
49  "not compatible with memory operand 'MakePointerVisible'");
50  }
51 
52  // The 'Aligned' memory operand requires an alignment literal to follow,
53  // which needs to be implemented on the level of op parsing and
54  // (de-)serialization.
55  // TODO: Consider adding support for this attribute value.
56  if (spirv::bitEnumContainsAll(memoryOperand.getValue(),
57  spirv::MemoryAccess::Aligned)) {
58  return op->emitOpError("has unhandled memory operand 'Aligned'");
59  }
60  }
61 
62  // TODO: Verify the memory object behind the pointer:
63  // > If the Shader capability was declared, Pointer must point into an array
64  // > and any ArrayStride decoration on Pointer is ignored.
65 
66  return success();
67 }
68 
69 //===----------------------------------------------------------------------===//
70 // spirv.KHR.CooperativeMatrixLoad
71 //===----------------------------------------------------------------------===//
72 
74  return verifyCoopMatrixAccess(*this, getPointer().getType(),
75  getResult().getType(), getMemoryOperandAttr());
76 }
77 
78 //===----------------------------------------------------------------------===//
79 // spirv.KHR.CooperativeMatrixStore
80 //===----------------------------------------------------------------------===//
81 
83  return verifyCoopMatrixAccess(*this, getPointer().getType(),
84  getObject().getType(), getMemoryOperandAttr());
85 }
86 
87 //===----------------------------------------------------------------------===//
88 // spirv.KHR.CooperativeMatrixMulAdd
89 //===----------------------------------------------------------------------===//
90 
92  auto typeA = cast<spirv::CooperativeMatrixType>(getA().getType());
93  auto typeB = cast<spirv::CooperativeMatrixType>(getB().getType());
94  auto typeC = cast<spirv::CooperativeMatrixType>(getC().getType());
95 
96  // Check element types. ODS enforces that `type(c) == type(result)`, so no
97  // need to check it here.
98 
99  // Check the 'use' part of the type against the operands and the result.
100  if (typeA.getUse() != CooperativeMatrixUseKHR::MatrixA)
101  return emitOpError("operand #0 must be of use 'MatrixA'");
102  if (typeB.getUse() != CooperativeMatrixUseKHR::MatrixB)
103  return emitOpError("operand #1 must be of use 'MatrixB'");
104  if (typeC.getUse() != CooperativeMatrixUseKHR::MatrixAcc)
105  return emitOpError("operand #2 must be of use 'MatrixAcc'");
106 
107  // Check the 'scope' part of the type.
108  if (!llvm::all_equal({typeA.getScope(), typeB.getScope(), typeC.getScope()}))
109  return emitOpError("matrix scope mismatch");
110 
111  // Check dimension sizes. We expect 'MxK * KxN + MxN -> MxN'.
112  if (typeA.getRows() != typeC.getRows())
113  return emitOpError("matrix size mismatch on dimension 'M'");
114  if (typeB.getColumns() != typeC.getColumns())
115  return emitOpError("matrix size mismatch on dimension 'N'");
116  if (typeA.getColumns() != typeB.getRows())
117  return emitOpError("matrix size mismatch on dimension 'K'");
118 
119  // The spec does not restrict the element types:
120  // > A, B, C, and Result Type need not necessarily have the same component
121  // > type, this is defined by the client API.
122 
123  // Check that if Cooperative Matrix Operands are provided, the element type
124  // is integer.
125  if (getMatrixOperands()) {
126  Type elementTypes[] = {typeA.getElementType(), typeB.getElementType(),
127  typeC.getElementType()};
128  if (!llvm::all_of(elementTypes,
129  [](Type ty) { return isa<IntegerType>(ty); })) {
130  return emitOpError("Matrix Operands require all matrix element types to "
131  "be Integer Types");
132  }
133  }
134 
135  // Any further requirements need to be checked against VCE.
136  return success();
137 }
138 
139 //===----------------------------------------------------------------------===//
140 // spirv.NV.CooperativeMatrixLength
141 //===----------------------------------------------------------------------===//
142 
143 LogicalResult NVCooperativeMatrixLengthOp::verify() {
144  if (!isa<CooperativeMatrixNVType>(getCooperativeMatrixType())) {
145  return emitOpError(
146  "type attribute must be a '!spirv.NV.coopmatrix' type, found ")
147  << getCooperativeMatrixType() << " instead";
148  }
149 
150  return success();
151 }
152 
153 //===----------------------------------------------------------------------===//
154 // spirv.NV.CooperativeMatrixLoad
155 //===----------------------------------------------------------------------===//
156 
157 ParseResult NVCooperativeMatrixLoadOp::parse(OpAsmParser &parser,
158  OperationState &result) {
159  SmallVector<OpAsmParser::UnresolvedOperand, 3> operandInfo;
160  Type strideType = parser.getBuilder().getIntegerType(32);
161  Type columnMajorType = parser.getBuilder().getIntegerType(1);
162  Type ptrType;
163  Type elementType;
164  if (parser.parseOperandList(operandInfo, 3) ||
165  parseMemoryAccessAttributes(parser, result) || parser.parseColon() ||
166  parser.parseType(ptrType) || parser.parseKeywordType("as", elementType)) {
167  return failure();
168  }
169  if (parser.resolveOperands(operandInfo,
170  {ptrType, strideType, columnMajorType},
171  parser.getNameLoc(), result.operands)) {
172  return failure();
173  }
174 
175  result.addTypes(elementType);
176  return success();
177 }
178 
179 void NVCooperativeMatrixLoadOp::print(OpAsmPrinter &printer) {
180  printer << " " << getPointer() << ", " << getStride() << ", "
181  << getColumnmajor();
182  // Print optional memory access attribute.
183  if (auto memAccess = getMemoryAccess())
184  printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
185  printer << " : " << getPointer().getType() << " as " << getType();
186 }
187 
188 static LogicalResult
190  Type pointeeType = llvm::cast<PointerType>(pointer).getPointeeType();
191  if (!llvm::isa<ScalarType>(pointeeType) &&
192  !llvm::isa<VectorType>(pointeeType))
193  return op->emitError(
194  "Pointer must point to a scalar or vector type but provided ")
195  << pointeeType;
196  StorageClass storage = llvm::cast<PointerType>(pointer).getStorageClass();
197  if (storage != StorageClass::Workgroup &&
198  storage != StorageClass::StorageBuffer &&
199  storage != StorageClass::PhysicalStorageBuffer)
200  return op->emitError(
201  "Pointer storage class must be Workgroup, StorageBuffer or "
202  "PhysicalStorageBufferEXT but provided ")
203  << stringifyStorageClass(storage);
204  return success();
205 }
206 
208  return verifyPointerAndCoopMatrixNVType(*this, getPointer().getType(),
209  getResult().getType());
210 }
211 
212 //===----------------------------------------------------------------------===//
213 // spirv.NV.CooperativeMatrixStore
214 //===----------------------------------------------------------------------===//
215 
216 ParseResult NVCooperativeMatrixStoreOp::parse(OpAsmParser &parser,
217  OperationState &result) {
218  SmallVector<OpAsmParser::UnresolvedOperand, 4> operandInfo;
219  Type strideType = parser.getBuilder().getIntegerType(32);
220  Type columnMajorType = parser.getBuilder().getIntegerType(1);
221  Type ptrType;
222  Type elementType;
223  if (parser.parseOperandList(operandInfo, 4) ||
224  parseMemoryAccessAttributes(parser, result) || parser.parseColon() ||
225  parser.parseType(ptrType) || parser.parseComma() ||
226  parser.parseType(elementType)) {
227  return failure();
228  }
229  if (parser.resolveOperands(
230  operandInfo, {ptrType, elementType, strideType, columnMajorType},
231  parser.getNameLoc(), result.operands)) {
232  return failure();
233  }
234 
235  return success();
236 }
237 
238 void NVCooperativeMatrixStoreOp::print(OpAsmPrinter &printer) {
239  printer << " " << getPointer() << ", " << getObject() << ", " << getStride()
240  << ", " << getColumnmajor();
241  // Print optional memory access attribute.
242  if (auto memAccess = getMemoryAccess())
243  printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
244  printer << " : " << getPointer().getType() << ", " << getOperand(1).getType();
245 }
246 
247 LogicalResult NVCooperativeMatrixStoreOp::verify() {
248  return verifyPointerAndCoopMatrixNVType(*this, getPointer().getType(),
249  getObject().getType());
250 }
251 
252 //===----------------------------------------------------------------------===//
253 // spirv.NV.CooperativeMatrixMulAdd
254 //===----------------------------------------------------------------------===//
255 
256 static LogicalResult verifyCoopMatrixMulAddNV(NVCooperativeMatrixMulAddOp op) {
257  if (op.getC().getType() != op.getResult().getType())
258  return op.emitOpError("result and third operand must have the same type");
259  auto typeA = llvm::cast<CooperativeMatrixNVType>(op.getA().getType());
260  auto typeB = llvm::cast<CooperativeMatrixNVType>(op.getB().getType());
261  auto typeC = llvm::cast<CooperativeMatrixNVType>(op.getC().getType());
262  auto typeR = llvm::cast<CooperativeMatrixNVType>(op.getResult().getType());
263  if (typeA.getRows() != typeR.getRows() ||
264  typeA.getColumns() != typeB.getRows() ||
265  typeB.getColumns() != typeR.getColumns())
266  return op.emitOpError("matrix size must match");
267  if (typeR.getScope() != typeA.getScope() ||
268  typeR.getScope() != typeB.getScope() ||
269  typeR.getScope() != typeC.getScope())
270  return op.emitOpError("matrix scope must match");
271  auto elementTypeA = typeA.getElementType();
272  auto elementTypeB = typeB.getElementType();
273  if (isa<IntegerType>(elementTypeA) && isa<IntegerType>(elementTypeB)) {
274  if (llvm::cast<IntegerType>(elementTypeA).getWidth() !=
275  llvm::cast<IntegerType>(elementTypeB).getWidth())
276  return op.emitOpError(
277  "matrix A and B integer element types must be the same bit width");
278  } else if (elementTypeA != elementTypeB) {
279  return op.emitOpError(
280  "matrix A and B non-integer element types must match");
281  }
282  if (typeR.getElementType() != typeC.getElementType())
283  return op.emitOpError("matrix accumulator element type must match");
284  return success();
285 }
286 
288  return verifyCoopMatrixMulAddNV(*this);
289 }
290 
291 } // namespace mlir::spirv
static void print(spirv::VerCapExtAttr triple, DialectAsmPrinter &printer)
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
@ Type
An inlay hint that for a type annotation.
QueryRef parse(llvm::StringRef line, const QuerySession &qs)
Definition: Query.cpp:19
ParseResult parseMemoryAccessAttributes(OpAsmParser &parser, OperationState &state, StringRef attrName)
Parses optional memory access (a.k.a.
static LogicalResult verifyCoopMatrixAccess(Operation *op, Type pointer, Type coopMatrix, spirv::MemoryAccessAttr memoryOperand)
static LogicalResult verifyCoopMatrixMulAddNV(NVCooperativeMatrixMulAddOp op)
static LogicalResult verifyPointerAndCoopMatrixNVType(Operation *op, Type pointer, Type coopMatrix)
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
LogicalResult verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs,...
Definition: Verifier.cpp:421
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26