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