[llvm] [AMDGPU] Fix alias handling in module splitting functionality (PR #187295)

[Joseph Huber via llvm-commits]

https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/187295

Summary:
The module splitting used for `-flto-partitions=8` support (which is
passed by default) did not correctly handle aliases. We mainly need to
do two things: keep the aliases in the they are used in and externalize
them. Internalize linkage needs to be handled conservatively.

This is needed because these aliases show up in PGO contexts.


>From 396b70778a849b11c8f7e35d20d423bcdee0641a Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Wed, 18 Mar 2026 10:04:18 -0500
Subject: [PATCH] [AMDGPU] Fix alias handling in module splitting functionality

Summary:
The module splitting used for `-flto-partitions=8` support (which is
passed by default) did not correctly handle aliases. We mainly need to
do two things: keep the aliases in the they are used in and externalize
them. Internalize linkage needs to be handled conservatively.

This is needed because these aliases show up in PGO contexts.
---
 llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp  | 22 ++++++++++++++-----
 .../llvm-split/AMDGPU/kernels-dependencies.ll | 11 ++++++++--
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp
index c155d85df1769..105ce9a245f5c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSplitModule.cpp
@@ -43,6 +43,7 @@
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
@@ -1285,7 +1286,9 @@ namespace {
 static bool needsConservativeImport(const GlobalValue *GV) {
   if (const auto *Var = dyn_cast<GlobalVariable>(GV))
     return Var->hasLocalLinkage();
-  return isa<GlobalAlias>(GV);
+  if (const auto *GA = dyn_cast<GlobalAlias>(GV))
+    return GA->hasLocalLinkage();
+  return false;
 }
 
 /// Prints a summary of the partition \p N, represented by module \p M, to \p
@@ -1394,8 +1397,6 @@ static void splitAMDGPUModule(
   // visible copy, then internalize all other copies" for some functions?
   if (!NoExternalizeOnAddrTaken) {
     for (auto &Fn : M) {
-      // TODO: Should aliases count? Probably not but they're so rare I'm not
-      // sure it's worth fixing.
       if (Fn.hasLocalLinkage() && Fn.hasAddressTaken()) {
         LLVM_DEBUG(dbgs() << "[externalize] "; Fn.printAsOperand(dbgs());
                    dbgs() << " because its address is taken\n");
@@ -1414,6 +1415,13 @@ static void splitAMDGPUModule(
     }
   }
 
+  for (auto &GA : M.aliases()) {
+    if (GA.hasLocalLinkage()) {
+      LLVM_DEBUG(dbgs() << "[externalize] alias " << GA.getName() << '\n');
+      externalize(GA);
+    }
+  }
+
   // Start by calculating the cost of every function in the module, as well as
   // the module's overall cost.
   FunctionsCostMap FnCosts;
@@ -1510,15 +1518,17 @@ static void splitAMDGPUModule(
             return false;
           }
 
+          // Aliases should not be separated from their underlying object.
+          if (const auto *GA = dyn_cast<GlobalAlias>(GV))
+            if (const auto *Fn = dyn_cast<Function>(GA->getAliaseeObject()))
+              return FnsInPart.contains(Fn);
+
           // Everything else goes in the first non-empty module we create.
           return ImportAllGVs || needsConservativeImport(GV);
         }));
 
     ImportAllGVs = false;
 
-    // FIXME: Aliases aren't seen often, and their handling isn't perfect so
-    // bugs are possible.
-
     // Clean-up conservatively imported GVs without any users.
     for (auto &GV : make_early_inc_range(MPart->global_values())) {
       if (needsConservativeImport(&GV) && GV.use_empty())
diff --git a/llvm/test/tools/llvm-split/AMDGPU/kernels-dependencies.ll b/llvm/test/tools/llvm-split/AMDGPU/kernels-dependencies.ll
index bea527f15bbaa..2e555ad0266eb 100644
--- a/llvm/test/tools/llvm-split/AMDGPU/kernels-dependencies.ll
+++ b/llvm/test/tools/llvm-split/AMDGPU/kernels-dependencies.ll
@@ -4,23 +4,30 @@
 ; RUN: llvm-dis -o - %t2 | FileCheck --check-prefix=CHECK2 %s
 
 ; 3 kernels with each their own dependencies should go into 3
-; distinct partitions.
+; distinct partitions. Aliases should follow their aliasee's partition
+; and not be cleaned up even if unused in that partition.
 
+; CHECK0-NOT: @HelperA_alias = {{.*}}alias
 ; CHECK0-NOT: define
 ; CHECK0: define amdgpu_kernel void @C
 ; CHECK0: define internal void @HelperC
+; CHECK0-NOT: @HelperA_alias = {{.*}}alias
 ; CHECK0-NOT: define
 
+; CHECK1-NOT: @HelperA_alias = {{.*}}alias
 ; CHECK1-NOT: define
 ; CHECK1: define amdgpu_kernel void @B
 ; CHECK1: define internal void @HelperB
+; CHECK1-NOT: @HelperA_alias = {{.*}}alias
 ; CHECK1-NOT: define
 
+; CHECK2: @HelperA_alias = hidden alias void (), ptr @HelperA
 ; CHECK2-NOT: define
 ; CHECK2: define amdgpu_kernel void @A
-; CHECK2: define internal void @HelperA
+; CHECK2: define hidden void @HelperA
 ; CHECK2-NOT: define
 
+ at HelperA_alias = internal alias void (), ptr @HelperA
 
 define amdgpu_kernel void @A() {
   call void @HelperA()