336 lines · cpp
1//===-- CUFGPUToLLVMConversion.cpp ----------------------------------------===//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#include "flang/Optimizer/Transforms/CUFGPUToLLVMConversion.h"10#include "flang/Optimizer/Builder/CUFCommon.h"11#include "flang/Optimizer/CodeGen/TypeConverter.h"12#include "flang/Optimizer/Dialect/CUF/CUFOps.h"13#include "flang/Optimizer/Support/DataLayout.h"14#include "flang/Runtime/CUDA/common.h"15#include "flang/Support/Fortran.h"16#include "mlir/Conversion/LLVMCommon/Pattern.h"17#include "mlir/Dialect/DLTI/DLTI.h"18#include "mlir/Dialect/GPU/IR/GPUDialect.h"19#include "mlir/Dialect/LLVMIR/NVVMDialect.h"20#include "mlir/Pass/Pass.h"21#include "mlir/Transforms/DialectConversion.h"22#include "mlir/Transforms/GreedyPatternRewriteDriver.h"23#include "llvm/Support/FormatVariadic.h"24 25namespace fir {26#define GEN_PASS_DEF_CUFGPUTOLLVMCONVERSION27#include "flang/Optimizer/Transforms/Passes.h.inc"28} // namespace fir29 30using namespace fir;31using namespace mlir;32using namespace Fortran::runtime;33 34namespace {35 36static mlir::Value createKernelArgArray(mlir::Location loc,37 mlir::ValueRange operands,38 mlir::PatternRewriter &rewriter) {39 40 auto *ctx = rewriter.getContext();41 llvm::SmallVector<mlir::Type> structTypes(operands.size(), nullptr);42 43 for (auto [i, arg] : llvm::enumerate(operands))44 structTypes[i] = arg.getType();45 46 auto structTy = mlir::LLVM::LLVMStructType::getLiteral(ctx, structTypes);47 auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());48 mlir::Type i32Ty = rewriter.getI32Type();49 auto zero = mlir::LLVM::ConstantOp::create(rewriter, loc, i32Ty,50 rewriter.getIntegerAttr(i32Ty, 0));51 auto one = mlir::LLVM::ConstantOp::create(rewriter, loc, i32Ty,52 rewriter.getIntegerAttr(i32Ty, 1));53 mlir::Value argStruct =54 mlir::LLVM::AllocaOp::create(rewriter, loc, ptrTy, structTy, one);55 auto size = mlir::LLVM::ConstantOp::create(56 rewriter, loc, i32Ty, rewriter.getIntegerAttr(i32Ty, structTypes.size()));57 mlir::Value argArray =58 mlir::LLVM::AllocaOp::create(rewriter, loc, ptrTy, ptrTy, size);59 60 for (auto [i, arg] : llvm::enumerate(operands)) {61 auto indice = mlir::LLVM::ConstantOp::create(62 rewriter, loc, i32Ty, rewriter.getIntegerAttr(i32Ty, i));63 mlir::Value structMember =64 LLVM::GEPOp::create(rewriter, loc, ptrTy, structTy, argStruct,65 mlir::ArrayRef<mlir::Value>({zero, indice}));66 LLVM::StoreOp::create(rewriter, loc, arg, structMember);67 mlir::Value arrayMember =68 LLVM::GEPOp::create(rewriter, loc, ptrTy, ptrTy, argArray,69 mlir::ArrayRef<mlir::Value>({indice}));70 LLVM::StoreOp::create(rewriter, loc, structMember, arrayMember);71 }72 return argArray;73}74 75struct GPULaunchKernelConversion76 : public mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp> {77 explicit GPULaunchKernelConversion(78 const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)79 : mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp>(typeConverter,80 benefit) {}81 82 using OpAdaptor = typename mlir::gpu::LaunchFuncOp::Adaptor;83 84 mlir::LogicalResult85 matchAndRewrite(mlir::gpu::LaunchFuncOp op, OpAdaptor adaptor,86 mlir::ConversionPatternRewriter &rewriter) const override {87 // Only convert gpu.launch_func for CUDA Fortran.88 if (!op.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(89 cuf::getProcAttrName()))90 return mlir::failure();91 92 mlir::Location loc = op.getLoc();93 auto *ctx = rewriter.getContext();94 mlir::ModuleOp mod = op->getParentOfType<mlir::ModuleOp>();95 mlir::Value dynamicMemorySize = op.getDynamicSharedMemorySize();96 mlir::Type i32Ty = rewriter.getI32Type();97 if (!dynamicMemorySize)98 dynamicMemorySize = mlir::LLVM::ConstantOp::create(99 rewriter, loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 0));100 101 mlir::Value kernelArgs =102 createKernelArgArray(loc, adaptor.getKernelOperands(), rewriter);103 104 auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());105 auto kernel = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(op.getKernelName());106 mlir::Value kernelPtr;107 if (!kernel) {108 auto funcOp = mod.lookupSymbol<mlir::func::FuncOp>(op.getKernelName());109 if (!funcOp)110 return mlir::failure();111 kernelPtr =112 LLVM::AddressOfOp::create(rewriter, loc, ptrTy, funcOp.getName());113 } else {114 kernelPtr =115 LLVM::AddressOfOp::create(rewriter, loc, ptrTy, kernel.getName());116 }117 118 auto llvmIntPtrType = mlir::IntegerType::get(119 ctx, this->getTypeConverter()->getPointerBitwidth(0));120 auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);121 122 mlir::Value nullPtr = LLVM::ZeroOp::create(rewriter, loc, ptrTy);123 124 if (op.hasClusterSize()) {125 auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(126 RTNAME_STRING(CUFLaunchClusterKernel));127 auto funcTy = mlir::LLVM::LLVMFunctionType::get(128 voidTy,129 {ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,130 llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,131 llvmIntPtrType, llvmIntPtrType, ptrTy, i32Ty, ptrTy, ptrTy},132 /*isVarArg=*/false);133 auto cufLaunchClusterKernel = mlir::SymbolRefAttr::get(134 mod.getContext(), RTNAME_STRING(CUFLaunchClusterKernel));135 if (!funcOp) {136 mlir::OpBuilder::InsertionGuard insertGuard(rewriter);137 rewriter.setInsertionPointToStart(mod.getBody());138 auto launchKernelFuncOp = mlir::LLVM::LLVMFuncOp::create(139 rewriter, loc, RTNAME_STRING(CUFLaunchClusterKernel), funcTy);140 launchKernelFuncOp.setVisibility(141 mlir::SymbolTable::Visibility::Private);142 }143 144 mlir::Value stream = nullPtr;145 if (!adaptor.getAsyncDependencies().empty()) {146 if (adaptor.getAsyncDependencies().size() != 1)147 return rewriter.notifyMatchFailure(148 op, "Can only convert with exactly one stream dependency.");149 stream = adaptor.getAsyncDependencies().front();150 }151 152 mlir::LLVM::CallOp::create(153 rewriter, loc, funcTy, cufLaunchClusterKernel,154 mlir::ValueRange{kernelPtr, adaptor.getClusterSizeX(),155 adaptor.getClusterSizeY(), adaptor.getClusterSizeZ(),156 adaptor.getGridSizeX(), adaptor.getGridSizeY(),157 adaptor.getGridSizeZ(), adaptor.getBlockSizeX(),158 adaptor.getBlockSizeY(), adaptor.getBlockSizeZ(),159 stream, dynamicMemorySize, kernelArgs, nullPtr});160 rewriter.eraseOp(op);161 } else {162 auto procAttr =163 op->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());164 bool isGridGlobal =165 procAttr && procAttr.getValue() == cuf::ProcAttribute::GridGlobal;166 llvm::StringRef fctName = isGridGlobal167 ? RTNAME_STRING(CUFLaunchCooperativeKernel)168 : RTNAME_STRING(CUFLaunchKernel);169 auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(fctName);170 auto funcTy = mlir::LLVM::LLVMFunctionType::get(171 voidTy,172 {ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,173 llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, ptrTy, i32Ty, ptrTy,174 ptrTy},175 /*isVarArg=*/false);176 auto cufLaunchKernel =177 mlir::SymbolRefAttr::get(mod.getContext(), fctName);178 if (!funcOp) {179 mlir::OpBuilder::InsertionGuard insertGuard(rewriter);180 rewriter.setInsertionPointToStart(mod.getBody());181 auto launchKernelFuncOp =182 mlir::LLVM::LLVMFuncOp::create(rewriter, loc, fctName, funcTy);183 launchKernelFuncOp.setVisibility(184 mlir::SymbolTable::Visibility::Private);185 }186 187 mlir::Value stream = nullPtr;188 if (!adaptor.getAsyncDependencies().empty()) {189 if (adaptor.getAsyncDependencies().size() != 1)190 return rewriter.notifyMatchFailure(191 op, "Can only convert with exactly one stream dependency.");192 stream = adaptor.getAsyncDependencies().front();193 }194 195 mlir::LLVM::CallOp::create(196 rewriter, loc, funcTy, cufLaunchKernel,197 mlir::ValueRange{kernelPtr, adaptor.getGridSizeX(),198 adaptor.getGridSizeY(), adaptor.getGridSizeZ(),199 adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),200 adaptor.getBlockSizeZ(), stream, dynamicMemorySize,201 kernelArgs, nullPtr});202 rewriter.eraseOp(op);203 }204 205 return mlir::success();206 }207};208 209static std::string getFuncName(cuf::SharedMemoryOp op) {210 if (auto gpuFuncOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())211 return gpuFuncOp.getName().str();212 if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>())213 return funcOp.getName().str();214 if (auto llvmFuncOp = op->getParentOfType<mlir::LLVM::LLVMFuncOp>())215 return llvmFuncOp.getSymName().str();216 return "";217}218 219static mlir::Value createAddressOfOp(mlir::ConversionPatternRewriter &rewriter,220 mlir::Location loc,221 gpu::GPUModuleOp gpuMod,222 std::string &sharedGlobalName) {223 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(224 rewriter.getContext(),225 static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Shared));226 if (auto g = gpuMod.lookupSymbol<fir::GlobalOp>(sharedGlobalName))227 return mlir::LLVM::AddressOfOp::create(rewriter, loc, llvmPtrTy,228 g.getSymName());229 if (auto g = gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(sharedGlobalName))230 return mlir::LLVM::AddressOfOp::create(rewriter, loc, llvmPtrTy,231 g.getSymName());232 return {};233}234 235struct CUFSharedMemoryOpConversion236 : public mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp> {237 explicit CUFSharedMemoryOpConversion(238 const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)239 : mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp>(typeConverter,240 benefit) {}241 using OpAdaptor = typename cuf::SharedMemoryOp::Adaptor;242 243 mlir::LogicalResult244 matchAndRewrite(cuf::SharedMemoryOp op, OpAdaptor adaptor,245 mlir::ConversionPatternRewriter &rewriter) const override {246 mlir::Location loc = op->getLoc();247 if (!op.getOffset())248 mlir::emitError(loc,249 "cuf.shared_memory must have an offset for code gen");250 251 auto gpuMod = op->getParentOfType<gpu::GPUModuleOp>();252 std::string sharedGlobalName =253 (getFuncName(op) + llvm::Twine(cudaSharedMemSuffix)).str();254 mlir::Value sharedGlobalAddr =255 createAddressOfOp(rewriter, loc, gpuMod, sharedGlobalName);256 257 if (!sharedGlobalAddr)258 mlir::emitError(loc, "Could not find the shared global operation\n");259 260 auto castPtr = mlir::LLVM::AddrSpaceCastOp::create(261 rewriter, loc, mlir::LLVM::LLVMPointerType::get(rewriter.getContext()),262 sharedGlobalAddr);263 mlir::Type baseType = castPtr->getResultTypes().front();264 llvm::SmallVector<mlir::LLVM::GEPArg> gepArgs = {op.getOffset()};265 mlir::Value shmemPtr = mlir::LLVM::GEPOp::create(266 rewriter, loc, baseType, rewriter.getI8Type(), castPtr, gepArgs);267 rewriter.replaceOp(op, {shmemPtr});268 return mlir::success();269 }270};271 272struct CUFStreamCastConversion273 : public mlir::ConvertOpToLLVMPattern<cuf::StreamCastOp> {274 explicit CUFStreamCastConversion(const fir::LLVMTypeConverter &typeConverter,275 mlir::PatternBenefit benefit)276 : mlir::ConvertOpToLLVMPattern<cuf::StreamCastOp>(typeConverter,277 benefit) {}278 using OpAdaptor = typename cuf::StreamCastOp::Adaptor;279 280 mlir::LogicalResult281 matchAndRewrite(cuf::StreamCastOp op, OpAdaptor adaptor,282 mlir::ConversionPatternRewriter &rewriter) const override {283 rewriter.replaceOp(op, adaptor.getStream());284 return mlir::success();285 }286};287 288class CUFGPUToLLVMConversion289 : public fir::impl::CUFGPUToLLVMConversionBase<CUFGPUToLLVMConversion> {290public:291 void runOnOperation() override {292 auto *ctx = &getContext();293 mlir::RewritePatternSet patterns(ctx);294 mlir::ConversionTarget target(*ctx);295 296 mlir::Operation *op = getOperation();297 mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(op);298 if (!module)299 return signalPassFailure();300 301 std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(302 module, /*allowDefaultLayout=*/false);303 fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,304 /*forceUnifiedTBAATree=*/false, *dl);305 cuf::populateCUFGPUToLLVMConversionPatterns(typeConverter, patterns);306 307 target.addDynamicallyLegalOp<mlir::gpu::LaunchFuncOp>(308 [&](mlir::gpu::LaunchFuncOp op) {309 if (op.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(310 cuf::getProcAttrName()))311 return false;312 return true;313 });314 315 target.addIllegalOp<cuf::SharedMemoryOp>();316 target.addLegalDialect<mlir::LLVM::LLVMDialect>();317 if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,318 std::move(patterns)))) {319 mlir::emitError(mlir::UnknownLoc::get(ctx),320 "error in CUF GPU op conversion\n");321 signalPassFailure();322 }323 }324};325} // namespace326 327void cuf::populateCUFGPUToLLVMConversionPatterns(328 fir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns,329 mlir::PatternBenefit benefit) {330 converter.addConversion([&converter](mlir::gpu::AsyncTokenType) -> Type {331 return mlir::LLVM::LLVMPointerType::get(&converter.getContext());332 });333 patterns.add<CUFSharedMemoryOpConversion, GPULaunchKernelConversion,334 CUFStreamCastConversion>(converter, benefit);335}336