flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp - llvm-project - Git at Google

 //===-- CUFAddConstructor.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/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/CUFCommon.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
 #include "flang/Optimizer/CodeGen/Target.h"
 #include "flang/Optimizer/CodeGen/TypeConverter.h"
 #include "flang/Optimizer/Dialect/CUF/CUFOps.h"
 #include "flang/Optimizer/Dialect/FIRAttr.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Support/DataLayout.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "flang/Runtime/CUDA/registration.h"
 #include "flang/Runtime/entry-names.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/SmallVector.h"

 namespace fir {
 #define GEN_PASS_DEF_CUFADDCONSTRUCTOR
 #include "flang/Optimizer/Transforms/Passes.h.inc"
 } // namespace fir

 using namespace Fortran::runtime::cuda;

 namespace {

 static constexpr llvm::StringRef cudaFortranCtorName{
     "__cudaFortranConstructor"};
 static constexpr llvm::StringRef managedPtrSuffix{".managed.ptr"};

 /// Create an 8-byte pointer global in the __nv_managed_data__ section.
 /// The CUDA runtime populates this pointer with the unified memory address
 /// when the module is initialized via __cudaInitModule.
 static fir::GlobalOp createManagedPointerGlobal(fir::FirOpBuilder &builder,
                                                 mlir::ModuleOp mod,
                                                 fir::GlobalOp globalOp) {
   mlir::MLIRContext *ctx = mod.getContext();
   std::string ptrGlobalName = (globalOp.getSymName() + managedPtrSuffix).str();
   auto ptrTy = fir::LLVMPointerType::get(ctx, mlir::IntegerType::get(ctx, 8));

   mlir::OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointAfter(globalOp);

   llvm::SmallVector<mlir::NamedAttribute> attrs;
   attrs.push_back(
       mlir::NamedAttribute(mlir::StringAttr::get(ctx, "section"),
                            mlir::StringAttr::get(ctx, "__nv_managed_data__")));

   mlir::DenseElementsAttr initAttr = {};
   auto ptrGlobal = fir::GlobalOp::create(
       builder, globalOp.getLoc(), ptrGlobalName, /*isConstant=*/false,
       /*isTarget=*/false, ptrTy, initAttr,
       /*linkName=*/builder.createInternalLinkage(), attrs);

   mlir::Region &region = ptrGlobal.getRegion();
   mlir::Block *block = builder.createBlock(&region);
   builder.setInsertionPointToStart(block);
   mlir::Value zero = fir::ZeroOp::create(builder, globalOp.getLoc(), ptrTy);
   fir::HasValueOp::create(builder, globalOp.getLoc(), zero);

   return ptrGlobal;
 }

 /// Return true if \p hostGlobal is a host module-scope global that has been
 /// mirrored in the GPU module as an external (no-body) declaration by the
 /// CUFDeviceGlobal pass under -gpu=mem:unified. Such globals must be
 /// registered with the CUDA driver via CUFRegisterExternalVariable so the
 /// device-side `.extern` symbol resolves to the host pointer at module-load
 /// time and HMM/ATS handles migration.
 static bool isCudaUnifiedExternalGlobal(fir::GlobalOp hostGlobal,
                                         mlir::SymbolTable &gpuSymTable) {
   if (hostGlobal.getDataAttrAttr())
     return false;
   if (hostGlobal.getConstant())
     return false;
   auto gpuGlobal = gpuSymTable.lookup<fir::GlobalOp>(hostGlobal.getSymName());
   if (!gpuGlobal)
     return false;
   return !gpuGlobal.isInitialized();
 }

 /// Build a C-style name literal (`<symname>\0`) for use as the deviceName
 /// argument of a CUF registration runtime call.
 static mlir::Value buildGlobalNameLiteral(fir::FirOpBuilder &builder,
                                           mlir::Location loc,
                                           fir::GlobalOp globalOp) {
   std::string nameStr = globalOp.getSymbol().getValue().str();
   nameStr += '\0';
   return fir::getBase(fir::factory::createStringLiteral(builder, loc, nameStr));
 }

 /// Compute the storage size in bytes of \p globalOp. For a box-typed
 /// allocatable global the size is the descriptor size (after type
 /// conversion); otherwise it's the size of the global's declared type.
 static mlir::Value computeGlobalSize(fir::FirOpBuilder &builder,
                                      mlir::Location loc, mlir::Type idxTy,
                                      const mlir::DataLayout &dl,
                                      const fir::KindMapping &kindMap,
                                      fir::LLVMTypeConverter &typeConverter,
                                      fir::GlobalOp globalOp) {
   std::optional<uint64_t> size;
   if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
     mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);
     size = dl.getTypeSizeInBits(structTy) / 8;
   }
   if (!size) {
     if (auto s =
             fir::getTypeSizeAndAlignment(loc, globalOp.getType(), dl, kindMap))
       size = s->first;
   }
   if (!size) {
     // A global embedding descriptor (allocatable/pointer) components has no
     // structural size; size it via its LLVM type, which inlines the
     // descriptors.
     mlir::Type llvmTy = typeConverter.convertType(globalOp.getType());
     if (llvmTy && mlir::isa<mlir::DataLayoutTypeInterface>(llvmTy))
       size = dl.getTypeSizeInBits(llvmTy) / 8;
   }
   if (!size) {
     size = fir::getTypeSizeAndAlignmentOrCrash(loc, globalOp.getType(), dl,
                                                kindMap)
                .first;
   }
   return builder.createIntegerConstant(loc, idxTy, *size);
 }

 /// Storage size in bytes of \p globalOp as a raw integer (see computeGlobalSize
 /// for the box vs declared-type handling).
 static uint64_t getGlobalSizeInBytes(mlir::Location loc,
                                      const mlir::DataLayout &dl,
                                      const fir::KindMapping &kindMap,
                                      fir::LLVMTypeConverter &typeConverter,
                                      fir::GlobalOp globalOp) {
   std::optional<uint64_t> size;
   if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
     mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);
     size = dl.getTypeSizeInBits(structTy) / 8;
   }
   if (!size) {
     if (auto s =
             fir::getTypeSizeAndAlignment(loc, globalOp.getType(), dl, kindMap))
       size = s->first;
   }
   if (!size) {
     // A global embedding descriptor (allocatable/pointer) components has no
     // structural size; size it via its LLVM type, which inlines the
     // descriptors.
     mlir::Type llvmTy = typeConverter.convertType(globalOp.getType());
     if (llvmTy && mlir::isa<mlir::DataLayoutTypeInterface>(llvmTy))
       size = dl.getTypeSizeInBits(llvmTy) / 8;
   }
   if (!size) {
     size = fir::getTypeSizeAndAlignmentOrCrash(loc, globalOp.getType(), dl,
                                                kindMap)
                .first;
   }
   return *size;
 }

 /// Emit a call to a CUF registration runtime function with the canonical
 /// (module, addr, name, size) signature, where addr is the address of \p
 /// addrGlobal taken via fir.address_of and name/size describe \p nameGlobal.
 /// Used both for CUFRegisterVariable / CUFRegisterManagedVariable / and
 /// CUFRegisterExternalVariable.
 static void
 emitCUFRegistrationCall(fir::FirOpBuilder &builder, mlir::Location loc,
                         mlir::Type idxTy, const mlir::DataLayout &dl,
                         const fir::KindMapping &kindMap,
                         fir::LLVMTypeConverter &typeConverter,
                         mlir::Value registeredMod, mlir::func::FuncOp func,
                         fir::GlobalOp addrGlobal, fir::GlobalOp nameGlobal) {
   mlir::Value gblName = buildGlobalNameLiteral(builder, loc, nameGlobal);
   mlir::Value sizeVal = computeGlobalSize(builder, loc, idxTy, dl, kindMap,
                                           typeConverter, nameGlobal);
   mlir::Value addr = fir::AddrOfOp::create(
       builder, loc, addrGlobal.resultType(), addrGlobal.getSymbol());
   llvm::SmallVector<mlir::Value> args{
       fir::runtime::createArguments(builder, loc, func.getFunctionType(),
                                     registeredMod, addr, gblName, sizeVal)};
   fir::CallOp::create(builder, loc, func, args);
 }

 static bool hasRegisteredGlobals(mlir::ModuleOp mod,
                                  mlir::SymbolTable gpuSymTable,
                                  bool cudaUnified) {
   for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
     auto attr = globalOp.getDataAttrAttr();
     if (!attr) {
       if (cudaUnified && isCudaUnifiedExternalGlobal(globalOp, gpuSymTable))
         return true;
       continue;
     }
     if (!gpuSymTable.lookup(globalOp.getSymName()))
       continue;
     if (attr.getValue() == cuf::DataAttribute::Managed &&
         !mlir::isa<fir::BaseBoxType>(globalOp.getType()))
       return true;
     switch (attr.getValue()) {
     case cuf::DataAttribute::Device:
     case cuf::DataAttribute::Constant:
     case cuf::DataAttribute::Managed: {
       return true;
     } break;
     default:
       break;
     }
   }
   return false;
 }

 static bool hasKernel(mlir::gpu::GPUModuleOp gpuMod) {
   for (auto func : gpuMod.getOps<mlir::gpu::GPUFuncOp>())
     if (func.isKernel())
       return true;
   return false;
 }

 struct CUFAddConstructor
     : public fir::impl::CUFAddConstructorBase<CUFAddConstructor> {

   using CUFAddConstructorBase::CUFAddConstructorBase;

   void runOnOperation() override {
     mlir::ModuleOp mod = getOperation();
     mlir::SymbolTable symTab(mod);
     mlir::OpBuilder opBuilder{mod.getBodyRegion()};
     fir::FirOpBuilder builder(opBuilder, mod);
     fir::KindMapping kindMap{fir::getKindMapping(mod)};
     builder.setInsertionPointToEnd(mod.getBody());
     mlir::Location loc = mod.getLoc();
     auto *ctx = mod.getContext();
     auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);
     auto idxTy = builder.getIndexType();
     auto funcTy =
         mlir::LLVM::LLVMFunctionType::get(voidTy, {}, /*isVarArg=*/false);
     std::optional<mlir::DataLayout> dl =
         fir::support::getOrSetMLIRDataLayout(mod, /*allowDefaultLayout=*/false);
     if (!dl) {
       mlir::emitError(mod.getLoc(),
                       "data layout attribute is required to perform " +
                           getName() + "pass");
     }

     bool needAllocatorRegistration = false;
     mod.walk([&](fir::DeclareOp declOp) {
       if (declOp.getFortranAttrs() &&
           fir::bitEnumContainsAny(*declOp.getFortranAttrs(),
                                   fir::FortranVariableFlagsEnum::allocatable |
                                       fir::FortranVariableFlagsEnum::pointer)) {
         needAllocatorRegistration = true;
         return mlir::WalkResult::interrupt();
       }
       return mlir::WalkResult::advance();
     });
     if (!needAllocatorRegistration) {
       mod.walk([&](fir::GlobalOp globalOp) {
         if (globalOp.getDataAttrAttr()) {
           if (auto baseBoxType =
                   mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
             if (baseBoxType.isPointerOrAllocatable()) {
               needAllocatorRegistration = true;
               return mlir::WalkResult::interrupt();
             }
           }
         }
         return mlir::WalkResult::advance();
       });
     }

     // Create the constructor function that call CUFRegisterAllocator.
     builder.setInsertionPointToEnd(mod.getBody());
     auto func = mlir::LLVM::LLVMFuncOp::create(builder, loc,
                                                cudaFortranCtorName, funcTy);
     func.setLinkage(mlir::LLVM::Linkage::Internal);
     auto entryBlock = func.addEntryBlock(builder);
     builder.setInsertionPointToStart(entryBlock);

     if (needAllocatorRegistration) {
       // Symbol reference to CUFRegisterAllocator.
       builder.setInsertionPointToEnd(mod.getBody());
       auto registerFuncOp = mlir::LLVM::LLVMFuncOp::create(
           builder, loc, RTNAME_STRING(CUFRegisterAllocator), funcTy);
       registerFuncOp.setVisibility(mlir::SymbolTable::Visibility::Private);
       auto cufRegisterAllocatorRef = mlir::SymbolRefAttr::get(
           mod.getContext(), RTNAME_STRING(CUFRegisterAllocator));
       builder.setInsertionPointToStart(entryBlock);
       mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);
     }

     auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName);
     if (gpuMod) {
       mlir::SymbolTable gpuSymTable(gpuMod);
       bool needsModuleRegistration =
           hasKernel(gpuMod) ||
           hasRegisteredGlobals(mod, gpuSymTable, cudaUnified);
       if (needsModuleRegistration) {
         auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(ctx);
         auto registeredMod = cuf::RegisterModuleOp::create(
             builder, loc, llvmPtrTy,
             mlir::SymbolRefAttr::get(ctx, gpuMod.getName()));

         fir::LLVMTypeConverter typeConverter(
             mod, /*applyTBAA=*/false, /*forceUnifiedTBAATree=*/false, *dl);
         // Register kernels
         for (auto func : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) {
           if (func.isKernel()) {
             auto kernelName = mlir::SymbolRefAttr::get(
                 builder.getStringAttr(cudaDeviceModuleName),
                 {mlir::SymbolRefAttr::get(builder.getContext(),
                                           func.getName())});
             cuf::RegisterKernelOp::create(builder, loc, kernelName,
                                           registeredMod);
           }
         }

         // Register variables
         bool hasNonAllocManagedGlobal = false;
         for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
           auto attr = globalOp.getDataAttrAttr();
           if (!attr)
             continue;
           if (!gpuSymTable.lookup(globalOp.getSymName()))
             continue;

           bool isNonAllocManagedGlobal =
               attr.getValue() == cuf::DataAttribute::Managed &&
               !mlir::isa<fir::BaseBoxType>(globalOp.getType());

           switch (attr.getValue()) {
           case cuf::DataAttribute::Device:
           case cuf::DataAttribute::Constant:
           case cuf::DataAttribute::Managed: {
             if (isNonAllocManagedGlobal) {
               hasNonAllocManagedGlobal = true;
               // Non-allocatable managed globals use pointer indirection:
               // a companion pointer in __nv_managed_data__ holds the unified
               // memory address, registered via __cudaRegisterManagedVar.
               fir::GlobalOp ptrGlobal =
                   createManagedPointerGlobal(builder, mod, globalOp);
               auto func = fir::runtime::getRuntimeFunc<mkRTKey(
                   CUFRegisterManagedVariable)>(loc, builder);
               emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
                                       typeConverter, registeredMod, func,
                                       /*addrGlobal=*/ptrGlobal,
                                       /*nameGlobal=*/globalOp);
             } else {
               auto func =
                   fir::runtime::getRuntimeFunc<mkRTKey(CUFRegisterVariable)>(
                       loc, builder);
               emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
                                       typeConverter, registeredMod, func,
                                       /*addrGlobal=*/globalOp,
                                       /*nameGlobal=*/globalOp);
               // Under -gpu=mem:unified, also register the global as
               // device-resident so a matching host symbol from another
               // translation unit is not treated as host memory.
               if (cudaUnified) {
                 uint64_t szBytes = getGlobalSizeInBytes(
                     loc, *dl, kindMap, typeConverter, globalOp);
                 cuf::RegisterVariableStaticOp::create(
                     builder, loc,
                     mlir::SymbolRefAttr::get(ctx, globalOp.getSymName()),
                     builder.getStringAttr(globalOp.getSymName()),
                     builder.getI64IntegerAttr(szBytes));
               }
             }
           } break;
           default:
             break;
           }
         }

         // Register externally-linked module globals under -gpu=mem:unified.
         // CUFDeviceGlobal cloned them into the GPU module with external
         // linkage so PTX emits .extern; the CUDA driver patches the device
         // reference to the host pointer at module-load time after this call.
         // Works uniformly for fixed-shape (e.g. fir.array<5xi32>) and
         // allocatable (fir.box<fir.heap<...>>) module globals.
         if (cudaUnified) {
           for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
             if (!isCudaUnifiedExternalGlobal(globalOp, gpuSymTable))
               continue;
             auto func = fir::runtime::getRuntimeFunc<mkRTKey(
                 CUFRegisterExternalVariable)>(loc, builder);
             emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
                                     typeConverter, registeredMod, func,
                                     /*addrGlobal=*/globalOp,
                                     /*nameGlobal=*/globalOp);
           }
         }

         if (hasNonAllocManagedGlobal) {
           // Initialize the module after all variables are registered so the
           // runtime populates managed variable unified memory pointers.
           mlir::func::FuncOp initFunc =
               fir::runtime::getRuntimeFunc<mkRTKey(CUFInitModule)>(loc,
                                                                    builder);
           mlir::FunctionType initFTy = initFunc.getFunctionType();
           llvm::SmallVector<mlir::Value> initArgs{fir::runtime::createArguments(
               builder, loc, initFTy, registeredMod)};
           fir::CallOp::create(builder, loc, initFunc, initArgs);
         }
       }
     }
     mlir::LLVM::ReturnOp::create(builder, loc, mlir::ValueRange{});

     // Create the llvm.global_ctor with the function.
     // TODO: We might want to have a utility that retrieve it if already
     // created and adds new functions.
     builder.setInsertionPointToEnd(mod.getBody());
     llvm::SmallVector<mlir::Attribute> funcs;
     funcs.push_back(
         mlir::FlatSymbolRefAttr::get(mod.getContext(), func.getSymName()));
     llvm::SmallVector<int> priorities;
     llvm::SmallVector<mlir::Attribute> data;
     priorities.push_back(0);
     data.push_back(mlir::LLVM::ZeroAttr::get(mod.getContext()));
     mlir::LLVM::GlobalCtorsOp::create(
         builder, mod.getLoc(), builder.getArrayAttr(funcs),
         builder.getI32ArrayAttr(priorities), builder.getArrayAttr(data));
   }
 };

 } // end anonymous namespace
	//===-- CUFAddConstructor.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/Builder/BoxValue.h"
	#include "flang/Optimizer/Builder/CUFCommon.h"
	#include "flang/Optimizer/Builder/FIRBuilder.h"
	#include "flang/Optimizer/Builder/Runtime/RTBuilder.h"
	#include "flang/Optimizer/CodeGen/Target.h"
	#include "flang/Optimizer/CodeGen/TypeConverter.h"
	#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
	#include "flang/Optimizer/Dialect/FIRAttr.h"
	#include "flang/Optimizer/Dialect/FIROps.h"
	#include "flang/Optimizer/Dialect/FIRType.h"
	#include "flang/Optimizer/Support/DataLayout.h"
	#include "flang/Optimizer/Transforms/Passes.h"
	#include "flang/Runtime/CUDA/registration.h"
	#include "flang/Runtime/entry-names.h"
	#include "mlir/Dialect/DLTI/DLTI.h"
	#include "mlir/Dialect/GPU/IR/GPUDialect.h"
	#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
	#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
	#include "mlir/IR/Value.h"
	#include "mlir/Pass/Pass.h"
	#include "llvm/ADT/SmallVector.h"

	namespace fir {
	#define GEN_PASS_DEF_CUFADDCONSTRUCTOR
	#include "flang/Optimizer/Transforms/Passes.h.inc"
	} // namespace fir

	using namespace Fortran::runtime::cuda;

	namespace {

	static constexpr llvm::StringRef cudaFortranCtorName{
	"__cudaFortranConstructor"};
	static constexpr llvm::StringRef managedPtrSuffix{".managed.ptr"};

	/// Create an 8-byte pointer global in the __nv_managed_data__ section.
	/// The CUDA runtime populates this pointer with the unified memory address
	/// when the module is initialized via __cudaInitModule.
	static fir::GlobalOp createManagedPointerGlobal(fir::FirOpBuilder &builder,
	mlir::ModuleOp mod,
	fir::GlobalOp globalOp) {
	mlir::MLIRContext *ctx = mod.getContext();
	std::string ptrGlobalName = (globalOp.getSymName() + managedPtrSuffix).str();
	auto ptrTy = fir::LLVMPointerType::get(ctx, mlir::IntegerType::get(ctx, 8));

	mlir::OpBuilder::InsertionGuard guard(builder);
	builder.setInsertionPointAfter(globalOp);

	llvm::SmallVector<mlir::NamedAttribute> attrs;
	attrs.push_back(
	mlir::NamedAttribute(mlir::StringAttr::get(ctx, "section"),
	mlir::StringAttr::get(ctx, "__nv_managed_data__")));

	mlir::DenseElementsAttr initAttr = {};
	auto ptrGlobal = fir::GlobalOp::create(
	builder, globalOp.getLoc(), ptrGlobalName, /isConstant=/false,
	/isTarget=/false, ptrTy, initAttr,
	/linkName=/builder.createInternalLinkage(), attrs);

	mlir::Region &region = ptrGlobal.getRegion();
	mlir::Block *block = builder.createBlock(&region);
	builder.setInsertionPointToStart(block);
	mlir::Value zero = fir::ZeroOp::create(builder, globalOp.getLoc(), ptrTy);
	fir::HasValueOp::create(builder, globalOp.getLoc(), zero);

	return ptrGlobal;
	}

	/// Return true if \p hostGlobal is a host module-scope global that has been
	/// mirrored in the GPU module as an external (no-body) declaration by the
	/// CUFDeviceGlobal pass under -gpu=mem:unified. Such globals must be
	/// registered with the CUDA driver via CUFRegisterExternalVariable so the
	/// device-side `.extern` symbol resolves to the host pointer at module-load
	/// time and HMM/ATS handles migration.
	static bool isCudaUnifiedExternalGlobal(fir::GlobalOp hostGlobal,
	mlir::SymbolTable &gpuSymTable) {
	if (hostGlobal.getDataAttrAttr())
	return false;
	if (hostGlobal.getConstant())
	return false;
	auto gpuGlobal = gpuSymTable.lookup<fir::GlobalOp>(hostGlobal.getSymName());
	if (!gpuGlobal)
	return false;
	return !gpuGlobal.isInitialized();
	}

	/// Build a C-style name literal (`<symname>\0`) for use as the deviceName
	/// argument of a CUF registration runtime call.
	static mlir::Value buildGlobalNameLiteral(fir::FirOpBuilder &builder,
	mlir::Location loc,
	fir::GlobalOp globalOp) {
	std::string nameStr = globalOp.getSymbol().getValue().str();
	nameStr += '\0';
	return fir::getBase(fir::factory::createStringLiteral(builder, loc, nameStr));
	}

	/// Compute the storage size in bytes of \p globalOp. For a box-typed
	/// allocatable global the size is the descriptor size (after type
	/// conversion); otherwise it's the size of the global's declared type.
	static mlir::Value computeGlobalSize(fir::FirOpBuilder &builder,
	mlir::Location loc, mlir::Type idxTy,
	const mlir::DataLayout &dl,
	const fir::KindMapping &kindMap,
	fir::LLVMTypeConverter &typeConverter,
	fir::GlobalOp globalOp) {
	std::optional<uint64_t> size;
	if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
	mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);
	size = dl.getTypeSizeInBits(structTy) / 8;
	}
	if (!size) {
	if (auto s =
	fir::getTypeSizeAndAlignment(loc, globalOp.getType(), dl, kindMap))
	size = s->first;
	}
	if (!size) {
	// A global embedding descriptor (allocatable/pointer) components has no
	// structural size; size it via its LLVM type, which inlines the
	// descriptors.
	mlir::Type llvmTy = typeConverter.convertType(globalOp.getType());
	if (llvmTy && mlir::isa<mlir::DataLayoutTypeInterface>(llvmTy))
	size = dl.getTypeSizeInBits(llvmTy) / 8;
	}
	if (!size) {
	size = fir::getTypeSizeAndAlignmentOrCrash(loc, globalOp.getType(), dl,
	kindMap)
	.first;
	}
	return builder.createIntegerConstant(loc, idxTy, *size);
	}

	/// Storage size in bytes of \p globalOp as a raw integer (see computeGlobalSize
	/// for the box vs declared-type handling).
	static uint64_t getGlobalSizeInBytes(mlir::Location loc,
	const mlir::DataLayout &dl,
	const fir::KindMapping &kindMap,
	fir::LLVMTypeConverter &typeConverter,
	fir::GlobalOp globalOp) {
	std::optional<uint64_t> size;
	if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
	mlir::Type structTy = typeConverter.convertBoxTypeAsStruct(boxTy);
	size = dl.getTypeSizeInBits(structTy) / 8;
	}
	if (!size) {
	if (auto s =
	fir::getTypeSizeAndAlignment(loc, globalOp.getType(), dl, kindMap))
	size = s->first;
	}
	if (!size) {
	// A global embedding descriptor (allocatable/pointer) components has no
	// structural size; size it via its LLVM type, which inlines the
	// descriptors.
	mlir::Type llvmTy = typeConverter.convertType(globalOp.getType());
	if (llvmTy && mlir::isa<mlir::DataLayoutTypeInterface>(llvmTy))
	size = dl.getTypeSizeInBits(llvmTy) / 8;
	}
	if (!size) {
	size = fir::getTypeSizeAndAlignmentOrCrash(loc, globalOp.getType(), dl,
	kindMap)
	.first;
	}
	return *size;
	}

	/// Emit a call to a CUF registration runtime function with the canonical
	/// (module, addr, name, size) signature, where addr is the address of \p
	/// addrGlobal taken via fir.address_of and name/size describe \p nameGlobal.
	/// Used both for CUFRegisterVariable / CUFRegisterManagedVariable / and
	/// CUFRegisterExternalVariable.
	static void
	emitCUFRegistrationCall(fir::FirOpBuilder &builder, mlir::Location loc,
	mlir::Type idxTy, const mlir::DataLayout &dl,
	const fir::KindMapping &kindMap,
	fir::LLVMTypeConverter &typeConverter,
	mlir::Value registeredMod, mlir::func::FuncOp func,
	fir::GlobalOp addrGlobal, fir::GlobalOp nameGlobal) {
	mlir::Value gblName = buildGlobalNameLiteral(builder, loc, nameGlobal);
	mlir::Value sizeVal = computeGlobalSize(builder, loc, idxTy, dl, kindMap,
	typeConverter, nameGlobal);
	mlir::Value addr = fir::AddrOfOp::create(
	builder, loc, addrGlobal.resultType(), addrGlobal.getSymbol());
	llvm::SmallVector<mlir::Value> args{
	fir::runtime::createArguments(builder, loc, func.getFunctionType(),
	registeredMod, addr, gblName, sizeVal)};
	fir::CallOp::create(builder, loc, func, args);
	}

	static bool hasRegisteredGlobals(mlir::ModuleOp mod,
	mlir::SymbolTable gpuSymTable,
	bool cudaUnified) {
	for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
	auto attr = globalOp.getDataAttrAttr();
	if (!attr) {
	if (cudaUnified && isCudaUnifiedExternalGlobal(globalOp, gpuSymTable))
	return true;
	continue;
	}
	if (!gpuSymTable.lookup(globalOp.getSymName()))
	continue;
	if (attr.getValue() == cuf::DataAttribute::Managed &&
	!mlir::isa<fir::BaseBoxType>(globalOp.getType()))
	return true;
	switch (attr.getValue()) {
	case cuf::DataAttribute::Device:
	case cuf::DataAttribute::Constant:
	case cuf::DataAttribute::Managed: {
	return true;
	} break;
	default:
	break;
	}
	}
	return false;
	}

	static bool hasKernel(mlir::gpu::GPUModuleOp gpuMod) {
	for (auto func : gpuMod.getOps<mlir::gpu::GPUFuncOp>())
	if (func.isKernel())
	return true;
	return false;
	}

	struct CUFAddConstructor
	: public fir::impl::CUFAddConstructorBase<CUFAddConstructor> {

	using CUFAddConstructorBase::CUFAddConstructorBase;

	void runOnOperation() override {
	mlir::ModuleOp mod = getOperation();
	mlir::SymbolTable symTab(mod);
	mlir::OpBuilder opBuilder{mod.getBodyRegion()};
	fir::FirOpBuilder builder(opBuilder, mod);
	fir::KindMapping kindMap{fir::getKindMapping(mod)};
	builder.setInsertionPointToEnd(mod.getBody());
	mlir::Location loc = mod.getLoc();
	auto *ctx = mod.getContext();
	auto voidTy = mlir::LLVM::LLVMVoidType::get(ctx);
	auto idxTy = builder.getIndexType();
	auto funcTy =
	mlir::LLVM::LLVMFunctionType::get(voidTy, {}, /isVarArg=/false);
	std::optional<mlir::DataLayout> dl =
	fir::support::getOrSetMLIRDataLayout(mod, /allowDefaultLayout=/false);
	if (!dl) {
	mlir::emitError(mod.getLoc(),
	"data layout attribute is required to perform " +
	getName() + "pass");
	}

	bool needAllocatorRegistration = false;
	mod.walk([&](fir::DeclareOp declOp) {
	if (declOp.getFortranAttrs() &&
	fir::bitEnumContainsAny(*declOp.getFortranAttrs(),
	fir::FortranVariableFlagsEnum::allocatable \|
	fir::FortranVariableFlagsEnum::pointer)) {
	needAllocatorRegistration = true;
	return mlir::WalkResult::interrupt();
	}
	return mlir::WalkResult::advance();
	});
	if (!needAllocatorRegistration) {
	mod.walk([&](fir::GlobalOp globalOp) {
	if (globalOp.getDataAttrAttr()) {
	if (auto baseBoxType =
	mlir::dyn_cast<fir::BaseBoxType>(globalOp.getType())) {
	if (baseBoxType.isPointerOrAllocatable()) {
	needAllocatorRegistration = true;
	return mlir::WalkResult::interrupt();
	}
	}
	}
	return mlir::WalkResult::advance();
	});
	}

	// Create the constructor function that call CUFRegisterAllocator.
	builder.setInsertionPointToEnd(mod.getBody());
	auto func = mlir::LLVM::LLVMFuncOp::create(builder, loc,
	cudaFortranCtorName, funcTy);
	func.setLinkage(mlir::LLVM::Linkage::Internal);
	auto entryBlock = func.addEntryBlock(builder);
	builder.setInsertionPointToStart(entryBlock);

	if (needAllocatorRegistration) {
	// Symbol reference to CUFRegisterAllocator.
	builder.setInsertionPointToEnd(mod.getBody());
	auto registerFuncOp = mlir::LLVM::LLVMFuncOp::create(
	builder, loc, RTNAME_STRING(CUFRegisterAllocator), funcTy);
	registerFuncOp.setVisibility(mlir::SymbolTable::Visibility::Private);
	auto cufRegisterAllocatorRef = mlir::SymbolRefAttr::get(
	mod.getContext(), RTNAME_STRING(CUFRegisterAllocator));
	builder.setInsertionPointToStart(entryBlock);
	mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);
	}

	auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName);
	if (gpuMod) {
	mlir::SymbolTable gpuSymTable(gpuMod);
	bool needsModuleRegistration =
	hasKernel(gpuMod) \|\|
	hasRegisteredGlobals(mod, gpuSymTable, cudaUnified);
	if (needsModuleRegistration) {
	auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(ctx);
	auto registeredMod = cuf::RegisterModuleOp::create(
	builder, loc, llvmPtrTy,
	mlir::SymbolRefAttr::get(ctx, gpuMod.getName()));

	fir::LLVMTypeConverter typeConverter(
	mod, /applyTBAA=/false, /forceUnifiedTBAATree=/false, *dl);
	// Register kernels
	for (auto func : gpuMod.getOps<mlir::gpu::GPUFuncOp>()) {
	if (func.isKernel()) {
	auto kernelName = mlir::SymbolRefAttr::get(
	builder.getStringAttr(cudaDeviceModuleName),
	{mlir::SymbolRefAttr::get(builder.getContext(),
	func.getName())});
	cuf::RegisterKernelOp::create(builder, loc, kernelName,
	registeredMod);
	}
	}

	// Register variables
	bool hasNonAllocManagedGlobal = false;
	for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
	auto attr = globalOp.getDataAttrAttr();
	if (!attr)
	continue;
	if (!gpuSymTable.lookup(globalOp.getSymName()))
	continue;

	bool isNonAllocManagedGlobal =
	attr.getValue() == cuf::DataAttribute::Managed &&
	!mlir::isa<fir::BaseBoxType>(globalOp.getType());

	switch (attr.getValue()) {
	case cuf::DataAttribute::Device:
	case cuf::DataAttribute::Constant:
	case cuf::DataAttribute::Managed: {
	if (isNonAllocManagedGlobal) {
	hasNonAllocManagedGlobal = true;
	// Non-allocatable managed globals use pointer indirection:
	// a companion pointer in __nv_managed_data__ holds the unified
	// memory address, registered via __cudaRegisterManagedVar.
	fir::GlobalOp ptrGlobal =
	createManagedPointerGlobal(builder, mod, globalOp);
	auto func = fir::runtime::getRuntimeFunc<mkRTKey(
	CUFRegisterManagedVariable)>(loc, builder);
	emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
	typeConverter, registeredMod, func,
	/addrGlobal=/ptrGlobal,
	/nameGlobal=/globalOp);
	} else {
	auto func =
	fir::runtime::getRuntimeFunc<mkRTKey(CUFRegisterVariable)>(
	loc, builder);
	emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
	typeConverter, registeredMod, func,
	/addrGlobal=/globalOp,
	/nameGlobal=/globalOp);
	// Under -gpu=mem:unified, also register the global as
	// device-resident so a matching host symbol from another
	// translation unit is not treated as host memory.
	if (cudaUnified) {
	uint64_t szBytes = getGlobalSizeInBytes(
	loc, *dl, kindMap, typeConverter, globalOp);
	cuf::RegisterVariableStaticOp::create(
	builder, loc,
	mlir::SymbolRefAttr::get(ctx, globalOp.getSymName()),
	builder.getStringAttr(globalOp.getSymName()),
	builder.getI64IntegerAttr(szBytes));
	}
	}
	} break;
	default:
	break;
	}
	}

	// Register externally-linked module globals under -gpu=mem:unified.
	// CUFDeviceGlobal cloned them into the GPU module with external
	// linkage so PTX emits .extern; the CUDA driver patches the device
	// reference to the host pointer at module-load time after this call.
	// Works uniformly for fixed-shape (e.g. fir.array<5xi32>) and
	// allocatable (fir.box<fir.heap<...>>) module globals.
	if (cudaUnified) {
	for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
	if (!isCudaUnifiedExternalGlobal(globalOp, gpuSymTable))
	continue;
	auto func = fir::runtime::getRuntimeFunc<mkRTKey(
	CUFRegisterExternalVariable)>(loc, builder);
	emitCUFRegistrationCall(builder, loc, idxTy, *dl, kindMap,
	typeConverter, registeredMod, func,
	/addrGlobal=/globalOp,
	/nameGlobal=/globalOp);
	}
	}

	if (hasNonAllocManagedGlobal) {
	// Initialize the module after all variables are registered so the
	// runtime populates managed variable unified memory pointers.
	mlir::func::FuncOp initFunc =
	fir::runtime::getRuntimeFunc<mkRTKey(CUFInitModule)>(loc,
	builder);
	mlir::FunctionType initFTy = initFunc.getFunctionType();
	llvm::SmallVector<mlir::Value> initArgs{fir::runtime::createArguments(
	builder, loc, initFTy, registeredMod)};
	fir::CallOp::create(builder, loc, initFunc, initArgs);
	}
	}
	}
	mlir::LLVM::ReturnOp::create(builder, loc, mlir::ValueRange{});

	// Create the llvm.global_ctor with the function.
	// TODO: We might want to have a utility that retrieve it if already
	// created and adds new functions.
	builder.setInsertionPointToEnd(mod.getBody());
	llvm::SmallVector<mlir::Attribute> funcs;
	funcs.push_back(
	mlir::FlatSymbolRefAttr::get(mod.getContext(), func.getSymName()));
	llvm::SmallVector<int> priorities;
	llvm::SmallVector<mlir::Attribute> data;
	priorities.push_back(0);
	data.push_back(mlir::LLVM::ZeroAttr::get(mod.getContext()));
	mlir::LLVM::GlobalCtorsOp::create(
	builder, mod.getLoc(), builder.getArrayAttr(funcs),
	builder.getI32ArrayAttr(priorities), builder.getArrayAttr(data));
	}
	};

	} // end anonymous namespace