[flang-commits] [flang] [flang][cuda] Register allocator only when needed (PR #208355)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval created https://github.com/llvm/llvm-project/pull/208355

Register allocator only when needed to avoid bringing in symbol when not strictly needed. 

>From 2c0f00887d75c82b3eb852c4e53a203a74bfb2cc Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 8 Jul 2026 17:00:02 -0700
Subject: [PATCH] [flang][cuda] Register allocator only when needed

---
 .../Transforms/CUDA/CUFAddConstructor.cpp     | 36 +++++++++++++------
 flang/test/Fir/CUDA/cuda-constructor-2.f90    |  5 ++-
 flang/test/Fir/CUDA/cuda-constructor.f90      |  6 ++--
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
index 5e8136af7991e..73129c88b0b19 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFAddConstructor.cpp
@@ -208,21 +208,37 @@ struct CUFAddConstructor
                           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();
+    });
 
     // 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 bb2a98c294558..addbb737748dc 100644
--- a/flang/test/Fir/CUDA/cuda-constructor-2.f90
+++ b/flang/test/Fir/CUDA/cuda-constructor-2.f90
@@ -49,7 +49,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
 
 // 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"} {
@@ -170,7 +170,6 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
 }
 
 // CHECK: llvm.func internal @__cudaFortranConstructor()
-// CHECK-NEXT: llvm.call @_FortranACUFRegisterAllocator()
 // CHECK-NEXT: llvm.return
 // CHECK: llvm.mlir.global_ctors ctors = [@__cudaFortranConstructor]
 
@@ -198,7 +197,7 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr, dense<
 }
 
 // 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 1ba3fdc4a6b73..e2f2ba88e9205 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]