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   Type i32 = builder.getIntegerType(32);
200   result.addAttribute(
201       "shape", MMAShapeAttr::get(builder.getIntegerAttr(i32, shape[0]),
202                                  builder.getIntegerAttr(i32, shape[1]),
203                                  builder.getIntegerAttr(i32, shape[2]), ctx));
204 
205   result.addOperands(operandA);
206   result.addOperands(operandB);
207   result.addOperands(operandC);
208 
209   if (multiplicandPtxTypes.hasValue()) {
210     result.addAttribute("multiplicandAPtxType",
211                         MMATypesAttr::get(ctx, (*multiplicandPtxTypes)[0]));
212     result.addAttribute("multiplicandBPtxType",
213                         MMATypesAttr::get(ctx, (*multiplicandPtxTypes)[1]));
214   } else {
215     if (auto res = inferOperandMMAType(operandA[0].getType(), false))
216       result.addAttribute("multiplicandAPtxType", MMATypesAttr::get(ctx, *res));
217     if (auto res = inferOperandMMAType(operandB[0].getType(), false))
218       result.addAttribute("multiplicandBPtxType", MMATypesAttr::get(ctx, *res));
219   }
220 
221   if (multiplicandLayouts.hasValue()) {
222     result.addAttribute("layoutA",
223                         MMALayoutAttr::get(ctx, (*multiplicandLayouts)[0]));
224     result.addAttribute("layoutB",
225                         MMALayoutAttr::get(ctx, (*multiplicandLayouts)[1]));
226   } else {
227     result.addAttribute("layoutA", MMALayoutAttr::get(ctx, MMALayout::row));
228     result.addAttribute("layoutB", MMALayoutAttr::get(ctx, MMALayout::col));
229   }
230 
231   if (intOverflow.hasValue())
232     result.addAttribute("intOverflowBehavior",
233                         MMAIntOverflowAttr::get(ctx, *intOverflow));
234   if (b1Op.hasValue())
235     result.addAttribute("b1Op", MMAB1OpAttr::get(ctx, *b1Op));
236 
237   result.addTypes(resultType);
238   result.addAttribute(
239       MmaOp::getOperandSegmentSizeAttr(),
240       builder.getI32VectorAttr({static_cast<int32_t>(operandA.size()),
241                                 static_cast<int32_t>(operandB.size()),
242                                 static_cast<int32_t>(operandC.size())}));
243 }
244 
245 // <operation> :=
246 //   A `[` $operandA `]` B `[` $operandB `]` C `[` $operandC `]`
247 //   attr-dict : (type($operandA[0]), type($operandB[0]), type($operandC[0]))
248 //     `->` type($res)
249 ParseResult MmaOp::parse(OpAsmParser &parser, OperationState &result) {
250   struct OperandFragment {
251     Optional<MMATypes> elemtype;
252     SmallVector<OpAsmParser::UnresolvedOperand, 4> regs;
253     SmallVector<Type> regTypes;
254   };
255 
256   Builder &builder = parser.getBuilder();
257   std::array<OperandFragment, 4> frags;
258 
259   NamedAttrList namedAttributes;
260 
261   // A helper to parse the operand segments.
262   auto parseMmaOperand = [&](StringRef operandName,
263                              OperandFragment &frag) -> LogicalResult {
264     if (parser.parseKeyword(operandName).failed())
265       return failure();
266     if (parser
267             .parseOperandList(frag.regs, OpAsmParser::Delimiter::OptionalSquare)
268             .failed())
269       return failure();
270     return success();
271   };
272 
273   // Parse the operand segments.
274   if (parseMmaOperand("A", frags[0]).failed())
275     return failure();
276   if (parseMmaOperand("B", frags[1]).failed())
277     return failure();
278   if (parseMmaOperand("C", frags[2]).failed())
279     return failure();
280 
281   if (parser.parseOptionalAttrDict(namedAttributes).failed())
282     return failure();
283 
284   // Parse the type specification and resolve operands.
285   SmallVector<Type, 3> operandTypes;
286   if (failed(parser.parseColon()))
287     return failure();
288   if (failed(parser.parseLParen()))
289     return failure();
290   if (failed(parser.parseTypeList(operandTypes)))
291     return failure();
292   if (failed(parser.parseRParen()))
293     if (operandTypes.size() != 3)
294       return parser.emitError(
295           parser.getNameLoc(),
296           "expected one type for each operand segment but got " +
297               Twine(operandTypes.size()) + " types");
298   for (const auto &iter : llvm::enumerate(operandTypes)) {
299     auto &frag = frags[iter.index()];
300     frag.regTypes.resize(frag.regs.size(), iter.value());
301     if (failed(parser.resolveOperands(frag.regs, frag.regTypes,
302                                       parser.getNameLoc(), result.operands)))
303       return failure();
304     frag.elemtype =
305         inferOperandMMAType(frag.regTypes[0], /*isAccum=*/iter.index() < 2);
306   }
307 
308   Type resultType;
309   if (parser.parseArrow() || parser.parseType(resultType))
310     return failure();
311   frags[3].elemtype = inferOperandMMAType(resultType, /*isAccum=*/true);
312 
313   std::array<StringRef, 2> names{"multiplicandAPtxType",
314                                  "multiplicandBPtxType"};
315   for (unsigned idx = 0; idx < names.size(); idx++) {
316     const auto &frag = frags[idx];
317     Optional<NamedAttribute> attr = namedAttributes.getNamed(names[idx]);
318     if (!frag.elemtype.hasValue() && !attr.hasValue()) {
319       return parser.emitError(
320           parser.getNameLoc(),
321           "attribute " + names[idx] +
322               " is not provided explicitly and cannot be inferred");
323     }
324     if (!attr.hasValue())
325       result.addAttribute(
326           names[idx], MMATypesAttr::get(parser.getContext(), *frag.elemtype));
327   }
328 
329   result.addTypes(resultType);
330   if (!namedAttributes.empty())
331     result.addAttributes(namedAttributes);
332   result.addAttribute(MmaOp::getOperandSegmentSizeAttr(),
333                       builder.getI32VectorAttr({
334                           static_cast<int32_t>(frags[0].regs.size()),
335                           static_cast<int32_t>(frags[1].regs.size()),
336                           static_cast<int32_t>(frags[2].regs.size()),
337                       }));
338   return success();
339 }
340 
341 LogicalResult MmaOp::verify() {
342   MLIRContext *context = getContext();
343   auto f16Ty = Float16Type::get(context);
344   auto i32Ty = IntegerType::get(context, 32);
345   auto f16x2Ty = LLVM::getFixedVectorType(f16Ty, 2);
346   auto f32Ty = Float32Type::get(context);
347   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
348       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
349 
350   auto s32x4StructTy =
351       LLVM::LLVMStructType::getLiteral(context, {i32Ty, i32Ty, i32Ty, i32Ty});
352   auto f32x8StructTy =
353       LLVM::LLVMStructType::getLiteral(context, SmallVector<Type>(8, f32Ty));
354   auto f16x2x2StructTy =
355       LLVM::LLVMStructType::getLiteral(context, {f16x2Ty, f16x2Ty});
356   auto f32x4StructTy =
357       LLVM::LLVMStructType::getLiteral(context, {f32Ty, f32Ty, f32Ty, f32Ty});
358   auto s32x2StructTy =
359       LLVM::LLVMStructType::getLiteral(context, {i32Ty, i32Ty});
360 
361   std::array<int64_t, 3> mmaShape{shapeAttr().m().getInt(),
362                                   shapeAttr().n().getInt(),
363                                   shapeAttr().k().getInt()};
364 
365   // These variables define the set of allowed data types for matrices A, B, C,
366   // and result.
367   using AllowedShapes = SmallVector<std::array<int64_t, 3>, 2>;
368   using AllowedTypes = SmallVector<SmallVector<Type, 4>, 2>;
369   AllowedShapes allowedShapes;
370   AllowedTypes expectedA;
371   AllowedTypes expectedB;
372   AllowedTypes expectedC;
373   SmallVector<Type> expectedResult;
374 
375   // When M = 16, we just need to calculate the number of 8xk tiles, where
376   // k is a factor that depends on the data type.
377   if (mmaShape[0] == 16) {
378     int64_t kFactor;
379     Type multiplicandFragType;
380     switch (multiplicandAPtxType().getValue()) {
381     case MMATypes::tf32:
382       kFactor = 4;
383       multiplicandFragType = i32Ty;
384       expectedResult.push_back(LLVM::LLVMStructType::getLiteral(
385           context, {f32Ty, f32Ty, f32Ty, f32Ty}));
386       break;
387     case MMATypes::f16:
388     case MMATypes::bf16:
389       kFactor = 8;
390       multiplicandFragType = f16x2Ty;
391       expectedResult.push_back(f16x2x2StructTy);
392       expectedResult.push_back(f32x4StructTy);
393       break;
394     case MMATypes::s4:
395     case MMATypes::u4:
396       kFactor = 32;
397       break;
398     case MMATypes::b1:
399       kFactor = 128;
400       break;
401     case MMATypes::s8:
402     case MMATypes::u8:
403       kFactor = 16;
404       break;
405     default:
406       return emitError("invalid shape or multiplicand type: " +
407                        stringifyEnum(multiplicandAPtxType().getValue()));
408     }
409 
410     if (isIntegerPtxType(multiplicandAPtxType().getValue())) {
411       expectedResult.push_back(s32x4StructTy);
412       expectedC.emplace_back(4, i32Ty);
413       multiplicandFragType = i32Ty;
414     } else {
415       expectedC.emplace_back(2, f16x2Ty);
416       expectedC.emplace_back(4, f32Ty);
417     }
418 
419     int64_t unitA = (mmaShape[0] / 8) * (mmaShape[2] / kFactor);
420     int64_t unitB = (mmaShape[1] / 8) * (mmaShape[2] / kFactor);
421     expectedA.emplace_back(unitA, multiplicandFragType);
422     expectedB.emplace_back(unitB, multiplicandFragType);
423     allowedShapes.push_back({16, 8, kFactor});
424     allowedShapes.push_back({16, 8, kFactor * 2});
425   }
426 
427   // In the M=8 case, there is only 1 possible case per data type.
428   if (mmaShape[0] == 8) {
429     if (multiplicandAPtxType().getValue() == MMATypes::f16) {
430       expectedA.emplace_back(2, f16x2Ty);
431       expectedB.emplace_back(2, f16x2Ty);
432       expectedResult.push_back(f16x2x4StructTy);
433       expectedResult.push_back(f32x8StructTy);
434       expectedC.emplace_back(4, f16x2Ty);
435       expectedC.emplace_back(8, f32Ty);
436       allowedShapes.push_back({8, 8, 4});
437     }
438     if (multiplicandAPtxType().getValue() == MMATypes::f64) {
439       Type f64Ty = Float64Type::get(context);
440       expectedA.emplace_back(1, f64Ty);
441       expectedB.emplace_back(1, f64Ty);
442       expectedC.emplace_back(2, f64Ty);
443       // expectedC.emplace_back(1, LLVM::getFixedVectorType(f64Ty, 2));
444       expectedResult.emplace_back(LLVM::LLVMStructType::getLiteral(
445           context, SmallVector<Type>(2, f64Ty)));
446       allowedShapes.push_back({8, 8, 4});
447     }
448     if (isIntegerPtxType(multiplicandAPtxType().getValue())) {
449       expectedA.push_back({i32Ty});
450       expectedB.push_back({i32Ty});
451       expectedC.push_back({i32Ty, i32Ty});
452       expectedResult.push_back(s32x2StructTy);
453       if (isInt4PtxType(multiplicandAPtxType().getValue()))
454         allowedShapes.push_back({8, 8, 32});
455       if (isInt8PtxType(multiplicandAPtxType().getValue()))
456         allowedShapes.push_back({8, 8, 16});
457       if (multiplicandAPtxType().getValue() == MMATypes::b1)
458         allowedShapes.push_back({8, 8, 128});
459     }
460   }
461 
462   std::string errorMessage;
463   llvm::raw_string_ostream errorStream(errorMessage);
464 
465   // Check that we matched an existing shape/dtype combination.
466   if (expectedA.empty() || expectedB.empty() || expectedC.empty() ||
467       !llvm::any_of(allowedShapes,
468                     [&](const auto &allowed) { return allowed == mmaShape; })) {
469     errorStream << "unimplemented variant for MMA shape <";
470     llvm::interleaveComma(mmaShape, errorStream);
471     errorStream << ">";
472     return emitOpError(errorMessage);
473   }
474 
475   // Verify the operand types for segments of A, B, and C operands.
476   std::array<StringRef, 3> operandNames{"A", "B", "C"};
477   for (const auto &iter : llvm::enumerate(
478            SmallVector<AllowedTypes, 3>{expectedA, expectedB, expectedC})) {
479     auto spec = this->getODSOperandIndexAndLength(iter.index());
480     SmallVector<Type, 4> operandTySeg(operand_type_begin() + spec.first,
481                                       operand_type_begin() + spec.first +
482                                           spec.second);
483     bool match =
484         llvm::any_of(iter.value(), [&](const SmallVector<Type, 4> &typeSet) {
485           return typeSet == operandTySeg;
486         });
487 
488     if (!match) {
489       errorStream << "Could not match types for the "
490                   << operandNames[iter.index()]
491                   << " operands; expected one of ";
492       for (const auto &x : iter.value()) {
493         errorStream << x.size() << "x" << x[0] << " ";
494       }
495       errorStream << "but got ";
496       llvm::interleaveComma(operandTySeg, errorStream);
497       return emitOpError(errorStream.str());
498     }
499   }
500 
501   // Check the result type
502   if (!llvm::any_of(expectedResult, [&](Type expectedResultType) {
503         return expectedResultType == getResult().getType();
504       })) {
505     errorStream
506         << "Could not match allowed types for the result; expected one of ";
507     llvm::interleaveComma(expectedResult, errorStream);
508     errorStream << " but got " << getResult().getType();
509     return emitOpError(errorStream.str());
510   }
511 
512   // Ensure that binary MMA variants have a b1 MMA operation defined.
513   if (multiplicandAPtxType() == MMATypes::b1 && !b1Op().hasValue()) {
514     return emitOpError("op requires " + b1OpAttrName().strref() + " attribute");
515   }
516 
517   // Ensure int4/int8 MMA variants specify the accum overflow behavior
518   // attribute.
519   if (isInt4PtxType(*multiplicandAPtxType()) ||
520       isInt8PtxType(*multiplicandAPtxType())) {
521     if (!intOverflowBehavior().hasValue())
522       return emitOpError("op requires " +
523                          intOverflowBehaviorAttrName().strref() + " attribute");
524   }
525 
526   return success();
527 }
528 
529 LogicalResult ShflOp::verify() {
530   if (!(*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid"))
531     return success();
532   auto type = getType().dyn_cast<LLVM::LLVMStructType>();
533   auto elementType = (type && type.getBody().size() == 2)
534                          ? type.getBody()[1].dyn_cast<IntegerType>()
535                          : nullptr;
536   if (!elementType || elementType.getWidth() != 1)
537     return emitError("expected return type to be a two-element struct with "
538                      "i1 as the second element");
539   return success();
540 }
541 
542 std::pair<mlir::Type, unsigned> NVVM::inferMMAType(NVVM::MMATypes type,
543                                                    NVVM::MMAFrag frag,
544                                                    MLIRContext *context) {
545   unsigned numberElements = 0;
546   Type elementType;
547   OpBuilder builder(context);
548   Type f16x2 = VectorType::get(2, builder.getF16Type());
549   if (type == NVVM::MMATypes::f16) {
550     elementType = f16x2;
551     if (frag == NVVM::MMAFrag::a || frag == NVVM::MMAFrag::b)
552       numberElements = 8;
553     else
554       numberElements = 4;
555   } else if (type == NVVM::MMATypes::f32) {
556     elementType = builder.getF32Type();
557     numberElements = 8;
558   } else if (type == NVVM::MMATypes::tf32) {
559     elementType = builder.getI32Type();
560     numberElements = 4;
561   }
562   assert(numberElements != 0 && elementType != nullptr);
563   return std::make_pair(elementType, numberElements);
564 }
565 
566 LogicalResult NVVM::WMMALoadOp::verify() {
567   unsigned addressSpace =
568       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
569   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
570     return emitOpError("expected source pointer in memory "
571                        "space 0, 1, 3");
572 
573   if (NVVM::WMMALoadOp::getIntrinsicID(m(), n(), k(), layout(), eltype(),
574                                        frag()) == 0)
575     return emitOpError() << "invalid attribute combination";
576   std::pair<Type, unsigned> typeInfo =
577       inferMMAType(eltype(), frag(), getContext());
578   Type dstType = LLVM::LLVMStructType::getLiteral(
579       getContext(), SmallVector<Type, 8>(typeInfo.second, typeInfo.first));
580   if (getType() != dstType)
581     return emitOpError("expected destination type is a structure of ")
582            << typeInfo.second << " elements of type " << typeInfo.first;
583   return success();
584 }
585 
586 LogicalResult NVVM::WMMAStoreOp::verify() {
587   unsigned addressSpace =
588       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
589   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
590     return emitOpError("expected operands to be a source pointer in memory "
591                        "space 0, 1, 3");
592 
593   if (NVVM::WMMAStoreOp::getIntrinsicID(m(), n(), k(), layout(), eltype()) == 0)
594     return emitOpError() << "invalid attribute combination";
595   std::pair<Type, unsigned> typeInfo =
596       inferMMAType(eltype(), NVVM::MMAFrag::c, getContext());
597   if (args().size() != typeInfo.second)
598     return emitOpError() << "expected " << typeInfo.second << " data operands";
599   if (llvm::any_of(args(), [&typeInfo](Value operands) {
600         return operands.getType() != typeInfo.first;
601       }))
602     return emitOpError() << "expected data operands of type " << typeInfo.first;
603   return success();
604 }
605 
606 LogicalResult NVVM::WMMAMmaOp::verify() {
607   if (NVVM::WMMAMmaOp::getIntrinsicID(m(), n(), k(), layoutA(), layoutB(),
608                                       eltypeA(), eltypeB()) == 0)
609     return emitOpError() << "invalid attribute combination";
610   std::pair<Type, unsigned> typeInfoA =
611       inferMMAType(eltypeA(), NVVM::MMAFrag::a, getContext());
612   std::pair<Type, unsigned> typeInfoB =
613       inferMMAType(eltypeA(), NVVM::MMAFrag::b, getContext());
614   std::pair<Type, unsigned> typeInfoC =
615       inferMMAType(eltypeB(), NVVM::MMAFrag::c, getContext());
616   SmallVector<Type, 32> arguments;
617   arguments.append(typeInfoA.second, typeInfoA.first);
618   arguments.append(typeInfoB.second, typeInfoB.first);
619   arguments.append(typeInfoC.second, typeInfoC.first);
620   unsigned numArgs = arguments.size();
621   if (args().size() != numArgs)
622     return emitOpError() << "expected " << numArgs << " arguments";
623   for (unsigned i = 0; i < numArgs; i++) {
624     if (args()[i].getType() != arguments[i])
625       return emitOpError() << "expected argument " << i << " to be of type "
626                            << arguments[i];
627   }
628   Type dstType = LLVM::LLVMStructType::getLiteral(
629       getContext(), SmallVector<Type, 8>(typeInfoC.second, typeInfoC.first));
630   if (getType() != dstType)
631     return emitOpError("expected destination type is a structure of ")
632            << typeInfoC.second << " elements of type " << typeInfoC.first;
633   return success();
634 }
635 
636 LogicalResult NVVM::LdMatrixOp::verify() {
637   unsigned addressSpace =
638       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
639   if (addressSpace != 3)
640     return emitOpError("expected source pointer in memory space 3");
641 
642   if (num() != 1 && num() != 2 && num() != 4)
643     return emitOpError("expected num attribute to be 1, 2 or 4");
644 
645   Type i32 = IntegerType::get(getContext(), 32);
646   if (num() == 1 && getType() != i32)
647     return emitOpError("expected destination type is i32");
648   if (num() == 2 || num() == 4) {
649     Type dstType = LLVM::LLVMStructType::getLiteral(
650         getContext(), SmallVector<Type>(num(), i32));
651     if (getType() != dstType)
652       return emitOpError("expected destination type is a structure of ")
653              << num() << " elements of type i32";
654   }
655   return success();
656 }
657 
658 //===----------------------------------------------------------------------===//
659 // NVVMDialect initialization, type parsing, and registration.
660 //===----------------------------------------------------------------------===//
661 
662 // TODO: This should be the llvm.nvvm dialect once this is supported.
663 void NVVMDialect::initialize() {
664   addOperations<
665 #define GET_OP_LIST
666 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
667       >();
668   addAttributes<
669 #define GET_ATTRDEF_LIST
670 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
671       >();
672 
673   // Support unknown operations because not all NVVM operations are
674   // registered.
675   allowUnknownOperations();
676 }
677 
678 LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
679                                                     NamedAttribute attr) {
680   // Kernel function attribute should be attached to functions.
681   if (attr.getName() == NVVMDialect::getKernelFuncAttrName()) {
682     if (!isa<LLVM::LLVMFuncOp>(op)) {
683       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
684                              << "' attribute attached to unexpected op";
685     }
686   }
687   return success();
688 }
689 
690 #define GET_OP_CLASSES
691 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
692 
693 #define GET_ATTRDEF_CLASSES
694 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
695