20 #include "llvm/IR/IRBuilder.h"
21 #include "llvm/IR/IntrinsicsAMDGPU.h"
22 #include "llvm/IR/MDBuilder.h"
23 #include "llvm/Support/raw_ostream.h"
32 auto *inst = llvm::cast<llvm::CallInst>(
37 apInts.push_back(llvm::APInt(32, i));
38 llvm::MDBuilder mdBuilder(builder.getContext());
39 llvm::MDNode *range = mdBuilder.createRange(apInts[0], apInts[1]);
40 inst->setMetadata(llvm::LLVMContext::MD_range, range);
50 StringRef fnName,
int parameter) {
51 llvm::Module *module = builder.GetInsertBlock()->getModule();
53 llvm::Type::getInt64Ty(module->getContext()),
54 llvm::Type::getInt32Ty(module->getContext()),
56 llvm::Function *fn = dyn_cast<llvm::Function>(
57 module->getOrInsertFunction(fnName, functionType).getCallee());
59 llvm::Type::getInt32Ty(module->getContext()), parameter);
66 class ROCDLDialectLLVMIRTranslationInterface
74 convertOperation(
Operation *op, llvm::IRBuilderBase &builder,
77 #include "mlir/Dialect/LLVMIR/ROCDLConversions.inc"
87 auto *dialect = dyn_cast<ROCDL::ROCDLDialect>(attribute.getNameDialect());
88 if (dialect->getKernelAttrHelper().getName() == attribute.getName()) {
89 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
92 " is only supported on `llvm.func` operations");
99 llvm::Function *llvmFunc =
100 moduleTranslation.lookupFunction(func.getName());
101 llvmFunc->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL);
102 if (!llvmFunc->hasFnAttribute(
"amdgpu-flat-work-group-size")) {
103 llvmFunc->addFnAttr(
"amdgpu-flat-work-group-size",
"1,256");
110 if (!llvmFunc->hasFnAttribute(
"uniform-work-group-size"))
111 llvmFunc->addFnAttr(
"uniform-work-group-size",
"true");
116 if (dialect->getMaxFlatWorkGroupSizeAttrHelper().getName() ==
117 attribute.getName()) {
118 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
120 return op->
emitOpError(Twine(attribute.getName()) +
121 " is only supported on `llvm.func` operations");
122 auto value = dyn_cast<IntegerAttr>(attribute.getValue());
124 return op->
emitOpError(Twine(attribute.getName()) +
125 " must be an integer");
127 llvm::Function *llvmFunc =
128 moduleTranslation.lookupFunction(func.getName());
130 llvm::raw_svector_ostream attrValueStream(llvmAttrValue);
131 attrValueStream <<
"1," << value.getInt();
132 llvmFunc->addFnAttr(
"amdgpu-flat-work-group-size", llvmAttrValue);
134 if (dialect->getFlatWorkGroupSizeAttrHelper().getName() ==
135 attribute.getName()) {
136 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
138 return op->
emitOpError(Twine(attribute.getName()) +
139 " is only supported on `llvm.func` operations");
140 auto value = dyn_cast<StringAttr>(attribute.getValue());
142 return op->
emitOpError(Twine(attribute.getName()) +
143 " must be a string");
145 llvm::Function *llvmFunc =
146 moduleTranslation.lookupFunction(func.getName());
148 llvmAttrValue.append(value.getValue());
149 llvmFunc->addFnAttr(
"amdgpu-flat-work-group-size", llvmAttrValue);
151 if (ROCDL::ROCDLDialect::getUniformWorkGroupSizeAttrName() ==
152 attribute.getName()) {
153 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
155 return op->
emitOpError(Twine(attribute.getName()) +
156 " is only supported on `llvm.func` operations");
157 auto value = dyn_cast<BoolAttr>(attribute.getValue());
159 return op->
emitOpError(Twine(attribute.getName()) +
160 " must be a boolean");
161 llvm::Function *llvmFunc =
162 moduleTranslation.lookupFunction(func.getName());
163 llvmFunc->addFnAttr(
"uniform-work-group-size",
164 value.getValue() ?
"true" :
"false");
167 if (dialect->getReqdWorkGroupSizeAttrHelper().getName() ==
168 attribute.getName()) {
169 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
171 return op->
emitOpError(Twine(attribute.getName()) +
172 " is only supported on `llvm.func` operations");
173 auto value = dyn_cast<DenseI32ArrayAttr>(attribute.getValue());
175 return op->
emitOpError(Twine(attribute.getName()) +
176 " must be a dense i32 array attribute");
177 llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
180 for (int32_t i : value.asArrayRef()) {
184 llvm::Function *llvmFunc =
185 moduleTranslation.lookupFunction(func.getName());
187 llvmFunc->setMetadata(
"reqd_work_group_size", node);
195 registry.
insert<ROCDL::ROCDLDialect>();
197 dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
static llvm::Value * createDeviceFunctionCall(llvm::IRBuilderBase &builder, StringRef fnName, int parameter)
static llvm::Value * createIntrinsicCallWithRange(llvm::IRBuilderBase &builder, llvm::Intrinsic::ID intrinsic, DenseI32ArrayAttr maybeRange)
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
Base class for dialect interfaces providing translation to LLVM IR.
LLVMTranslationDialectInterface(Dialect *dialect)
Implementation class for module translation.
MLIRContext is the top-level object for a collection of MLIR operations.
void appendDialectRegistry(const DialectRegistry ®istry)
Append the contents of the given dialect registry to the registry associated with this context.
NamedAttribute represents a combination of a name and an Attribute value.
Operation is the basic unit of execution within MLIR.
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
ArrayRef< T > asArrayRef() const
llvm::CallInst * createIntrinsicCall(llvm::IRBuilderBase &builder, llvm::Intrinsic::ID intrinsic, ArrayRef< llvm::Value * > args={}, ArrayRef< llvm::Type * > tys={})
Creates a call to an LLVM IR intrinsic function with the given arguments.
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
void registerROCDLDialectTranslation(DialectRegistry ®istry)
Register the ROCDL dialect and the translation from it to the LLVM IR in the given registry;.
This class represents an efficient way to signal success or failure.