[flang-commits] [flang] 98ba607 - [flang][cuda] initialize the CUDA module after all units register managed variables (#230350)
via flang-commits
flang-commits at lists.llvm.org
Fri Oct 9 11:53:38 PDT 2026
Author: Valentin Clement (バレンタイン クレメン)
Date: 2026-10-09T11:53:31-07:00
New Revision: 98ba6072a735aa6d9ea0d475c613506f293a513c
URL: https://github.com/llvm/llvm-project/commit/98ba6072a735aa6d9ea0d475c613506f293a513c
DIFF: https://github.com/llvm/llvm-project/commit/98ba6072a735aa6d9ea0d475c613506f293a513c.diff
LOG: [flang][cuda] initialize the CUDA module after all units register managed variables (#230350)
With relocatable device code, every unit registers its managed variables
on the same CUDA module, and the runtime does not fill variables
registered after the module is initialized. Move CUFInitModule to a
second constructor with the next priority value, so that it runs after
the registration constructors of every unit in the executable or shared
library.
This will fix segmentation fault in example like:
```
module m1
integer, managed :: x1
contains
attributes(global) subroutine k1()
x1 = 1
end subroutine
end module
```
```
module m2
integer, managed :: x2
contains
attributes(global) subroutine k2()
x2 = 2
end subroutine
end module
program p
use cudafor
use m1
use m2
integer :: s
call k1<<<1,1>>>()
call k2<<<1,1>>>()
s = cudaDeviceSynchronize()
print *, x1, x2
end program
```
Added:
Modified:
flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
flang/test/Fir/CUDA/cuda-constructor-2.f90
flang/test/Fir/CUDA/cuda-managed-pointer-linkage.cuf
Removed:
################################################################################
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
index 3acc1de16c638..3fa91e7214e0c 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
@@ -80,6 +80,10 @@ namespace {
static constexpr llvm::StringRef cudaFortranCtorName{
"__cudaFortranConstructor"};
+static constexpr llvm::StringRef cudaFortranInitCtorName{
+ "__cudaFortranInitConstructor"};
+static constexpr llvm::StringRef cudaModuleHandleName{
+ "__cudaFortranModuleHandle"};
static constexpr llvm::StringRef managedPtrSuffix{".managed.ptr"};
static constexpr llvm::StringRef cudaCompiledSymbolName{"Mcuda_compiled"};
@@ -392,6 +396,7 @@ struct CUFAddConstructor
func.setLinkage(mlir::LLVM::Linkage::Internal);
auto entryBlock = func.addEntryBlock(builder);
builder.setInsertionPointToStart(entryBlock);
+ mlir::LLVM::LLVMFuncOp initCtor;
if (needAllocatorRegistration) {
llvm::StringRef allocatorRegistrationFunctionName =
@@ -530,14 +535,47 @@ struct CUFAddConstructor
if (hasNonAllocManagedGlobal) {
// Initialize the module after all variables are registered so the
- // runtime populates managed variable unified memory pointers.
+ // runtime populates managed variable unified memory pointers. With
+ // relocatable device code, every unit registers its variables with
+ // the same module, and the runtime ignores the ones registered after
+ // the module is initialized. The initialization therefore runs in a
+ // second constructor with the next priority value, after the
+ // registration constructors of every unit in the same executable or
+ // shared library.
+ mlir::OpBuilder::InsertionGuard guard(builder);
+ builder.setInsertionPointToEnd(mod.getBody());
+ auto handleGlobal = mlir::LLVM::GlobalOp::create(
+ builder, loc, llvmPtrTy, /*isConstant=*/false,
+ mlir::LLVM::Linkage::Internal, cudaModuleHandleName,
+ mlir::Attribute{});
+ builder.createBlock(&handleGlobal.getInitializerRegion());
+ mlir::Value nullHandle =
+ mlir::LLVM::ZeroOp::create(builder, loc, llvmPtrTy);
+ mlir::LLVM::ReturnOp::create(builder, loc, nullHandle);
+
+ builder.setInsertionPointToEnd(mod.getBody());
+ initCtor = mlir::LLVM::LLVMFuncOp::create(
+ builder, loc, cudaFortranInitCtorName, funcTy);
+ initCtor.setLinkage(mlir::LLVM::Linkage::Internal);
+ mlir::Block *initBlock = initCtor.addEntryBlock(builder);
+
+ builder.setInsertionPointToEnd(entryBlock);
+ mlir::LLVM::StoreOp::create(
+ builder, loc, registeredMod,
+ mlir::LLVM::AddressOfOp::create(builder, loc, handleGlobal));
+
+ builder.setInsertionPointToStart(initBlock);
+ mlir::Value handle = mlir::LLVM::LoadOp::create(
+ builder, loc, llvmPtrTy,
+ mlir::LLVM::AddressOfOp::create(builder, loc, handleGlobal));
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)};
+ llvm::SmallVector<mlir::Value> initArgs{
+ fir::runtime::createArguments(builder, loc, initFTy, handle)};
fir::CallOp::create(builder, loc, initFunc, initArgs);
+ mlir::LLVM::ReturnOp::create(builder, loc, mlir::ValueRange{});
}
}
}
@@ -582,6 +620,12 @@ struct CUFAddConstructor
llvm::SmallVector<mlir::Attribute> data;
priorities.push_back(priority);
data.push_back(mlir::LLVM::ZeroAttr::get(mod.getContext()));
+ if (initCtor) {
+ funcs.push_back(mlir::FlatSymbolRefAttr::get(mod.getContext(),
+ initCtor.getSymName()));
+ priorities.push_back(priority + 1);
+ 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));
diff --git a/flang/test/Fir/CUDA/cuda-constructor-2.f90 b/flang/test/Fir/CUDA/cuda-constructor-2.f90
index 6465a217a7572..70a3066eedbe6 100644
--- a/flang/test/Fir/CUDA/cuda-constructor-2.f90
+++ b/flang/test/Fir/CUDA/cuda-constructor-2.f90
@@ -151,11 +151,20 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
// CHECK: fir.global external @_QMtestEmanx.managed.ptr <{section = "__nv_managed_data__"}> : !fir.llvm_ptr<i8>
// CHECK: fir.zero_bits !fir.llvm_ptr<i8>
-// Constructor should register with CUFRegisterManagedVariable then init module.
+// The constructor registers with CUFRegisterManagedVariable. The module is
+// initialized by a second constructor, which runs after the registration
+// constructors of all the units.
// CHECK: llvm.func internal @__cudaFortranConstructor()
// CHECK: fir.address_of(@_QMtestEmanx.managed.ptr) : !fir.ref<!fir.llvm_ptr<i8>>
// CHECK: fir.call @_FortranACUFRegisterManagedVariable
+// CHECK-NOT: fir.call @_FortranACUFInitModule
+// CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm.ptr, !llvm.ptr
+// CHECK: llvm.return
+// CHECK: llvm.func internal @__cudaFortranInitConstructor()
+// CHECK: llvm.load
// CHECK: fir.call @_FortranACUFInitModule
+// CHECK: llvm.return
+// CHECK: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor, @__cudaFortranInitConstructor], priorities = [0 : i32, 1 : i32]
// -----
diff --git a/flang/test/Fir/CUDA/cuda-managed-pointer-linkage.cuf b/flang/test/Fir/CUDA/cuda-managed-pointer-linkage.cuf
index ec474bf92d803..db73b75a96544 100644
--- a/flang/test/Fir/CUDA/cuda-managed-pointer-linkage.cuf
+++ b/flang/test/Fir/CUDA/cuda-managed-pointer-linkage.cuf
@@ -33,6 +33,7 @@ end module
! DEFINING: llvm.func internal @__cudaFortranConstructor()
! DEFINING: fir.address_of(@_QMmEg.managed.ptr) : !fir.ref<!fir.llvm_ptr<i8>>
! DEFINING: fir.call @_FortranACUFRegisterManagedVariable
+! DEFINING: llvm.func internal @__cudaFortranInitConstructor()
! DEFINING: fir.call @_FortranACUFInitModule
!--- useonly.cuf
@@ -51,3 +52,4 @@ end subroutine
! USEONLY: llvm.func internal @__cudaFortranConstructor()
! USEONLY-NOT: fir.call @_FortranACUFRegisterManagedVariable
! USEONLY-NOT: fir.call @_FortranACUFInitModule
+! USEONLY-NOT: __cudaFortranInitConstructor
More information about the flang-commits
mailing list