1//===- ROCDLDialect.cpp - ROCDL 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 ROCDL IR dialect in
10// MLIR, and the LLVM IR dialect. It also registers the dialect.
11//
12// The ROCDL dialect only contains GPU specific additions on top of the general
13// LLVM dialect.
14//
15//===----------------------------------------------------------------------===//
16
17#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
18
19#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
20#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
21#include "mlir/IR/Builders.h"
22#include "mlir/IR/BuiltinTypes.h"
23#include "mlir/IR/DialectImplementation.h"
24#include "mlir/IR/MLIRContext.h"
25#include "mlir/IR/Operation.h"
26#include "llvm/ADT/TypeSwitch.h"
27#include "llvm/AsmParser/Parser.h"
28#include "llvm/IR/Attributes.h"
29#include "llvm/IR/Function.h"
30#include "llvm/IR/Type.h"
31#include "llvm/Support/SourceMgr.h"
32
33using namespace mlir;
34using namespace ROCDL;
35
36#include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.cpp.inc"
37
38//===----------------------------------------------------------------------===//
39// Parsing for ROCDL ops
40//===----------------------------------------------------------------------===//
41
42// <operation> ::=
43// `llvm.amdgcn.buffer.load.* %rsrc, %vindex, %offset, %glc, %slc :
44// result_type`
45ParseResult MubufLoadOp::parse(OpAsmParser &parser, OperationState &result) {
46 SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
47 Type type;
48 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type) ||
49 parser.addTypeToList(type, result.types))
50 return failure();
51
52 MLIRContext *context = parser.getContext();
53 auto int32Ty = IntegerType::get(context, 32);
54 auto int1Ty = IntegerType::get(context, 1);
55 auto i32x4Ty = LLVM::getFixedVectorType(int32Ty, 4);
56 return parser.resolveOperands(ops,
57 {i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
58 parser.getNameLoc(), result.operands);
59}
60
61void MubufLoadOp::print(OpAsmPrinter &p) {
62 p << " " << getOperands() << " : " << (*this)->getResultTypes();
63}
64
65// <operation> ::=
66// `llvm.amdgcn.buffer.store.* %vdata, %rsrc, %vindex, %offset, %glc, %slc :
67// result_type`
68ParseResult MubufStoreOp::parse(OpAsmParser &parser, OperationState &result) {
69 SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
70 Type type;
71 if (parser.parseOperandList(ops, 6) || parser.parseColonType(type))
72 return failure();
73
74 MLIRContext *context = parser.getContext();
75 auto int32Ty = IntegerType::get(context, 32);
76 auto int1Ty = IntegerType::get(context, 1);
77 auto i32x4Ty = LLVM::getFixedVectorType(int32Ty, 4);
78
79 if (parser.resolveOperands(ops,
80 {type, i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
81 parser.getNameLoc(), result.operands))
82 return failure();
83 return success();
84}
85
86void MubufStoreOp::print(OpAsmPrinter &p) {
87 p << " " << getOperands() << " : " << getVdata().getType();
88}
89
90// <operation> ::=
91// `llvm.amdgcn.raw.buffer.load.* %rsrc, %offset, %soffset, %aux
92// : result_type`
93ParseResult RawBufferLoadOp::parse(OpAsmParser &parser,
94 OperationState &result) {
95 SmallVector<OpAsmParser::UnresolvedOperand, 4> ops;
96 Type type;
97 if (parser.parseOperandList(ops, 4) || parser.parseColonType(type) ||
98 parser.addTypeToList(type, result.types))
99 return failure();
100
101 auto bldr = parser.getBuilder();
102 auto int32Ty = bldr.getI32Type();
103 auto i32x4Ty = VectorType::get({4}, int32Ty);
104 return parser.resolveOperands(ops, {i32x4Ty, int32Ty, int32Ty, int32Ty},
105 parser.getNameLoc(), result.operands);
106}
107
108void RawBufferLoadOp::print(OpAsmPrinter &p) {
109 p << " " << getOperands() << " : " << getRes().getType();
110}
111
112// <operation> ::=
113// `llvm.amdgcn.raw.buffer.store.* %vdata, %rsrc, %offset,
114// %soffset, %aux : result_type`
115ParseResult RawBufferStoreOp::parse(OpAsmParser &parser,
116 OperationState &result) {
117 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
118 Type type;
119 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
120 return failure();
121
122 auto bldr = parser.getBuilder();
123 auto int32Ty = bldr.getI32Type();
124 auto i32x4Ty = VectorType::get({4}, int32Ty);
125
126 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
127 parser.getNameLoc(), result.operands))
128 return failure();
129 return success();
130}
131
132void RawBufferStoreOp::print(OpAsmPrinter &p) {
133 p << " " << getOperands() << " : " << getVdata().getType();
134}
135
136// <operation> ::=
137// `llvm.amdgcn.raw.buffer.atomic.fadd.* %vdata, %rsrc, %offset,
138// %soffset, %aux : result_type`
139ParseResult RawBufferAtomicFAddOp::parse(OpAsmParser &parser,
140 OperationState &result) {
141 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
142 Type type;
143 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
144 return failure();
145
146 auto bldr = parser.getBuilder();
147 auto int32Ty = bldr.getI32Type();
148 auto i32x4Ty = VectorType::get({4}, int32Ty);
149
150 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
151 parser.getNameLoc(), result.operands))
152 return failure();
153 return success();
154}
155
156void RawBufferAtomicFAddOp::print(mlir::OpAsmPrinter &p) {
157 p << " " << getOperands() << " : " << getVdata().getType();
158}
159
160// <operation> ::=
161// `llvm.amdgcn.raw.buffer.atomic.fmax.* %vdata, %rsrc, %offset,
162// %soffset, %aux : result_type`
163ParseResult RawBufferAtomicFMaxOp::parse(OpAsmParser &parser,
164 OperationState &result) {
165 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
166 Type type;
167 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
168 return failure();
169
170 auto bldr = parser.getBuilder();
171 auto int32Ty = bldr.getI32Type();
172 auto i32x4Ty = VectorType::get({4}, int32Ty);
173
174 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
175 parser.getNameLoc(), result.operands))
176 return failure();
177 return success();
178}
179
180void RawBufferAtomicFMaxOp::print(mlir::OpAsmPrinter &p) {
181 p << " " << getOperands() << " : " << getVdata().getType();
182}
183
184// <operation> ::=
185// `llvm.amdgcn.raw.buffer.atomic.smax.* %vdata, %rsrc, %offset,
186// %soffset, %aux : result_type`
187ParseResult RawBufferAtomicSMaxOp::parse(OpAsmParser &parser,
188 OperationState &result) {
189 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
190 Type type;
191 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
192 return failure();
193
194 auto bldr = parser.getBuilder();
195 auto int32Ty = bldr.getI32Type();
196 auto i32x4Ty = VectorType::get({4}, int32Ty);
197
198 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
199 parser.getNameLoc(), result.operands))
200 return failure();
201 return success();
202}
203
204void RawBufferAtomicSMaxOp::print(mlir::OpAsmPrinter &p) {
205 p << " " << getOperands() << " : " << getVdata().getType();
206}
207
208// <operation> ::=
209// `llvm.amdgcn.raw.buffer.atomic.umin.* %vdata, %rsrc, %offset,
210// %soffset, %aux : result_type`
211ParseResult RawBufferAtomicUMinOp::parse(OpAsmParser &parser,
212 OperationState &result) {
213 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
214 Type type;
215 if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
216 return failure();
217
218 auto bldr = parser.getBuilder();
219 auto int32Ty = bldr.getI32Type();
220 auto i32x4Ty = VectorType::get({4}, int32Ty);
221
222 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
223 parser.getNameLoc(), result.operands))
224 return failure();
225 return success();
226}
227
228void RawBufferAtomicUMinOp::print(mlir::OpAsmPrinter &p) {
229 p << " " << getOperands() << " : " << getVdata().getType();
230}
231
232//===----------------------------------------------------------------------===//
233// ROCDLDialect initialization, type parsing, and registration.
234//===----------------------------------------------------------------------===//
235
236// TODO: This should be the llvm.rocdl dialect once this is supported.
237void ROCDLDialect::initialize() {
238 addOperations<
239#define GET_OP_LIST
240#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
241 >();
242
243 addAttributes<
244#define GET_ATTRDEF_LIST
245#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
246 >();
247
248 // Support unknown operations because not all ROCDL operations are registered.
249 allowUnknownOperations();
250 declarePromisedInterface<gpu::TargetAttrInterface, ROCDLTargetAttr>();
251}
252
253LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
254 NamedAttribute attr) {
255 // Kernel function attribute should be attached to functions.
256 if (kernelAttrName.getName() == attr.getName()) {
257 if (!isa<LLVM::LLVMFuncOp>(op)) {
258 return op->emitError() << "'" << kernelAttrName.getName()
259 << "' attribute attached to unexpected op";
260 }
261 }
262 return success();
263}
264
265//===----------------------------------------------------------------------===//
266// ROCDL target attribute.
267//===----------------------------------------------------------------------===//
268LogicalResult
269ROCDLTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
270 int optLevel, StringRef triple, StringRef chip,
271 StringRef features, StringRef abiVersion,
272 DictionaryAttr flags, ArrayAttr files) {
273 if (optLevel < 0 || optLevel > 3) {
274 emitError() << "The optimization level must be a number between 0 and 3.";
275 return failure();
276 }
277 if (triple.empty()) {
278 emitError() << "The target triple cannot be empty.";
279 return failure();
280 }
281 if (chip.empty()) {
282 emitError() << "The target chip cannot be empty.";
283 return failure();
284 }
285 if (abiVersion != "400" && abiVersion != "500") {
286 emitError() << "Invalid ABI version, it must be either `400` or `500`.";
287 return failure();
288 }
289 if (files && !llvm::all_of(files, [](::mlir::Attribute attr) {
290 return attr && mlir::isa<StringAttr>(attr);
291 })) {
292 emitError() << "All the elements in the `link` array must be strings.";
293 return failure();
294 }
295 return success();
296}
297
298#define GET_OP_CLASSES
299#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
300
301#define GET_ATTRDEF_CLASSES
302#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
303

source code of mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp