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