20 #include "llvm/IR/IRBuilder.h"
21 #include "llvm/IR/IntrinsicsNVPTX.h"
28 NVVM::ReduxKind kind) {
29 if (!resultType->isIntegerTy(32))
30 llvm_unreachable(
"unsupported data type for redux");
33 case NVVM::ReduxKind::ADD:
34 return llvm::Intrinsic::nvvm_redux_sync_add;
35 case NVVM::ReduxKind::UMAX:
36 return llvm::Intrinsic::nvvm_redux_sync_umax;
37 case NVVM::ReduxKind::UMIN:
38 return llvm::Intrinsic::nvvm_redux_sync_umin;
39 case NVVM::ReduxKind::AND:
40 return llvm::Intrinsic::nvvm_redux_sync_and;
41 case NVVM::ReduxKind::OR:
42 return llvm::Intrinsic::nvvm_redux_sync_or;
43 case NVVM::ReduxKind::XOR:
44 return llvm::Intrinsic::nvvm_redux_sync_xor;
45 case NVVM::ReduxKind::MAX:
46 return llvm::Intrinsic::nvvm_redux_sync_max;
47 case NVVM::ReduxKind::MIN:
48 return llvm::Intrinsic::nvvm_redux_sync_min;
50 llvm_unreachable(
"unknown redux kind");
58 resultType = cast<llvm::StructType>(resultType)->getElementType(0);
60 case NVVM::ShflKind::bfly:
61 return resultType->isFloatTy()
62 ? llvm::Intrinsic::nvvm_shfl_sync_bfly_f32p
63 : llvm::Intrinsic::nvvm_shfl_sync_bfly_i32p;
64 case NVVM::ShflKind::up:
65 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_up_f32p
66 : llvm::Intrinsic::nvvm_shfl_sync_up_i32p;
67 case NVVM::ShflKind::down:
68 return resultType->isFloatTy()
69 ? llvm::Intrinsic::nvvm_shfl_sync_down_f32p
70 : llvm::Intrinsic::nvvm_shfl_sync_down_i32p;
71 case NVVM::ShflKind::idx:
72 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_idx_f32p
73 : llvm::Intrinsic::nvvm_shfl_sync_idx_i32p;
77 case NVVM::ShflKind::bfly:
78 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_bfly_f32
79 : llvm::Intrinsic::nvvm_shfl_sync_bfly_i32;
80 case NVVM::ShflKind::up:
81 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_up_f32
82 : llvm::Intrinsic::nvvm_shfl_sync_up_i32;
83 case NVVM::ShflKind::down:
84 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_down_f32
85 : llvm::Intrinsic::nvvm_shfl_sync_down_i32;
86 case NVVM::ShflKind::idx:
87 return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_idx_f32
88 : llvm::Intrinsic::nvvm_shfl_sync_idx_i32;
91 llvm_unreachable(
"unknown shuffle kind");
97 if (layout == NVVM::MMALayout::row) {
100 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x1_b16;
102 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x2_b16;
104 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x4_b16;
106 llvm_unreachable(
"unsupported number of matrix");
112 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16;
114 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16;
116 return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16;
118 llvm_unreachable(
"unsupported number of matrix");
126 class NVVMDialectLLVMIRTranslationInterface
134 convertOperation(
Operation *op, llvm::IRBuilderBase &builder,
137 #include "mlir/Dialect/LLVMIR/NVVMConversions.inc"
147 auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
150 llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
151 llvm::Function *llvmFunc = moduleTranslation.lookupFunction(func.getName());
153 auto generateMetadata = [&](
int dim, StringRef name) {
154 llvm::Metadata *llvmMetadata[] = {
158 llvm::Type::getInt32Ty(llvmContext), dim))};
159 llvm::MDNode *llvmMetadataNode =
161 moduleTranslation.getOrInsertNamedModuleMetadata(
"nvvm.annotations")
162 ->addOperand(llvmMetadataNode);
164 if (attribute.getName() == NVVM::NVVMDialect::getMaxntidAttrName()) {
165 if (!dyn_cast<DenseI32ArrayAttr>(attribute.getValue()))
167 auto values = cast<DenseI32ArrayAttr>(attribute.getValue());
168 generateMetadata(values[0], NVVM::NVVMDialect::getMaxntidXName());
169 if (values.size() > 1)
170 generateMetadata(values[1], NVVM::NVVMDialect::getMaxntidYName());
171 if (values.size() > 2)
172 generateMetadata(values[2], NVVM::NVVMDialect::getMaxntidZName());
173 }
else if (attribute.getName() == NVVM::NVVMDialect::getReqntidAttrName()) {
174 if (!dyn_cast<DenseI32ArrayAttr>(attribute.getValue()))
176 auto values = cast<DenseI32ArrayAttr>(attribute.getValue());
177 generateMetadata(values[0], NVVM::NVVMDialect::getReqntidXName());
178 if (values.size() > 1)
179 generateMetadata(values[1], NVVM::NVVMDialect::getReqntidYName());
180 if (values.size() > 2)
181 generateMetadata(values[2], NVVM::NVVMDialect::getReqntidZName());
182 }
else if (attribute.getName() ==
183 NVVM::NVVMDialect::getMinctasmAttrName()) {
184 auto value = dyn_cast<IntegerAttr>(attribute.getValue());
185 generateMetadata(value.getInt(),
"minctasm");
186 }
else if (attribute.getName() == NVVM::NVVMDialect::getMaxnregAttrName()) {
187 auto value = dyn_cast<IntegerAttr>(attribute.getValue());
188 generateMetadata(value.getInt(),
"maxnreg");
189 }
else if (attribute.getName() ==
190 NVVM::NVVMDialect::getKernelFuncAttrName()) {
191 llvm::Metadata *llvmMetadataKernel[] = {
196 llvm::MDNode *llvmMetadataNode =
198 moduleTranslation.getOrInsertNamedModuleMetadata(
"nvvm.annotations")
199 ->addOperand(llvmMetadataNode);
205 convertParameterAttr(LLVMFuncOp funcOp,
int argIdx,
NamedAttribute attribute,
208 llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
209 llvm::Function *llvmFunc =
210 moduleTranslation.lookupFunction(funcOp.getName());
211 llvm::NamedMDNode *nvvmAnnotations =
212 moduleTranslation.getOrInsertNamedModuleMetadata(
"nvvm.annotations");
214 if (attribute.getName() == NVVM::NVVMDialect::getGridConstantAttrName()) {
215 llvm::MDNode *gridConstantMetaData =
nullptr;
218 for (llvm::MDNode *opnd : llvm::reverse(nvvmAnnotations->operands())) {
219 if (opnd->getNumOperands() == 3 &&
221 opnd->getOperand(1) ==
223 gridConstantMetaData = opnd;
233 if (gridConstantMetaData ==
nullptr) {
236 llvm::ValueAsMetadata::getConstant(
238 llvm::Metadata *llvmMetadata[] = {
242 llvm::MDNode *llvmMetadataNode =
244 nvvmAnnotations->addOperand(llvmMetadataNode);
248 dyn_cast<llvm::MDTuple>(gridConstantMetaData->getOperand(2))) {
249 llvm::TempMDTuple clonedArgList = argList->clone();
250 clonedArgList->push_back((llvm::ValueAsMetadata::getConstant(
252 gridConstantMetaData->replaceOperandWith(
253 2, llvm::MDNode::replaceWithUniqued(std::move(clonedArgList)));
263 registry.
insert<NVVM::NVVMDialect>();
265 dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
static llvm::Intrinsic::ID getReduxIntrinsicId(llvm::Type *resultType, NVVM::ReduxKind kind)
static llvm::Intrinsic::ID getShflIntrinsicId(llvm::Type *resultType, NVVM::ShflKind kind, bool withPredicate)
static llvm::Intrinsic::ID getLdMatrixIntrinsicId(NVVM::MMALayout layout, int32_t num)
Return the intrinsic ID associated with ldmatrix for the given paramters.
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.
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.
void registerNVVMDialectTranslation(DialectRegistry ®istry)
Register the NVVM dialect and the translation from it to the LLVM IR in the given registry;.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...