1 //===- NVVMDialect.cpp - NVVM IR Ops and Dialect registration -------------===//
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 // This file defines the types and operation details for the NVVM IR dialect in
10 // MLIR, and the LLVM IR dialect.  It also registers the dialect.
11 //
12 // The NVVM dialect only contains GPU specific additions on top of the general
13 // LLVM dialect.
14 //
15 //===----------------------------------------------------------------------===//
16 
17 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
18 
19 #include "mlir/IR/Builders.h"
20 #include "mlir/IR/BuiltinTypes.h"
21 #include "mlir/IR/DialectImplementation.h"
22 #include "mlir/IR/MLIRContext.h"
23 #include "mlir/IR/Operation.h"
24 #include "mlir/IR/OperationSupport.h"
25 #include "llvm/ADT/TypeSwitch.h"
26 #include "llvm/AsmParser/Parser.h"
27 #include "llvm/IR/Attributes.h"
28 #include "llvm/IR/Function.h"
29 #include "llvm/IR/Type.h"
30 #include "llvm/Support/SourceMgr.h"
31 
32 using namespace mlir;
33 using namespace NVVM;
34 
35 #include "mlir/Dialect/LLVMIR/NVVMOpsDialect.cpp.inc"
36 #include "mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc"
37 #include "mlir/Dialect/LLVMIR/NVVMOpsStructs.cpp.inc"
38 
39 //===----------------------------------------------------------------------===//
40 // Printing/parsing for NVVM ops
41 //===----------------------------------------------------------------------===//
42 
43 static void printNVVMIntrinsicOp(OpAsmPrinter &p, Operation *op) {
44   p << " " << op->getOperands();
45   if (op->getNumResults() > 0)
46     p << " : " << op->getResultTypes();
47 }
48 
49 // <operation> ::= `llvm.nvvm.vote.ballot.sync %mask, %pred` : result_type
50 ParseResult VoteBallotOp::parse(OpAsmParser &parser, OperationState &result) {
51   MLIRContext *context = parser.getContext();
52   auto int32Ty = IntegerType::get(context, 32);
53   auto int1Ty = IntegerType::get(context, 1);
54 
55   SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
56   Type type;
57   return failure(parser.parseOperandList(ops) ||
58                  parser.parseOptionalAttrDict(result.attributes) ||
59                  parser.parseColonType(type) ||
60                  parser.addTypeToList(type, result.types) ||
61                  parser.resolveOperands(ops, {int32Ty, int1Ty},
62                                         parser.getNameLoc(), result.operands));
63 }
64 
65 void VoteBallotOp::print(OpAsmPrinter &p) { printNVVMIntrinsicOp(p, *this); }
66 
67 LogicalResult CpAsyncOp::verify() {
68   if (size() != 4 && size() != 8 && size() != 16)
69     return emitError("expected byte size to be either 4, 8 or 16.");
70   return success();
71 }
72 
73 // Given the element type of an operand and whether or not it is an accumulator,
74 // this function returns the PTX type (`NVVM::MMATypes`) that corresponds to the
75 // operand's element type.
76 Optional<mlir::NVVM::MMATypes> MmaOp::inferOperandMMAType(Type operandElType,
77                                                           bool isAccumulator) {
78   auto half2Type =
79       LLVM::getFixedVectorType(Float16Type::get(operandElType.getContext()), 2);
80   if (operandElType.isF64())
81     return NVVM::MMATypes::f64;
82   if (operandElType.isF16() || operandElType == half2Type)
83     return NVVM::MMATypes::f16;
84   if (operandElType.isF32())
85     return NVVM::MMATypes::f32;
86   if (operandElType.isa<IntegerType>()) {
87     if (isAccumulator)
88       return NVVM::MMATypes::s32;
89     return llvm::None;
90   }
91 
92   if (auto structType = operandElType.dyn_cast<LLVM::LLVMStructType>()) {
93     if (structType.getBody().empty())
94       return llvm::None;
95     return inferOperandMMAType(structType.getBody()[0], isAccumulator);
96   }
97 
98   return llvm::None;
99 }
100 
101 static bool isInt4PtxType(MMATypes type) {
102   return (type == MMATypes::u4 || type == MMATypes::s4);
103 }
104 
105 static bool isInt8PtxType(MMATypes type) {
106   return (type == MMATypes::u8 || type == MMATypes::s8);
107 }
108 
109 static bool isIntegerPtxType(MMATypes type) {
110   return isInt4PtxType(type) || isInt8PtxType(type) || type == MMATypes::b1 ||
111          type == MMATypes::s32;
112 }
113 
114 MMATypes MmaOp::accumPtxType() {
115   Optional<mlir::NVVM::MMATypes> val = inferOperandMMAType(
116       getODSOperands(2).getTypes().front(), /*isAccum=*/true);
117   assert(val.hasValue() && "accumulator PTX type should always be inferrable");
118   return val.getValue();
119 }
120 
121 MMATypes MmaOp::resultPtxType() {
122   Optional<mlir::NVVM::MMATypes> val =
123       inferOperandMMAType(getResult().getType(), /*isAccum=*/true);
124   assert(val.hasValue() && "result PTX type should always be inferrable");
125   return val.getValue();
126 }
127 
128 void MmaOp::print(OpAsmPrinter &p) {
129   SmallVector<Type, 4> regTypes;
130   struct OperandFragment {
131     StringRef operandName;
132     StringRef ptxTypeAttr;
133     SmallVector<Value, 4> regs;
134     explicit OperandFragment(StringRef name, StringRef ptxTypeName)
135         : operandName(name), ptxTypeAttr(ptxTypeName) {}
136   };
137 
138   std::array<OperandFragment, 3> frags{
139       OperandFragment("A", multiplicandAPtxTypeAttrName()),
140       OperandFragment("B", multiplicandBPtxTypeAttrName()),
141       OperandFragment("C", "")};
142   SmallVector<StringRef, 4> ignoreAttrNames{
143       mlir::NVVM::MmaOp::getOperandSegmentSizeAttr()};
144 
145   for (unsigned fragIdx = 0; fragIdx < frags.size(); fragIdx++) {
146     auto &frag = frags[fragIdx];
147     auto varOperandSpec = getODSOperandIndexAndLength(fragIdx);
148     for (auto operandIdx = varOperandSpec.first;
149          operandIdx < varOperandSpec.first + varOperandSpec.second;
150          operandIdx++) {
151       frag.regs.push_back(this->getOperand(operandIdx));
152       if (operandIdx == 0) {
153         regTypes.push_back(this->getOperand(operandIdx).getType());
154       }
155     }
156     Optional<MMATypes> inferredType =
157         inferOperandMMAType(regTypes.back(), /*isAccum=*/fragIdx >= 2);
158     if (inferredType)
159       ignoreAttrNames.push_back(frag.ptxTypeAttr);
160   }
161 
162   auto printMmaOperand = [&](const OperandFragment &frag) -> void {
163     p << " " << frag.operandName;
164     p << "[";
165     p.printOperands(frag.regs);
166     p << "] ";
167   };
168 
169   for (const auto &frag : frags) {
170     printMmaOperand(frag);
171   }
172 
173   p.printOptionalAttrDict(this->getOperation()->getAttrs(), ignoreAttrNames);
174 
175   // Print the types of the operands and result.
176   p << " : "
177     << "(";
178   llvm::interleaveComma(SmallVector<Type, 3>{frags[0].regs[0].getType(),
179                                              frags[1].regs[0].getType(),
180                                              frags[2].regs[0].getType()},
181                         p);
182   p << ")";
183   p.printArrowTypeList(TypeRange{this->res().getType()});
184 }
185 
186 void MmaOp::build(OpBuilder &builder, OperationState &result, Type resultType,
187                   ValueRange operandA, ValueRange operandB, ValueRange operandC,
188                   ArrayRef<int64_t> shape, Optional<MMAB1Op> b1Op,
189                   Optional<MMAIntOverflow> intOverflow,
190                   Optional<std::array<MMATypes, 2>> multiplicandPtxTypes,
191                   Optional<std::array<MMALayout, 2>> multiplicandLayouts) {
192 
193   assert(shape.size() == 3 && "expected shape to have size 3 (m, n, k)");
194   MLIRContext *ctx = builder.getContext();
195   Type i32 = builder.getIntegerType(32);
196   result.addAttribute(
197       "shape", MMAShapeAttr::get(builder.getIntegerAttr(i32, shape[0]),
198                                  builder.getIntegerAttr(i32, shape[1]),
199                                  builder.getIntegerAttr(i32, shape[2]), ctx));
200 
201   result.addOperands(operandA);
202   result.addOperands(operandB);
203   result.addOperands(operandC);
204 
205   if (multiplicandPtxTypes.hasValue()) {
206     result.addAttribute("multiplicandAPtxType",
207                         MMATypesAttr::get(ctx, (*multiplicandPtxTypes)[0]));
208     result.addAttribute("multiplicandBPtxType",
209                         MMATypesAttr::get(ctx, (*multiplicandPtxTypes)[1]));
210   } else {
211     if (auto res = inferOperandMMAType(operandA[0].getType(), false))
212       result.addAttribute("multiplicandAPtxType", MMATypesAttr::get(ctx, *res));
213     if (auto res = inferOperandMMAType(operandB[0].getType(), false))
214       result.addAttribute("multiplicandBPtxType", MMATypesAttr::get(ctx, *res));
215   }
216 
217   if (multiplicandLayouts.hasValue()) {
218     result.addAttribute("layoutA",
219                         MMALayoutAttr::get(ctx, (*multiplicandLayouts)[0]));
220     result.addAttribute("layoutB",
221                         MMALayoutAttr::get(ctx, (*multiplicandLayouts)[1]));
222   } else {
223     result.addAttribute("layoutA", MMALayoutAttr::get(ctx, MMALayout::row));
224     result.addAttribute("layoutB", MMALayoutAttr::get(ctx, MMALayout::col));
225   }
226 
227   if (intOverflow.hasValue())
228     result.addAttribute("intOverflowBehavior",
229                         MMAIntOverflowAttr::get(ctx, *intOverflow));
230   if (b1Op.hasValue())
231     result.addAttribute("b1Op", MMAB1OpAttr::get(ctx, *b1Op));
232 
233   result.addTypes(resultType);
234   result.addAttribute(
235       MmaOp::getOperandSegmentSizeAttr(),
236       builder.getI32VectorAttr({static_cast<int32_t>(operandA.size()),
237                                 static_cast<int32_t>(operandB.size()),
238                                 static_cast<int32_t>(operandC.size())}));
239 }
240 
241 // <operation> :=
242 //   A `[` $operandA `]` B `[` $operandB `]` C `[` $operandC `]`
243 //   attr-dict : (type($operandA[0]), type($operandB[0]), type($operandC[0]))
244 //     `->` type($res)
245 ParseResult MmaOp::parse(OpAsmParser &parser, OperationState &result) {
246   struct OperandFragment {
247     Optional<MMATypes> elemtype;
248     SmallVector<OpAsmParser::UnresolvedOperand, 4> regs;
249     SmallVector<Type> regTypes;
250   };
251 
252   Builder &builder = parser.getBuilder();
253   std::array<OperandFragment, 4> frags;
254 
255   NamedAttrList namedAttributes;
256 
257   // A helper to parse the operand segments.
258   auto parseMmaOperand = [&](StringRef operandName,
259                              OperandFragment &frag) -> LogicalResult {
260     if (parser.parseKeyword(operandName).failed())
261       return failure();
262     if (parser
263             .parseOperandList(frag.regs, OpAsmParser::Delimiter::OptionalSquare)
264             .failed())
265       return failure();
266     return success();
267   };
268 
269   // Parse the operand segments.
270   if (parseMmaOperand("A", frags[0]).failed())
271     return failure();
272   if (parseMmaOperand("B", frags[1]).failed())
273     return failure();
274   if (parseMmaOperand("C", frags[2]).failed())
275     return failure();
276 
277   if (parser.parseOptionalAttrDict(namedAttributes).failed())
278     return failure();
279 
280   // Parse the type specification and resolve operands.
281   SmallVector<Type, 3> operandTypes;
282   if (failed(parser.parseColon()))
283     return failure();
284   if (failed(parser.parseLParen()))
285     return failure();
286   if (failed(parser.parseTypeList(operandTypes)))
287     return failure();
288   if (failed(parser.parseRParen()))
289     if (operandTypes.size() != 3)
290       return parser.emitError(
291           parser.getNameLoc(),
292           "expected one type for each operand segment but got " +
293               Twine(operandTypes.size()) + " types");
294   for (const auto& iter : llvm::enumerate(operandTypes)) {
295     auto &frag = frags[iter.index()];
296     frag.regTypes.resize(frag.regs.size(), iter.value());
297     if (failed(parser.resolveOperands(frag.regs, frag.regTypes,
298                                       parser.getNameLoc(), result.operands)))
299       return failure();
300     frag.elemtype =
301         inferOperandMMAType(frag.regTypes[0], /*isAccum=*/iter.index() < 2);
302   }
303 
304   Type resultType;
305   parser.parseArrow();
306   parser.parseType(resultType);
307   frags[3].elemtype = inferOperandMMAType(resultType, /*isAccum=*/true);
308 
309   std::array<StringRef, 2> names{"multiplicandAPtxType",
310                                  "multiplicandBPtxType"};
311   for (unsigned idx = 0; idx < names.size(); idx++) {
312     const auto &frag = frags[idx];
313     Optional<NamedAttribute> attr = namedAttributes.getNamed(names[idx]);
314     if (!frag.elemtype.hasValue() && !attr.hasValue()) {
315       return parser.emitError(
316           parser.getNameLoc(),
317           "attribute " + names[idx] +
318               " is not provided explicitly and cannot be inferred");
319     }
320     if (!attr.hasValue())
321       result.addAttribute(
322           names[idx], MMATypesAttr::get(parser.getContext(), *frag.elemtype));
323   }
324 
325   result.addTypes(resultType);
326   if (!namedAttributes.empty())
327     result.addAttributes(namedAttributes);
328   result.addAttribute(MmaOp::getOperandSegmentSizeAttr(),
329                       builder.getI32VectorAttr({
330                           static_cast<int32_t>(frags[0].regs.size()),
331                           static_cast<int32_t>(frags[1].regs.size()),
332                           static_cast<int32_t>(frags[2].regs.size()),
333                       }));
334   return success();
335 }
336 
337 LogicalResult MmaOp::verify() {
338   MLIRContext *context = getContext();
339   auto f16Ty = Float16Type::get(context);
340   auto i32Ty = IntegerType::get(context, 32);
341   auto f16x2Ty = LLVM::getFixedVectorType(f16Ty, 2);
342   auto f32Ty = Float32Type::get(context);
343   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
344       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
345 
346   auto s32x4StructTy =
347       LLVM::LLVMStructType::getLiteral(context, {i32Ty, i32Ty, i32Ty, i32Ty});
348   auto f32x8StructTy =
349       LLVM::LLVMStructType::getLiteral(context, SmallVector<Type>(8, f32Ty));
350   auto f16x2x2StructTy =
351       LLVM::LLVMStructType::getLiteral(context, {f16x2Ty, f16x2Ty});
352   auto f32x4StructTy =
353       LLVM::LLVMStructType::getLiteral(context, {f32Ty, f32Ty, f32Ty, f32Ty});
354   auto s32x2StructTy =
355       LLVM::LLVMStructType::getLiteral(context, {i32Ty, i32Ty});
356 
357   std::array<int64_t, 3> mmaShape{shapeAttr().m().getInt(),
358                                   shapeAttr().n().getInt(),
359                                   shapeAttr().k().getInt()};
360 
361   // These variables define the set of allowed data types for matrices A, B, C,
362   // and result.
363   using AllowedShapes = SmallVector<std::array<int64_t, 3>, 2>;
364   using AllowedTypes = SmallVector<SmallVector<Type, 4>, 2>;
365   AllowedShapes allowedShapes;
366   AllowedTypes expectedA;
367   AllowedTypes expectedB;
368   AllowedTypes expectedC;
369   SmallVector<Type> expectedResult;
370 
371   // When M = 16, we just need to calculate the number of 8xk tiles, where
372   // k is a factor that depends on the data type.
373   if (mmaShape[0] == 16) {
374     int64_t kFactor;
375     Type multiplicandFragType;
376     switch (multiplicandAPtxType().getValue()) {
377     case MMATypes::tf32:
378       kFactor = 4;
379       expectedResult.push_back(LLVM::LLVMStructType::getLiteral(
380           context, {i32Ty, i32Ty, i32Ty, i32Ty}));
381       break;
382     case MMATypes::f16:
383     case MMATypes::bf16:
384       kFactor = 8;
385       multiplicandFragType = f16x2Ty;
386       expectedResult.push_back(f16x2x2StructTy);
387       expectedResult.push_back(f32x4StructTy);
388       break;
389     case MMATypes::s4:
390     case MMATypes::u4:
391       kFactor = 32;
392       break;
393     case MMATypes::b1:
394       kFactor = 128;
395       break;
396     case MMATypes::s8:
397     case MMATypes::u8:
398       kFactor = 16;
399       break;
400     default:
401       return emitError("invalid shape or multiplicand type: " +
402                        stringifyEnum(multiplicandAPtxType().getValue()));
403     }
404 
405     if (isIntegerPtxType(multiplicandAPtxType().getValue())) {
406       expectedResult.push_back(s32x4StructTy);
407       expectedC.emplace_back(4, i32Ty);
408       multiplicandFragType = i32Ty;
409     } else {
410       expectedC.emplace_back(2, f16x2Ty);
411       expectedC.emplace_back(4, f32Ty);
412     }
413 
414     int64_t unitA = (mmaShape[0] / 8) * (mmaShape[2] / kFactor);
415     int64_t unitB = (mmaShape[1] / 8) * (mmaShape[2] / kFactor);
416     expectedA.emplace_back(unitA, multiplicandFragType);
417     expectedB.emplace_back(unitB, multiplicandFragType);
418     allowedShapes.push_back({16, 8, kFactor});
419     allowedShapes.push_back({16, 8, kFactor * 2});
420   }
421 
422   // In the M=8 case, there is only 1 possible case per data type.
423   if (mmaShape[0] == 8) {
424     if (multiplicandAPtxType().getValue() == MMATypes::f16) {
425       expectedA.emplace_back(2, f16x2Ty);
426       expectedB.emplace_back(2, f16x2Ty);
427       expectedResult.push_back(f16x2x4StructTy);
428       expectedResult.push_back(f32x8StructTy);
429       expectedC.emplace_back(4, f16x2Ty);
430       expectedC.emplace_back(8, f32Ty);
431       allowedShapes.push_back({8, 8, 4});
432     }
433     if (multiplicandAPtxType().getValue() == MMATypes::f64) {
434       Type f64Ty = Float64Type::get(context);
435       expectedA.emplace_back(1, f64Ty);
436       expectedB.emplace_back(1, f64Ty);
437       expectedC.emplace_back(2, f64Ty);
438       // expectedC.emplace_back(1, LLVM::getFixedVectorType(f64Ty, 2));
439       expectedResult.emplace_back(LLVM::LLVMStructType::getLiteral(
440           context, SmallVector<Type>(2, f64Ty)));
441       allowedShapes.push_back({8, 8, 4});
442     }
443     if (isIntegerPtxType(multiplicandAPtxType().getValue())) {
444       expectedA.push_back({i32Ty});
445       expectedB.push_back({i32Ty});
446       expectedC.push_back({i32Ty, i32Ty});
447       expectedResult.push_back(s32x2StructTy);
448       if (isInt4PtxType(multiplicandAPtxType().getValue()))
449         allowedShapes.push_back({8, 8, 32});
450       if (isInt8PtxType(multiplicandAPtxType().getValue()))
451         allowedShapes.push_back({8, 8, 16});
452       if (multiplicandAPtxType().getValue() == MMATypes::b1)
453         allowedShapes.push_back({8, 8, 128});
454     }
455   }
456 
457   std::string errorMessage;
458   llvm::raw_string_ostream errorStream(errorMessage);
459 
460   // Check that we matched an existing shape/dtype combination.
461   if (expectedA.empty() || expectedB.empty() || expectedC.empty() ||
462       !llvm::any_of(allowedShapes,
463                     [&](const auto &allowed) { return allowed == mmaShape; })) {
464     errorStream << "unimplemented variant for MMA shape <";
465     llvm::interleaveComma(mmaShape, errorStream);
466     errorStream << ">";
467     return emitOpError(errorMessage);
468   }
469 
470   // Verify the operand types for segments of A, B, and C operands.
471   std::array<StringRef, 3> operandNames{"A", "B", "C"};
472   for (const auto &iter : llvm::enumerate(
473            SmallVector<AllowedTypes, 3>{expectedA, expectedB, expectedC})) {
474     auto spec = this->getODSOperandIndexAndLength(iter.index());
475     SmallVector<Type, 4> operandTySeg(operand_type_begin() + spec.first,
476                                       operand_type_begin() + spec.first +
477                                           spec.second);
478     bool match =
479         llvm::any_of(iter.value(), [&](const SmallVector<Type, 4> &typeSet) {
480           return typeSet == operandTySeg;
481         });
482 
483     if (!match) {
484       errorStream << "Could not match types for the "
485                   << operandNames[iter.index()]
486                   << " operands; expected one of ";
487       for (const auto &x : iter.value()) {
488         errorStream << x.size() << "x" << x[0] << " ";
489       }
490       errorStream << "but got ";
491       llvm::interleaveComma(operandTySeg, errorStream);
492       return emitOpError(errorStream.str());
493     }
494   }
495 
496   // Check the result type
497   if (!llvm::any_of(expectedResult, [&](Type expectedResultType) {
498         return expectedResultType == getResult().getType();
499       })) {
500     errorStream
501         << "Could not match allowed types for the result; expected one of ";
502     llvm::interleaveComma(expectedResult, errorStream);
503     errorStream << " but got " << getResult().getType();
504     return emitOpError(errorStream.str());
505   }
506 
507   // Ensure that binary MMA variants have a b1 MMA operation defined.
508   if (multiplicandAPtxType() == MMATypes::b1 && !b1Op().hasValue()) {
509     return emitOpError("op requires " + b1OpAttrName().strref() + " attribute");
510   }
511 
512   // Ensure int4/int8 MMA variants specify the accum overflow behavior
513   // attribute.
514   if (isInt4PtxType(*multiplicandAPtxType()) ||
515       isInt8PtxType(*multiplicandAPtxType())) {
516     if (!intOverflowBehavior().hasValue())
517       return emitOpError("op requires " +
518                          intOverflowBehaviorAttrName().strref() + " attribute");
519   }
520 
521   return success();
522 }
523 
524 LogicalResult ShflOp::verify() {
525   if (!(*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid"))
526     return success();
527   auto type = getType().dyn_cast<LLVM::LLVMStructType>();
528   auto elementType = (type && type.getBody().size() == 2)
529                          ? type.getBody()[1].dyn_cast<IntegerType>()
530                          : nullptr;
531   if (!elementType || elementType.getWidth() != 1)
532     return emitError("expected return type to be a two-element struct with "
533                      "i1 as the second element");
534   return success();
535 }
536 
537 std::pair<mlir::Type, unsigned> NVVM::inferMMAType(NVVM::MMATypes type,
538                                                    NVVM::MMAFrag frag,
539                                                    MLIRContext *context) {
540   unsigned numberElements = 0;
541   Type elementType;
542   OpBuilder builder(context);
543   Type f16x2 = VectorType::get(2, builder.getF16Type());
544   if (type == NVVM::MMATypes::f16) {
545     elementType = f16x2;
546     if (frag == NVVM::MMAFrag::a || frag == NVVM::MMAFrag::b)
547       numberElements = 8;
548     else
549       numberElements = 4;
550   } else if (type == NVVM::MMATypes::f32) {
551     elementType = builder.getF32Type();
552     numberElements = 8;
553   } else if (type == NVVM::MMATypes::tf32) {
554     elementType = builder.getI32Type();
555     numberElements = 4;
556   }
557   assert(numberElements != 0 && elementType != nullptr);
558   return std::make_pair(elementType, numberElements);
559 }
560 
561 LogicalResult NVVM::WMMALoadOp::verify() {
562   unsigned addressSpace =
563       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
564   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
565     return emitOpError("expected source pointer in memory "
566                        "space 0, 1, 3");
567 
568   if (NVVM::WMMALoadOp::getIntrinsicID(m(), n(), k(), layout(), eltype(),
569                                        frag()) == 0)
570     return emitOpError() << "invalid attribute combination";
571   std::pair<Type, unsigned> typeInfo =
572       inferMMAType(eltype(), frag(), getContext());
573   Type dstType = LLVM::LLVMStructType::getLiteral(
574       getContext(), SmallVector<Type, 8>(typeInfo.second, typeInfo.first));
575   if (getType() != dstType)
576     return emitOpError("expected destination type is a structure of ")
577            << typeInfo.second << " elements of type " << typeInfo.first;
578   return success();
579 }
580 
581 LogicalResult NVVM::WMMAStoreOp::verify() {
582   unsigned addressSpace =
583       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
584   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
585     return emitOpError("expected operands to be a source pointer in memory "
586                        "space 0, 1, 3");
587 
588   if (NVVM::WMMAStoreOp::getIntrinsicID(m(), n(), k(), layout(), eltype()) == 0)
589     return emitOpError() << "invalid attribute combination";
590   std::pair<Type, unsigned> typeInfo =
591       inferMMAType(eltype(), NVVM::MMAFrag::c, getContext());
592   if (args().size() != typeInfo.second)
593     return emitOpError() << "expected " << typeInfo.second << " data operands";
594   if (llvm::any_of(args(), [&typeInfo](Value operands) {
595         return operands.getType() != typeInfo.first;
596       }))
597     return emitOpError() << "expected data operands of type " << typeInfo.first;
598   return success();
599 }
600 
601 LogicalResult NVVM::WMMAMmaOp::verify() {
602   if (NVVM::WMMAMmaOp::getIntrinsicID(m(), n(), k(), layoutA(), layoutB(),
603                                       eltypeA(), eltypeB()) == 0)
604     return emitOpError() << "invalid attribute combination";
605   std::pair<Type, unsigned> typeInfoA =
606       inferMMAType(eltypeA(), NVVM::MMAFrag::a, getContext());
607   std::pair<Type, unsigned> typeInfoB =
608       inferMMAType(eltypeA(), NVVM::MMAFrag::b, getContext());
609   std::pair<Type, unsigned> typeInfoC =
610       inferMMAType(eltypeB(), NVVM::MMAFrag::c, getContext());
611   SmallVector<Type, 32> arguments;
612   arguments.append(typeInfoA.second, typeInfoA.first);
613   arguments.append(typeInfoB.second, typeInfoB.first);
614   arguments.append(typeInfoC.second, typeInfoC.first);
615   unsigned numArgs = arguments.size();
616   if (args().size() != numArgs)
617     return emitOpError() << "expected " << numArgs << " arguments";
618   for (unsigned i = 0; i < numArgs; i++) {
619     if (args()[i].getType() != arguments[i])
620       return emitOpError() << "expected argument " << i << " to be of type "
621                            << arguments[i];
622   }
623   Type dstType = LLVM::LLVMStructType::getLiteral(
624       getContext(), SmallVector<Type, 8>(typeInfoC.second, typeInfoC.first));
625   if (getType() != dstType)
626     return emitOpError("expected destination type is a structure of ")
627            << typeInfoC.second << " elements of type " << typeInfoC.first;
628   return success();
629 }
630 
631 LogicalResult NVVM::LdMatrixOp::verify() {
632   unsigned addressSpace =
633       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
634   if (addressSpace != 3)
635     return emitOpError("expected source pointer in memory space 3");
636 
637   if (num() != 1 && num() != 2 && num() != 4)
638     return emitOpError("expected num attribute to be 1, 2 or 4");
639 
640   Type i32 = IntegerType::get(getContext(), 32);
641   if (num() == 1 && getType() != i32)
642     return emitOpError("expected destination type is i32");
643   if (num() == 2 || num() == 4) {
644     Type dstType = LLVM::LLVMStructType::getLiteral(
645         getContext(), SmallVector<Type>(num(), i32));
646     if (getType() != dstType)
647       return emitOpError("expected destination type is a structure of ")
648              << num() << " elements of type i32";
649   }
650   return success();
651 }
652 
653 //===----------------------------------------------------------------------===//
654 // NVVMDialect initialization, type parsing, and registration.
655 //===----------------------------------------------------------------------===//
656 
657 // TODO: This should be the llvm.nvvm dialect once this is supported.
658 void NVVMDialect::initialize() {
659   addOperations<
660 #define GET_OP_LIST
661 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
662       >();
663   addAttributes<
664 #define GET_ATTRDEF_LIST
665 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
666       >();
667 
668   // Support unknown operations because not all NVVM operations are
669   // registered.
670   allowUnknownOperations();
671 }
672 
673 LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
674                                                     NamedAttribute attr) {
675   // Kernel function attribute should be attached to functions.
676   if (attr.getName() == NVVMDialect::getKernelFuncAttrName()) {
677     if (!isa<LLVM::LLVMFuncOp>(op)) {
678       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
679                              << "' attribute attached to unexpected op";
680     }
681   }
682   return success();
683 }
684 
685 #define GET_OP_CLASSES
686 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
687 
688 #define GET_ATTRDEF_CLASSES
689 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
690