177 lines · cpp
1//===-- CUFAddConstructor.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/BoxValue.h"10#include "flang/Optimizer/Builder/CUFCommon.h"11#include "flang/Optimizer/Builder/FIRBuilder.h"12#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"13#include "flang/Optimizer/Builder/Todo.h"14#include "flang/Optimizer/CodeGen/Target.h"15#include "flang/Optimizer/CodeGen/TypeConverter.h"16#include "flang/Optimizer/Dialect/CUF/CUFOps.h"17#include "flang/Optimizer/Dialect/FIRAttr.h"18#include "flang/Optimizer/Dialect/FIRDialect.h"19#include "flang/Optimizer/Dialect/FIROps.h"20#include "flang/Optimizer/Dialect/FIROpsSupport.h"21#include "flang/Optimizer/Dialect/FIRType.h"22#include "flang/Optimizer/Support/DataLayout.h"23#include "flang/Runtime/CUDA/registration.h"24#include "flang/Runtime/entry-names.h"25#include "mlir/Dialect/DLTI/DLTI.h"26#include "mlir/Dialect/GPU/IR/GPUDialect.h"27#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"28#include "mlir/Dialect/LLVMIR/LLVMDialect.h"29#include "mlir/IR/Value.h"30#include "mlir/Pass/Pass.h"31#include "llvm/ADT/SmallVector.h"32 33namespace fir {34#define GEN_PASS_DEF_CUFADDCONSTRUCTOR35#include "flang/Optimizer/Transforms/Passes.h.inc"36} // namespace fir37 38using namespace Fortran::runtime::cuda;39 40namespace {41 42static constexpr llvm::StringRef cudaFortranCtorName{43 "__cudaFortranConstructor"};44 45struct CUFAddConstructor46 : public fir::impl::CUFAddConstructorBase<CUFAddConstructor> {47 48 void runOnOperation() override {49 mlir::ModuleOp mod = getOperation();50 mlir::SymbolTable symTab(mod);51 mlir::OpBuilder opBuilder{mod.getBodyRegion()};52 fir::FirOpBuilder builder(opBuilder, mod);53 fir::KindMapping kindMap{fir::getKindMapping(mod)};54 builder.setInsertionPointToEnd(mod.getBody());55 mlir::Location loc = mod.getLoc();56 auto *ctx = mod.getContext();57 auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);58 auto idxTy = builder.getIndexType();59 auto funcTy =60 mlir::LLVM::LLVMFunctionType::get(voidTy, {}, /*isVarArg=*/false);61 std::optional<mlir::DataLayout> dl =62 fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/false);63 if (!dl) {64 mlir::emitError(mod.getLoc(),65 "data layout attribute is required to perform " +66 getName() + "pass");67 }68 69 // Symbol reference to CUFRegisterAllocator.70 builder.setInsertionPointToEnd(mod.getBody());71 auto registerFuncOp = mlir::LLVM::LLVMFuncOp::create(72 builder, loc, RTNAME_STRING(CUFRegisterAllocator), funcTy);73 registerFuncOp.setVisibility(mlir::SymbolTable::Visibility::Private);74 auto cufRegisterAllocatorRef = mlir::SymbolRefAttr::get(75 mod.getContext(), RTNAME_STRING(CUFRegisterAllocator));76 builder.setInsertionPointToEnd(mod.getBody());77 78 // Create the constructor function that call CUFRegisterAllocator.79 auto func = mlir::LLVM::LLVMFuncOp::create(builder, loc,80 cudaFortranCtorName, funcTy);81 func.setLinkage(mlir::LLVM::Linkage::Internal);82 builder.setInsertionPointToStart(func.addEntryBlock(builder));83 mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);84 85 auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName);86 if (gpuMod) {87 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(ctx);88 auto registeredMod = cuf::RegisterModuleOp::create(89 builder, loc, llvmPtrTy,90 mlir::SymbolRefAttr::get(ctx, gpuMod.getName()));91 92 fir::LLVMTypeConverter typeConverter(mod, /*applyTBAA=*/false,93 /*forceUnifiedTBAATree=*/false, *dl);94 // Register kernels95 for (auto func : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) {96 if (func.isKernel()) {97 auto kernelName = mlir::SymbolRefAttr::get(98 builder.getStringAttr(cudaDeviceModuleName),99 {mlir::SymbolRefAttr::get(builder.getContext(), func.getName())});100 cuf::RegisterKernelOp::create(builder, loc, kernelName,101 registeredMod);102 }103 }104 105 // Register variables106 for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {107 auto attr = globalOp.getDataAttrAttr();108 if (!attr)109 continue;110 111 if (attr.getValue() == cuf::DataAttribute::Managed &&112 !mlir::isa<fir::BaseBoxType>(globalOp.getType()))113 TODO(loc, "registration of non-allocatable managed variables");114 115 mlir::func::FuncOp func;116 switch (attr.getValue()) {117 case cuf::DataAttribute::Device:118 case cuf::DataAttribute::Constant:119 case cuf::DataAttribute::Managed: {120 func = fir::runtime::getRuntimeFunc<mkRTKey(CUFRegisterVariable)>(121 loc, builder);122 auto fTy = func.getFunctionType();123 124 // Global variable name125 std::string gblNameStr = globalOp.getSymbol().getValue().str();126 gblNameStr += '\0';127 mlir::Value gblName = fir::getBase(128 fir::factory::createStringLiteral(builder, loc, gblNameStr));129 130 // Global variable size131 std::optional<uint64_t> size;132 if (auto boxTy =133 mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {134 mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);135 size = dl->getTypeSizeInBits(structTy) / 8;136 }137 if (!size) {138 size = fir::getTypeSizeAndAlignmentOrCrash(loc, globalOp.getType(),139 *dl, kindMap)140 .first;141 }142 auto sizeVal = builder.createIntegerConstant(loc, idxTy, *size);143 144 // Global variable address145 mlir::Value addr = fir::AddrOfOp::create(146 builder, loc, globalOp.resultType(), globalOp.getSymbol());147 148 llvm::SmallVector<mlir::Value> args{fir::runtime::createArguments(149 builder, loc, fTy, registeredMod, addr, gblName, sizeVal)};150 fir::CallOp::create(builder, loc, func, args);151 } break;152 default:153 break;154 }155 }156 }157 mlir::LLVM::ReturnOp::create(builder, loc, mlir::ValueRange{});158 159 // Create the llvm.global_ctor with the function.160 // TODO: We might want to have a utility that retrieve it if already161 // created and adds new functions.162 builder.setInsertionPointToEnd(mod.getBody());163 llvm::SmallVector<mlir::Attribute> funcs;164 funcs.push_back(165 mlir::FlatSymbolRefAttr::get(mod.getContext(), func.getSymName()));166 llvm::SmallVector<int> priorities;167 llvm::SmallVector<mlir::Attribute> data;168 priorities.push_back(0);169 data.push_back(mlir::LLVM::ZeroAttr::get(mod.getContext()));170 mlir::LLVM::GlobalCtorsOp::create(171 builder, mod.getLoc(), builder.getArrayAttr(funcs),172 builder.getI32ArrayAttr(priorities), builder.getArrayAttr(data));173 }174};175 176} // end anonymous namespace177