Add support for moving target specific AA before basic AA

Author: Chengjunp

llvm/include/llvm/Analysis/AliasAnalysis.h

@@ -993,6 +993,8 @@ struct ExternalAAWrapperPass : ImmutablePass {
 
   explicit ExternalAAWrapperPass(CallbackT CB);
 
+  virtual bool runEarly() { return false; }
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }

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 &) {}

llvm/lib/Analysis/AliasAnalysis.cpp

@@ -752,6 +752,11 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
   AAR.reset(
       new AAResults(getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F)));
 
+  // Add any target-specific alias analyses that should be run early.
+  auto *ExtWrapperPass = getAnalysisIfAvailable<ExternalAAWrapperPass>();
+  if (ExtWrapperPass && ExtWrapperPass->runEarly() && 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.
   // FIXME: TBAA should have an explicit mode to support this and then we
@@ -771,9 +776,8 @@ bool AAResultsWrapperPass::runOnFunction(Function &F) {
 
   // 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 (ExtWrapperPass && !ExtWrapperPass->runEarly() && ExtWrapperPass->CB)
+    ExtWrapperPass->CB(*this, F, *AAR);
 
   // Analyses don't mutate the IR, so return false.
   return false;

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -2202,6 +2202,10 @@ AAManager PassBuilder::buildDefaultAAPipeline() {
   // The order in which these are registered determines their priority when
   // being queried.
 
+  // Add any target-specific alias analyses that should be run early.
+  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.

llvm/lib/Target/NVPTX/NVPTXAliasAnalysis.h

@@ -79,10 +79,13 @@ class NVPTXAAWrapperPass : public ImmutablePass {
 
 // Wrapper around ExternalAAWrapperPass so that the default
 // constructor gets the callback.
+// Note that NVPTXAA will run before BasicAA to for compile time considerations.
 class NVPTXExternalAAWrapper : public ExternalAAWrapperPass {
 public:
   static char ID;
 
+  bool runEarly() override { return true; }
+
   NVPTXExternalAAWrapper()
       : ExternalAAWrapperPass([](Pass &P, Function &, AAResults &AAR) {
           if (auto *WrapperPass =

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>();
 }
 

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;
 

llvm/test/Analysis/NVPTXAA/NVPTXAA_before_BasicAA.ll

@@ -0,0 +1,15 @@
+; RUN: opt -aa-pipeline=default -passes='require<aa>' -debug-pass-manager -disable-output -S < %s 2>&1 | FileCheck %s
+; RUN: opt -O2 -debug-aa -disable-output -S < %s 2>&1 | FileCheck %s --check-prefix=LEGACY
+; LEGACY: NVPTX AA
+; LEGACY-NEXT: Basic Alias Analysis (stateless AA impl)
+
+; 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
+}
+