[llvm] [ElimAvailExtern] Add an option to allow to convert global variables in a specified address space to local (PR #144287)

[Shilei Tian via llvm-commits]

https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/144287

>From 354afb9ce54da0e7addcc7d855d4949248b80b4d Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Mon, 16 Jun 2025 22:05:53 -0400
Subject: [PATCH] [ElimAvailExtern] Add an option to allow to convert global
 variables in a specified address space to local

Currently, the `EliminateAvailableExternallyPass` only converts certain
available externally functions to local. For global variables, it only drops
their initializers.

This PR adds an option to allow the pass to convert global variables in a
specified address space to local. The motivation for this change is to correctly
support lowering of LDS variables (`__shared__` variables, in more generic
terminology) when ThinLTO is enabled for AMDGPU.

A `__shared__` variable is lowered to a hidden global variable in a particular
address space by the frontend, which is roughly same as a `static` local
variable. To properly lower it in the backend, the compiler needs to check all
its uses. Enabling ThinLTO currently breaks this when a function containing a
`__shared__` variable is imported from another module. Even though the global
variable is imported along with its associated function, and the function is
privatized by the `EliminateAvailableExternallyPass`, the global variable itself
is not.

It's safe to privatize such global variables, because they're _local_ to their
associated functions. If the function itself is privatized, its associated
global variables should also be privatized accordingly.
---
 llvm/lib/Transforms/IPO/ElimAvailExtern.cpp   | 36 ++++++++++++++++---
 .../convert-global-variables-to-local.ll      | 21 +++++++++++
 2 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/EliminateAvailableExternally/convert-global-variables-to-local.ll

diff --git a/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp b/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
index 718452fc02764..bc98f994f490c 100644
--- a/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
+++ b/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
@@ -35,8 +35,15 @@ static cl::opt<bool> ConvertToLocal(
     cl::desc("Convert available_externally into locals, renaming them "
              "to avoid link-time clashes."));
 
+static cl::opt<unsigned> ConvertGlobalVariableInAddrSpace(
+    "avail-extern-gv-in-addrspace-to-local", cl::Hidden,
+    cl::desc(
+        "Convert available_externally global variables into locals if they are "
+        "in specificed addrspace, renaming them to avoid link-time clashes."));
+
 STATISTIC(NumRemovals, "Number of functions removed");
-STATISTIC(NumConversions, "Number of functions converted");
+STATISTIC(NumFunctionsConverted, "Number of functions converted");
+STATISTIC(NumGlobalVariablesConverted, "Number of global variables converted");
 STATISTIC(NumVariables, "Number of global variables removed");
 
 void deleteFunction(Function &F) {
@@ -45,6 +52,10 @@ void deleteFunction(Function &F) {
   ++NumRemovals;
 }
 
+static std::string getNewName(Module &M, const GlobalValue &GV) {
+  return GV.getName().str() + ".__uniq" + getUniqueModuleId(&M);
+}
+
 /// Create a copy of the thinlto import, mark it local, and redirect direct
 /// calls to the copy. Only direct calls are replaced, so that e.g. indirect
 /// call function pointer tests would use the global identity of the function.
@@ -68,7 +79,7 @@ static void convertToLocalCopy(Module &M, Function &F) {
   // functions with the same name, but that just creates more trouble than
   // necessary e.g. distinguishing profiles or debugging. Instead, we append the
   // module identifier.
-  auto NewName = OrigName + ".__uniq" + getUniqueModuleId(&M);
+  std::string NewName = getNewName(M, F);
   F.setName(NewName);
   if (auto *SP = F.getSubprogram())
     SP->replaceLinkageName(MDString::get(F.getParent()->getContext(), NewName));
@@ -85,16 +96,33 @@ static void convertToLocalCopy(Module &M, Function &F) {
                        F.getAddressSpace(), OrigName, F.getParent());
   F.replaceUsesWithIf(Decl,
                       [&](Use &U) { return !isa<CallBase>(U.getUser()); });
-  ++NumConversions;
+  ++NumFunctionsConverted;
+}
+
+/// Similar to the function above, this is to convert an externally available
+/// global variable to local.
+static void convertToLocalCopy(Module &M, GlobalVariable &GV) {
+  assert(GV.hasAvailableExternallyLinkage());
+  GV.setName(getNewName(M, GV));
+  GV.setLinkage(GlobalValue::InternalLinkage);
+  ++NumGlobalVariablesConverted;
 }
 
 static bool eliminateAvailableExternally(Module &M, bool Convert) {
   bool Changed = false;
 
-  // Drop initializers of available externally global variables.
+  // If a global variable is available externally and in the specified address
+  // space, convert it to local linkage; otherwise, drop its initializer.
   for (GlobalVariable &GV : M.globals()) {
     if (!GV.hasAvailableExternallyLinkage())
       continue;
+    if (ConvertGlobalVariableInAddrSpace.getNumOccurrences() &&
+        GV.getAddressSpace() == ConvertGlobalVariableInAddrSpace &&
+        !GV.use_empty()) {
+      convertToLocalCopy(M, GV);
+      Changed = true;
+      continue;
+    }
     if (GV.hasInitializer()) {
       Constant *Init = GV.getInitializer();
       GV.setInitializer(nullptr);
diff --git a/llvm/test/Transforms/EliminateAvailableExternally/convert-global-variables-to-local.ll b/llvm/test/Transforms/EliminateAvailableExternally/convert-global-variables-to-local.ll
new file mode 100644
index 0000000000000..6995b97e79887
--- /dev/null
+++ b/llvm/test/Transforms/EliminateAvailableExternally/convert-global-variables-to-local.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
+; RUN: opt -S -passes=elim-avail-extern -avail-extern-gv-in-addrspace-to-local=3 %s -o - | FileCheck %s
+
+ at shared = internal addrspace(3) global i32 undef, align 4
+ at shared.imported = available_externally hidden unnamed_addr addrspace(3) global i32 undef, align 4
+
+;.
+; CHECK: @shared = internal addrspace(3) global i32 undef, align 4
+; CHECK: @shared.imported.__uniq.[[UUID:.*]] = internal unnamed_addr addrspace(3) global i32 undef, align 4
+;.
+define void @foo(i32 %v) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i32 [[V:%.*]]) {
+; CHECK-NEXT:    store i32 [[V]], ptr addrspace(3) @shared, align 4
+; CHECK-NEXT:    store i32 [[V]], ptr addrspace(3) @shared.imported.__uniq.[[UUID]], align 4
+; CHECK-NEXT:    ret void
+;
+  store i32 %v, ptr addrspace(3) @shared, align 4
+  store i32 %v, ptr addrspace(3) @shared.imported, align 4
+  ret void
+}