[flang-commits] [flang] [flang][cuda] share runtime type info between host and device under managed memory (PR #229213)
Valentin Clement バレンタイン クレメン via flang-commits
flang-commits at lists.llvm.org
Fri Oct 9 11:55:58 PDT 2026
https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/229213
>From 6753693a627daef40ed7a3f6691fab4293b37de9 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 8 Oct 2026 20:43:08 -0700
Subject: [PATCH] [flang][cuda] share runtime type info between host and device
under managed memory
---
.../flang/Optimizer/Builder/CUFCommon.h | 20 +++
.../flang/Optimizer/Transforms/Passes.td | 28 ++++
flang/lib/Optimizer/Builder/CUFCommon.cpp | 27 ++++
flang/lib/Optimizer/CodeGen/CodeGen.cpp | 58 ++++++++
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/CUDA/CUFAddConstructor.cpp | 88 +++++++++++-
.../Transforms/CUDA/CUFSharedTypeInfo.cpp | 131 ++++++++++++++++++
.../CUDA/cuda-shared-type-info-codegen.mlir | 71 ++++++++++
.../CUDA/cuda-shared-type-info-rename.mlir | 61 ++++++++
.../test/Fir/CUDA/cuda-shared-type-info.mlir | 105 ++++++++++++++
10 files changed, 589 insertions(+), 1 deletion(-)
create mode 100644 flang/lib/Optimizer/Transforms/CUDA/CUFSharedTypeInfo.cpp
create mode 100644 flang/test/Fir/CUDA/cuda-shared-type-info-codegen.mlir
create mode 100644 flang/test/Fir/CUDA/cuda-shared-type-info-rename.mlir
create mode 100644 flang/test/Fir/CUDA/cuda-shared-type-info.mlir
diff --git a/flang/include/flang/Optimizer/Builder/CUFCommon.h b/flang/include/flang/Optimizer/Builder/CUFCommon.h
index e1f8ae454b68e3..6c2afc4ec0e4db 100644
--- a/flang/include/flang/Optimizer/Builder/CUFCommon.h
+++ b/flang/include/flang/Optimizer/Builder/CUFCommon.h
@@ -22,6 +22,18 @@ static constexpr llvm::StringRef cudaSharedMemSuffix = "__shared_mem__";
/// the host module, where it cannot share the symbol of the original. The dot
/// keeps it clear of any Fortran or C identifier.
static constexpr llvm::StringRef cudaDeviceCopySuffix = ".device";
+/// Section holding the host runtime type information shared with device code
+/// under -gpu=mem:managed. It must be a valid C identifier so the linker
+/// defines the __start_ and __stop_ symbols bounding it.
+static constexpr llvm::StringRef cudaSharedTypeInfoSection = "__nv_type_info";
+/// Attribute on the GPU module mapping the name of each shared type
+/// descriptor to the managed pointer holding its host address.
+static constexpr llvm::StringRef cudaSharedTypeDescsAttrName =
+ "cuf.shared_type_descs";
+/// Attribute on the host copy of such a managed pointer, referencing the type
+/// descriptor whose address it holds.
+static constexpr llvm::StringRef cudaHostTypeDescAttrName =
+ "cuf.host_type_desc";
namespace fir {
class FirOpBuilder;
@@ -47,6 +59,14 @@ bool isExecutingOnDevice(mlir::Operation *op);
bool isRegisteredDeviceGlobal(fir::GlobalOp op);
bool isRegisteredDeviceAttr(std::optional<cuf::DataAttribute> attr);
+/// True when \p op holds runtime type information generated for a derived
+/// type (type descriptor, binding table, component table, names, ...).
+bool isTypeInfoGlobal(fir::GlobalOp op);
+
+/// True when \p op is the type descriptor of a derived type, the type-info
+/// global whose address is stored in descriptors.
+bool isTypeDescriptorGlobal(fir::GlobalOp op);
+
/// True for procedures that have a device side: attributes(device), (global),
/// (grid_global) and (host,device). Unlike isCUDADeviceContext, host_device
/// counts, since its body is compiled for the device as well.
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 195bbe014f0063..a161af28260799 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -647,6 +647,29 @@ def CUFDeviceGlobal :
];
}
+def CUFSharedTypeInfo : Pass<"cuf-shared-type-info", "mlir::ModuleOp"> {
+ let summary = "Share host runtime type information with device code";
+ let description = [{
+ Under -gpu=mem:managed, a descriptor in managed memory can be written by
+ device code and read by host code, so the type descriptor address it holds
+ must be valid on both sides.
+
+ This pass makes the host runtime type information writable and places it
+ in the `__nv_type_info` section, which CUFAddConstructor registers with the
+ CUDA runtime so device code can read it through its host address. Each
+ type descriptor present in the GPU module gets a managed pointer, which
+ the CUDA Fortran constructor sets to the host address of the type
+ descriptor. The GPU module attribute `cuf.shared_type_descs` maps type
+ descriptor names to these pointers, and code generation loads the type
+ descriptor address from them in device code. The `acc.declare` attribute
+ is dropped so OpenACC does not copy the type information to the device.
+
+ Device code then sees the host binding tables, so type-bound procedures
+ cannot be called from device code.
+ }];
+ let dependentDialects = ["mlir::gpu::GPUDialect", "cuf::CUFDialect"];
+}
+
def CUFAddConstructor : Pass<"cuf-add-constructor", "mlir::ModuleOp"> {
let summary = "Add constructor to register CUDA Fortran allocators";
let dependentDialects = [
@@ -664,6 +687,11 @@ def CUFAddConstructor : Pass<"cuf-add-constructor", "mlir::ModuleOp"> {
"GPU module via cuf.register_variable_static, so the CUDA driver "
"maps the device-side symbol to the host pointer at module-load "
"time and HMM/ATS handles migration.">,
+ Option<"cudaManagedTypeInfo", "cuda-managed-type-info", "bool",
+ /*default=*/"false",
+ "Register the runtime type information shared with device code "
+ "under -gpu=mem:managed (see cuf-shared-type-info): register the "
+ "pages of its section with the CUDA runtime.">,
Option<"priority", "priority", "int32_t", /*default=*/"0",
"Priority for the CUDA Fortran global constructor in "
"llvm.mlir.global_ctors.">,
diff --git a/flang/lib/Optimizer/Builder/CUFCommon.cpp b/flang/lib/Optimizer/Builder/CUFCommon.cpp
index eb2e64158b10bf..389a46de600370 100644
--- a/flang/lib/Optimizer/Builder/CUFCommon.cpp
+++ b/flang/lib/Optimizer/Builder/CUFCommon.cpp
@@ -102,6 +102,33 @@ bool cuf::isRegisteredDeviceGlobal(fir::GlobalOp op) {
return isRegisteredDeviceAttr(op.getDataAttr());
}
+bool cuf::isTypeInfoGlobal(fir::GlobalOp op) {
+ if (op.getDataAttrAttr())
+ return false;
+ auto [nameKind, deconstructed] =
+ fir::NameUniquer::deconstruct(op.getSymName());
+ // Literals and runtime tables live in the _QQ namespace; type information
+ // is mangled as an ordinary variable whose name starts with a separator
+ // marker (e.g. _QMmodE.dt.t), or 'X' once the markers have been renamed.
+ if (nameKind == fir::NameUniquer::NameKind::NOT_UNIQUED ||
+ nameKind == fir::NameUniquer::NameKind::GENERATED)
+ return false;
+ return fir::NameUniquer::isSpecialSymbol(deconstructed.name);
+}
+
+bool cuf::isTypeDescriptorGlobal(fir::GlobalOp op) {
+ if (!isTypeInfoGlobal(op))
+ return false;
+ auto [nameKind, deconstructed] =
+ fir::NameUniquer::deconstruct(op.getSymName());
+ llvm::StringRef name = deconstructed.name;
+ static const std::string renamedSeparator =
+ fir::NameUniquer::replaceSpecialSymbols(
+ fir::kTypeDescriptorSeparator.str());
+ return name.starts_with(fir::kTypeDescriptorSeparator) ||
+ name.starts_with(renamedSeparator);
+}
+
void cuf::genPointerSync(const mlir::Value box, fir::FirOpBuilder &builder) {
if (auto declareOp = box.getDefiningOp<hlfir::DeclareOp>()) {
if (auto addrOfOp = declareOp.getMemref().getDefiningOp<fir::AddrOfOp>()) {
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 4923cfe2e18787..3eec423b5669b5 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -235,6 +235,41 @@ mlir::Value replaceWithAddrOfOrASCast(mlir::ConversionPatternRewriter &rewriter,
return mlir::LLVM::AddressOfOp::create(rewriter, loc, type, symName);
}
+/// In device code under -gpu=mem:managed, load the host address of the type
+/// descriptor \p name from the managed pointer created for it by the
+/// CUFSharedTypeInfo pass. \p name may be in source or assembly form; the
+/// mapping is keyed by the assembly form. Return a null value when the type
+/// descriptor is not shared, or when the code is emitted in a global
+/// initializer, which needs a constant address and keeps the device copy.
+static mlir::Value
+loadSharedTypeDescriptor(mlir::gpu::GPUModuleOp gpuMod,
+ mlir::ConversionPatternRewriter &rewriter,
+ mlir::Location loc, llvm::StringRef name) {
+ if (mlir::Block *block = rewriter.getInsertionBlock()) {
+ mlir::Operation *parent = block->getParentOp();
+ if (mlir::isa<fir::GlobalOp, mlir::LLVM::GlobalOp>(parent) ||
+ parent->getParentOfType<fir::GlobalOp>() ||
+ parent->getParentOfType<mlir::LLVM::GlobalOp>())
+ return {};
+ }
+ auto typeDescs =
+ gpuMod->getAttrOfType<mlir::DictionaryAttr>(cudaSharedTypeDescsAttrName);
+ if (!typeDescs)
+ return {};
+ auto ptrSym = typeDescs.getAs<mlir::FlatSymbolRefAttr>(
+ fir::NameUniquer::replaceSpecialSymbols(name.str()));
+ if (!ptrSym)
+ return {};
+ mlir::Type llvmPtrTy = getLlvmPtrType(gpuMod.getContext());
+ mlir::DataLayout dataLayout(gpuMod);
+ // Managed globals live in the global address space on the device.
+ mlir::Value ptrAddr = replaceWithAddrOfOrASCast(
+ rewriter, loc, static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Global),
+ fir::factory::getProgramAddressSpace(&dataLayout), ptrSym.getValue(),
+ llvmPtrTy);
+ return mlir::LLVM::LoadOp::create(rewriter, loc, llvmPtrTy, ptrAddr);
+}
+
/// Return the NVVM address space implied by a CUF data attribute on a
/// fir::GlobalOp that has not yet been converted to llvm.mlir.global.
/// Returns std::nullopt if no CUF-specific address space applies.
@@ -270,6 +305,12 @@ struct AddrOfOpConversion : public fir::FIROpConversion<fir::AddrOfOp> {
mlir::ConversionPatternRewriter &rewriter) const override {
if (auto gpuMod = addr->getParentOfType<mlir::gpu::GPUModuleOp>()) {
+ if (mlir::Value typeDesc = loadSharedTypeDescriptor(
+ gpuMod, rewriter, addr.getLoc(),
+ addr.getSymbol().getRootReference().getValue())) {
+ rewriter.replaceOp(addr, typeDesc);
+ return mlir::success();
+ }
auto global = gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol());
replaceWithAddrOfOrASCast(
rewriter, addr->getLoc(),
@@ -1818,6 +1859,11 @@ getTypeDescriptor(ModOpTy mod, mlir::ConversionPatternRewriter &rewriter,
options.typeDescriptorsRenamedForAssembly
? fir::NameUniquer::getTypeDescriptorAssemblyName(recType.getName())
: fir::NameUniquer::getTypeDescriptorName(recType.getName());
+ if constexpr (std::is_same_v<ModOpTy, mlir::gpu::GPUModuleOp>) {
+ if (mlir::Value typeDesc =
+ loadSharedTypeDescriptor(mod, rewriter, loc, name))
+ return typeDesc;
+ }
mlir::Type llvmPtrTy = ::getLlvmPtrType(mod.getContext());
mlir::DataLayout dataLayout(mod);
if (auto global = mod.template lookupSymbol<fir::GlobalOp>(name))
@@ -3731,6 +3777,18 @@ struct TypeDescOpConversion : public fir::FIROpConversion<fir::TypeDescOp> {
mlir::Type inTy = typeDescOp.getInType();
assert(mlir::isa<fir::RecordType>(inTy) && "expecting fir.type");
auto recordType = mlir::dyn_cast<fir::RecordType>(inTy);
+ if (auto gpuMod = typeDescOp->getParentOfType<mlir::gpu::GPUModuleOp>()) {
+ std::string name =
+ this->options.typeDescriptorsRenamedForAssembly
+ ? fir::NameUniquer::getTypeDescriptorAssemblyName(
+ recordType.getName())
+ : fir::NameUniquer::getTypeDescriptorName(recordType.getName());
+ if (mlir::Value typeDesc = loadSharedTypeDescriptor(
+ gpuMod, rewriter, typeDescOp.getLoc(), name)) {
+ rewriter.replaceOp(typeDescOp, typeDesc);
+ return mlir::success();
+ }
+ }
auto module = typeDescOp.getOperation()->getParentOfType<mlir::ModuleOp>();
mlir::Value typeDesc = getTypeDescriptor(
module, rewriter, typeDescOp.getLoc(), recordType, this->options);
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index ba63cbd564a97b..b7526e2cec289f 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -21,6 +21,7 @@ add_flang_library(FIRTransforms
CUDA/CUFOpConversion.cpp
CUDA/CUFOpConversionLate.cpp
CUDA/CUFPredefinedVarToGPU.cpp
+ CUDA/CUFSharedTypeInfo.cpp
CharacterConversion.cpp
CompilerGeneratedNames.cpp
ConstantArgumentGlobalisation.cpp
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
index 3fa91e7214e0c4..85aa87a626bcca 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
@@ -271,6 +271,16 @@ emitCUFRegistrationCall(fir::FirOpBuilder &builder, mlir::Location loc,
fir::CallOp::create(builder, loc, func, args);
}
+/// Return true if \p mod defines runtime type information placed in the
+/// section shared with device code.
+static bool definesSharedTypeInfo(mlir::ModuleOp mod) {
+ for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>())
+ if (definesGlobal(globalOp) &&
+ globalOp.getSection() == cudaSharedTypeInfoSection)
+ return true;
+ return false;
+}
+
static bool hasRegisteredGlobals(mlir::ModuleOp mod,
mlir::SymbolTable gpuSymTable,
bool cudaUnified) {
@@ -414,6 +424,44 @@ struct CUFAddConstructor
mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);
}
+ // Device code reads the runtime type information defined here through its
+ // host address, so its pages must be registered with the CUDA runtime.
+ // The section bounds are local to the executable or shared library being
+ // linked, and every unit defining type information registers the same
+ // range.
+ if (cudaManagedTypeInfo && definesSharedTypeInfo(mod)) {
+ auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(ctx);
+ auto getSectionBound = [&](llvm::StringRef prefix) -> mlir::Value {
+ std::string name = (prefix + cudaSharedTypeInfoSection).str();
+ auto bound = mod.lookupSymbol<mlir::LLVM::GlobalOp>(name);
+ if (!bound) {
+ mlir::OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToEnd(mod.getBody());
+ bound = mlir::LLVM::GlobalOp::create(
+ builder, loc, mlir::IntegerType::get(ctx, 8),
+ /*isConstant=*/false, mlir::LLVM::Linkage::ExternWeak, name,
+ mlir::Attribute{});
+ bound.setVisibility_(mlir::LLVM::Visibility::Hidden);
+ }
+ return mlir::LLVM::AddressOfOp::create(builder, loc, bound);
+ };
+ mlir::Value begin = getSectionBound("__start_");
+ mlir::Value end = getSectionBound("__stop_");
+ auto rangeFuncTy = mlir::LLVM::LLVMFunctionType::get(
+ voidTy, {llvmPtrTy, llvmPtrTy}, /*isVarArg=*/false);
+ llvm::StringRef rangeFuncName = RTNAME_STRING(CUFRegisterHostMemoryRange);
+ if (!mod.lookupSymbol(rangeFuncName)) {
+ mlir::OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToEnd(mod.getBody());
+ auto rangeFunc = mlir::LLVM::LLVMFuncOp::create(
+ builder, loc, rangeFuncName, rangeFuncTy);
+ rangeFunc.setVisibility(mlir::SymbolTable::Visibility::Private);
+ }
+ mlir::LLVM::CallOp::create(builder, loc, rangeFuncTy,
+ mlir::SymbolRefAttr::get(ctx, rangeFuncName),
+ mlir::ValueRange{begin, end});
+ }
+
auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName);
if (gpuMod) {
mlir::SymbolTable gpuSymTable(gpuMod);
@@ -532,7 +580,6 @@ struct CUFAddConstructor
builder.getI64IntegerAttr(szBytes));
}
}
-
if (hasNonAllocManagedGlobal) {
// Initialize the module after all variables are registered so the
// runtime populates managed variable unified memory pointers. With
@@ -575,6 +622,45 @@ struct CUFAddConstructor
llvm::SmallVector<mlir::Value> initArgs{
fir::runtime::createArguments(builder, loc, initFTy, handle)};
fir::CallOp::create(builder, loc, initFunc, initArgs);
+
+ // Device code reads each shared type descriptor through a managed
+ // pointer (see CUFSharedTypeInfo). Now that the managed memory is
+ // allocated, store the host address of the type descriptor in it.
+ // The runtime may leave a pointer that no device code references
+ // null; nothing reads it, so the store is skipped.
+ for (fir::GlobalOp globalOp : mod.getOps<fir::GlobalOp>()) {
+ auto typeDescRef = globalOp->getAttrOfType<mlir::FlatSymbolRefAttr>(
+ cudaHostTypeDescAttrName);
+ if (!typeDescRef)
+ continue;
+ auto typeDesc =
+ symTab.lookup<fir::GlobalOp>(typeDescRef.getValue());
+ auto ptrGlobal = mod.lookupSymbol<fir::GlobalOp>(
+ (globalOp.getSymName() + managedPtrSuffix).str());
+ if (!typeDesc || !ptrGlobal)
+ continue;
+ mlir::Value ptrRef = fir::AddrOfOp::create(
+ builder, loc, ptrGlobal.resultType(), ptrGlobal.getSymbol());
+ mlir::Value managedAddr = fir::LoadOp::create(builder, loc, ptrRef);
+ mlir::Value isSet = builder.genIsNotNullAddr(loc, managedAddr);
+ mlir::Block *currentBlock = builder.getInsertionBlock();
+ mlir::Region *region = currentBlock->getParent();
+ mlir::Block *storeBlock = builder.createBlock(region, region->end());
+ mlir::Block *nextBlock = builder.createBlock(region, region->end());
+ builder.setInsertionPointToEnd(currentBlock);
+ mlir::LLVM::CondBrOp::create(builder, loc, isSet, storeBlock,
+ nextBlock);
+ builder.setInsertionPointToStart(storeBlock);
+ mlir::Value dest = builder.createConvert(
+ loc, fir::ReferenceType::get(globalOp.getType()), managedAddr);
+ mlir::Value typeDescAddr = fir::AddrOfOp::create(
+ builder, loc, typeDesc.resultType(), typeDesc.getSymbol());
+ mlir::Value value =
+ builder.createConvert(loc, globalOp.getType(), typeDescAddr);
+ fir::StoreOp::create(builder, loc, value, dest);
+ mlir::LLVM::BrOp::create(builder, loc, nextBlock);
+ builder.setInsertionPointToStart(nextBlock);
+ }
mlir::LLVM::ReturnOp::create(builder, loc, mlir::ValueRange{});
}
}
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFSharedTypeInfo.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFSharedTypeInfo.cpp
new file mode 100644
index 00000000000000..a4cee32c27a293
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFSharedTypeInfo.cpp
@@ -0,0 +1,131 @@
+//===-- CUFSharedTypeInfo.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/CUFCommon.h"
+#include "flang/Optimizer/Dialect/CUF/Attributes/CUFAttr.h"
+#include "flang/Optimizer/Dialect/CUF/CUFOps.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Support/InternalNames.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/xxhash.h"
+
+namespace fir {
+#define GEN_PASS_DEF_CUFSHAREDTYPEINFO
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+
+/// Return a tag identifying the translation unit of \p mod.
+///
+/// Without relocatable device code, every unit has its own device module, and
+/// the CUDA runtime only populates the first managed variable registered under
+/// a given name. The managed pointers created here must therefore have a name
+/// unique to their unit.
+static std::string getUnitTag(mlir::ModuleOp mod) {
+ std::string key;
+ llvm::raw_string_ostream os(key);
+ if (auto fileLoc = mlir::dyn_cast<mlir::FileLineColLoc>(mod.getLoc()))
+ os << fileLoc.getFilename().getValue();
+ for (mlir::Operation &op : mod.getBody()->getOperations())
+ if (auto sym = mlir::dyn_cast<mlir::SymbolOpInterface>(op))
+ os << ';' << sym.getName();
+ return llvm::utohexstr(llvm::xxh3_64bits(key), /*LowerCase=*/true);
+}
+
+class CUFSharedTypeInfo
+ : public fir::impl::CUFSharedTypeInfoBase<CUFSharedTypeInfo> {
+public:
+ using CUFSharedTypeInfoBase::CUFSharedTypeInfoBase;
+
+ void runOnOperation() override {
+ mlir::ModuleOp mod = getOperation();
+ mlir::MLIRContext *ctx = mod.getContext();
+ mlir::SymbolTable symTab(mod);
+ mlir::StringAttr declareAttrName =
+ mlir::StringAttr::get(ctx, mlir::acc::getDeclareAttrName());
+ mlir::StringAttr section =
+ mlir::StringAttr::get(ctx, cudaSharedTypeInfoSection);
+
+ llvm::SmallVector<fir::GlobalOp> typeDescs;
+ for (fir::GlobalOp global : mod.getOps<fir::GlobalOp>()) {
+ if (!cuf::isTypeInfoGlobal(global))
+ continue;
+ // The type information is no longer copied to the device.
+ global->removeAttr(declareAttrName);
+ // Constant data would be placed in a read-only section, which not every
+ // GPU can register. Every type-info global must land in the same
+ // section, and a section cannot mix read-only and writable data.
+ if (global.isInitialized()) {
+ global.removeConstantAttr();
+ global.setSectionAttr(section);
+ }
+ if (cuf::isTypeDescriptorGlobal(global))
+ typeDescs.push_back(global);
+ }
+
+ auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(cudaDeviceModuleName);
+ if (!gpuMod)
+ return;
+ mlir::SymbolTable gpuSymTab(gpuMod);
+ for (fir::GlobalOp global : gpuMod.getOps<fir::GlobalOp>())
+ if (cuf::isTypeInfoGlobal(global))
+ global->removeAttr(declareAttrName);
+
+ // Device code reads each type descriptor through a managed pointer that
+ // the CUDA Fortran constructor sets to the host address. A single managed
+ // allocation is valid on every device.
+ //
+ // The pointer names and the dictionary keys use the assembly form of the
+ // type descriptor name. CompilerGeneratedNamesConversion does not update
+ // names nested in the dictionary, and leaves names in this form unchanged.
+ std::string unitTag = getUnitTag(mod);
+ mlir::OpBuilder builder(ctx);
+ auto ptrTy = fir::LLVMPointerType::get(ctx, mlir::IntegerType::get(ctx, 8));
+ auto managed =
+ cuf::DataAttributeAttr::get(ctx, cuf::DataAttribute::Managed);
+ llvm::SmallVector<mlir::NamedAttribute> sharedTypeDescs;
+ for (fir::GlobalOp typeDesc : typeDescs) {
+ if (!gpuSymTab.lookup(typeDesc.getSymName()))
+ continue;
+ std::string typeDescName =
+ fir::NameUniquer::replaceSpecialSymbols(typeDesc.getSymName().str());
+ std::string ptrName = typeDescName + "Xhostaddr" + unitTag;
+ builder.setInsertionPointAfter(typeDesc);
+ auto ptrGlobal = fir::GlobalOp::create(
+ builder, typeDesc.getLoc(), ptrName, /*isConstant=*/false,
+ /*isTarget=*/false, ptrTy, mlir::Attribute{},
+ /*linkage=*/fir::LinkageAttr{});
+ ptrGlobal.setDataAttrAttr(managed);
+ mlir::Block *block = builder.createBlock(&ptrGlobal.getRegion());
+ builder.setInsertionPointToStart(block);
+ mlir::Value zero = fir::ZeroOp::create(builder, typeDesc.getLoc(), ptrTy);
+ fir::HasValueOp::create(builder, typeDesc.getLoc(), zero);
+
+ gpuSymTab.insert(ptrGlobal->clone());
+ ptrGlobal->setAttr(
+ cudaHostTypeDescAttrName,
+ mlir::FlatSymbolRefAttr::get(typeDesc.getSymNameAttr()));
+ sharedTypeDescs.emplace_back(
+ mlir::StringAttr::get(ctx, typeDescName),
+ mlir::FlatSymbolRefAttr::get(ptrGlobal.getSymNameAttr()));
+ }
+ if (!sharedTypeDescs.empty())
+ gpuMod->setAttr(cudaSharedTypeDescsAttrName,
+ mlir::DictionaryAttr::get(ctx, sharedTypeDescs));
+ }
+};
+
+} // namespace
diff --git a/flang/test/Fir/CUDA/cuda-shared-type-info-codegen.mlir b/flang/test/Fir/CUDA/cuda-shared-type-info-codegen.mlir
new file mode 100644
index 00000000000000..bd50560c2ac13c
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-shared-type-info-codegen.mlir
@@ -0,0 +1,71 @@
+// Under -gpu=mem:managed, device code loads the address of a shared type
+// descriptor from the managed pointer recorded in the GPU module attribute
+// cuf.shared_type_descs (see cuda-shared-type-info.mlir).
+
+// RUN: fir-opt --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu" %s | FileCheck %s
+
+module attributes {gpu.container_module} {
+ gpu.module @cuda_device_mod attributes {cuf.shared_type_descs = {_QMmEXdtXt = @_QMmEXdtXtXhostaddr0}} {
+ fir.global linkonce_odr @_QMmE.dt.t constant : i8
+ fir.global @_QMmEXdtXtXhostaddr0 <{data_attr = #cuf.cuda<managed>}> : !fir.llvm_ptr<i8> {
+ %0 = fir.zero_bits !fir.llvm_ptr<i8>
+ fir.has_value %0 : !fir.llvm_ptr<i8>
+ }
+ // An initializer needs a constant address and keeps the device copy.
+ fir.global @_QMmEholder constant : !fir.llvm_ptr<i8> {
+ %0 = fir.address_of(@_QMmE.dt.t) : !fir.ref<i8>
+ %1 = fir.convert %0 : (!fir.ref<i8>) -> !fir.llvm_ptr<i8>
+ fir.has_value %1 : !fir.llvm_ptr<i8>
+ }
+ fir.global @_QMmEboxed : !fir.box<!fir.ptr<!fir.type<_QMmTt{i:i32}>>> {
+ %0 = fir.zero_bits !fir.ptr<!fir.type<_QMmTt{i:i32}>>
+ %1 = fir.embox %0 : (!fir.ptr<!fir.type<_QMmTt{i:i32}>>) -> !fir.box<!fir.ptr<!fir.type<_QMmTt{i:i32}>>>
+ fir.has_value %1 : !fir.box<!fir.ptr<!fir.type<_QMmTt{i:i32}>>>
+ }
+
+ func.func @embox(%arg0: !fir.ref<!fir.type<_QMmTt{i:i32}>>) -> !fir.box<!fir.type<_QMmTt{i:i32}>> {
+ %0 = fir.embox %arg0 : (!fir.ref<!fir.type<_QMmTt{i:i32}>>) -> !fir.box<!fir.type<_QMmTt{i:i32}>>
+ return %0 : !fir.box<!fir.type<_QMmTt{i:i32}>>
+ }
+
+ func.func @type_desc() -> !fir.tdesc<!fir.type<_QMmTt{i:i32}>> {
+ %0 = fir.type_desc !fir.type<_QMmTt{i:i32}>
+ return %0 : !fir.tdesc<!fir.type<_QMmTt{i:i32}>>
+ }
+
+ func.func @address_of() -> !fir.ref<i8> {
+ %0 = fir.address_of(@_QMmE.dt.t) : !fir.ref<i8>
+ return %0 : !fir.ref<i8>
+ }
+ }
+}
+
+// CHECK-LABEL: gpu.module @cuda_device_mod
+// CHECK: llvm.mlir.global external @_QMmEXdtXtXhostaddr0() {{.*}}nvvm.managed
+
+// CHECK: llvm.mlir.global {{.*}}@_QMmEholder()
+// CHECK: llvm.mlir.addressof @_QMmE.dt.t : !llvm.ptr
+
+// CHECK: llvm.mlir.global {{.*}}@_QMmEboxed()
+// CHECK-NOT: llvm.load
+// CHECK: llvm.mlir.addressof @_QMmE.dt.t : !llvm.ptr
+// CHECK-NOT: llvm.load
+// CHECK: llvm.return
+
+// CHECK-LABEL: llvm.func @embox
+// CHECK: %[[PTR:.*]] = llvm.mlir.addressof @_QMmEXdtXtXhostaddr0 : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[PTR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.insertvalue %[[TYPEDESC]], %{{.*}}[{{[0-9]+}}] : !llvm.struct<
+
+// CHECK-LABEL: llvm.func @type_desc
+// CHECK: %[[PTR:.*]] = llvm.mlir.addressof @_QMmEXdtXtXhostaddr0 : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[PTR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.return %[[TYPEDESC]] : !llvm.ptr
+
+// CHECK-LABEL: llvm.func @address_of
+// CHECK: %[[PTR:.*]] = llvm.mlir.addressof @_QMmEXdtXtXhostaddr0 : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[PTR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.return %[[TYPEDESC]] : !llvm.ptr
diff --git a/flang/test/Fir/CUDA/cuda-shared-type-info-rename.mlir b/flang/test/Fir/CUDA/cuda-shared-type-info-rename.mlir
new file mode 100644
index 00000000000000..088239e1a5ec6e
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-shared-type-info-rename.mlir
@@ -0,0 +1,61 @@
+// Shared type descriptors survive the conversion of compiler-generated names:
+// device code still loads the type descriptor address from the managed
+// pointer when type descriptors are renamed for assembly.
+
+// RUN: fir-opt --cuf-shared-type-info --compiler-generated-names --fir-to-llvm-ir="target=x86_64-unknown-linux-gnu type-descriptors-renamed-for-assembly=true" %s | FileCheck %s
+
+module attributes {gpu.container_module} {
+ fir.global linkonce_odr @_QMmE.dt.t constant : i8 {
+ %0 = arith.constant 0 : i8
+ fir.has_value %0 : i8
+ }
+
+ gpu.module @cuda_device_mod {
+ fir.global linkonce_odr @_QMmE.dt.t constant : i8 {
+ %0 = arith.constant 0 : i8
+ fir.has_value %0 : i8
+ }
+
+ func.func @embox(%arg0: !fir.ref<!fir.type<_QMmTt{i:i32}>>) -> !fir.box<!fir.type<_QMmTt{i:i32}>> {
+ %0 = fir.embox %arg0 : (!fir.ref<!fir.type<_QMmTt{i:i32}>>) -> !fir.box<!fir.type<_QMmTt{i:i32}>>
+ return %0 : !fir.box<!fir.type<_QMmTt{i:i32}>>
+ }
+
+ func.func @type_desc() -> !fir.tdesc<!fir.type<_QMmTt{i:i32}>> {
+ %0 = fir.type_desc !fir.type<_QMmTt{i:i32}>
+ return %0 : !fir.tdesc<!fir.type<_QMmTt{i:i32}>>
+ }
+
+ func.func @address_of() -> !fir.ref<i8> {
+ %0 = fir.address_of(@_QMmE.dt.t) : !fir.ref<i8>
+ return %0 : !fir.ref<i8>
+ }
+ }
+}
+
+// CHECK: llvm.mlir.global linkonce_odr @_QMmEXdtXt()
+// CHECK: llvm.mlir.global external @[[PTR:_QMmEXdtXtXhostaddr[0-9a-f]+]]()
+
+// CHECK-LABEL: gpu.module @cuda_device_mod
+// CHECK-SAME: cuf.shared_type_descs = {_QMmEXdtXt = @[[PTR]]}
+// CHECK: llvm.mlir.global linkonce_odr constant @_QMmEXdtXt()
+
+// CHECK-LABEL: llvm.func @embox
+// CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @[[PTR]] : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[ADDR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.insertvalue %[[TYPEDESC]], %{{.*}}[{{[0-9]+}}] : !llvm.struct<
+
+// CHECK-LABEL: llvm.func @type_desc
+// CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @[[PTR]] : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[ADDR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.return %[[TYPEDESC]] : !llvm.ptr
+
+// CHECK-LABEL: llvm.func @address_of
+// CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @[[PTR]] : !llvm.ptr<1>
+// CHECK: %[[CAST:.*]] = llvm.addrspacecast %[[ADDR]] : !llvm.ptr<1> to !llvm.ptr
+// CHECK: %[[TYPEDESC:.*]] = llvm.load %[[CAST]] : !llvm.ptr -> !llvm.ptr
+// CHECK: llvm.return %[[TYPEDESC]] : !llvm.ptr
+
+// CHECK: llvm.mlir.global external @[[PTR]]() {{.*}}nvvm.managed
diff --git a/flang/test/Fir/CUDA/cuda-shared-type-info.mlir b/flang/test/Fir/CUDA/cuda-shared-type-info.mlir
new file mode 100644
index 00000000000000..220b146802d2ae
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-shared-type-info.mlir
@@ -0,0 +1,105 @@
+// Under -gpu=mem:managed, runtime type information is shared between host and
+// device so a descriptor in managed memory holds a type descriptor address
+// valid on both sides.
+
+// RUN: fir-opt --cuf-shared-type-info %s | FileCheck %s
+// RUN: fir-opt --cuf-shared-type-info %s | FileCheck %s --check-prefix=NODECLARE
+// RUN: fir-opt --cuf-shared-type-info --cuf-add-constructor="cuda-managed-type-info=true" %s | FileCheck %s --check-prefix=CTOR
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<i8 = dense<8> : vector<2xi64>, i16 = dense<16> : vector<2xi64>, i1 = dense<8> : vector<2xi64>, !llvm.ptr = dense<64> : vector<4xi64>, f80 = dense<128> : vector<2xi64>, i128 = dense<128> : vector<2xi64>, i64 = dense<64> : vector<2xi64>, !llvm.ptr<271> = dense<32> : vector<4xi64>, !llvm.ptr<272> = dense<64> : vector<4xi64>, f128 = dense<128> : vector<2xi64>, !llvm.ptr<270> = dense<32> : vector<4xi64>, f16 = dense<16> : vector<2xi64>, f64 = dense<64> : vector<2xi64>, i32 = dense<32> : vector<2xi64>, "dlti.stack_alignment" = 128 : i64, "dlti.endianness" = "little">, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", gpu.container_module, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu"} {
+ fir.global linkonce_odr @_QMmE.n.t {acc.declare = #acc.declare<dataClause = acc_copyin>} constant : !fir.char<1> {
+ %0 = fir.string_lit "t"(1) : !fir.char<1>
+ fir.has_value %0 : !fir.char<1>
+ }
+ fir.global linkonce_odr @_QMmE.dt.t {acc.declare = #acc.declare<dataClause = acc_copyin>} constant : !fir.llvm_ptr<i8> {
+ %0 = fir.address_of(@_QMmE.n.t) : !fir.ref<!fir.char<1>>
+ %1 = fir.convert %0 : (!fir.ref<!fir.char<1>>) -> !fir.llvm_ptr<i8>
+ fir.has_value %1 : !fir.llvm_ptr<i8>
+ }
+ // Type information defined in another unit, not used by device code.
+ fir.global @_QMotherE.dt.u {acc.declare = #acc.declare<dataClause = acc_copyin>} constant : !fir.char<1>
+ // Not type information.
+ fir.global linkonce @_QQclX74 constant : !fir.char<1> {
+ %0 = fir.string_lit "t"(1) : !fir.char<1>
+ fir.has_value %0 : !fir.char<1>
+ }
+ fir.global @_QMmEx : i32 {
+ %0 = arith.constant 0 : i32
+ fir.has_value %0 : i32
+ }
+
+ gpu.module @cuda_device_mod {
+ fir.global linkonce_odr @_QMmE.n.t {acc.declare = #acc.declare<dataClause = acc_copyin>} constant : !fir.char<1> {
+ %0 = fir.string_lit "t"(1) : !fir.char<1>
+ fir.has_value %0 : !fir.char<1>
+ }
+ fir.global linkonce_odr @_QMmE.dt.t {acc.declare = #acc.declare<dataClause = acc_copyin>} constant : !fir.llvm_ptr<i8> {
+ %0 = fir.address_of(@_QMmE.n.t) : !fir.ref<!fir.char<1>>
+ %1 = fir.convert %0 : (!fir.ref<!fir.char<1>>) -> !fir.llvm_ptr<i8>
+ fir.has_value %1 : !fir.llvm_ptr<i8>
+ }
+ gpu.func @_QMmPk() kernel {
+ gpu.return
+ }
+ }
+}
+
+// Host type information is writable and placed in the shared section. Each
+// type descriptor used on the device gets a managed pointer.
+// CHECK: fir.global linkonce_odr @_QMmE.n.t <{section = "__nv_type_info"}> : !fir.char<1> {
+// CHECK: fir.global linkonce_odr @_QMmE.dt.t <{section = "__nv_type_info"}> : !fir.llvm_ptr<i8> {
+// CHECK: fir.global @_QMmEXdtXtXhostaddr[[TAG:[0-9a-f]+]] <{data_attr = #cuf.cuda<managed>}> {cuf.host_type_desc = @_QMmE.dt.t} : !fir.llvm_ptr<i8> {
+// CHECK: fir.zero_bits !fir.llvm_ptr<i8>
+// CHECK: fir.global @_QMotherE.dt.u constant : !fir.char<1>{{$}}
+// CHECK: fir.global linkonce @_QQclX74 constant : !fir.char<1> {
+// CHECK: fir.global @_QMmEx : i32 {
+
+// The device copies stay, and the GPU module maps each shared type descriptor
+// to its pointer.
+// CHECK: gpu.module @cuda_device_mod
+// CHECK-SAME: cuf.shared_type_descs
+// CHECK-SAME: _QMmEXdtXt = @_QMmEXdtXtXhostaddr[[TAG]]
+// CHECK-DAG: fir.global linkonce_odr @_QMmE.n.t constant : !fir.char<1> {
+// CHECK-DAG: fir.global linkonce_odr @_QMmE.dt.t constant : !fir.llvm_ptr<i8> {
+// CHECK-DAG: fir.global @_QMmEXdtXtXhostaddr[[TAG]] <{data_attr = #cuf.cuda<managed>}> : !fir.llvm_ptr<i8> {
+
+// NODECLARE-NOT: acc.declare
+
+// The constructor registers the pages of the section and the managed pointer,
+// and saves the module handle. A second constructor, which runs after the
+// registration constructors of all the units, initializes the module and then
+// stores the host address of the type descriptor in the managed pointer,
+// unless the runtime left the pointer null.
+// CTOR-LABEL: llvm.func internal @__cudaFortranConstructor()
+// CTOR: %[[START:.*]] = llvm.mlir.addressof @__start___nv_type_info : !llvm.ptr
+// CTOR: %[[STOP:.*]] = llvm.mlir.addressof @__stop___nv_type_info : !llvm.ptr
+// CTOR: llvm.call @_FortranACUFRegisterHostMemoryRange(%[[START]], %[[STOP]]) : (!llvm.ptr, !llvm.ptr) -> ()
+// CTOR: %[[MOD:.*]] = cuf.register_module @cuda_device_mod -> !llvm.ptr
+// CTOR: fir.call @_FortranACUFRegisterManagedVariable
+// CTOR-NOT: fir.call @_FortranACUFInitModule
+// CTOR: %[[HANDLEADDR:.*]] = llvm.mlir.addressof @__cudaFortranModuleHandle : !llvm.ptr
+// CTOR: llvm.store %[[MOD]], %[[HANDLEADDR]] : !llvm.ptr, !llvm.ptr
+// CTOR: llvm.return
+// CTOR-DAG: llvm.mlir.global extern_weak hidden @__start___nv_type_info() {{.*}}: i8
+// CTOR-DAG: llvm.mlir.global extern_weak hidden @__stop___nv_type_info() {{.*}}: i8
+// CTOR-DAG: llvm.mlir.global internal @__cudaFortranModuleHandle() {{.*}}: !llvm.ptr
+
+// CTOR-LABEL: llvm.func internal @__cudaFortranInitConstructor()
+// CTOR: %[[HANDLEADDR2:.*]] = llvm.mlir.addressof @__cudaFortranModuleHandle : !llvm.ptr
+// CTOR: %[[HANDLE:.*]] = llvm.load %[[HANDLEADDR2]] : !llvm.ptr -> !llvm.ptr
+// CTOR: fir.convert %[[HANDLE]]
+// CTOR: fir.call @_FortranACUFInitModule
+// CTOR: %[[PTRREF:.*]] = fir.address_of(@_QMmEXdtXtXhostaddr{{[0-9a-f]+}}.managed.ptr) : !fir.ref<!fir.llvm_ptr<i8>>
+// CTOR: %[[MANAGED:.*]] = fir.load %[[PTRREF]] : !fir.ref<!fir.llvm_ptr<i8>>
+// CTOR: %[[NOTNULL:.*]] = arith.cmpi ne
+// CTOR: llvm.cond_br %[[NOTNULL]], ^[[STORE:bb[0-9]+]], ^[[NEXT:bb[0-9]+]]
+// CTOR: ^[[STORE]]:
+// CTOR: %[[DEST:.*]] = fir.convert %[[MANAGED]] : (!fir.llvm_ptr<i8>) -> !fir.ref<!fir.llvm_ptr<i8>>
+// CTOR: %[[TYPEDESC:.*]] = fir.address_of(@_QMmE.dt.t) : !fir.ref<!fir.llvm_ptr<i8>>
+// CTOR: %[[VAL:.*]] = fir.convert %[[TYPEDESC]] : (!fir.ref<!fir.llvm_ptr<i8>>) -> !fir.llvm_ptr<i8>
+// CTOR: fir.store %[[VAL]] to %[[DEST]] : !fir.ref<!fir.llvm_ptr<i8>>
+// CTOR: llvm.br ^[[NEXT]]
+// CTOR: ^[[NEXT]]:
+// CTOR: llvm.return
+
+// CTOR: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor, @__cudaFortranInitConstructor], priorities = [0 : i32, 1 : i32]
More information about the flang-commits
mailing list