291 lines
13 KiB
C++
291 lines
13 KiB
C++
//===-- CUFGPUToLLVMConversion.cpp ----------------------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "flang/Optimizer/Transforms/CUFGPUToLLVMConversion.h"
|
|
#include "flang/Optimizer/Builder/CUFCommon.h"
|
|
#include "flang/Optimizer/CodeGen/TypeConverter.h"
|
|
#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
|
|
#include "flang/Optimizer/Support/DataLayout.h"
|
|
#include "flang/Runtime/CUDA/common.h"
|
|
#include "flang/Support/Fortran.h"
|
|
#include "mlir/Conversion/LLVMCommon/Pattern.h"
|
|
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
|
|
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "llvm/Support/FormatVariadic.h"
|
|
|
|
namespace fir {
|
|
#define GEN_PASS_DEF_CUFGPUTOLLVMCONVERSION
|
|
#include "flang/Optimizer/Transforms/Passes.h.inc"
|
|
} // namespace fir
|
|
|
|
using namespace fir;
|
|
using namespace mlir;
|
|
using namespace Fortran::runtime;
|
|
|
|
namespace {
|
|
|
|
static mlir::Value createKernelArgArray(mlir::Location loc,
|
|
mlir::ValueRange operands,
|
|
mlir::PatternRewriter &rewriter) {
|
|
|
|
auto *ctx = rewriter.getContext();
|
|
llvm::SmallVector<mlir::Type> structTypes(operands.size(), nullptr);
|
|
|
|
for (auto [i, arg] : llvm::enumerate(operands))
|
|
structTypes[i] = arg.getType();
|
|
|
|
auto structTy = mlir::LLVM::LLVMStructType::getLiteral(ctx, structTypes);
|
|
auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());
|
|
mlir::Type i32Ty = rewriter.getI32Type();
|
|
auto zero = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 0));
|
|
auto one = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 1));
|
|
mlir::Value argStruct =
|
|
rewriter.create<mlir::LLVM::AllocaOp>(loc, ptrTy, structTy, one);
|
|
auto size = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, i32Ty, rewriter.getIntegerAttr(i32Ty, structTypes.size()));
|
|
mlir::Value argArray =
|
|
rewriter.create<mlir::LLVM::AllocaOp>(loc, ptrTy, ptrTy, size);
|
|
|
|
for (auto [i, arg] : llvm::enumerate(operands)) {
|
|
auto indice = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, i32Ty, rewriter.getIntegerAttr(i32Ty, i));
|
|
mlir::Value structMember = rewriter.create<LLVM::GEPOp>(
|
|
loc, ptrTy, structTy, argStruct,
|
|
mlir::ArrayRef<mlir::Value>({zero, indice}));
|
|
rewriter.create<LLVM::StoreOp>(loc, arg, structMember);
|
|
mlir::Value arrayMember = rewriter.create<LLVM::GEPOp>(
|
|
loc, ptrTy, ptrTy, argArray, mlir::ArrayRef<mlir::Value>({indice}));
|
|
rewriter.create<LLVM::StoreOp>(loc, structMember, arrayMember);
|
|
}
|
|
return argArray;
|
|
}
|
|
|
|
struct GPULaunchKernelConversion
|
|
: public mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp> {
|
|
explicit GPULaunchKernelConversion(
|
|
const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)
|
|
: mlir::ConvertOpToLLVMPattern<mlir::gpu::LaunchFuncOp>(typeConverter,
|
|
benefit) {}
|
|
|
|
using OpAdaptor = typename mlir::gpu::LaunchFuncOp::Adaptor;
|
|
|
|
mlir::LogicalResult
|
|
matchAndRewrite(mlir::gpu::LaunchFuncOp op, OpAdaptor adaptor,
|
|
mlir::ConversionPatternRewriter &rewriter) const override {
|
|
mlir::Location loc = op.getLoc();
|
|
auto *ctx = rewriter.getContext();
|
|
mlir::ModuleOp mod = op->getParentOfType<mlir::ModuleOp>();
|
|
mlir::Value dynamicMemorySize = op.getDynamicSharedMemorySize();
|
|
mlir::Type i32Ty = rewriter.getI32Type();
|
|
if (!dynamicMemorySize)
|
|
dynamicMemorySize = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, i32Ty, rewriter.getIntegerAttr(i32Ty, 0));
|
|
|
|
mlir::Value kernelArgs =
|
|
createKernelArgArray(loc, adaptor.getKernelOperands(), rewriter);
|
|
|
|
auto ptrTy = mlir::LLVM::LLVMPointerType::get(rewriter.getContext());
|
|
auto kernel = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(op.getKernelName());
|
|
mlir::Value kernelPtr;
|
|
if (!kernel) {
|
|
auto funcOp = mod.lookupSymbol<mlir::func::FuncOp>(op.getKernelName());
|
|
if (!funcOp)
|
|
return mlir::failure();
|
|
kernelPtr =
|
|
rewriter.create<LLVM::AddressOfOp>(loc, ptrTy, funcOp.getName());
|
|
} else {
|
|
kernelPtr =
|
|
rewriter.create<LLVM::AddressOfOp>(loc, ptrTy, kernel.getName());
|
|
}
|
|
|
|
auto llvmIntPtrType = mlir::IntegerType::get(
|
|
ctx, this->getTypeConverter()->getPointerBitwidth(0));
|
|
auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);
|
|
|
|
mlir::Value nullPtr = rewriter.create<LLVM::ZeroOp>(loc, ptrTy);
|
|
|
|
if (op.hasClusterSize()) {
|
|
auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(
|
|
RTNAME_STRING(CUFLaunchClusterKernel));
|
|
auto funcTy = mlir::LLVM::LLVMFunctionType::get(
|
|
voidTy,
|
|
{ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
|
|
llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
|
|
llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, i32Ty, ptrTy, ptrTy},
|
|
/*isVarArg=*/false);
|
|
auto cufLaunchClusterKernel = mlir::SymbolRefAttr::get(
|
|
mod.getContext(), RTNAME_STRING(CUFLaunchClusterKernel));
|
|
if (!funcOp) {
|
|
mlir::OpBuilder::InsertionGuard insertGuard(rewriter);
|
|
rewriter.setInsertionPointToStart(mod.getBody());
|
|
auto launchKernelFuncOp = rewriter.create<mlir::LLVM::LLVMFuncOp>(
|
|
loc, RTNAME_STRING(CUFLaunchClusterKernel), funcTy);
|
|
launchKernelFuncOp.setVisibility(
|
|
mlir::SymbolTable::Visibility::Private);
|
|
}
|
|
mlir::Value stream = adaptor.getAsyncObject();
|
|
if (!stream)
|
|
stream = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, llvmIntPtrType, rewriter.getIntegerAttr(llvmIntPtrType, -1));
|
|
rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
|
|
op, funcTy, cufLaunchClusterKernel,
|
|
mlir::ValueRange{kernelPtr, adaptor.getClusterSizeX(),
|
|
adaptor.getClusterSizeY(), adaptor.getClusterSizeZ(),
|
|
adaptor.getGridSizeX(), adaptor.getGridSizeY(),
|
|
adaptor.getGridSizeZ(), adaptor.getBlockSizeX(),
|
|
adaptor.getBlockSizeY(), adaptor.getBlockSizeZ(),
|
|
stream, dynamicMemorySize, kernelArgs, nullPtr});
|
|
} else {
|
|
auto procAttr =
|
|
op->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());
|
|
bool isGridGlobal =
|
|
procAttr && procAttr.getValue() == cuf::ProcAttribute::GridGlobal;
|
|
llvm::StringRef fctName = isGridGlobal
|
|
? RTNAME_STRING(CUFLaunchCooperativeKernel)
|
|
: RTNAME_STRING(CUFLaunchKernel);
|
|
auto funcOp = mod.lookupSymbol<mlir::LLVM::LLVMFuncOp>(fctName);
|
|
auto funcTy = mlir::LLVM::LLVMFunctionType::get(
|
|
voidTy,
|
|
{ptrTy, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
|
|
llvmIntPtrType, llvmIntPtrType, llvmIntPtrType, llvmIntPtrType,
|
|
i32Ty, ptrTy, ptrTy},
|
|
/*isVarArg=*/false);
|
|
auto cufLaunchKernel =
|
|
mlir::SymbolRefAttr::get(mod.getContext(), fctName);
|
|
if (!funcOp) {
|
|
mlir::OpBuilder::InsertionGuard insertGuard(rewriter);
|
|
rewriter.setInsertionPointToStart(mod.getBody());
|
|
auto launchKernelFuncOp =
|
|
rewriter.create<mlir::LLVM::LLVMFuncOp>(loc, fctName, funcTy);
|
|
launchKernelFuncOp.setVisibility(
|
|
mlir::SymbolTable::Visibility::Private);
|
|
}
|
|
|
|
mlir::Value stream = adaptor.getAsyncObject();
|
|
if (!stream)
|
|
stream = rewriter.create<mlir::LLVM::ConstantOp>(
|
|
loc, llvmIntPtrType, rewriter.getIntegerAttr(llvmIntPtrType, -1));
|
|
|
|
rewriter.replaceOpWithNewOp<mlir::LLVM::CallOp>(
|
|
op, funcTy, cufLaunchKernel,
|
|
mlir::ValueRange{kernelPtr, adaptor.getGridSizeX(),
|
|
adaptor.getGridSizeY(), adaptor.getGridSizeZ(),
|
|
adaptor.getBlockSizeX(), adaptor.getBlockSizeY(),
|
|
adaptor.getBlockSizeZ(), stream, dynamicMemorySize,
|
|
kernelArgs, nullPtr});
|
|
}
|
|
|
|
return mlir::success();
|
|
}
|
|
};
|
|
|
|
static std::string getFuncName(cuf::SharedMemoryOp op) {
|
|
if (auto gpuFuncOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())
|
|
return gpuFuncOp.getName().str();
|
|
if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>())
|
|
return funcOp.getName().str();
|
|
if (auto llvmFuncOp = op->getParentOfType<mlir::LLVM::LLVMFuncOp>())
|
|
return llvmFuncOp.getSymName().str();
|
|
return "";
|
|
}
|
|
|
|
static mlir::Value createAddressOfOp(mlir::ConversionPatternRewriter &rewriter,
|
|
mlir::Location loc,
|
|
gpu::GPUModuleOp gpuMod,
|
|
std::string &sharedGlobalName) {
|
|
auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(
|
|
rewriter.getContext(), mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace);
|
|
if (auto g = gpuMod.lookupSymbol<fir::GlobalOp>(sharedGlobalName))
|
|
return rewriter.create<mlir::LLVM::AddressOfOp>(loc, llvmPtrTy,
|
|
g.getSymName());
|
|
if (auto g = gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(sharedGlobalName))
|
|
return rewriter.create<mlir::LLVM::AddressOfOp>(loc, llvmPtrTy,
|
|
g.getSymName());
|
|
return {};
|
|
}
|
|
|
|
struct CUFSharedMemoryOpConversion
|
|
: public mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp> {
|
|
explicit CUFSharedMemoryOpConversion(
|
|
const fir::LLVMTypeConverter &typeConverter, mlir::PatternBenefit benefit)
|
|
: mlir::ConvertOpToLLVMPattern<cuf::SharedMemoryOp>(typeConverter,
|
|
benefit) {}
|
|
using OpAdaptor = typename cuf::SharedMemoryOp::Adaptor;
|
|
|
|
mlir::LogicalResult
|
|
matchAndRewrite(cuf::SharedMemoryOp op, OpAdaptor adaptor,
|
|
mlir::ConversionPatternRewriter &rewriter) const override {
|
|
mlir::Location loc = op->getLoc();
|
|
if (!op.getOffset())
|
|
mlir::emitError(loc,
|
|
"cuf.shared_memory must have an offset for code gen");
|
|
|
|
auto gpuMod = op->getParentOfType<gpu::GPUModuleOp>();
|
|
std::string sharedGlobalName =
|
|
(getFuncName(op) + llvm::Twine(cudaSharedMemSuffix)).str();
|
|
mlir::Value sharedGlobalAddr =
|
|
createAddressOfOp(rewriter, loc, gpuMod, sharedGlobalName);
|
|
|
|
if (!sharedGlobalAddr)
|
|
mlir::emitError(loc, "Could not find the shared global operation\n");
|
|
|
|
auto castPtr = rewriter.create<mlir::LLVM::AddrSpaceCastOp>(
|
|
loc, mlir::LLVM::LLVMPointerType::get(rewriter.getContext()),
|
|
sharedGlobalAddr);
|
|
mlir::Type baseType = castPtr->getResultTypes().front();
|
|
llvm::SmallVector<mlir::LLVM::GEPArg> gepArgs = {op.getOffset()};
|
|
mlir::Value shmemPtr = rewriter.create<mlir::LLVM::GEPOp>(
|
|
loc, baseType, rewriter.getI8Type(), castPtr, gepArgs);
|
|
rewriter.replaceOp(op, {shmemPtr});
|
|
return mlir::success();
|
|
}
|
|
};
|
|
|
|
class CUFGPUToLLVMConversion
|
|
: public fir::impl::CUFGPUToLLVMConversionBase<CUFGPUToLLVMConversion> {
|
|
public:
|
|
void runOnOperation() override {
|
|
auto *ctx = &getContext();
|
|
mlir::RewritePatternSet patterns(ctx);
|
|
mlir::ConversionTarget target(*ctx);
|
|
|
|
mlir::Operation *op = getOperation();
|
|
mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(op);
|
|
if (!module)
|
|
return signalPassFailure();
|
|
|
|
std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(
|
|
module, /*allowDefaultLayout=*/false);
|
|
fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,
|
|
/*forceUnifiedTBAATree=*/false, *dl);
|
|
cuf::populateCUFGPUToLLVMConversionPatterns(typeConverter, patterns);
|
|
target.addIllegalOp<mlir::gpu::LaunchFuncOp>();
|
|
target.addIllegalOp<cuf::SharedMemoryOp>();
|
|
target.addLegalDialect<mlir::LLVM::LLVMDialect>();
|
|
if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
|
|
std::move(patterns)))) {
|
|
mlir::emitError(mlir::UnknownLoc::get(ctx),
|
|
"error in CUF GPU op conversion\n");
|
|
signalPassFailure();
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void cuf::populateCUFGPUToLLVMConversionPatterns(
|
|
const fir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns,
|
|
mlir::PatternBenefit benefit) {
|
|
patterns.add<CUFSharedMemoryOpConversion, GPULaunchKernelConversion>(
|
|
converter, benefit);
|
|
}
|