469 lines · cpp
1//===-- CUFAllocationConversion.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/CUDA/CUFAllocationConversion.h"10#include "flang/Optimizer/Builder/CUFCommon.h"11#include "flang/Optimizer/Builder/FIRBuilder.h"12#include "flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h"13#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"14#include "flang/Optimizer/CodeGen/TypeConverter.h"15#include "flang/Optimizer/Dialect/CUF/CUFOps.h"16#include "flang/Optimizer/Dialect/FIRDialect.h"17#include "flang/Optimizer/Dialect/FIROps.h"18#include "flang/Optimizer/HLFIR/HLFIROps.h"19#include "flang/Optimizer/Support/DataLayout.h"20#include "flang/Runtime/CUDA/allocatable.h"21#include "flang/Runtime/CUDA/common.h"22#include "flang/Runtime/CUDA/descriptor.h"23#include "flang/Runtime/CUDA/memory.h"24#include "flang/Runtime/CUDA/pointer.h"25#include "flang/Runtime/allocatable.h"26#include "flang/Runtime/allocator-registry-consts.h"27#include "flang/Support/Fortran.h"28#include "mlir/Dialect/Func/IR/FuncOps.h"29#include "mlir/IR/Matchers.h"30#include "mlir/Pass/Pass.h"31#include "mlir/Transforms/DialectConversion.h"32#include "mlir/Transforms/GreedyPatternRewriteDriver.h"33 34namespace fir {35#define GEN_PASS_DEF_CUFALLOCATIONCONVERSION36#include "flang/Optimizer/Transforms/Passes.h.inc"37} // namespace fir38 39using namespace fir;40using namespace mlir;41using namespace Fortran::runtime;42using namespace Fortran::runtime::cuda;43 44namespace {45 46template <typename OpTy>47static bool isPinned(OpTy op) {48 if (op.getDataAttr() && *op.getDataAttr() == cuf::DataAttribute::Pinned)49 return true;50 return false;51}52 53static inline unsigned getMemType(cuf::DataAttribute attr) {54 if (attr == cuf::DataAttribute::Device)55 return kMemTypeDevice;56 if (attr == cuf::DataAttribute::Managed)57 return kMemTypeManaged;58 if (attr == cuf::DataAttribute::Pinned)59 return kMemTypePinned;60 if (attr == cuf::DataAttribute::Unified)61 return kMemTypeUnified;62 llvm_unreachable("unsupported memory type");63}64 65template <typename OpTy>66static bool hasDoubleDescriptors(OpTy op) {67 if (auto declareOp =68 mlir::dyn_cast_or_null<fir::DeclareOp>(op.getBox().getDefiningOp())) {69 if (mlir::isa_and_nonnull<fir::AddrOfOp>(70 declareOp.getMemref().getDefiningOp())) {71 if (isPinned(declareOp))72 return false;73 return true;74 }75 } else if (auto declareOp = mlir::dyn_cast_or_null<hlfir::DeclareOp>(76 op.getBox().getDefiningOp())) {77 if (mlir::isa_and_nonnull<fir::AddrOfOp>(78 declareOp.getMemref().getDefiningOp())) {79 if (isPinned(declareOp))80 return false;81 return true;82 }83 }84 return false;85}86 87static bool inDeviceContext(mlir::Operation *op) {88 if (op->getParentOfType<cuf::KernelOp>())89 return true;90 if (auto funcOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())91 return true;92 if (auto funcOp = op->getParentOfType<mlir::gpu::LaunchOp>())93 return true;94 if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>()) {95 if (auto cudaProcAttr =96 funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(97 cuf::getProcAttrName())) {98 return cudaProcAttr.getValue() != cuf::ProcAttribute::Host &&99 cudaProcAttr.getValue() != cuf::ProcAttribute::HostDevice;100 }101 }102 return false;103}104 105template <typename OpTy>106static mlir::LogicalResult convertOpToCall(OpTy op,107 mlir::PatternRewriter &rewriter,108 mlir::func::FuncOp func) {109 auto mod = op->template getParentOfType<mlir::ModuleOp>();110 fir::FirOpBuilder builder(rewriter, mod);111 mlir::Location loc = op.getLoc();112 auto fTy = func.getFunctionType();113 114 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);115 mlir::Value sourceLine;116 if constexpr (std::is_same_v<OpTy, cuf::AllocateOp>)117 sourceLine = fir::factory::locationToLineNo(118 builder, loc, op.getSource() ? fTy.getInput(7) : fTy.getInput(6));119 else120 sourceLine = fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));121 122 mlir::Value hasStat = op.getHasStat() ? builder.createBool(loc, true)123 : builder.createBool(loc, false);124 125 mlir::Value errmsg;126 if (op.getErrmsg()) {127 errmsg = op.getErrmsg();128 } else {129 mlir::Type boxNoneTy = fir::BoxType::get(builder.getNoneType());130 errmsg = fir::AbsentOp::create(builder, loc, boxNoneTy).getResult();131 }132 llvm::SmallVector<mlir::Value> args;133 if constexpr (std::is_same_v<OpTy, cuf::AllocateOp>) {134 mlir::Value pinned =135 op.getPinned()136 ? op.getPinned()137 : builder.createNullConstant(138 loc, fir::ReferenceType::get(139 mlir::IntegerType::get(op.getContext(), 1)));140 if (op.getSource()) {141 mlir::Value stream =142 op.getStream() ? op.getStream()143 : builder.createNullConstant(loc, fTy.getInput(2));144 args = fir::runtime::createArguments(145 builder, loc, fTy, op.getBox(), op.getSource(), stream, pinned,146 hasStat, errmsg, sourceFile, sourceLine);147 } else {148 mlir::Value stream =149 op.getStream() ? op.getStream()150 : builder.createNullConstant(loc, fTy.getInput(1));151 args = fir::runtime::createArguments(builder, loc, fTy, op.getBox(),152 stream, pinned, hasStat, errmsg,153 sourceFile, sourceLine);154 }155 } else {156 args =157 fir::runtime::createArguments(builder, loc, fTy, op.getBox(), hasStat,158 errmsg, sourceFile, sourceLine);159 }160 auto callOp = fir::CallOp::create(builder, loc, func, args);161 rewriter.replaceOp(op, callOp);162 return mlir::success();163}164 165struct CUFAllocOpConversion : public mlir::OpRewritePattern<cuf::AllocOp> {166 using OpRewritePattern::OpRewritePattern;167 168 CUFAllocOpConversion(mlir::MLIRContext *context, mlir::DataLayout *dl,169 const fir::LLVMTypeConverter *typeConverter)170 : OpRewritePattern(context), dl{dl}, typeConverter{typeConverter} {}171 172 mlir::LogicalResult173 matchAndRewrite(cuf::AllocOp op,174 mlir::PatternRewriter &rewriter) const override {175 176 mlir::Location loc = op.getLoc();177 178 if (inDeviceContext(op.getOperation())) {179 // In device context just replace the cuf.alloc operation with a fir.alloc180 // the cuf.free will be removed.181 auto allocaOp =182 fir::AllocaOp::create(rewriter, loc, op.getInType(),183 op.getUniqName() ? *op.getUniqName() : "",184 op.getBindcName() ? *op.getBindcName() : "",185 op.getTypeparams(), op.getShape());186 allocaOp->setAttr(cuf::getDataAttrName(), op.getDataAttrAttr());187 rewriter.replaceOp(op, allocaOp);188 return mlir::success();189 }190 191 auto mod = op->getParentOfType<mlir::ModuleOp>();192 fir::FirOpBuilder builder(rewriter, mod);193 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);194 195 if (!mlir::dyn_cast_or_null<fir::BaseBoxType>(op.getInType())) {196 // Convert scalar and known size array allocations.197 mlir::Value bytes;198 fir::KindMapping kindMap{fir::getKindMapping(mod)};199 if (fir::isa_trivial(op.getInType())) {200 int width = cuf::computeElementByteSize(loc, op.getInType(), kindMap);201 bytes =202 builder.createIntegerConstant(loc, builder.getIndexType(), width);203 } else if (auto seqTy = mlir::dyn_cast_or_null<fir::SequenceType>(204 op.getInType())) {205 std::size_t size = 0;206 if (fir::isa_derived(seqTy.getEleTy())) {207 mlir::Type structTy = typeConverter->convertType(seqTy.getEleTy());208 size = dl->getTypeSizeInBits(structTy) / 8;209 } else {210 size = cuf::computeElementByteSize(loc, seqTy.getEleTy(), kindMap);211 }212 mlir::Value width =213 builder.createIntegerConstant(loc, builder.getIndexType(), size);214 mlir::Value nbElem;215 if (fir::sequenceWithNonConstantShape(seqTy)) {216 assert(!op.getShape().empty() && "expect shape with dynamic arrays");217 nbElem = builder.loadIfRef(loc, op.getShape()[0]);218 for (unsigned i = 1; i < op.getShape().size(); ++i) {219 nbElem = mlir::arith::MulIOp::create(220 rewriter, loc, nbElem,221 builder.loadIfRef(loc, op.getShape()[i]));222 }223 } else {224 nbElem = builder.createIntegerConstant(loc, builder.getIndexType(),225 seqTy.getConstantArraySize());226 }227 bytes = mlir::arith::MulIOp::create(rewriter, loc, nbElem, width);228 } else if (fir::isa_derived(op.getInType())) {229 mlir::Type structTy = typeConverter->convertType(op.getInType());230 std::size_t structSize = dl->getTypeSizeInBits(structTy) / 8;231 bytes = builder.createIntegerConstant(loc, builder.getIndexType(),232 structSize);233 } else if (fir::isa_char(op.getInType())) {234 mlir::Type charTy = typeConverter->convertType(op.getInType());235 std::size_t charSize = dl->getTypeSizeInBits(charTy) / 8;236 bytes = builder.createIntegerConstant(loc, builder.getIndexType(),237 charSize);238 } else {239 mlir::emitError(loc, "unsupported type in cuf.alloc\n");240 }241 mlir::func::FuncOp func =242 fir::runtime::getRuntimeFunc<mkRTKey(CUFMemAlloc)>(loc, builder);243 auto fTy = func.getFunctionType();244 mlir::Value sourceLine =245 fir::factory::locationToLineNo(builder, loc, fTy.getInput(3));246 mlir::Value memTy = builder.createIntegerConstant(247 loc, builder.getI32Type(), getMemType(op.getDataAttr()));248 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(249 builder, loc, fTy, bytes, memTy, sourceFile, sourceLine)};250 auto callOp = fir::CallOp::create(builder, loc, func, args);251 callOp->setAttr(cuf::getDataAttrName(), op.getDataAttrAttr());252 auto convOp = builder.createConvert(loc, op.getResult().getType(),253 callOp.getResult(0));254 rewriter.replaceOp(op, convOp);255 return mlir::success();256 }257 258 // Convert descriptor allocations to function call.259 auto boxTy = mlir::dyn_cast_or_null<fir::BaseBoxType>(op.getInType());260 mlir::func::FuncOp func =261 fir::runtime::getRuntimeFunc<mkRTKey(CUFAllocDescriptor)>(loc, builder);262 auto fTy = func.getFunctionType();263 mlir::Value sourceLine =264 fir::factory::locationToLineNo(builder, loc, fTy.getInput(2));265 266 mlir::Type structTy = typeConverter->convertBoxTypeAsStruct(boxTy);267 std::size_t boxSize = dl->getTypeSizeInBits(structTy) / 8;268 mlir::Value sizeInBytes =269 builder.createIntegerConstant(loc, builder.getIndexType(), boxSize);270 271 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(272 builder, loc, fTy, sizeInBytes, sourceFile, sourceLine)};273 auto callOp = fir::CallOp::create(builder, loc, func, args);274 callOp->setAttr(cuf::getDataAttrName(), op.getDataAttrAttr());275 auto convOp = builder.createConvert(loc, op.getResult().getType(),276 callOp.getResult(0));277 rewriter.replaceOp(op, convOp);278 return mlir::success();279 }280 281private:282 mlir::DataLayout *dl;283 const fir::LLVMTypeConverter *typeConverter;284};285 286struct CUFFreeOpConversion : public mlir::OpRewritePattern<cuf::FreeOp> {287 using OpRewritePattern::OpRewritePattern;288 289 mlir::LogicalResult290 matchAndRewrite(cuf::FreeOp op,291 mlir::PatternRewriter &rewriter) const override {292 if (inDeviceContext(op.getOperation())) {293 rewriter.eraseOp(op);294 return mlir::success();295 }296 297 if (!mlir::isa<fir::ReferenceType>(op.getDevptr().getType()))298 return failure();299 300 auto mod = op->getParentOfType<mlir::ModuleOp>();301 fir::FirOpBuilder builder(rewriter, mod);302 mlir::Location loc = op.getLoc();303 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);304 305 auto refTy = mlir::dyn_cast<fir::ReferenceType>(op.getDevptr().getType());306 if (!mlir::isa<fir::BaseBoxType>(refTy.getEleTy())) {307 mlir::func::FuncOp func =308 fir::runtime::getRuntimeFunc<mkRTKey(CUFMemFree)>(loc, builder);309 auto fTy = func.getFunctionType();310 mlir::Value sourceLine =311 fir::factory::locationToLineNo(builder, loc, fTy.getInput(3));312 mlir::Value memTy = builder.createIntegerConstant(313 loc, builder.getI32Type(), getMemType(op.getDataAttr()));314 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(315 builder, loc, fTy, op.getDevptr(), memTy, sourceFile, sourceLine)};316 fir::CallOp::create(builder, loc, func, args);317 rewriter.eraseOp(op);318 return mlir::success();319 }320 321 // Convert cuf.free on descriptors.322 mlir::func::FuncOp func =323 fir::runtime::getRuntimeFunc<mkRTKey(CUFFreeDescriptor)>(loc, builder);324 auto fTy = func.getFunctionType();325 mlir::Value sourceLine =326 fir::factory::locationToLineNo(builder, loc, fTy.getInput(2));327 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(328 builder, loc, fTy, op.getDevptr(), sourceFile, sourceLine)};329 auto callOp = fir::CallOp::create(builder, loc, func, args);330 callOp->setAttr(cuf::getDataAttrName(), op.getDataAttrAttr());331 rewriter.eraseOp(op);332 return mlir::success();333 }334};335 336struct CUFAllocateOpConversion337 : public mlir::OpRewritePattern<cuf::AllocateOp> {338 using OpRewritePattern::OpRewritePattern;339 340 mlir::LogicalResult341 matchAndRewrite(cuf::AllocateOp op,342 mlir::PatternRewriter &rewriter) const override {343 auto mod = op->getParentOfType<mlir::ModuleOp>();344 fir::FirOpBuilder builder(rewriter, mod);345 mlir::Location loc = op.getLoc();346 347 bool isPointer = false;348 349 if (auto declareOp =350 mlir::dyn_cast_or_null<fir::DeclareOp>(op.getBox().getDefiningOp()))351 if (declareOp.getFortranAttrs() &&352 bitEnumContainsAny(*declareOp.getFortranAttrs(),353 fir::FortranVariableFlagsEnum::pointer))354 isPointer = true;355 356 if (hasDoubleDescriptors(op)) {357 // Allocation for module variable are done with custom runtime entry point358 // so the descriptors can be synchronized.359 mlir::func::FuncOp func;360 if (op.getSource()) {361 func = isPointer ? fir::runtime::getRuntimeFunc<mkRTKey(362 CUFPointerAllocateSourceSync)>(loc, builder)363 : fir::runtime::getRuntimeFunc<mkRTKey(364 CUFAllocatableAllocateSourceSync)>(loc, builder);365 } else {366 func =367 isPointer368 ? fir::runtime::getRuntimeFunc<mkRTKey(CUFPointerAllocateSync)>(369 loc, builder)370 : fir::runtime::getRuntimeFunc<mkRTKey(371 CUFAllocatableAllocateSync)>(loc, builder);372 }373 return convertOpToCall<cuf::AllocateOp>(op, rewriter, func);374 }375 376 mlir::func::FuncOp func;377 if (op.getSource()) {378 func =379 isPointer380 ? fir::runtime::getRuntimeFunc<mkRTKey(CUFPointerAllocateSource)>(381 loc, builder)382 : fir::runtime::getRuntimeFunc<mkRTKey(383 CUFAllocatableAllocateSource)>(loc, builder);384 } else {385 func =386 isPointer387 ? fir::runtime::getRuntimeFunc<mkRTKey(CUFPointerAllocate)>(388 loc, builder)389 : fir::runtime::getRuntimeFunc<mkRTKey(CUFAllocatableAllocate)>(390 loc, builder);391 }392 393 return convertOpToCall<cuf::AllocateOp>(op, rewriter, func);394 }395};396 397struct CUFDeallocateOpConversion398 : public mlir::OpRewritePattern<cuf::DeallocateOp> {399 using OpRewritePattern::OpRewritePattern;400 401 mlir::LogicalResult402 matchAndRewrite(cuf::DeallocateOp op,403 mlir::PatternRewriter &rewriter) const override {404 405 auto mod = op->getParentOfType<mlir::ModuleOp>();406 fir::FirOpBuilder builder(rewriter, mod);407 mlir::Location loc = op.getLoc();408 409 if (hasDoubleDescriptors(op)) {410 // Deallocation for module variable are done with custom runtime entry411 // point so the descriptors can be synchronized.412 mlir::func::FuncOp func =413 fir::runtime::getRuntimeFunc<mkRTKey(CUFAllocatableDeallocate)>(414 loc, builder);415 return convertOpToCall<cuf::DeallocateOp>(op, rewriter, func);416 }417 418 // Deallocation for local descriptor falls back on the standard runtime419 // AllocatableDeallocate as the dedicated deallocator is set in the420 // descriptor before the call.421 mlir::func::FuncOp func =422 fir::runtime::getRuntimeFunc<mkRTKey(AllocatableDeallocate)>(loc,423 builder);424 return convertOpToCall<cuf::DeallocateOp>(op, rewriter, func);425 }426};427 428class CUFAllocationConversion429 : public fir::impl::CUFAllocationConversionBase<CUFAllocationConversion> {430public:431 void runOnOperation() override {432 auto *ctx = &getContext();433 mlir::RewritePatternSet patterns(ctx);434 mlir::ConversionTarget target(*ctx);435 436 mlir::Operation *op = getOperation();437 mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(op);438 if (!module)439 return signalPassFailure();440 mlir::SymbolTable symtab(module);441 442 std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(443 module, /*allowDefaultLayout=*/false);444 fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,445 /*forceUnifiedTBAATree=*/false, *dl);446 target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,447 mlir::gpu::GPUDialect>();448 target.addLegalOp<cuf::StreamCastOp>();449 cuf::populateCUFAllocationConversionPatterns(typeConverter, *dl, symtab,450 patterns);451 if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,452 std::move(patterns)))) {453 mlir::emitError(mlir::UnknownLoc::get(ctx),454 "error in CUF allocation conversion\n");455 signalPassFailure();456 }457 }458};459 460} // namespace461 462void cuf::populateCUFAllocationConversionPatterns(463 const fir::LLVMTypeConverter &converter, mlir::DataLayout &dl,464 const mlir::SymbolTable &symtab, mlir::RewritePatternSet &patterns) {465 patterns.insert<CUFAllocOpConversion>(patterns.getContext(), &dl, &converter);466 patterns.insert<CUFFreeOpConversion, CUFAllocateOpConversion,467 CUFDeallocateOpConversion>(patterns.getContext());468}469