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/IR/Type.h"
28
29using namespace mlir;
30using namespace ROCDL;
31
32#include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.cpp.inc"
33
34//===----------------------------------------------------------------------===//
35// Parsing for ROCDL ops
36//===----------------------------------------------------------------------===//
37
38// <operation> ::=
39// `llvm.amdgcn.raw.buffer.load.* %rsrc, %offset, %soffset, %aux
40// : result_type`
41ParseResult RawBufferLoadOp::parse(OpAsmParser &parser,
42 OperationState &result) {
43 SmallVector<OpAsmParser::UnresolvedOperand, 4> ops;
44 Type type;
45 if (parser.parseOperandList(result&: ops, requiredOperandCount: 4) || parser.parseColonType(result&: type) ||
46 parser.addTypeToList(type, result&: result.types))
47 return failure();
48
49 auto bldr = parser.getBuilder();
50 auto int32Ty = bldr.getI32Type();
51 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
52 return parser.resolveOperands(operands&: ops, types: {i32x4Ty, int32Ty, int32Ty, int32Ty},
53 loc: parser.getNameLoc(), result&: result.operands);
54}
55
56void RawBufferLoadOp::print(OpAsmPrinter &p) {
57 p << " " << getOperands() << " : " << getRes().getType();
58}
59
60// <operation> ::=
61// `llvm.amdgcn.raw.buffer.store.* %vdata, %rsrc, %offset,
62// %soffset, %aux : result_type`
63ParseResult RawBufferStoreOp::parse(OpAsmParser &parser,
64 OperationState &result) {
65 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
66 Type type;
67 if (parser.parseOperandList(result&: ops, requiredOperandCount: 5) || parser.parseColonType(result&: type))
68 return failure();
69
70 auto bldr = parser.getBuilder();
71 auto int32Ty = bldr.getI32Type();
72 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
73
74 if (parser.resolveOperands(operands&: ops, types: {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
75 loc: parser.getNameLoc(), result&: result.operands))
76 return failure();
77 return success();
78}
79
80void RawBufferStoreOp::print(OpAsmPrinter &p) {
81 p << " " << getOperands() << " : " << getVdata().getType();
82}
83
84// <operation> ::=
85// `llvm.amdgcn.raw.buffer.atomic.fadd.* %vdata, %rsrc, %offset,
86// %soffset, %aux : result_type`
87ParseResult RawBufferAtomicFAddOp::parse(OpAsmParser &parser,
88 OperationState &result) {
89 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
90 Type type;
91 if (parser.parseOperandList(result&: ops, requiredOperandCount: 5) || parser.parseColonType(result&: type))
92 return failure();
93
94 auto bldr = parser.getBuilder();
95 auto int32Ty = bldr.getI32Type();
96 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
97
98 if (parser.resolveOperands(operands&: ops, types: {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
99 loc: parser.getNameLoc(), result&: result.operands))
100 return failure();
101 return success();
102}
103
104void RawBufferAtomicFAddOp::print(mlir::OpAsmPrinter &p) {
105 p << " " << getOperands() << " : " << getVdata().getType();
106}
107
108// <operation> ::=
109// `llvm.amdgcn.raw.buffer.atomic.fmax.* %vdata, %rsrc, %offset,
110// %soffset, %aux : result_type`
111ParseResult RawBufferAtomicFMaxOp::parse(OpAsmParser &parser,
112 OperationState &result) {
113 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
114 Type type;
115 if (parser.parseOperandList(result&: ops, requiredOperandCount: 5) || parser.parseColonType(result&: type))
116 return failure();
117
118 auto bldr = parser.getBuilder();
119 auto int32Ty = bldr.getI32Type();
120 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
121
122 if (parser.resolveOperands(operands&: ops, types: {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
123 loc: parser.getNameLoc(), result&: result.operands))
124 return failure();
125 return success();
126}
127
128void RawBufferAtomicFMaxOp::print(mlir::OpAsmPrinter &p) {
129 p << " " << getOperands() << " : " << getVdata().getType();
130}
131
132// <operation> ::=
133// `llvm.amdgcn.raw.buffer.atomic.smax.* %vdata, %rsrc, %offset,
134// %soffset, %aux : result_type`
135ParseResult RawBufferAtomicSMaxOp::parse(OpAsmParser &parser,
136 OperationState &result) {
137 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
138 Type type;
139 if (parser.parseOperandList(result&: ops, requiredOperandCount: 5) || parser.parseColonType(result&: type))
140 return failure();
141
142 auto bldr = parser.getBuilder();
143 auto int32Ty = bldr.getI32Type();
144 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
145
146 if (parser.resolveOperands(operands&: ops, types: {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
147 loc: parser.getNameLoc(), result&: result.operands))
148 return failure();
149 return success();
150}
151
152void RawBufferAtomicSMaxOp::print(mlir::OpAsmPrinter &p) {
153 p << " " << getOperands() << " : " << getVdata().getType();
154}
155
156// <operation> ::=
157// `llvm.amdgcn.raw.buffer.atomic.umin.* %vdata, %rsrc, %offset,
158// %soffset, %aux : result_type`
159ParseResult RawBufferAtomicUMinOp::parse(OpAsmParser &parser,
160 OperationState &result) {
161 SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
162 Type type;
163 if (parser.parseOperandList(result&: ops, requiredOperandCount: 5) || parser.parseColonType(result&: type))
164 return failure();
165
166 auto bldr = parser.getBuilder();
167 auto int32Ty = bldr.getI32Type();
168 auto i32x4Ty = VectorType::get(shape: {4}, elementType: int32Ty);
169
170 if (parser.resolveOperands(operands&: ops, types: {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
171 loc: parser.getNameLoc(), result&: result.operands))
172 return failure();
173 return success();
174}
175
176void RawBufferAtomicUMinOp::print(mlir::OpAsmPrinter &p) {
177 p << " " << getOperands() << " : " << getVdata().getType();
178}
179
180//===----------------------------------------------------------------------===//
181// ROCDLDialect initialization, type parsing, and registration.
182//===----------------------------------------------------------------------===//
183
184// TODO: This should be the llvm.rocdl dialect once this is supported.
185void ROCDLDialect::initialize() {
186 addOperations<
187#define GET_OP_LIST
188#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
189 >();
190
191 addAttributes<
192#define GET_ATTRDEF_LIST
193#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
194 >();
195
196 // Support unknown operations because not all ROCDL operations are registered.
197 allowUnknownOperations();
198 declarePromisedInterface<gpu::TargetAttrInterface, ROCDLTargetAttr>();
199}
200
201LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
202 NamedAttribute attr) {
203 // Kernel function attribute should be attached to functions.
204 if (kernelAttrName.getName() == attr.getName()) {
205 if (!isa<LLVM::LLVMFuncOp>(Val: op)) {
206 return op->emitError() << "'" << kernelAttrName.getName()
207 << "' attribute attached to unexpected op";
208 }
209 }
210 return success();
211}
212
213//===----------------------------------------------------------------------===//
214// ROCDL target attribute.
215//===----------------------------------------------------------------------===//
216LogicalResult
217ROCDLTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
218 int optLevel, StringRef triple, StringRef chip,
219 StringRef features, StringRef abiVersion,
220 DictionaryAttr flags, ArrayAttr files) {
221 if (optLevel < 0 || optLevel > 3) {
222 emitError() << "The optimization level must be a number between 0 and 3.";
223 return failure();
224 }
225 if (triple.empty()) {
226 emitError() << "The target triple cannot be empty.";
227 return failure();
228 }
229 if (chip.empty()) {
230 emitError() << "The target chip cannot be empty.";
231 return failure();
232 }
233 if (abiVersion != "400" && abiVersion != "500" && abiVersion != "600") {
234 emitError() << "Invalid ABI version, it must be `400`, `500` or '600'.";
235 return failure();
236 }
237 if (files && !llvm::all_of(Range&: files, P: [](::mlir::Attribute attr) {
238 return attr && mlir::isa<StringAttr>(Val: attr);
239 })) {
240 emitError() << "All the elements in the `link` array must be strings.";
241 return failure();
242 }
243 return success();
244}
245
246#define GET_OP_CLASSES
247#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
248
249#define GET_ATTRDEF_CLASSES
250#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
251

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