[flang-commits] [flang] [llvm] [flang][cuda] Use cuf.register_variable_static instead of CUFRegisterExternalVariable (PR #211886)
Valentin Clement バレンタイン クレメン via flang-commits
flang-commits at lists.llvm.org
Fri Jul 24 11:51:23 PDT 2026
https://github.com/clementval created https://github.com/llvm/llvm-project/pull/211886
This gives more flexibility to the backend to select a different registration mechanism
>From f7c3e2ec27b99c07c6219fcc322acadfd7b97a5d Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Fri, 24 Jul 2026 11:48:43 -0700
Subject: [PATCH] [flang][cuda] Use cuf.register_variable_static instead of
CUFRegisterExternalVariable
---
flang-rt/lib/cuda/registration.cpp | 8 --------
.../flang/Optimizer/Transforms/Passes.td | 2 +-
.../include/flang/Runtime/CUDA/registration.h | 7 -------
.../Transforms/CUDA/CUFAddConstructor.cpp | 18 +++++++++---------
.../Transforms/CUDA/CUFDeviceGlobal.cpp | 2 +-
flang/test/Fir/CUDA/cuda-constructor-2.f90 | 12 ++++++------
.../Fir/CUDA/cuda-unified-module-global.f90 | 5 ++---
7 files changed, 19 insertions(+), 35 deletions(-)
diff --git a/flang-rt/lib/cuda/registration.cpp b/flang-rt/lib/cuda/registration.cpp
index a35284a2460dc..1910b149653a9 100644
--- a/flang-rt/lib/cuda/registration.cpp
+++ b/flang-rt/lib/cuda/registration.cpp
@@ -48,14 +48,6 @@ void RTDEF(CUFRegisterVariable)(
__cudaRegisterVar(module, varSym, varName, varName, 0, size, 0, 0);
}
-void RTDEF(CUFRegisterExternalVariable)(
- void **module, char *varSym, const char *varName, int64_t size) {
- // Tell the CUDA driver to bind the device-side global <varName> to the
- // host-resident storage at <varSym>. Kernel accesses to <varName> then go
- // through the host address; HMM/ATS handles migration.
- __cudaRegisterHostVar(module, varName, varSym, size);
-}
-
void RTDEF(CUFRegisterManagedVariable)(
void **module, void **varSym, char *varName, int64_t size) {
__cudaRegisterManagedVar(module, varSym, varName, varName, 0, size, 0, 0);
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index db02fa0180ae5..8573c0e4f3f00 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -556,7 +556,7 @@ def CUFAddConstructor : Pass<"cuf-add-constructor", "mlir::ModuleOp"> {
Option<"cudaUnified", "cuda-unified", "bool", /*default=*/"false",
"Treat host module globals as unified memory (-gpu=mem:unified): "
"register every plain host module global that is mirrored in the "
- "GPU module via CUFRegisterExternalVariable, so the CUDA driver "
+ "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.">
];
diff --git a/flang/include/flang/Runtime/CUDA/registration.h b/flang/include/flang/Runtime/CUDA/registration.h
index 48a74dc1d6d03..74dbf9e189076 100644
--- a/flang/include/flang/Runtime/CUDA/registration.h
+++ b/flang/include/flang/Runtime/CUDA/registration.h
@@ -28,13 +28,6 @@ void RTDECL(CUFRegisterFunction)(
void RTDECL(CUFRegisterVariable)(
void **module, char *varSym, const char *varName, int64_t size);
-/// Register a module-scope variable as host-resident under -gpu=mem:unified,
-/// so that the device-side symbol of the same name is mapped to the host
-/// pointer at module-load time. Wraps __cudaRegisterHostVar. Kernel accesses
-/// to the variable then reach the host storage directly via HMM/ATS.
-void RTDECL(CUFRegisterExternalVariable)(
- void **module, char *varSym, const char *varName, int64_t size);
-
/// Register a managed variable.
void RTDECL(CUFRegisterManagedVariable)(
void **module, void **varSym, char *varName, int64_t size);
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
index 6f2c17db5eb1d..c8f3a6f738e34 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
@@ -77,7 +77,7 @@ static fir::GlobalOp createManagedPointerGlobal(fir::FirOpBuilder &builder,
/// 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
+/// registered with the CUDA driver via cuf.register_variable_static 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,
@@ -173,8 +173,7 @@ static uint64_t getGlobalSizeInBytes(mlir::Location loc,
/// 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.
+/// Used both for CUFRegisterVariable / CUFRegisterManagedVariable.
static void
emitCUFRegistrationCall(fir::FirOpBuilder &builder, mlir::Location loc,
mlir::Type idxTy, const mlir::DataLayout &dl,
@@ -410,12 +409,13 @@ struct CUFAddConstructor
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);
+ 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));
}
}
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceGlobal.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceGlobal.cpp
index ba0a6aec9992a..2e1dad7034621 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceGlobal.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFDeviceGlobal.cpp
@@ -249,7 +249,7 @@ class CUFDeviceGlobal : public fir::impl::CUFDeviceGlobalBase<CUFDeviceGlobal> {
// External linkage (see convertLinkage in CodeGen.cpp), so an
// initializer-less global emits as `.extern .global ...` in PTX.
// The host-side definition stays. CUFAddConstructor will emit
- // CUFRegisterExternalVariable (= __cudaRegisterHostVar) so the CUDA
+ // cuf.register_variable_static so the CUDA
// runtime maps the device extern to the host pointer at module-load
// time, and HMM/ATS handles migration.
if (cudaUnified && !globalOp.getConstant() &&
diff --git a/flang/test/Fir/CUDA/cuda-constructor-2.f90 b/flang/test/Fir/CUDA/cuda-constructor-2.f90
index b1695ac7af72f..57bb1244bb364 100644
--- a/flang/test/Fir/CUDA/cuda-constructor-2.f90
+++ b/flang/test/Fir/CUDA/cuda-constructor-2.f90
@@ -211,7 +211,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
// Under -gpu=mem:unified, a non-allocatable host module global referenced from
// device code (mirrored as an external-linkage clone in the GPU module by the
// CUFDeviceGlobal pass) must be registered with
-// _FortranACUFRegisterExternalVariable so the device-side `.extern` symbol
+// cuf.register_variable_static so the device-side `.extern` symbol
// resolves to the host pointer at module-load time.
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"} {
@@ -229,14 +229,13 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<i8 = dense<8> : vector<2xi64>, i
// UNIFIED: llvm.func internal @__cudaFortranConstructor()
// UNIFIED: cuf.register_module @cuda_device_mod -> !llvm.ptr
-// UNIFIED: fir.address_of(@_QMmtestsEm) : !fir.ref<!fir.array<5xi32>>
-// UNIFIED: fir.call @_FortranACUFRegisterExternalVariable
+// UNIFIED: cuf.register_variable_static @_QMmtestsEm
// UNIFIED-NOT: fir.call @_FortranACUFInitModule
// -----
// Under -gpu=mem:unified, an allocatable host module global also gets
-// registered via _FortranACUFRegisterExternalVariable; the registered symbol
+// registered via cuf.register_variable_static; the registered symbol
// is the descriptor (fir.box<fir.heap<...>>). The host runtime allocates the
// data buffer in HMM/ATS-accessible memory.
@@ -257,12 +256,12 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<i8 = dense<8> : vector<2xi64>, i
// UNIFIED: llvm.func internal @__cudaFortranConstructor()
// UNIFIED: cuf.register_module @cuda_device_mod -> !llvm.ptr
-// UNIFIED: fir.address_of(@_QMmtestsEma) : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
-// UNIFIED: fir.call @_FortranACUFRegisterExternalVariable
+// UNIFIED: cuf.register_variable_static @_QMmtestsEma
// UNIFIED-NOT: fir.call @_FortranACUFInitModule
// -----
+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"} {
func.func @_QPsub1() {
%0 = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFsub1Ea"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
%4:2 = hlfir.declare %0 {data_attr = #cuf.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub1Ea"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
@@ -274,6 +273,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<i8 = dense<8> : vector<2xi64>, i
cuf.free %4#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>}
return
}
+}
// CHECK: llvm.func internal @__cudaFortranConstructor()
// CHECK: llvm.call @_FortranACUFRegisterAllocator()
diff --git a/flang/test/Fir/CUDA/cuda-unified-module-global.f90 b/flang/test/Fir/CUDA/cuda-unified-module-global.f90
index fc34c9c4686e6..4ea4500322c77 100644
--- a/flang/test/Fir/CUDA/cuda-unified-module-global.f90
+++ b/flang/test/Fir/CUDA/cuda-unified-module-global.f90
@@ -3,7 +3,7 @@
// 1. is mirrored into the GPU module by CUFDeviceGlobal as a no-body
// external declaration (so PTX gets `.extern .global ...`); and
// 2. is registered with the CUDA driver via
-// _FortranACUFRegisterExternalVariable (= __cudaRegisterHostVar) from
+// cuf.register_variable_static from
// __cudaFortranConstructor, so the device-side symbol is mapped to
// the host-resident storage at module-load time and HMM/ATS handles
// migration.
@@ -44,6 +44,5 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<i8 = dense<8> : vector<2xi64>, i
// Constructor registers the host pointer.
// CHECK: llvm.func internal @__cudaFortranConstructor()
// CHECK: cuf.register_module @cuda_device_mod -> !llvm.ptr
-// CHECK: fir.address_of(@_QMmtestsEm) : !fir.ref<!fir.array<5xi32>>
-// CHECK: fir.call @_FortranACUFRegisterExternalVariable
+// CHECK: cuf.register_variable_static @_QMmtestsEm
// CHECK-NOT: fir.call @_FortranACUFInitModule
More information about the flang-commits
mailing list