[flang][cuda] Register allocator only when needed (#208355)
Register allocator only when needed to avoid bringing in symbol when not
strictly needed.
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
index 44654ad..47f4acd 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
@@ -254,21 +254,51 @@
getName() + "pass");
}
- // 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.setInsertionPointToEnd(mod.getBody());
+ 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);
- builder.setInsertionPointToStart(func.addEntryBlock(builder));
- mlir::LLVM::CallOp::create(builder, loc, funcTy, cufRegisterAllocatorRef);
+ 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) {
diff --git a/flang/test/Fir/CUDA/cuda-constructor-2.f90 b/flang/test/Fir/CUDA/cuda-constructor-2.f90
index 26ebcf7..c2a1839 100644
--- a/flang/test/Fir/CUDA/cuda-constructor-2.f90
+++ b/flang/test/Fir/CUDA/cuda-constructor-2.f90
@@ -53,7 +53,7 @@
// Checking that constant global variables are not registered
-// CHECK: @_FortranACUFRegisterAllocator
+// CHECK: @__cudaFortranConstructor()
// CHECK-NOT: fir.call @_FortranACUFRegisterVariable
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.ident = "flang version 20.0.0 (https://github.com/llvm/llvm-project.git 3372303188df0f7f8ac26e7ab610cf8b0f716d42)", llvm.target_triple = "x86_64-unknown-linux-gnu"} {
@@ -174,7 +174,6 @@
}
// CHECK: llvm.func internal @__cudaFortranConstructor()
-// CHECK-NEXT: llvm.call @_FortranACUFRegisterAllocator()
// CHECK-NEXT: llvm.return
// CHECK: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor]
@@ -202,7 +201,7 @@
}
// CHECK: llvm.func internal @__cudaFortranConstructor()
-// CHECK: llvm.call @_FortranACUFRegisterAllocator()
+// CHECK-NOT: llvm.call @_FortranACUFRegisterAllocator()
// CHECK: cuf.register_module @cuda_device_mod -> !llvm.ptr
// CHECK: fir.address_of(@_QMkernels_mEdev_var) : !fir.ref<f32>
// CHECK: fir.call @_FortranACUFRegisterVariable(%3, %4, %5, %6) : (!fir.ref<!fir.llvm_ptr<i8>>, !fir.ref<i8>, !fir.ref<i8>, i64) -> ()
diff --git a/flang/test/Fir/CUDA/cuda-constructor.f90 b/flang/test/Fir/CUDA/cuda-constructor.f90
index 1ba3fdc..e2f2ba8 100644
--- a/flang/test/Fir/CUDA/cuda-constructor.f90
+++ b/flang/test/Fir/CUDA/cuda-constructor.f90
@@ -1,12 +1,12 @@
-! RUN: bbc -fcuda -emit-hlfir %s -o - | fir-opt --cuf-add-constructor | FileCheck %s
+! RUN: bbc -fcuda -emit-fir %s -o - | fir-opt --cuf-add-constructor | FileCheck %s
program main
- real, device :: ahost(10)
+ real, allocatable, device :: ahost(:)
end
-! CHECK: llvm.func @_FortranACUFRegisterAllocator() attributes {sym_visibility = "private"}
! CHECK-LABEL: llvm.func internal @__cudaFortranConstructor() {
! CHECK: llvm.call @_FortranACUFRegisterAllocator() : () -> ()
! CHECK: llvm.return
! CHECK: }
+! CHECK: llvm.func @_FortranACUFRegisterAllocator() attributes {sym_visibility = "private"}
! CHECK: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor], priorities = [0 : i32], data = [#llvm.zero]
diff --git a/flang/test/Fir/CUDA/cuda-register-allocator-0.f90 b/flang/test/Fir/CUDA/cuda-register-allocator-0.f90
new file mode 100644
index 0000000..cb40525
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-register-allocator-0.f90
@@ -0,0 +1,8 @@
+! RUN: bbc -fcuda -emit-fir %s -o - | fir-opt --cuf-add-constructor | FileCheck %s
+
+module m0
+ integer, device, allocatable :: a(:)
+end module
+
+! CHECK-LABEL: llvm.func internal @__cudaFortranConstructor() {
+! CHECK: llvm.call @_FortranACUFRegisterAllocator()
diff --git a/flang/test/Fir/CUDA/cuda-register-allocator-1.f90 b/flang/test/Fir/CUDA/cuda-register-allocator-1.f90
new file mode 100644
index 0000000..321154c
--- /dev/null
+++ b/flang/test/Fir/CUDA/cuda-register-allocator-1.f90
@@ -0,0 +1,8 @@
+! RUN: bbc -fcuda -emit-fir %s -o - | fir-opt --cuf-add-constructor | FileCheck %s
+
+module m0
+ integer, allocatable :: a(:)
+end module
+
+! CHECK-LABEL: llvm.func internal @__cudaFortranConstructor() {
+! CHECK-NOT: llvm.call @_FortranACUFRegisterAllocator()
