155 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/Builder/CUFCommon.h"10#include "flang/Optimizer/Dialect/CUF/CUFOps.h"11#include "flang/Optimizer/Dialect/FIRDialect.h"12#include "flang/Optimizer/Dialect/FIROps.h"13#include "flang/Optimizer/HLFIR/HLFIROps.h"14#include "flang/Optimizer/Support/InternalNames.h"15#include "flang/Runtime/CUDA/common.h"16#include "flang/Runtime/allocatable.h"17#include "flang/Support/Fortran.h"18#include "mlir/Dialect/LLVMIR/NVVMDialect.h"19#include "mlir/IR/SymbolTable.h"20#include "mlir/Pass/Pass.h"21#include "mlir/Transforms/DialectConversion.h"22#include "llvm/ADT/DenseSet.h"23 24namespace fir {25#define GEN_PASS_DEF_CUFDEVICEGLOBAL26#include "flang/Optimizer/Transforms/Passes.h.inc"27} // namespace fir28 29namespace {30 31static void processAddrOfOp(fir::AddrOfOp addrOfOp,32 mlir::SymbolTable &symbolTable,33 llvm::DenseSet<fir::GlobalOp> &candidates,34 bool recurseInGlobal) {35 36 // Check if there is a real use of the global.37 if (addrOfOp.getOperation()->hasOneUse()) {38 mlir::OpOperand &addrUse = *addrOfOp.getOperation()->getUses().begin();39 if (mlir::isa<fir::DeclareOp>(addrUse.getOwner()) &&40 addrUse.getOwner()->use_empty())41 return;42 }43 44 if (auto globalOp = symbolTable.lookup<fir::GlobalOp>(45 addrOfOp.getSymbol().getRootReference().getValue())) {46 // TO DO: limit candidates to non-scalars. Scalars appear to have been47 // folded in already.48 if (recurseInGlobal)49 globalOp.walk([&](fir::AddrOfOp op) {50 processAddrOfOp(op, symbolTable, candidates, recurseInGlobal);51 });52 candidates.insert(globalOp);53 }54}55 56static void processTypeDescriptor(fir::RecordType recTy,57 mlir::SymbolTable &symbolTable,58 llvm::DenseSet<fir::GlobalOp> &candidates) {59 if (auto globalOp = symbolTable.lookup<fir::GlobalOp>(60 fir::NameUniquer::getTypeDescriptorName(recTy.getName()))) {61 if (!candidates.contains(globalOp)) {62 globalOp.walk([&](fir::AddrOfOp op) {63 processAddrOfOp(op, symbolTable, candidates,64 /*recurseInGlobal=*/true);65 });66 candidates.insert(globalOp);67 }68 }69}70 71static void processEmboxOp(fir::EmboxOp emboxOp, mlir::SymbolTable &symbolTable,72 llvm::DenseSet<fir::GlobalOp> &candidates) {73 if (auto recTy = mlir::dyn_cast<fir::RecordType>(74 fir::unwrapRefType(emboxOp.getMemref().getType())))75 processTypeDescriptor(recTy, symbolTable, candidates);76}77 78static void79prepareImplicitDeviceGlobals(mlir::func::FuncOp funcOp,80 mlir::SymbolTable &symbolTable,81 llvm::DenseSet<fir::GlobalOp> &candidates) {82 auto cudaProcAttr{83 funcOp->getAttrOfType<cuf::ProcAttributeAttr>(cuf::getProcAttrName())};84 if (cudaProcAttr && cudaProcAttr.getValue() != cuf::ProcAttribute::Host) {85 funcOp.walk([&](fir::AddrOfOp op) {86 processAddrOfOp(op, symbolTable, candidates, /*recurseInGlobal=*/false);87 });88 funcOp.walk(89 [&](fir::EmboxOp op) { processEmboxOp(op, symbolTable, candidates); });90 }91}92 93static void94processPotentialTypeDescriptor(mlir::Type candidateType,95 mlir::SymbolTable &symbolTable,96 llvm::DenseSet<fir::GlobalOp> &candidates) {97 if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(candidateType))98 candidateType = boxTy.getEleTy();99 candidateType = fir::unwrapSequenceType(fir::unwrapRefType(candidateType));100 if (auto recTy = mlir::dyn_cast<fir::RecordType>(candidateType))101 processTypeDescriptor(recTy, symbolTable, candidates);102}103 104class CUFDeviceGlobal : public fir::impl::CUFDeviceGlobalBase<CUFDeviceGlobal> {105public:106 void runOnOperation() override {107 mlir::Operation *op = getOperation();108 mlir::ModuleOp mod = mlir::dyn_cast<mlir::ModuleOp>(op);109 if (!mod)110 return signalPassFailure();111 112 llvm::DenseSet<fir::GlobalOp> candidates;113 mlir::SymbolTable symTable(mod);114 mod.walk([&](mlir::func::FuncOp funcOp) {115 prepareImplicitDeviceGlobals(funcOp, symTable, candidates);116 return mlir::WalkResult::advance();117 });118 mod.walk([&](cuf::KernelOp kernelOp) {119 kernelOp.walk([&](fir::AddrOfOp addrOfOp) {120 processAddrOfOp(addrOfOp, symTable, candidates,121 /*recurseInGlobal=*/false);122 });123 });124 125 // Copying the device global variable into the gpu module126 mlir::SymbolTable parentSymTable(mod);127 auto gpuMod = cuf::getOrCreateGPUModule(mod, parentSymTable);128 if (!gpuMod)129 return signalPassFailure();130 mlir::SymbolTable gpuSymTable(gpuMod);131 for (auto globalOp : mod.getOps<fir::GlobalOp>()) {132 if (cuf::isRegisteredDeviceGlobal(globalOp)) {133 candidates.insert(globalOp);134 processPotentialTypeDescriptor(globalOp.getType(), parentSymTable,135 candidates);136 } else if (globalOp.getConstant() &&137 mlir::isa<fir::SequenceType>(138 fir::unwrapRefType(globalOp.resultType()))) {139 mlir::Attribute initAttr =140 globalOp.getInitVal().value_or(mlir::Attribute());141 if (initAttr && mlir::dyn_cast<mlir::DenseElementsAttr>(initAttr))142 candidates.insert(globalOp);143 }144 }145 for (auto globalOp : candidates) {146 auto globalName{globalOp.getSymbol().getValue()};147 if (gpuSymTable.lookup<fir::GlobalOp>(globalName)) {148 break;149 }150 gpuSymTable.insert(globalOp->clone());151 }152 }153};154} // namespace155