614 lines · cpp
1//===-- CUFDeviceGlobal.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/CUFOpConversion.h"10#include "flang/Optimizer/Builder/CUFCommon.h"11#include "flang/Optimizer/Builder/Runtime/CUDA/Descriptor.h"12#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"13#include "flang/Optimizer/CodeGen/TypeConverter.h"14#include "flang/Optimizer/Dialect/CUF/CUFOps.h"15#include "flang/Optimizer/Dialect/FIRDialect.h"16#include "flang/Optimizer/Dialect/FIROps.h"17#include "flang/Optimizer/HLFIR/HLFIROps.h"18#include "flang/Optimizer/Support/DataLayout.h"19#include "flang/Optimizer/Transforms/CUDA/CUFAllocationConversion.h"20#include "flang/Optimizer/Transforms/Passes.h"21#include "flang/Runtime/CUDA/allocatable.h"22#include "flang/Runtime/CUDA/common.h"23#include "flang/Runtime/CUDA/descriptor.h"24#include "flang/Runtime/CUDA/memory.h"25#include "flang/Runtime/CUDA/pointer.h"26#include "flang/Runtime/allocatable.h"27#include "flang/Runtime/allocator-registry-consts.h"28#include "flang/Support/Fortran.h"29#include "mlir/Conversion/LLVMCommon/Pattern.h"30#include "mlir/Dialect/DLTI/DLTI.h"31#include "mlir/Dialect/GPU/IR/GPUDialect.h"32#include "mlir/IR/Matchers.h"33#include "mlir/Pass/Pass.h"34#include "mlir/Transforms/DialectConversion.h"35#include "mlir/Transforms/GreedyPatternRewriteDriver.h"36 37namespace fir {38#define GEN_PASS_DEF_CUFOPCONVERSION39#include "flang/Optimizer/Transforms/Passes.h.inc"40} // namespace fir41 42using namespace fir;43using namespace mlir;44using namespace Fortran::runtime;45using namespace Fortran::runtime::cuda;46 47namespace {48 49static bool inDeviceContext(mlir::Operation *op) {50 if (op->getParentOfType<cuf::KernelOp>())51 return true;52 if (auto funcOp = op->getParentOfType<mlir::gpu::GPUFuncOp>())53 return true;54 if (auto funcOp = op->getParentOfType<mlir::gpu::LaunchOp>())55 return true;56 if (auto funcOp = op->getParentOfType<mlir::func::FuncOp>()) {57 if (auto cudaProcAttr =58 funcOp.getOperation()->getAttrOfType<cuf::ProcAttributeAttr>(59 cuf::getProcAttrName())) {60 return cudaProcAttr.getValue() != cuf::ProcAttribute::Host &&61 cudaProcAttr.getValue() != cuf::ProcAttribute::HostDevice;62 }63 }64 return false;65}66 67static mlir::Value createConvertOp(mlir::PatternRewriter &rewriter,68 mlir::Location loc, mlir::Type toTy,69 mlir::Value val) {70 if (val.getType() != toTy)71 return fir::ConvertOp::create(rewriter, loc, toTy, val);72 return val;73}74 75struct CUFDeviceAddressOpConversion76 : public mlir::OpRewritePattern<cuf::DeviceAddressOp> {77 using OpRewritePattern::OpRewritePattern;78 79 CUFDeviceAddressOpConversion(mlir::MLIRContext *context,80 const mlir::SymbolTable &symtab)81 : OpRewritePattern(context), symTab{symtab} {}82 83 mlir::LogicalResult84 matchAndRewrite(cuf::DeviceAddressOp op,85 mlir::PatternRewriter &rewriter) const override {86 if (auto global = symTab.lookup<fir::GlobalOp>(87 op.getHostSymbol().getRootReference().getValue())) {88 auto mod = op->getParentOfType<mlir::ModuleOp>();89 mlir::Location loc = op.getLoc();90 auto hostAddr = fir::AddrOfOp::create(91 rewriter, loc, fir::ReferenceType::get(global.getType()),92 op.getHostSymbol());93 fir::FirOpBuilder builder(rewriter, mod);94 mlir::func::FuncOp callee =95 fir::runtime::getRuntimeFunc<mkRTKey(CUFGetDeviceAddress)>(loc,96 builder);97 auto fTy = callee.getFunctionType();98 mlir::Value conv =99 createConvertOp(rewriter, loc, fTy.getInput(0), hostAddr);100 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);101 mlir::Value sourceLine =102 fir::factory::locationToLineNo(builder, loc, fTy.getInput(2));103 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(104 builder, loc, fTy, conv, sourceFile, sourceLine)};105 auto call = fir::CallOp::create(rewriter, loc, callee, args);106 mlir::Value addr = createConvertOp(rewriter, loc, hostAddr.getType(),107 call->getResult(0));108 rewriter.replaceOp(op, addr.getDefiningOp());109 return success();110 }111 return failure();112 }113 114private:115 const mlir::SymbolTable &symTab;116};117 118struct DeclareOpConversion : public mlir::OpRewritePattern<fir::DeclareOp> {119 using OpRewritePattern::OpRewritePattern;120 121 DeclareOpConversion(mlir::MLIRContext *context,122 const mlir::SymbolTable &symtab)123 : OpRewritePattern(context), symTab{symtab} {}124 125 mlir::LogicalResult126 matchAndRewrite(fir::DeclareOp op,127 mlir::PatternRewriter &rewriter) const override {128 if (op.getResult().getUsers().empty())129 return success();130 if (auto addrOfOp = op.getMemref().getDefiningOp<fir::AddrOfOp>()) {131 if (auto global = symTab.lookup<fir::GlobalOp>(132 addrOfOp.getSymbol().getRootReference().getValue())) {133 if (cuf::isRegisteredDeviceGlobal(global)) {134 rewriter.setInsertionPointAfter(addrOfOp);135 mlir::Value devAddr = cuf::DeviceAddressOp::create(136 rewriter, op.getLoc(), addrOfOp.getType(), addrOfOp.getSymbol());137 rewriter.startOpModification(op);138 op.getMemrefMutable().assign(devAddr);139 rewriter.finalizeOpModification(op);140 return success();141 }142 }143 }144 return failure();145 }146 147private:148 const mlir::SymbolTable &symTab;149};150 151static bool isDstGlobal(cuf::DataTransferOp op) {152 if (auto declareOp = op.getDst().getDefiningOp<fir::DeclareOp>())153 if (declareOp.getMemref().getDefiningOp<fir::AddrOfOp>())154 return true;155 if (auto declareOp = op.getDst().getDefiningOp<hlfir::DeclareOp>())156 if (declareOp.getMemref().getDefiningOp<fir::AddrOfOp>())157 return true;158 return false;159}160 161static mlir::Value getShapeFromDecl(mlir::Value src) {162 if (auto declareOp = src.getDefiningOp<fir::DeclareOp>())163 return declareOp.getShape();164 if (auto declareOp = src.getDefiningOp<hlfir::DeclareOp>())165 return declareOp.getShape();166 return mlir::Value{};167}168 169static mlir::Value emboxSrc(mlir::PatternRewriter &rewriter,170 cuf::DataTransferOp op,171 const mlir::SymbolTable &symtab,172 mlir::Type dstEleTy = nullptr) {173 auto mod = op->getParentOfType<mlir::ModuleOp>();174 mlir::Location loc = op.getLoc();175 fir::FirOpBuilder builder(rewriter, mod);176 mlir::Value addr;177 mlir::Type srcTy = fir::unwrapRefType(op.getSrc().getType());178 if (fir::isa_trivial(srcTy) &&179 mlir::matchPattern(op.getSrc().getDefiningOp(), mlir::m_Constant())) {180 mlir::Value src = op.getSrc();181 if (srcTy.isInteger(1)) {182 // i1 is not a supported type in the descriptor and it is actually coming183 // from a LOGICAL constant. Use the destination type to avoid mismatch.184 assert(dstEleTy && "expect dst element type to be set");185 srcTy = dstEleTy;186 src = createConvertOp(rewriter, loc, srcTy, src);187 addr = builder.createTemporary(loc, srcTy);188 fir::StoreOp::create(builder, loc, src, addr);189 } else {190 if (dstEleTy && fir::isa_trivial(dstEleTy) && srcTy != dstEleTy) {191 // Use dstEleTy and convert to avoid assign mismatch.192 addr = builder.createTemporary(loc, dstEleTy);193 auto conv = fir::ConvertOp::create(builder, loc, dstEleTy, src);194 fir::StoreOp::create(builder, loc, conv, addr);195 srcTy = dstEleTy;196 } else {197 // Put constant in memory if it is not.198 addr = builder.createTemporary(loc, srcTy);199 fir::StoreOp::create(builder, loc, src, addr);200 }201 }202 } else {203 addr = op.getSrc();204 }205 llvm::SmallVector<mlir::Value> lenParams;206 mlir::Type boxTy = fir::BoxType::get(srcTy);207 mlir::Value box =208 builder.createBox(loc, boxTy, addr, getShapeFromDecl(op.getSrc()),209 /*slice=*/nullptr, lenParams,210 /*tdesc=*/nullptr);211 mlir::Value src = builder.createTemporary(loc, box.getType());212 fir::StoreOp::create(builder, loc, box, src);213 return src;214}215 216static mlir::Value emboxDst(mlir::PatternRewriter &rewriter,217 cuf::DataTransferOp op,218 const mlir::SymbolTable &symtab) {219 auto mod = op->getParentOfType<mlir::ModuleOp>();220 mlir::Location loc = op.getLoc();221 fir::FirOpBuilder builder(rewriter, mod);222 mlir::Type dstTy = fir::unwrapRefType(op.getDst().getType());223 mlir::Value dstAddr = op.getDst();224 mlir::Type dstBoxTy = fir::BoxType::get(dstTy);225 llvm::SmallVector<mlir::Value> lenParams;226 mlir::Value dstBox =227 builder.createBox(loc, dstBoxTy, dstAddr, getShapeFromDecl(op.getDst()),228 /*slice=*/nullptr, lenParams,229 /*tdesc=*/nullptr);230 mlir::Value dst = builder.createTemporary(loc, dstBox.getType());231 fir::StoreOp::create(builder, loc, dstBox, dst);232 return dst;233}234 235struct CUFDataTransferOpConversion236 : public mlir::OpRewritePattern<cuf::DataTransferOp> {237 using OpRewritePattern::OpRewritePattern;238 239 CUFDataTransferOpConversion(mlir::MLIRContext *context,240 const mlir::SymbolTable &symtab,241 mlir::DataLayout *dl,242 const fir::LLVMTypeConverter *typeConverter)243 : OpRewritePattern(context), symtab{symtab}, dl{dl},244 typeConverter{typeConverter} {}245 246 mlir::LogicalResult247 matchAndRewrite(cuf::DataTransferOp op,248 mlir::PatternRewriter &rewriter) const override {249 250 mlir::Type srcTy = fir::unwrapRefType(op.getSrc().getType());251 mlir::Type dstTy = fir::unwrapRefType(op.getDst().getType());252 253 mlir::Location loc = op.getLoc();254 unsigned mode = 0;255 if (op.getTransferKind() == cuf::DataTransferKind::HostDevice) {256 mode = kHostToDevice;257 } else if (op.getTransferKind() == cuf::DataTransferKind::DeviceHost) {258 mode = kDeviceToHost;259 } else if (op.getTransferKind() == cuf::DataTransferKind::DeviceDevice) {260 mode = kDeviceToDevice;261 } else {262 mlir::emitError(loc, "unsupported transfer kind\n");263 }264 265 auto mod = op->getParentOfType<mlir::ModuleOp>();266 fir::FirOpBuilder builder(rewriter, mod);267 fir::KindMapping kindMap{fir::getKindMapping(mod)};268 mlir::Value modeValue =269 builder.createIntegerConstant(loc, builder.getI32Type(), mode);270 271 // Convert data transfer without any descriptor.272 if (!mlir::isa<fir::BaseBoxType>(srcTy) &&273 !mlir::isa<fir::BaseBoxType>(dstTy)) {274 275 if (fir::isa_trivial(srcTy) && !fir::isa_trivial(dstTy)) {276 // Initialization of an array from a scalar value should be implemented277 // via a kernel launch. Use the flang runtime via the Assign function278 // until we have more infrastructure.279 mlir::Type dstEleTy = fir::unwrapInnerType(fir::unwrapRefType(dstTy));280 mlir::Value src = emboxSrc(rewriter, op, symtab, dstEleTy);281 mlir::Value dst = emboxDst(rewriter, op, symtab);282 mlir::func::FuncOp func =283 fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferCstDesc)>(284 loc, builder);285 auto fTy = func.getFunctionType();286 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);287 mlir::Value sourceLine =288 fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));289 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(290 builder, loc, fTy, dst, src, modeValue, sourceFile, sourceLine)};291 fir::CallOp::create(builder, loc, func, args);292 rewriter.eraseOp(op);293 return mlir::success();294 }295 296 mlir::Type i64Ty = builder.getI64Type();297 mlir::Value nbElement =298 cuf::computeElementCount(rewriter, loc, op.getShape(), dstTy, i64Ty);299 unsigned width = 0;300 if (fir::isa_derived(fir::unwrapSequenceType(dstTy))) {301 mlir::Type structTy =302 typeConverter->convertType(fir::unwrapSequenceType(dstTy));303 width = dl->getTypeSizeInBits(structTy) / 8;304 } else {305 width = cuf::computeElementByteSize(loc, dstTy, kindMap);306 }307 mlir::Value widthValue = mlir::arith::ConstantOp::create(308 rewriter, loc, i64Ty, rewriter.getIntegerAttr(i64Ty, width));309 mlir::Value bytes = nbElement ? mlir::arith::MulIOp::create(310 rewriter, loc, nbElement, widthValue)311 : widthValue;312 313 mlir::func::FuncOp func =314 fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferPtrPtr)>(loc,315 builder);316 auto fTy = func.getFunctionType();317 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);318 mlir::Value sourceLine =319 fir::factory::locationToLineNo(builder, loc, fTy.getInput(5));320 321 mlir::Value dst = op.getDst();322 mlir::Value src = op.getSrc();323 // Materialize the src if constant.324 if (matchPattern(src.getDefiningOp(), mlir::m_Constant())) {325 mlir::Value temp = builder.createTemporary(loc, srcTy);326 fir::StoreOp::create(builder, loc, src, temp);327 src = temp;328 }329 llvm::SmallVector<mlir::Value> args{330 fir::runtime::createArguments(builder, loc, fTy, dst, src, bytes,331 modeValue, sourceFile, sourceLine)};332 fir::CallOp::create(builder, loc, func, args);333 rewriter.eraseOp(op);334 return mlir::success();335 }336 337 auto materializeBoxIfNeeded = [&](mlir::Value val) -> mlir::Value {338 if (mlir::isa<fir::EmboxOp, fir::ReboxOp>(val.getDefiningOp())) {339 // Materialize the box to memory to be able to call the runtime.340 mlir::Value box = builder.createTemporary(loc, val.getType());341 fir::StoreOp::create(builder, loc, val, box);342 return box;343 }344 if (mlir::isa<fir::BaseBoxType>(val.getType()))345 if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(val.getDefiningOp()))346 return loadOp.getMemref();347 return val;348 };349 350 // Conversion of data transfer involving at least one descriptor.351 if (auto dstBoxTy = mlir::dyn_cast<fir::BaseBoxType>(dstTy)) {352 // Transfer to a descriptor.353 mlir::func::FuncOp func =354 isDstGlobal(op)355 ? fir::runtime::getRuntimeFunc<mkRTKey(356 CUFDataTransferGlobalDescDesc)>(loc, builder)357 : fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferDescDesc)>(358 loc, builder);359 mlir::Value dst = op.getDst();360 mlir::Value src = op.getSrc();361 if (!mlir::isa<fir::BaseBoxType>(srcTy)) {362 mlir::Type dstEleTy = fir::unwrapInnerType(dstBoxTy.getEleTy());363 src = emboxSrc(rewriter, op, symtab, dstEleTy);364 if (fir::isa_trivial(srcTy))365 func = fir::runtime::getRuntimeFunc<mkRTKey(CUFDataTransferCstDesc)>(366 loc, builder);367 }368 369 src = materializeBoxIfNeeded(src);370 dst = materializeBoxIfNeeded(dst);371 372 auto fTy = func.getFunctionType();373 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);374 mlir::Value sourceLine =375 fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));376 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(377 builder, loc, fTy, dst, src, modeValue, sourceFile, sourceLine)};378 fir::CallOp::create(builder, loc, func, args);379 rewriter.eraseOp(op);380 } else {381 // Transfer from a descriptor.382 mlir::Value dst = emboxDst(rewriter, op, symtab);383 mlir::Value src = materializeBoxIfNeeded(op.getSrc());384 385 mlir::func::FuncOp func = fir::runtime::getRuntimeFunc<mkRTKey(386 CUFDataTransferDescDescNoRealloc)>(loc, builder);387 388 auto fTy = func.getFunctionType();389 mlir::Value sourceFile = fir::factory::locationToFilename(builder, loc);390 mlir::Value sourceLine =391 fir::factory::locationToLineNo(builder, loc, fTy.getInput(4));392 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(393 builder, loc, fTy, dst, src, modeValue, sourceFile, sourceLine)};394 fir::CallOp::create(builder, loc, func, args);395 rewriter.eraseOp(op);396 }397 return mlir::success();398 }399 400private:401 const mlir::SymbolTable &symtab;402 mlir::DataLayout *dl;403 const fir::LLVMTypeConverter *typeConverter;404};405 406struct CUFLaunchOpConversion407 : public mlir::OpRewritePattern<cuf::KernelLaunchOp> {408public:409 using OpRewritePattern::OpRewritePattern;410 411 CUFLaunchOpConversion(mlir::MLIRContext *context,412 const mlir::SymbolTable &symTab)413 : OpRewritePattern(context), symTab{symTab} {}414 415 mlir::LogicalResult416 matchAndRewrite(cuf::KernelLaunchOp op,417 mlir::PatternRewriter &rewriter) const override {418 mlir::Location loc = op.getLoc();419 auto idxTy = mlir::IndexType::get(op.getContext());420 mlir::Value zero = mlir::arith::ConstantOp::create(421 rewriter, loc, rewriter.getIntegerType(32),422 rewriter.getI32IntegerAttr(0));423 auto gridSizeX =424 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getGridX());425 auto gridSizeY =426 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getGridY());427 auto gridSizeZ =428 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getGridZ());429 auto blockSizeX =430 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getBlockX());431 auto blockSizeY =432 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getBlockY());433 auto blockSizeZ =434 mlir::arith::IndexCastOp::create(rewriter, loc, idxTy, op.getBlockZ());435 auto kernelName = mlir::SymbolRefAttr::get(436 rewriter.getStringAttr(cudaDeviceModuleName),437 {mlir::SymbolRefAttr::get(438 rewriter.getContext(),439 op.getCallee().getLeafReference().getValue())});440 mlir::Value clusterDimX, clusterDimY, clusterDimZ;441 cuf::ProcAttributeAttr procAttr;442 if (auto funcOp = symTab.lookup<mlir::func::FuncOp>(443 op.getCallee().getLeafReference())) {444 if (auto clusterDimsAttr = funcOp->getAttrOfType<cuf::ClusterDimsAttr>(445 cuf::getClusterDimsAttrName())) {446 clusterDimX = mlir::arith::ConstantIndexOp::create(447 rewriter, loc, clusterDimsAttr.getX().getInt());448 clusterDimY = mlir::arith::ConstantIndexOp::create(449 rewriter, loc, clusterDimsAttr.getY().getInt());450 clusterDimZ = mlir::arith::ConstantIndexOp::create(451 rewriter, loc, clusterDimsAttr.getZ().getInt());452 }453 procAttr =454 funcOp->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName());455 }456 llvm::SmallVector<mlir::Value> args;457 for (mlir::Value arg : op.getArgs()) {458 // If the argument is a global descriptor, make sure we pass the device459 // copy of this descriptor and not the host one.460 if (mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(arg.getType()))) {461 if (auto declareOp =462 mlir::dyn_cast_or_null<fir::DeclareOp>(arg.getDefiningOp())) {463 if (auto addrOfOp = mlir::dyn_cast_or_null<fir::AddrOfOp>(464 declareOp.getMemref().getDefiningOp())) {465 if (auto global = symTab.lookup<fir::GlobalOp>(466 addrOfOp.getSymbol().getRootReference().getValue())) {467 if (cuf::isRegisteredDeviceGlobal(global)) {468 arg = cuf::DeviceAddressOp::create(rewriter, op.getLoc(),469 addrOfOp.getType(),470 addrOfOp.getSymbol())471 .getResult();472 }473 }474 }475 }476 }477 args.push_back(arg);478 }479 mlir::Value dynamicShmemSize = op.getBytes() ? op.getBytes() : zero;480 auto gpuLaunchOp = mlir::gpu::LaunchFuncOp::create(481 rewriter, loc, kernelName,482 mlir::gpu::KernelDim3{gridSizeX, gridSizeY, gridSizeZ},483 mlir::gpu::KernelDim3{blockSizeX, blockSizeY, blockSizeZ},484 dynamicShmemSize, args);485 if (clusterDimX && clusterDimY && clusterDimZ) {486 gpuLaunchOp.getClusterSizeXMutable().assign(clusterDimX);487 gpuLaunchOp.getClusterSizeYMutable().assign(clusterDimY);488 gpuLaunchOp.getClusterSizeZMutable().assign(clusterDimZ);489 }490 if (op.getStream()) {491 mlir::OpBuilder::InsertionGuard guard(rewriter);492 rewriter.setInsertionPoint(gpuLaunchOp);493 mlir::Value stream =494 cuf::StreamCastOp::create(rewriter, loc, op.getStream());495 gpuLaunchOp.getAsyncDependenciesMutable().append(stream);496 }497 if (procAttr)498 gpuLaunchOp->setAttr(cuf::getProcAttrName(), procAttr);499 else500 // Set default global attribute of the original was not found.501 gpuLaunchOp->setAttr(cuf::getProcAttrName(),502 cuf::ProcAttributeAttr::get(503 op.getContext(), cuf::ProcAttribute::Global));504 rewriter.replaceOp(op, gpuLaunchOp);505 return mlir::success();506 }507 508private:509 const mlir::SymbolTable &symTab;510};511 512struct CUFSyncDescriptorOpConversion513 : public mlir::OpRewritePattern<cuf::SyncDescriptorOp> {514 using OpRewritePattern::OpRewritePattern;515 516 mlir::LogicalResult517 matchAndRewrite(cuf::SyncDescriptorOp op,518 mlir::PatternRewriter &rewriter) const override {519 auto mod = op->getParentOfType<mlir::ModuleOp>();520 fir::FirOpBuilder builder(rewriter, mod);521 mlir::Location loc = op.getLoc();522 523 auto globalOp = mod.lookupSymbol<fir::GlobalOp>(op.getGlobalName());524 if (!globalOp)525 return mlir::failure();526 527 auto hostAddr = fir::AddrOfOp::create(528 builder, loc, fir::ReferenceType::get(globalOp.getType()),529 op.getGlobalName());530 fir::runtime::cuda::genSyncGlobalDescriptor(builder, loc, hostAddr);531 op.erase();532 return mlir::success();533 }534};535 536class CUFOpConversion : public fir::impl::CUFOpConversionBase<CUFOpConversion> {537 using CUFOpConversionBase::CUFOpConversionBase;538 539public:540 void runOnOperation() override {541 auto *ctx = &getContext();542 mlir::RewritePatternSet patterns(ctx);543 mlir::ConversionTarget target(*ctx);544 545 mlir::Operation *op = getOperation();546 mlir::ModuleOp module = mlir::dyn_cast<mlir::ModuleOp>(op);547 if (!module)548 return signalPassFailure();549 mlir::SymbolTable symtab(module);550 551 std::optional<mlir::DataLayout> dl = fir::support::getOrSetMLIRDataLayout(552 module, /*allowDefaultLayout=*/false);553 fir::LLVMTypeConverter typeConverter(module, /*applyTBAA=*/false,554 /*forceUnifiedTBAATree=*/false, *dl);555 target.addLegalDialect<fir::FIROpsDialect, mlir::arith::ArithDialect,556 mlir::gpu::GPUDialect>();557 target.addLegalOp<cuf::StreamCastOp>();558 cuf::populateCUFToFIRConversionPatterns(typeConverter, *dl, symtab,559 patterns);560 if (allocationConversion)561 cuf::populateCUFAllocationConversionPatterns(typeConverter, *dl, symtab,562 patterns);563 if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,564 std::move(patterns)))) {565 mlir::emitError(mlir::UnknownLoc::get(ctx),566 "error in CUF op conversion\n");567 signalPassFailure();568 }569 570 target.addDynamicallyLegalOp<fir::DeclareOp>([&](fir::DeclareOp op) {571 if (op.getResult().getUsers().empty())572 return true;573 if (inDeviceContext(op))574 return true;575 if (auto addrOfOp = op.getMemref().getDefiningOp<fir::AddrOfOp>()) {576 if (auto global = symtab.lookup<fir::GlobalOp>(577 addrOfOp.getSymbol().getRootReference().getValue())) {578 if (mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(global.getType())))579 return true;580 if (cuf::isRegisteredDeviceGlobal(global))581 return false;582 }583 }584 return true;585 });586 587 patterns.clear();588 cuf::populateFIRCUFConversionPatterns(symtab, patterns);589 if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,590 std::move(patterns)))) {591 mlir::emitError(mlir::UnknownLoc::get(ctx),592 "error in CUF op conversion\n");593 signalPassFailure();594 }595 }596};597} // namespace598 599void cuf::populateCUFToFIRConversionPatterns(600 const fir::LLVMTypeConverter &converter, mlir::DataLayout &dl,601 const mlir::SymbolTable &symtab, mlir::RewritePatternSet &patterns) {602 patterns.insert<CUFSyncDescriptorOpConversion>(patterns.getContext());603 patterns.insert<CUFDataTransferOpConversion>(patterns.getContext(), symtab,604 &dl, &converter);605 patterns.insert<CUFLaunchOpConversion, CUFDeviceAddressOpConversion>(606 patterns.getContext(), symtab);607}608 609void cuf::populateFIRCUFConversionPatterns(const mlir::SymbolTable &symtab,610 mlir::RewritePatternSet &patterns) {611 patterns.insert<DeclareOpConversion, CUFDeviceAddressOpConversion>(612 patterns.getContext(), symtab);613}614