[llvm] [AMDGPU][Attributor] Add `ThinOrFullLTOPhase` as an argument (PR #123994)

[Shilei Tian via llvm-commits]

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

>From c3685eb50f76fd91057941c9bdbba38da4633dd2 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 22 Jan 2025 13:17:36 -0500
Subject: [PATCH] [AMDGPU][Attributor] Add `ThinOrFullLTOPhase` as an argument

---
 llvm/lib/Target/AMDGPU/AMDGPU.h               |  7 ++--
 llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp   | 34 +++++++++++++++----
 .../lib/Target/AMDGPU/AMDGPUTargetMachine.cpp |  9 +++--
 .../LTO/AMDGPU/closed-world-assumption.ll     |  8 +++--
 4 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPU.h b/llvm/lib/Target/AMDGPU/AMDGPU.h
index eebc33aea7a86..84995136e3dff 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPU.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPU.h
@@ -339,9 +339,12 @@ class AMDGPUAttributorPass : public PassInfoMixin<AMDGPUAttributorPass> {
 
   AMDGPUAttributorOptions Options;
 
+  const ThinOrFullLTOPhase LTOPhase;
+
 public:
-  AMDGPUAttributorPass(TargetMachine &TM, AMDGPUAttributorOptions Options = {})
-      : TM(TM), Options(Options) {};
+  AMDGPUAttributorPass(TargetMachine &TM, AMDGPUAttributorOptions Options,
+                       ThinOrFullLTOPhase LTOPhase = ThinOrFullLTOPhase::None)
+      : TM(TM), Options(Options), LTOPhase(LTOPhase) {};
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
index 0cee3c3cb5e92..97c8ca15a39b8 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAttributor.cpp
@@ -1337,7 +1337,8 @@ static void addPreloadKernArgHint(Function &F, TargetMachine &TM) {
 }
 
 static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
-                    AMDGPUAttributorOptions Options) {
+                    AMDGPUAttributorOptions Options,
+                    ThinOrFullLTOPhase LTOPhase) {
   SetVector<Function *> Functions;
   for (Function &F : M) {
     if (!F.isIntrinsic())
@@ -1372,9 +1373,27 @@ static bool runImpl(Module &M, AnalysisGetter &AG, TargetMachine &TM,
 
   Attributor A(Functions, InfoCache, AC);
 
-  LLVM_DEBUG(dbgs() << "[AMDGPUAttributor] Module " << M.getName() << " is "
-                    << (AC.IsClosedWorldModule ? "" : "not ")
-                    << "assumed to be a closed world.\n");
+  LLVM_DEBUG({
+    auto PhaseToString = [](ThinOrFullLTOPhase LTOPhase) -> StringRef {
+      switch (LTOPhase) {
+      case ThinOrFullLTOPhase::None:
+        return "None";
+      case ThinOrFullLTOPhase::ThinLTOPreLink:
+        return "ThinLTOPreLink";
+      case ThinOrFullLTOPhase::ThinLTOPostLink:
+        return "ThinLTOPostLink";
+      case ThinOrFullLTOPhase::FullLTOPreLink:
+        return "FullLTOPreLink";
+      case ThinOrFullLTOPhase::FullLTOPostLink:
+        return "FullLTOPostLink";
+      }
+    };
+    StringRef LTOPhaseStr = PhaseToString(LTOPhase);
+    dbgs() << "[AMDGPUAttributor] Running at phase " << LTOPhaseStr << '\n'
+           << "[AMDGPUAttributor] Module " << M.getName() << " is "
+           << (AC.IsClosedWorldModule ? "" : "not ")
+           << "assumed to be a closed world.\n";
+  });
 
   for (auto *F : Functions) {
     A.getOrCreateAAFor<AAAMDAttributes>(IRPosition::function(*F));
@@ -1427,7 +1446,8 @@ class AMDGPUAttributorLegacy : public ModulePass {
 
   bool runOnModule(Module &M) override {
     AnalysisGetter AG(this);
-    return runImpl(M, AG, *TM, /*Options=*/{});
+    return runImpl(M, AG, *TM, /*Options=*/{},
+                   /*LTOPhase=*/ThinOrFullLTOPhase::None);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -1448,8 +1468,8 @@ PreservedAnalyses llvm::AMDGPUAttributorPass::run(Module &M,
   AnalysisGetter AG(FAM);
 
   // TODO: Probably preserves CFG
-  return runImpl(M, AG, TM, Options) ? PreservedAnalyses::none()
-                                     : PreservedAnalyses::all();
+  return runImpl(M, AG, TM, Options, LTOPhase) ? PreservedAnalyses::none()
+                                               : PreservedAnalyses::all();
 }
 
 char AMDGPUAttributorLegacy::ID = 0;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 827216f8fde59..caa7c3e35397c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -876,8 +876,10 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
                                             OptimizationLevel Level,
                                             ThinOrFullLTOPhase Phase) {
     if (Level != OptimizationLevel::O0) {
-      if (!isLTOPreLink(Phase))
-        MPM.addPass(AMDGPUAttributorPass(*this));
+      if (!isLTOPreLink(Phase)) {
+        AMDGPUAttributorOptions Opts;
+        MPM.addPass(AMDGPUAttributorPass(*this, Opts, Phase));
+      }
     }
   });
 
@@ -900,7 +902,8 @@ void AMDGPUTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
             AMDGPUAttributorOptions Opt;
             if (HasClosedWorldAssumption)
               Opt.IsClosedWorld = true;
-            PM.addPass(AMDGPUAttributorPass(*this, Opt));
+            PM.addPass(AMDGPUAttributorPass(
+                *this, Opt, ThinOrFullLTOPhase::FullLTOPostLink));
           }
         }
         if (!NoKernelInfoEndLTO) {
diff --git a/llvm/test/LTO/AMDGPU/closed-world-assumption.ll b/llvm/test/LTO/AMDGPU/closed-world-assumption.ll
index dd084e7f3d9ed..de1e614944d6b 100644
--- a/llvm/test/LTO/AMDGPU/closed-world-assumption.ll
+++ b/llvm/test/LTO/AMDGPU/closed-world-assumption.ll
@@ -1,10 +1,12 @@
-; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -O3 -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefix=NO-CW
-; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefix=NO-CW
-; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -amdgpu-link-time-closed-world=1 -o - %s 2>&1 | FileCheck %s --check-prefix=CW
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -O3 -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefixes=NO-CW,NO-LTO
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -o - %s 2>&1 | FileCheck %s --check-prefixes=NO-CW,LTO
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes="lto<O3>" -debug-only=amdgpu-attributor -amdgpu-link-time-closed-world=1 -o - %s 2>&1 | FileCheck %s --check-prefixes=CW,LTO
 
 ; REQUIRES: amdgpu-registered-target
 ; REQUIRES: asserts
 
+; NO-LTO: Running at phase None
+; LTO: Running at phase FullLTOPostLink
 ; NO-CW: Module {{.*}} is not assumed to be a closed world.
 ; CW: Module {{.*}} is assumed to be a closed world.
 define hidden noundef i32 @_Z3foov() {