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 
38 //===----------------------------------------------------------------------===//
39 // Printing/parsing for NVVM ops
40 //===----------------------------------------------------------------------===//
41 
42 static void printNVVMIntrinsicOp(OpAsmPrinter &p, Operation *op) {
43   p << " " << op->getOperands();
44   if (op->getNumResults() > 0)
45     p << " : " << op->getResultTypes();
46 }
47 
48 // <operation> ::= `llvm.nvvm.vote.ballot.sync %mask, %pred` : result_type
49 ParseResult VoteBallotOp::parse(OpAsmParser &parser, OperationState &result) {
50   MLIRContext *context = parser.getContext();
51   auto int32Ty = IntegerType::get(context, 32);
52   auto int1Ty = IntegerType::get(context, 1);
53 
54   SmallVector<OpAsmParser::OperandType, 8> ops;
55   Type type;
56   return failure(parser.parseOperandList(ops) ||
57                  parser.parseOptionalAttrDict(result.attributes) ||
58                  parser.parseColonType(type) ||
59                  parser.addTypeToList(type, result.types) ||
60                  parser.resolveOperands(ops, {int32Ty, int1Ty},
61                                         parser.getNameLoc(), result.operands));
62 }
63 
64 void VoteBallotOp::print(OpAsmPrinter &p) { printNVVMIntrinsicOp(p, *this); }
65 
66 LogicalResult CpAsyncOp::verify() {
67   if (size() != 4 && size() != 8 && size() != 16)
68     return emitError("expected byte size to be either 4, 8 or 16.");
69   return success();
70 }
71 
72 LogicalResult MmaOp::verify() {
73   MLIRContext *context = getContext();
74   auto f16Ty = Float16Type::get(context);
75   auto f16x2Ty = LLVM::getFixedVectorType(f16Ty, 2);
76   auto f32Ty = Float32Type::get(context);
77   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
78       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
79   auto f32x8StructTy = LLVM::LLVMStructType::getLiteral(
80       context, {f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty});
81 
82   auto operandTypes = getOperandTypes();
83   if (operandTypes != SmallVector<Type, 8>(8, f16x2Ty) &&
84       operandTypes != ArrayRef<Type>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f32Ty,
85                                      f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
86                                      f32Ty}) {
87     return emitOpError("expected operands to be 4 <halfx2>s followed by either "
88                        "4 <halfx2>s or 8 floats");
89   }
90   if (getType() != f32x8StructTy && getType() != f16x2x4StructTy) {
91     return emitOpError("expected result type to be a struct of either 4 "
92                        "<halfx2>s or 8 floats");
93   }
94 
95   auto alayout = (*this)->getAttrOfType<StringAttr>("alayout");
96   auto blayout = (*this)->getAttrOfType<StringAttr>("blayout");
97 
98   if (!(alayout && blayout) ||
99       !(alayout.getValue() == "row" || alayout.getValue() == "col") ||
100       !(blayout.getValue() == "row" || blayout.getValue() == "col")) {
101     return emitOpError("alayout and blayout attributes must be set to either "
102                        "\"row\" or \"col\"");
103   }
104 
105   if (operandTypes == ArrayRef<Type>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f32Ty,
106                                      f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
107                                      f32Ty} &&
108       getType() == f32x8StructTy && alayout.getValue() == "row" &&
109       blayout.getValue() == "col") {
110     return success();
111   }
112   return emitOpError("unimplemented mma.sync variant");
113 }
114 
115 LogicalResult ShflOp::verify() {
116   if (!(*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid"))
117     return success();
118   auto type = getType().dyn_cast<LLVM::LLVMStructType>();
119   auto elementType = (type && type.getBody().size() == 2)
120                          ? type.getBody()[1].dyn_cast<IntegerType>()
121                          : nullptr;
122   if (!elementType || elementType.getWidth() != 1)
123     return emitError("expected return type to be a two-element struct with "
124                      "i1 as the second element");
125   return success();
126 }
127 
128 std::pair<mlir::Type, unsigned> NVVM::inferMMAType(NVVM::MMATypes type,
129                                                    NVVM::MMAFrag frag,
130                                                    MLIRContext *context) {
131   unsigned numberElements = 0;
132   Type elementType;
133   OpBuilder builder(context);
134   Type f16x2 = VectorType::get(2, builder.getF16Type());
135   if (type == NVVM::MMATypes::f16) {
136     elementType = f16x2;
137     if (frag == NVVM::MMAFrag::a || frag == NVVM::MMAFrag::b)
138       numberElements = 8;
139     else
140       numberElements = 4;
141   } else if (type == NVVM::MMATypes::f32) {
142     elementType = builder.getF32Type();
143     numberElements = 8;
144   } else if (type == NVVM::MMATypes::tf32) {
145     elementType = builder.getI32Type();
146     numberElements = 4;
147   }
148   assert(numberElements != 0 && elementType != nullptr);
149   return std::make_pair(elementType, numberElements);
150 }
151 
152 LogicalResult NVVM::WMMALoadOp::verify() {
153   unsigned addressSpace =
154       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
155   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
156     return emitOpError("expected source pointer in memory "
157                        "space 0, 1, 3");
158 
159   if (NVVM::WMMALoadOp::getIntrinsicID(m(), n(), k(), layout(), eltype(),
160                                        frag()) == 0)
161     return emitOpError() << "invalid attribute combination";
162   std::pair<Type, unsigned> typeInfo =
163       inferMMAType(eltype(), frag(), getContext());
164   Type dstType = LLVM::LLVMStructType::getLiteral(
165       getContext(), SmallVector<Type, 8>(typeInfo.second, typeInfo.first));
166   if (getType() != dstType)
167     return emitOpError("expected destination type is a structure of ")
168            << typeInfo.second << " elements of type " << typeInfo.first;
169   return success();
170 }
171 
172 LogicalResult NVVM::WMMAStoreOp::verify() {
173   unsigned addressSpace =
174       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
175   if (addressSpace != 0 && addressSpace != 1 && addressSpace != 3)
176     return emitOpError("expected operands to be a source pointer in memory "
177                        "space 0, 1, 3");
178 
179   if (NVVM::WMMAStoreOp::getIntrinsicID(m(), n(), k(), layout(), eltype()) == 0)
180     return emitOpError() << "invalid attribute combination";
181   std::pair<Type, unsigned> typeInfo =
182       inferMMAType(eltype(), NVVM::MMAFrag::c, getContext());
183   if (args().size() != typeInfo.second)
184     return emitOpError() << "expected " << typeInfo.second << " data operands";
185   if (llvm::any_of(args(), [&typeInfo](Value operands) {
186         return operands.getType() != typeInfo.first;
187       }))
188     return emitOpError() << "expected data operands of type " << typeInfo.first;
189   return success();
190 }
191 
192 LogicalResult NVVM::WMMAMmaOp::verify() {
193   if (NVVM::WMMAMmaOp::getIntrinsicID(m(), n(), k(), layoutA(), layoutB(),
194                                       eltypeA(), eltypeB()) == 0)
195     return emitOpError() << "invalid attribute combination";
196   std::pair<Type, unsigned> typeInfoA =
197       inferMMAType(eltypeA(), NVVM::MMAFrag::a, getContext());
198   std::pair<Type, unsigned> typeInfoB =
199       inferMMAType(eltypeA(), NVVM::MMAFrag::b, getContext());
200   std::pair<Type, unsigned> typeInfoC =
201       inferMMAType(eltypeB(), NVVM::MMAFrag::c, getContext());
202   SmallVector<Type, 32> arguments;
203   arguments.append(typeInfoA.second, typeInfoA.first);
204   arguments.append(typeInfoB.second, typeInfoB.first);
205   arguments.append(typeInfoC.second, typeInfoC.first);
206   unsigned numArgs = arguments.size();
207   if (args().size() != numArgs)
208     return emitOpError() << "expected " << numArgs << " arguments";
209   for (unsigned i = 0; i < numArgs; i++) {
210     if (args()[i].getType() != arguments[i])
211       return emitOpError() << "expected argument " << i << " to be of type "
212                            << arguments[i];
213   }
214   Type dstType = LLVM::LLVMStructType::getLiteral(
215       getContext(), SmallVector<Type, 8>(typeInfoC.second, typeInfoC.first));
216   if (getType() != dstType)
217     return emitOpError("expected destination type is a structure of ")
218            << typeInfoC.second << " elements of type " << typeInfoC.first;
219   return success();
220 }
221 
222 LogicalResult NVVM::LdMatrixOp::verify() {
223   unsigned addressSpace =
224       ptr().getType().cast<LLVM::LLVMPointerType>().getAddressSpace();
225   if (addressSpace != 3)
226     return emitOpError("expected source pointer in memory space 3");
227 
228   if (num() != 1 && num() != 2 && num() != 4)
229     return emitOpError("expected num attribute to be 1, 2 or 4");
230 
231   Type i32 = IntegerType::get(getContext(), 32);
232   if (num() == 1 && getType() != i32)
233     return emitOpError("expected destination type is i32");
234   if (num() == 2 || num() == 4) {
235     Type dstType = LLVM::LLVMStructType::getLiteral(
236         getContext(), SmallVector<Type>(num(), i32));
237     if (getType() != dstType)
238       return emitOpError("expected destination type is a structure of ")
239              << num() << " elements of type i32";
240   }
241   return success();
242 }
243 
244 //===----------------------------------------------------------------------===//
245 // NVVMDialect initialization, type parsing, and registration.
246 //===----------------------------------------------------------------------===//
247 
248 // TODO: This should be the llvm.nvvm dialect once this is supported.
249 void NVVMDialect::initialize() {
250   addOperations<
251 #define GET_OP_LIST
252 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
253       >();
254   addAttributes<
255 #define GET_ATTRDEF_LIST
256 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
257       >();
258 
259   // Support unknown operations because not all NVVM operations are
260   // registered.
261   allowUnknownOperations();
262 }
263 
264 LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
265                                                     NamedAttribute attr) {
266   // Kernel function attribute should be attached to functions.
267   if (attr.getName() == NVVMDialect::getKernelFuncAttrName()) {
268     if (!isa<LLVM::LLVMFuncOp>(op)) {
269       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
270                              << "' attribute attached to unexpected op";
271     }
272   }
273   return success();
274 }
275 
276 #define GET_OP_CLASSES
277 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
278 
279 #define GET_ATTRDEF_CLASSES
280 #include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
281