[llvm] [AMDGPU][Attributor] Add an option to turn on internalization (PR #108420)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/108420

None

>From decc6851c102705d3306a9b9099c735651e225cc Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 12 Sep 2024 12:21:31 -0400
Subject: [PATCH] [AMDGPU][Attributor] Add an option to turn on internalization

---
 llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp | 37 ++++++++++++++++++---
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
index ffeec31bb930ab..b00474081d6dee 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
@@ -38,6 +38,11 @@ static cl::opt<unsigned> IndirectCallSpecializationThreshold(
         "A threshold controls whether an indirect call will be specialized"),
     cl::init(3));
 
+static cl::opt<bool>
+    DisableInternalization("amdgpu-disable-internalization",
+                           cl::desc("Disable function internalization."),
+                           cl::Hidden, cl::init(true));
+
 #define AMDGPU_ATTRIBUTE(Name, Str) Name##_POS,
 
 enum ImplicitArgumentPositions {
@@ -1031,9 +1036,33 @@ static void addPreloadKernArgHint(Function &F, TargetMachine &TM) {
 
 static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
                     AMDGPUAttributorOptions Options) {
+  bool Changed = false;
+
+  DenseMap<Function *, Function *> InternalizedMap;
+  if (!DisableInternalization) {
+    auto IsCalled = [](Function &F) {
+      for (const User *U : F.users())
+        if (!isa<BlockAddress>(U))
+          return true;
+      return false;
+    };
+
+    SmallPtrSet<Function *, 16> InternalizeFns;
+    for (Function &F : M) {
+      if (!F.isDeclaration() && AMDGPU::isEntryFunctionCC(F.getCallingConv()) &&
+          IsCalled(F)) {
+        if (Attributor::isInternalizable(F))
+          InternalizeFns.insert(&F);
+      }
+    }
+
+    Changed |=
+        Attributor::internalizeFunctions(InternalizeFns, InternalizedMap);
+  }
+
   SetVector<Function *> Functions;
   for (Function &F : M) {
-    if (!F.isIntrinsic())
+    if (!F.isIntrinsic() && !F.isDeclaration() && !InternalizedMap.lookup(&F))
       Functions.insert(&F);
   }
 
@@ -1066,7 +1095,7 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
   Attributor A(Functions, InfoCache, AC);
 
   for (Function &F : M) {
-    if (F.isIntrinsic())
+    if (F.isIntrinsic() || InternalizedMap.lookup(&F))
       continue;
 
     A.getOrCreateAAFor<AAAMDAttributes>(IRPosition::function(F));
@@ -1097,8 +1126,8 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
     }
   }
 
-  ChangeStatus Change = A.run();
-  return Change == ChangeStatus::CHANGED;
+  Changed |= (A.run() == ChangeStatus::CHANGED);
+  return Changed;
 }
 
 class AMDGPUAttributorLegacy : public ModulePass {