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/MLIRContext.h"
22 #include "mlir/IR/Operation.h"
23 #include "mlir/IR/OperationSupport.h"
24 #include "llvm/AsmParser/Parser.h"
25 #include "llvm/IR/Attributes.h"
26 #include "llvm/IR/Function.h"
27 #include "llvm/IR/Type.h"
28 #include "llvm/Support/SourceMgr.h"
29 
30 using namespace mlir;
31 using namespace NVVM;
32 
33 //===----------------------------------------------------------------------===//
34 // Printing/parsing for NVVM ops
35 //===----------------------------------------------------------------------===//
36 
37 static void printNVVMIntrinsicOp(OpAsmPrinter &p, Operation *op) {
38   p << op->getName() << " " << op->getOperands();
39   if (op->getNumResults() > 0)
40     p << " : " << op->getResultTypes();
41 }
42 
43 // <operation> ::=
44 //     `llvm.nvvm.shfl.sync.bfly %dst, %val, %offset, %clamp_and_mask`
45 //      ({return_value_and_is_valid})? : result_type
46 static ParseResult parseNVVMShflSyncBflyOp(OpAsmParser &parser,
47                                            OperationState &result) {
48   SmallVector<OpAsmParser::OperandType, 8> ops;
49   Type resultType;
50   if (parser.parseOperandList(ops) ||
51       parser.parseOptionalAttrDict(result.attributes) ||
52       parser.parseColonType(resultType) ||
53       parser.addTypeToList(resultType, result.types))
54     return failure();
55 
56   for (auto &attr : result.attributes) {
57     if (attr.first != "return_value_and_is_valid")
58       continue;
59     auto structType = resultType.dyn_cast<LLVM::LLVMStructType>();
60     if (structType && !structType.getBody().empty())
61       resultType = structType.getBody()[0];
62     break;
63   }
64 
65   auto int32Ty = IntegerType::get(parser.getBuilder().getContext(), 32);
66   return parser.resolveOperands(ops, {int32Ty, resultType, int32Ty, int32Ty},
67                                 parser.getNameLoc(), result.operands);
68 }
69 
70 // <operation> ::= `llvm.nvvm.vote.ballot.sync %mask, %pred` : result_type
71 static ParseResult parseNVVMVoteBallotOp(OpAsmParser &parser,
72                                          OperationState &result) {
73   MLIRContext *context = parser.getBuilder().getContext();
74   auto int32Ty = IntegerType::get(context, 32);
75   auto int1Ty = IntegerType::get(context, 1);
76 
77   SmallVector<OpAsmParser::OperandType, 8> ops;
78   Type type;
79   return failure(parser.parseOperandList(ops) ||
80                  parser.parseOptionalAttrDict(result.attributes) ||
81                  parser.parseColonType(type) ||
82                  parser.addTypeToList(type, result.types) ||
83                  parser.resolveOperands(ops, {int32Ty, int1Ty},
84                                         parser.getNameLoc(), result.operands));
85 }
86 
87 static LogicalResult verify(MmaOp op) {
88   MLIRContext *context = op.getContext();
89   auto f16Ty = Float16Type::get(context);
90   auto f16x2Ty = LLVM::getFixedVectorType(f16Ty, 2);
91   auto f32Ty = Float32Type::get(context);
92   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
93       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
94   auto f32x8StructTy = LLVM::LLVMStructType::getLiteral(
95       context, {f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty});
96 
97   SmallVector<Type, 12> operandTypes(op.getOperandTypes().begin(),
98                                      op.getOperandTypes().end());
99   if (operandTypes != SmallVector<Type, 8>(8, f16x2Ty) &&
100       operandTypes != SmallVector<Type, 12>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty,
101                                             f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
102                                             f32Ty, f32Ty, f32Ty}) {
103     return op.emitOpError(
104         "expected operands to be 4 <halfx2>s followed by either "
105         "4 <halfx2>s or 8 floats");
106   }
107   if (op.getType() != f32x8StructTy && op.getType() != f16x2x4StructTy) {
108     return op.emitOpError("expected result type to be a struct of either 4 "
109                           "<halfx2>s or 8 floats");
110   }
111 
112   auto alayout = op->getAttrOfType<StringAttr>("alayout");
113   auto blayout = op->getAttrOfType<StringAttr>("blayout");
114 
115   if (!(alayout && blayout) ||
116       !(alayout.getValue() == "row" || alayout.getValue() == "col") ||
117       !(blayout.getValue() == "row" || blayout.getValue() == "col")) {
118     return op.emitOpError(
119         "alayout and blayout attributes must be set to either "
120         "\"row\" or \"col\"");
121   }
122 
123   if (operandTypes == SmallVector<Type, 12>{f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty,
124                                             f32Ty, f32Ty, f32Ty, f32Ty, f32Ty,
125                                             f32Ty, f32Ty, f32Ty} &&
126       op.getType() == f32x8StructTy && alayout.getValue() == "row" &&
127       blayout.getValue() == "col") {
128     return success();
129   }
130   return op.emitOpError("unimplemented mma.sync variant");
131 }
132 
133 template <typename T>
134 static LogicalResult verifyWMMALoadOp(T op, StringRef operand) {
135   MLIRContext *context = op.getContext();
136   auto i32Ty = IntegerType::get(context, 32);
137   auto i32Ptr1Ty = LLVM::LLVMPointerType::get(i32Ty, 1);
138   auto i32Ptr3Ty = LLVM::LLVMPointerType::get(i32Ty, 3);
139   auto i32Ptr0Ty = LLVM::LLVMPointerType::get(i32Ty, 0);
140   auto f16Ty = FloatType::getF16(context);
141   auto f32Ty = FloatType::getF32(context);
142   auto f16x2Ty = VectorType::get(2, f16Ty);
143   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
144       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
145   auto f16x2x8StructTy = LLVM::LLVMStructType::getLiteral(
146       context,
147       {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
148   auto f32x8StructTy = LLVM::LLVMStructType::getLiteral(
149       context, {f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty});
150 
151   SmallVector<Type, 2> operandTypes(op.getOperandTypes().begin(),
152                                     op.getOperandTypes().end());
153   if (operandTypes != SmallVector<Type, 2>{i32Ptr1Ty, i32Ty} &&
154       operandTypes != SmallVector<Type, 2>{i32Ptr3Ty, i32Ty} &&
155       operandTypes != SmallVector<Type, 2>{i32Ptr0Ty, i32Ty}) {
156     return op.emitOpError("expected operands to be a source pointer in memory "
157                           "space 0, 1, 3 followed by ldm of the source");
158   }
159 
160   if (operand.equals("AOp") || operand.equals("BOp")) {
161     if (op.getType() != f16x2x8StructTy) {
162       return op.emitOpError("expected result type of loadAOp and loadBOp to be "
163                             "a struct of 8 <halfx2>s");
164     }
165   } else if (operand.equals("COp")) {
166     if (op.getType() != f16x2x4StructTy && op.getType() != f32x8StructTy) {
167       return op.emitOpError("expected result type of loadCOp to be a struct of "
168                             "4 <halfx2>s or 8 f32s");
169     }
170   }
171 
172   return success();
173 }
174 
175 static LogicalResult verify(WMMALoadAM16N16K16Op op) {
176   return verifyWMMALoadOp(op, "AOp");
177 }
178 
179 static LogicalResult verify(WMMALoadBM16N16K16Op op) {
180   return verifyWMMALoadOp(op, "BOp");
181 }
182 
183 static LogicalResult verify(WMMALoadCF16M16N16K16Op op) {
184   return verifyWMMALoadOp(op, "COp");
185 }
186 
187 static LogicalResult verify(WMMALoadCF32M16N16K16Op op) {
188   return verifyWMMALoadOp(op, "COp");
189 }
190 
191 template <typename T>
192 static bool verifyWMMAStoreOp(T op, SmallVector<Type> &containedElems) {
193   SmallVector<Type> operandTypes(op.getOperandTypes().begin(),
194                                  op.getOperandTypes().end());
195   if (operandTypes == containedElems)
196     return true;
197 
198   return false;
199 }
200 
201 static LogicalResult verify(WMMAStoreF16M16N16K16Op op) {
202   MLIRContext *context = op.getContext();
203   auto i32Ty = IntegerType::get(context, 32);
204   auto i32Ptr1Ty = LLVM::LLVMPointerType::get(i32Ty, 1);
205   auto i32Ptr3Ty = LLVM::LLVMPointerType::get(i32Ty, 3);
206   auto i32Ptr0Ty = LLVM::LLVMPointerType::get(i32Ty, 0);
207   auto f16Ty = FloatType::getF16(context);
208   auto f16x2Ty = VectorType::get(2, f16Ty);
209   SmallVector<Type> type1{i32Ptr1Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, i32Ty};
210   SmallVector<Type> type0{i32Ptr0Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, i32Ty};
211   SmallVector<Type> type3{i32Ptr3Ty, f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty, i32Ty};
212   if (verifyWMMAStoreOp(op, type1) || verifyWMMAStoreOp(op, type0) ||
213       verifyWMMAStoreOp(op, type3))
214     return success();
215 
216   return op.emitOpError("expected operands to be a source pointer in memory"
217                         "space 0, 1, 3 followed by ldm of the source");
218 }
219 
220 static LogicalResult verify(WMMAStoreF32M16N16K16Op op) {
221   MLIRContext *context = op.getContext();
222   auto i32Ty = IntegerType::get(context, 32);
223   auto i32Ptr1Ty = LLVM::LLVMPointerType::get(i32Ty, 1);
224   auto i32Ptr3Ty = LLVM::LLVMPointerType::get(i32Ty, 3);
225   auto i32Ptr0Ty = LLVM::LLVMPointerType::get(i32Ty, 0);
226   auto f32Ty = FloatType::getF32(context);
227 
228   SmallVector<Type> type1{i32Ptr1Ty, f32Ty, f32Ty, f32Ty, f32Ty,
229                           f32Ty,     f32Ty, f32Ty, f32Ty, i32Ty};
230   SmallVector<Type> type0{i32Ptr0Ty, f32Ty, f32Ty, f32Ty, f32Ty,
231                           f32Ty,     f32Ty, f32Ty, f32Ty, i32Ty};
232   SmallVector<Type> type3{i32Ptr3Ty, f32Ty, f32Ty, f32Ty, f32Ty,
233                           f32Ty,     f32Ty, f32Ty, f32Ty, i32Ty};
234   if (verifyWMMAStoreOp(op, type0) || verifyWMMAStoreOp(op, type1) ||
235       verifyWMMAStoreOp(op, type3))
236     return success();
237 
238   return op.emitOpError("expected operands to be a source pointer in memory"
239                         "space 0, 1, 3 followed by ldm of the source");
240 }
241 
242 static LogicalResult verify(WMMAMmaF16F16M16N16K16Op op) {
243   MLIRContext *context = op.getContext();
244   auto f16Ty = FloatType::getF16(context);
245   auto f16x2Ty = VectorType::get(2, f16Ty);
246   auto f16x2x4StructTy = LLVM::LLVMStructType::getLiteral(
247       context, {f16x2Ty, f16x2Ty, f16x2Ty, f16x2Ty});
248 
249   SmallVector<Type, 2> operandTypes(op.getOperandTypes().begin(),
250                                     op.getOperandTypes().end());
251   if (operandTypes != SmallVector<Type, 20>(20, f16x2Ty))
252     return op.emitOpError("expected 20 <halfx2>s as operands");
253 
254   if (op.getResult().getType() != f16x2x4StructTy)
255     return op.emitOpError("expected result type to be a struct of 4 <halfx2>s");
256 
257   return success();
258 }
259 
260 static LogicalResult parseWMMAMmaF16F16M16N16K16Op(OpAsmParser &parser,
261                                                    OperationState &result) {
262   SmallVector<OpAsmParser::OperandType, 4> operands;
263   ::llvm::SMLoc operandsLoc;
264   Type operandType;
265   Type resType;
266 
267   operandsLoc = parser.getCurrentLocation();
268   if (parser.parseOperandList(operands) ||
269       parser.parseOptionalAttrDict(result.attributes) || parser.parseColon() ||
270       parser.parseType(operandType) || parser.parseArrow())
271     return failure();
272 
273   unsigned numOperands = operands.size();
274   SmallVector<Type> operandTypes(numOperands, operandType);
275   if (parser.parseType(resType))
276     return failure();
277   result.addTypes(resType);
278   if (parser.resolveOperands(operands, operandTypes, operandsLoc,
279                              result.operands))
280     return failure();
281   return success();
282 }
283 
284 static void printWMMAMmaF16F16M16N16K16Op(OpAsmPrinter &p,
285                                           WMMAMmaF16F16M16N16K16Op &op) {
286   p << op.getOperationName();
287   p << ' ';
288   p << op.args();
289   p.printOptionalAttrDict(op->getAttrs(), {});
290   p << " : ";
291   p << op->getOperand(0).getType();
292   p << ' ' << "->";
293   p << ' ';
294   p << ::llvm::ArrayRef<::mlir::Type>(op.res().getType());
295 }
296 
297 static LogicalResult verify(WMMAMmaF32F32M16N16K16Op op) {
298   unsigned numABOperands = 16;
299   unsigned numCOperands = 8;
300   MLIRContext *context = op.getContext();
301   auto f16Ty = FloatType::getF16(context);
302   auto f32Ty = FloatType::getF32(context);
303   auto f16x2Ty = VectorType::get(2, f16Ty);
304   auto f32x8StructTy = LLVM::LLVMStructType::getLiteral(
305       context, {f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty, f32Ty});
306 
307   SmallVector<Type> abOpTypes;
308   SmallVector<Type> bOpTypes;
309   SmallVector<Type> cOpTypes;
310 
311   for (auto operand : op->getOperands().take_front(numABOperands)) {
312     abOpTypes.push_back(operand.getType());
313   }
314 
315   for (auto operand :
316        op->getOperands().drop_front(numABOperands).take_front(numCOperands)) {
317     cOpTypes.push_back(operand.getType());
318   }
319 
320   if (abOpTypes != SmallVector<Type>(16, f16x2Ty))
321     return op.emitOpError("expected 16 <halfx2>s for `a` and `b` operand");
322 
323   if (cOpTypes != SmallVector<Type>(8, f32Ty))
324     return op.emitOpError("expected 8 f32s for `c` operand");
325 
326   if (op.getResult().getType() != f32x8StructTy)
327     return op.emitOpError("expected result type to be a struct of 8 f32s");
328 
329   return success();
330 }
331 
332 //===----------------------------------------------------------------------===//
333 // NVVMDialect initialization, type parsing, and registration.
334 //===----------------------------------------------------------------------===//
335 
336 // TODO: This should be the llvm.nvvm dialect once this is supported.
337 void NVVMDialect::initialize() {
338   addOperations<
339 #define GET_OP_LIST
340 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
341       >();
342 
343   // Support unknown operations because not all NVVM operations are
344   // registered.
345   allowUnknownOperations();
346 }
347 
348 LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
349                                                     NamedAttribute attr) {
350   // Kernel function attribute should be attached to functions.
351   if (attr.first == NVVMDialect::getKernelFuncAttrName()) {
352     if (!isa<LLVM::LLVMFuncOp>(op)) {
353       return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
354                              << "' attribute attached to unexpected op";
355     }
356   }
357   return success();
358 }
359 
360 #define GET_OP_CLASSES
361 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
362