[llvm] [AA] Support Running Target Specific AA before BasicAA (PR #125965)

via llvm-commits llvm-commits at lists.llvm.org
Thu Feb 6 15:17:15 PST 2025


https://github.com/Chengjunp updated https://github.com/llvm/llvm-project/pull/125965

>From 613d8fa1b7bd5a33258ee770a9666f8f345461fc Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Thu, 6 Feb 2025 00:09:21 +0000
Subject: [PATCH 1/3] Initial commit

---
 llvm/include/llvm/Analysis/AliasAnalysis.h    |  7 ++++--
 llvm/include/llvm/Target/TargetMachine.h      |  5 ++++
 llvm/lib/Analysis/AliasAnalysis.cpp           | 23 +++++++++++--------
 llvm/lib/Passes/PassBuilderPipelines.cpp      |  5 ++++
 llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h    |  4 +++-
 llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp  |  4 ++--
 llvm/lib/Target/NVPTX/NVPTXTargetMachine.h    |  2 +-
 .../NVPTXAA/NVPTXAA_before_BasicAA.ll         | 11 +++++++++
 8 files changed, 46 insertions(+), 15 deletions(-)
 create mode 100644 llvm/test/Analysis/NVPTXAA/NVPTXAA_before_BasicAA.ll

diff --git a/llvm/include/llvm/Analysis/AliasAnalysis.h b/llvm/include/llvm/Analysis/AliasAnalysis.h
index 9c6084d2d9deeb1..33be0d6bdee3c02 100644
--- a/llvm/include/llvm/Analysis/AliasAnalysis.h
+++ b/llvm/include/llvm/Analysis/AliasAnalysis.h
@@ -987,11 +987,13 @@ struct ExternalAAWrapperPass : ImmutablePass {
 
   CallbackT CB;
 
+  bool Early; // Run before BasicAA or not
+
   static char ID;
 
   ExternalAAWrapperPass();
 
-  explicit ExternalAAWrapperPass(CallbackT CB);
+  explicit ExternalAAWrapperPass(CallbackT CB, bool Early = false);
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
@@ -1006,7 +1008,8 @@ struct ExternalAAWrapperPass : ImmutablePass {
 /// function, and the AAResults object to populate. This should be used when
 /// setting up a custom pass pipeline to inject a hook into the AA results.
 ImmutablePass *createExternalAAWrapperPass(
-    std::function<void(Pass &, Function &, AAResults &)> Callback);
+    std::function<void(Pass &, Function &, AAResults &)> Callback,
+    bool RunEarly = false);
 
 } // end namespace llvm
 
diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index c3e9d41315f617d..588025525d8feab 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -371,6 +371,11 @@ class TargetMachine {
   // TODO: Populate all pass names by using <Target>PassRegistry.def.
   virtual void registerPassBuilderCallbacks(PassBuilder &) {}
 
+  /// Allow the target to register early alias analyses(AA before basicAA) with
+  /// the AAManager for use with the new pass manager. Only affects the
+  /// "default" AAManager.
+  virtual void registerEarlyDefaultAliasAnalyses(AAManager &) {}
+
   /// Allow the target to register alias analyses with the AAManager for use
   /// with the new pass manager. Only affects the "default" AAManager.
   virtual void registerDefaultAliasAnalyses(AAManager &) {}
diff --git a/llvm/lib/Analysis/AliasAnalysis.cpp b/llvm/lib/Analysis/AliasAnalysis.cpp
index 20cdbb632032247..8cba8dffeba7be4 100644
--- a/llvm/lib/Analysis/AliasAnalysis.cpp
+++ b/llvm/lib/Analysis/AliasAnalysis.cpp
@@ -702,8 +702,8 @@ ExternalAAWrapperPass::ExternalAAWrapperPass() : ImmutablePass(ID) {
   initializeExternalAAWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
-ExternalAAWrapperPass::ExternalAAWrapperPass(CallbackT CB)
-    : ImmutablePass(ID), CB(std::move(CB)) {
+ExternalAAWrapperPass::ExternalAAWrapperPass(CallbackT CB, bool Early)
+    : ImmutablePass(ID), CB(std::move(CB)), Early(Early) {
   initializeExternalAAWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
@@ -713,8 +713,8 @@ INITIALIZE_PASS(ExternalAAWrapperPass, "external-aa", "External Alias Analysis",
                 false, true)
 
 ImmutablePass *
-llvm::createExternalAAWrapperPass(ExternalAAWrapperPass::CallbackT Callback) {
-  return new ExternalAAWrapperPass(std::move(Callback));
+llvm::createExternalAAWrapperPass(ExternalAAWrapperPass::CallbackT Callback, bool RunEarly) {
+  return new ExternalAAWrapperPass(std::move(Callback), RunEarly);
 }
 
 AAResultsWrapperPass::AAResultsWrapperPass() : FunctionPass(ID) {
@@ -751,6 +751,12 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
   // registering new results.
   AAR.reset(
       new AAResults(getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F)));
+  
+  // For the target need run early target-specific AA, we register it
+  // before basicAA.
+  auto *ExtWrapperPass = getAnalysisIfAvailable<ExternalAAWrapperPass>();
+  if (ExtWrapperPass && ExtWrapperPass->Early && ExtWrapperPass->CB)
+    ExtWrapperPass->CB(*this, F, *AAR);
 
   // BasicAA is always available for function analyses. Also, we add it first
   // so that it can trump TBAA results when it proves MustAlias.
@@ -769,11 +775,10 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
   if (auto *WrapperPass = getAnalysisIfAvailable<SCEVAAWrapperPass>())
     AAR->addAAResult(WrapperPass->getResult());
 
-  // If available, run an external AA providing callback over the results as
-  // well.
-  if (auto *WrapperPass = getAnalysisIfAvailable<ExternalAAWrapperPass>())
-    if (WrapperPass->CB)
-      WrapperPass->CB(*this, F, *AAR);
+  // If available, run an external non-early AA providing callback over the
+  // results as well.
+  if (ExtWrapperPass && !ExtWrapperPass->Early && ExtWrapperPass->CB)
+    ExtWrapperPass->CB(*this, F, *AAR);
 
   // Analyses don't mutate the IR, so return false.
   return false;
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 17710eb94b6dedb..19c040873bd1e95 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -2202,6 +2202,11 @@ AAManager PassBuilder::buildDefaultAAPipeline() {
   // The order in which these are registered determines their priority when
   // being queried.
 
+  // For the target need run early target-specific AA, we register it
+  // before basicAA.
+  if (TM)
+    TM->registerEarlyDefaultAliasAnalyses(AA);
+
   // First we register the basic alias analysis that provides the majority of
   // per-function local AA logic. This is a stateless, on-demand local set of
   // AA techniques.
diff --git a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
index 2d204979eb6cef4..490af0200a01cde 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
+++ b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
@@ -79,6 +79,8 @@ class NVPTXAAWrapperPass : public ImmutablePass {
 
 // Wrapper around ExternalAAWrapperPass so that the default
 // constructor gets the callback.
+// Note that NVPTXAA will run before BasicAA to improve the compile
+// time.
 class NVPTXExternalAAWrapper : public ExternalAAWrapperPass {
 public:
   static char ID;
@@ -88,7 +90,7 @@ class NVPTXExternalAAWrapper : public ExternalAAWrapperPass {
           if (auto *WrapperPass =
                   P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
             AAR.addAAResult(WrapperPass->getResult());
-        }) {}
+        }, true) {}
 };
 
 ImmutablePass *createNVPTXAAWrapperPass();
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index 7d04cf3dc51e673..b50ecb41d8b1147 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -221,7 +221,7 @@ MachineFunctionInfo *NVPTXTargetMachine::createMachineFunctionInfo(
                                                                     F, STI);
 }
 
-void NVPTXTargetMachine::registerDefaultAliasAnalyses(AAManager &AAM) {
+void NVPTXTargetMachine::registerEarlyDefaultAliasAnalyses(AAManager &AAM) {
   AAM.registerFunctionAnalysis<NVPTXAA>();
 }
 
@@ -320,7 +320,7 @@ void NVPTXPassConfig::addIRPasses() {
   addPass(createExternalAAWrapperPass([](Pass &P, Function &, AAResults &AAR) {
     if (auto *WrapperPass = P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
       AAR.addAAResult(WrapperPass->getResult());
-  }));
+  }, true));
 
   // NVVMReflectPass is added in addEarlyAsPossiblePasses, so hopefully running
   // it here does nothing.  But since we need it for correctness when lowering
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h
index 2b88da67a50f95a..bb65a14a8d79259 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.h
@@ -64,7 +64,7 @@ class NVPTXTargetMachine : public LLVMTargetMachine {
   createMachineFunctionInfo(BumpPtrAllocator &Allocator, const Function &F,
                             const TargetSubtargetInfo *STI) const override;
 
-  void registerDefaultAliasAnalyses(AAManager &AAM) override;
+  void registerEarlyDefaultAliasAnalyses(AAManager &AAM) override;
 
   void registerPassBuilderCallbacks(PassBuilder &PB) override;
 
diff --git a/llvm/test/Analysis/NVPTXAA/NVPTXAA_before_BasicAA.ll b/llvm/test/Analysis/NVPTXAA/NVPTXAA_before_BasicAA.ll
new file mode 100644
index 000000000000000..c1b900ce16d97d5
--- /dev/null
+++ b/llvm/test/Analysis/NVPTXAA/NVPTXAA_before_BasicAA.ll
@@ -0,0 +1,11 @@
+; RUN: opt -aa-pipeline=default -passes='require<aa>' -debug-pass-manager -disable-output -S < %s 2>&1 | FileCheck %s
+
+; In default AA pipeline, NVPTXAA should run before BasicAA to reduce compile time for NVPTX backend
+target triple = "nvptx64-nvidia-cuda"
+
+; CHECK: Running analysis: NVPTXAA on foo
+; CHECK-NEXT: Running analysis: BasicAA on foo
+define void @foo(){
+entry:
+  ret void
+}

>From c52aaa259f6d495db32d338b4395efe5d1519026 Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Thu, 6 Feb 2025 00:30:40 +0000
Subject: [PATCH 2/3] Format

---
 llvm/lib/Analysis/AliasAnalysis.cpp          |  5 +++--
 llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h   | 12 +++++++-----
 llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp | 10 ++++++----
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Analysis/AliasAnalysis.cpp b/llvm/lib/Analysis/AliasAnalysis.cpp
index 8cba8dffeba7be4..a54c505220c4385 100644
--- a/llvm/lib/Analysis/AliasAnalysis.cpp
+++ b/llvm/lib/Analysis/AliasAnalysis.cpp
@@ -713,7 +713,8 @@ INITIALIZE_PASS(ExternalAAWrapperPass, "external-aa", "External Alias Analysis",
                 false, true)
 
 ImmutablePass *
-llvm::createExternalAAWrapperPass(ExternalAAWrapperPass::CallbackT Callback, bool RunEarly) {
+llvm::createExternalAAWrapperPass(ExternalAAWrapperPass::CallbackT Callback,
+                                  bool RunEarly) {
   return new ExternalAAWrapperPass(std::move(Callback), RunEarly);
 }
 
@@ -751,7 +752,7 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
   // registering new results.
   AAR.reset(
       new AAResults(getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F)));
-  
+
   // For the target need run early target-specific AA, we register it
   // before basicAA.
   auto *ExtWrapperPass = getAnalysisIfAvailable<ExternalAAWrapperPass>();
diff --git a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
index 490af0200a01cde..4b521718edf0657 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
+++ b/llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h
@@ -86,11 +86,13 @@ class NVPTXExternalAAWrapper : public ExternalAAWrapperPass {
   static char ID;
 
   NVPTXExternalAAWrapper()
-      : ExternalAAWrapperPass([](Pass &P, Function &, AAResults &AAR) {
-          if (auto *WrapperPass =
-                  P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
-            AAR.addAAResult(WrapperPass->getResult());
-        }, true) {}
+      : ExternalAAWrapperPass(
+            [](Pass &P, Function &, AAResults &AAR) {
+              if (auto *WrapperPass =
+                      P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
+                AAR.addAAResult(WrapperPass->getResult());
+            },
+            true) {}
 };
 
 ImmutablePass *createNVPTXAAWrapperPass();
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index b50ecb41d8b1147..d02d9d31dff7905 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -317,10 +317,12 @@ void NVPTXPassConfig::addIRPasses() {
   disablePass(&RemoveLoadsIntoFakeUsesID);
 
   addPass(createNVPTXAAWrapperPass());
-  addPass(createExternalAAWrapperPass([](Pass &P, Function &, AAResults &AAR) {
-    if (auto *WrapperPass = P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
-      AAR.addAAResult(WrapperPass->getResult());
-  }, true));
+  addPass(createExternalAAWrapperPass(
+      [](Pass &P, Function &, AAResults &AAR) {
+        if (auto *WrapperPass = P.getAnalysisIfAvailable<NVPTXAAWrapperPass>())
+          AAR.addAAResult(WrapperPass->getResult());
+      },
+      true));
 
   // NVVMReflectPass is added in addEarlyAsPossiblePasses, so hopefully running
   // it here does nothing.  But since we need it for correctness when lowering

>From a199a8fce8319cefd2e9699182f13e94cb6bc06b Mon Sep 17 00:00:00 2001
From: chengjunp <chengjunp at nvidia.com>
Date: Thu, 6 Feb 2025 23:23:31 +0000
Subject: [PATCH 3/3] Update some comments

---
 llvm/include/llvm/Target/TargetMachine.h | 2 +-
 llvm/lib/Analysis/AliasAnalysis.cpp      | 3 +--
 llvm/lib/Passes/PassBuilderPipelines.cpp | 3 +--
 3 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/Target/TargetMachine.h b/llvm/include/llvm/Target/TargetMachine.h
index 588025525d8feab..e6ed4c306140326 100644
--- a/llvm/include/llvm/Target/TargetMachine.h
+++ b/llvm/include/llvm/Target/TargetMachine.h
@@ -371,7 +371,7 @@ class TargetMachine {
   // TODO: Populate all pass names by using <Target>PassRegistry.def.
   virtual void registerPassBuilderCallbacks(PassBuilder &) {}
 
-  /// Allow the target to register early alias analyses(AA before basicAA) with
+  /// Allow the target to register early alias analyses (before basicAA) with
   /// the AAManager for use with the new pass manager. Only affects the
   /// "default" AAManager.
   virtual void registerEarlyDefaultAliasAnalyses(AAManager &) {}
diff --git a/llvm/lib/Analysis/AliasAnalysis.cpp b/llvm/lib/Analysis/AliasAnalysis.cpp
index a54c505220c4385..21fbcdbfd031158 100644
--- a/llvm/lib/Analysis/AliasAnalysis.cpp
+++ b/llvm/lib/Analysis/AliasAnalysis.cpp
@@ -753,8 +753,7 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
   AAR.reset(
       new AAResults(getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F)));
 
-  // For the target need run early target-specific AA, we register it
-  // before basicAA.
+  // Add any target-specific alias analyses that should be run early.
   auto *ExtWrapperPass = getAnalysisIfAvailable<ExternalAAWrapperPass>();
   if (ExtWrapperPass && ExtWrapperPass->Early && ExtWrapperPass->CB)
     ExtWrapperPass->CB(*this, F, *AAR);
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 19c040873bd1e95..0afcf992dcde718 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -2202,8 +2202,7 @@ AAManager PassBuilder::buildDefaultAAPipeline() {
   // The order in which these are registered determines their priority when
   // being queried.
 
-  // For the target need run early target-specific AA, we register it
-  // before basicAA.
+  // Add any target-specific alias analyses that should be run early.
   if (TM)
     TM->registerEarlyDefaultAliasAnalyses(AA);
 



More information about the llvm-commits mailing list