261 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file defines the types and operation details for the ROCDL IR dialect in10// 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 general13// 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 "mlir/Transforms/InliningUtils.h"27#include "llvm/ADT/TypeSwitch.h"28 29using namespace mlir;30using namespace ROCDL;31 32#include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.cpp.inc"33 34//===----------------------------------------------------------------------===//35// Parsing for ROCDL ops36//===----------------------------------------------------------------------===//37 38// <operation> ::=39// `llvm.amdgcn.raw.buffer.load.* %rsrc, %offset, %soffset, %aux40// : result_type`41ParseResult RawBufferLoadOp::parse(OpAsmParser &parser,42 OperationState &result) {43 SmallVector<OpAsmParser::UnresolvedOperand, 4> ops;44 Type type;45 if (parser.parseOperandList(ops, 4) || parser.parseColonType(type) ||46 parser.addTypeToList(type, result.types))47 return failure();48 49 auto bldr = parser.getBuilder();50 auto int32Ty = bldr.getI32Type();51 auto i32x4Ty = VectorType::get({4}, int32Ty);52 return parser.resolveOperands(ops, {i32x4Ty, int32Ty, int32Ty, int32Ty},53 parser.getNameLoc(), 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(ops, 5) || parser.parseColonType(type))68 return failure();69 70 auto bldr = parser.getBuilder();71 auto int32Ty = bldr.getI32Type();72 auto i32x4Ty = VectorType::get({4}, int32Ty);73 74 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},75 parser.getNameLoc(), 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(ops, 5) || parser.parseColonType(type))92 return failure();93 94 auto bldr = parser.getBuilder();95 auto int32Ty = bldr.getI32Type();96 auto i32x4Ty = VectorType::get({4}, int32Ty);97 98 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},99 parser.getNameLoc(), 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(ops, 5) || parser.parseColonType(type))116 return failure();117 118 auto bldr = parser.getBuilder();119 auto int32Ty = bldr.getI32Type();120 auto i32x4Ty = VectorType::get({4}, int32Ty);121 122 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},123 parser.getNameLoc(), 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(ops, 5) || parser.parseColonType(type))140 return failure();141 142 auto bldr = parser.getBuilder();143 auto int32Ty = bldr.getI32Type();144 auto i32x4Ty = VectorType::get({4}, int32Ty);145 146 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},147 parser.getNameLoc(), 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(ops, 5) || parser.parseColonType(type))164 return failure();165 166 auto bldr = parser.getBuilder();167 auto int32Ty = bldr.getI32Type();168 auto i32x4Ty = VectorType::get({4}, int32Ty);169 170 if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},171 parser.getNameLoc(), 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 184namespace {185struct ROCDLInlinerInterface final : DialectInlinerInterface {186 using DialectInlinerInterface::DialectInlinerInterface;187 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {188 return true;189 }190};191} // namespace192 193// TODO: This should be the llvm.rocdl dialect once this is supported.194void ROCDLDialect::initialize() {195 addOperations<196#define GET_OP_LIST197#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"198 >();199 200 addAttributes<201#define GET_ATTRDEF_LIST202#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"203 >();204 205 // Support unknown operations because not all ROCDL operations are registered.206 allowUnknownOperations();207 addInterfaces<ROCDLInlinerInterface>();208 declarePromisedInterface<gpu::TargetAttrInterface, ROCDLTargetAttr>();209}210 211LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,212 NamedAttribute attr) {213 // Kernel function attribute should be attached to functions.214 if (kernelAttrName.getName() == attr.getName()) {215 if (!isa<LLVM::LLVMFuncOp>(op)) {216 return op->emitError() << "'" << kernelAttrName.getName()217 << "' attribute attached to unexpected op";218 }219 }220 return success();221}222 223//===----------------------------------------------------------------------===//224// ROCDL target attribute.225//===----------------------------------------------------------------------===//226LogicalResult227ROCDLTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,228 int optLevel, StringRef triple, StringRef chip,229 StringRef features, StringRef abiVersion,230 DictionaryAttr flags, ArrayAttr files) {231 if (optLevel < 0 || optLevel > 3) {232 emitError() << "The optimization level must be a number between 0 and 3.";233 return failure();234 }235 if (triple.empty()) {236 emitError() << "The target triple cannot be empty.";237 return failure();238 }239 if (chip.empty()) {240 emitError() << "The target chip cannot be empty.";241 return failure();242 }243 if (abiVersion != "400" && abiVersion != "500" && abiVersion != "600") {244 emitError() << "Invalid ABI version, it must be `400`, `500` or '600'.";245 return failure();246 }247 if (files && !llvm::all_of(files, [](::mlir::Attribute attr) {248 return attr && mlir::isa<StringAttr>(attr);249 })) {250 emitError() << "All the elements in the `link` array must be strings.";251 return failure();252 }253 return success();254}255 256#define GET_OP_CLASSES257#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"258 259#define GET_ATTRDEF_CLASSES260#include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"261