[AA] Enhance the new AliasAnalysis infrastructure with an optional "external" AA wrapper pass.

[chandlerc]

Summary

This is a generic hook that can be used to thread custom code into the
primary AAResultsWrapperPass for the legacy pass manager in order to
allow it to merge external AA results into the AA results it is
building. It does this by threading in a raw callback and so it is
*very* powerful and should serve almost any use case I have come up with
for extending the set of alias analyses used. The only thing not well
supported here is using a *different order* of alias analyses. That form
of extension *is* supportable with the new pass manager, and I can make
the callback structure here more elaborate to support it in the legacy
pass manager if this is a critical use case that people are already
depending on, but the only use cases I have heard of thus far should be
reasonably satisfied by this simpler extension mechanism.

It is hard to test this using normal facilities (the built-in AAs don't
use this for obvious reasons) so I've written a fairly extensive set of
custom passes in the alias analysis unit test that should be an
excellent test case because it models the out-of-tree users: it adds
a totally custom AA to the system. This should also serve as
a reasonably good example and guide for out-of-tree users to follow in
order to rig up their existing alias analyses.

No support in opt for commandline control is provided here however. I'm
really unhappy with the kind of contortions that would be required to
support that. It would fully re-introduce the analysis group
self-recursion kind of patterns. =/

Hopefully this is sufficient for the existing out-of-tree users to be
unblocked and able to use the AA infrastructure after the removal of
analysis group functionality.

[gottesmm]

I just prepared all of the code that I need internally to work with this and it enables me to work through/verify all my use cases. This LGTM! Thanks Chandler!

llvm/trunk/include/llvm/Analysis/AliasAnalysis.h

public:
 
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 };
 
 FunctionPass *createAAResultsWrapperPass();
 
+/// A wrapper pass around a callback which can be used to populate the
+/// AAResults in the AAResultsWrapperPass from an external AA.
+///
+/// The callback provided here will be used each time we prepare an AAResults
+/// object, and will receive a reference to the function wrapper pass, the
+/// 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);
+
 /// A helper for the legacy pass manager to create a \c AAResults
 /// object populated to the best of our ability for a particular function when
 /// inside of a \c ModulePass or a \c CallGraphSCCPass.
 AAResults createLegacyPMAAResults(Pass &P, Function &F, BasicAAResult &BAR);
 
 } // End llvm namespace
 
 #endif

llvm/trunk/include/llvm/InitializePasses.h

 void initializeBreakCriticalEdgesPass(PassRegistry&);
 void initializeCallGraphPrinterPass(PassRegistry&);
 void initializeCallGraphViewerPass(PassRegistry&);
 void initializeCFGOnlyPrinterPass(PassRegistry&);
 void initializeCFGOnlyViewerPass(PassRegistry&);
 void initializeCFGPrinterPass(PassRegistry&);
 void initializeCFGSimplifyPassPass(PassRegistry&);
 void initializeCFLAAWrapperPassPass(PassRegistry&);
+void initializeExternalAAWrapperPassPass(PassRegistry&);
 void initializeForwardControlFlowIntegrityPass(PassRegistry&);
 void initializeFlattenCFGPassPass(PassRegistry&);
 void initializeStructurizeCFGPass(PassRegistry&);
 void initializeCFGViewerPass(PassRegistry&);
 void initializeConstantHoistingPass(PassRegistry&);
 void initializeCodeGenPreparePass(PassRegistry&);
 void initializeConstantMergePass(PassRegistry&);
 void initializeConstantPropagationPass(PassRegistry&);

llvm/trunk/lib/Analysis/AliasAnalysis.cpp

for (; I != E; ++I) // Check every instruction in range
     if (getModRefInfo(&*I, Loc) & Mode)
       return true;
   return false;
 }
 
 // Provide a definition for the root virtual destructor.
 AAResults::Concept::~Concept() {}
 
+namespace {
+/// A wrapper pass for external alias analyses. This just squirrels away the
+/// callback used to run any analyses and register their results.
+struct ExternalAAWrapperPass : ImmutablePass {
+  typedef std::function<void(Pass &, Function &, AAResults &)> CallbackT;
+
+  CallbackT CB;
+
+  static char ID;
+
+  ExternalAAWrapperPass() : ImmutablePass(ID) {
+    initializeExternalAAWrapperPassPass(*PassRegistry::getPassRegistry());
+  }
+  explicit ExternalAAWrapperPass(CallbackT CB)
+      : ImmutablePass(ID), CB(std::move(CB)) {
+    initializeExternalAAWrapperPassPass(*PassRegistry::getPassRegistry());
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+  }
+};
+}
+
+char ExternalAAWrapperPass::ID = 0;
+INITIALIZE_PASS(ExternalAAWrapperPass, "external-aa", "External Alias Analysis",
+                false, true)
+
+ImmutablePass *
+llvm::createExternalAAWrapperPass(ExternalAAWrapperPass::CallbackT Callback) {
+  return new ExternalAAWrapperPass(std::move(Callback));
+}
+
 AAResultsWrapperPass::AAResultsWrapperPass() : FunctionPass(ID) {
   initializeAAResultsWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
 char AAResultsWrapperPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AAResultsWrapperPass, "aa",
                       "Function Alias Analysis Results", false, true)
 INITIALIZE_PASS_DEPENDENCY(BasicAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(CFLAAWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ExternalAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(GlobalsAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ObjCARCAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(SCEVAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScopedNoAliasAAWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TypeBasedAAWrapperPass)
 INITIALIZE_PASS_END(AAResultsWrapperPass, "aa",
                     "Function Alias Analysis Results", false, true)
 
if (auto *WrapperPass =
     AAR->addAAResult(WrapperPass->getResult());
   if (auto *WrapperPass = getAnalysisIfAvailable<GlobalsAAWrapperPass>())
     AAR->addAAResult(WrapperPass->getResult());
   if (auto *WrapperPass = getAnalysisIfAvailable<SCEVAAWrapperPass>())
     AAR->addAAResult(WrapperPass->getResult());
   if (auto *WrapperPass = getAnalysisIfAvailable<CFLAAWrapperPass>())
     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);
+
   // Analyses don't mutate the IR, so return false.
   return false;
 }
 
 void AAResultsWrapperPass::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<BasicAAWrapperPass>();
 

llvm/trunk/unittests/Analysis/AliasAnalysisTest.cpp

 //===--- AliasAnalysisTest.cpp - Mixed TBAA unit tests --------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
+using namespace llvm;
+
+// Set up some test passes.
 namespace llvm {
+void initializeAATestPassPass(PassRegistry&);
+void initializeTestCustomAAWrapperPassPass(PassRegistry&);
+}
+
+namespace {
+struct AATestPass : FunctionPass {
+  static char ID;
+  AATestPass() : FunctionPass(ID) {
+    initializeAATestPassPass(*PassRegistry::getPassRegistry());
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AAResultsWrapperPass>();
+    AU.setPreservesAll();
+  }
+
+  bool runOnFunction(Function &F) override {
+    AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
+
+    SetVector<Value *> Pointers;
+    for (Argument &A : F.args())
+      if (A.getType()->isPointerTy())
+        Pointers.insert(&A);
+    for (Instruction &I : instructions(F))
+      if (I.getType()->isPointerTy())
+        Pointers.insert(&I);
+
+    for (Value *P1 : Pointers)
+      for (Value *P2 : Pointers)
+        (void)AA.alias(P1, MemoryLocation::UnknownSize, P2,
+                       MemoryLocation::UnknownSize);
+
+    return false;
+  }
+};
+}
+
+char AATestPass::ID = 0;
+INITIALIZE_PASS_BEGIN(AATestPass, "aa-test-pas", "Alias Analysis Test Pass",
+                      false, true)
+INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+INITIALIZE_PASS_END(AATestPass, "aa-test-pass", "Alias Analysis Test Pass",
+                    false, true)
+
+namespace {
+/// A test customizable AA result. It merely accepts a callback to run whenever
+/// it receives an alias query. Useful for testing that a particular AA result
+/// is reached.
+struct TestCustomAAResult : AAResultBase<TestCustomAAResult> {
+  friend AAResultBase<TestCustomAAResult>;
+
+  std::function<void()> CB;
+
+  explicit TestCustomAAResult(const TargetLibraryInfo &TLI,
+                              std::function<void()> CB)
+      : AAResultBase(TLI), CB(std::move(CB)) {}
+  TestCustomAAResult(TestCustomAAResult &&Arg)
+      : AAResultBase(std::move(Arg)), CB(std::move(Arg.CB)) {}
+
+  bool invalidate(Function &, const PreservedAnalyses &) { return false; }
+
+  AliasResult alias(const MemoryLocation &LocA, const MemoryLocation &LocB) {
+    CB();
+    return MayAlias;
+  }
+};
+}
+
+namespace {
+/// A wrapper pass for the legacy pass manager to use with the above custom AA
+/// result.
+class TestCustomAAWrapperPass : public ImmutablePass {
+  std::function<void()> CB;
+  std::unique_ptr<TestCustomAAResult> Result;
+
+public:
+  static char ID;
+
+  explicit TestCustomAAWrapperPass(
+      std::function<void()> CB = std::function<void()>())
+      : ImmutablePass(ID), CB(std::move(CB)) {
+    initializeTestCustomAAWrapperPassPass(*PassRegistry::getPassRegistry());
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    AU.addRequired<TargetLibraryInfoWrapperPass>();
+  }
+
+  bool doInitialization(Module &M) override {
+    Result.reset(new TestCustomAAResult(
+        getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(), std::move(CB)));
+    return true;
+  }
+
+  bool doFinalization(Module &M) override {
+    Result.reset();
+    return true;
+  }
+
+  TestCustomAAResult &getResult() { return *Result; }
+  const TestCustomAAResult &getResult() const { return *Result; }
+};
+}
+
+char TestCustomAAWrapperPass::ID = 0;
+INITIALIZE_PASS_BEGIN(TestCustomAAWrapperPass, "test-custom-aa",
+                "Test Custom AA Wrapper Pass", false, true)
+INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_END(TestCustomAAWrapperPass, "test-custom-aa",
+                "Test Custom AA Wrapper Pass", false, true)
+
 namespace {
 
 class AliasAnalysisTest : public testing::Test {
 protected:
   LLVMContext C;
   Module M;
   TargetLibraryInfoImpl TLII;
   TargetLibraryInfo TLI;
TEST_F(AliasAnalysisTest, getModRefInfo) {
   EXPECT_EQ(AA.getModRefInfo(VAArg1, MemoryLocation()), MRI_ModRef);
   EXPECT_EQ(AA.getModRefInfo(VAArg1), MRI_ModRef);
   EXPECT_EQ(AA.getModRefInfo(CmpXChg1, MemoryLocation()), MRI_ModRef);
   EXPECT_EQ(AA.getModRefInfo(CmpXChg1), MRI_ModRef);
   EXPECT_EQ(AA.getModRefInfo(AtomicRMW, MemoryLocation()), MRI_ModRef);
   EXPECT_EQ(AA.getModRefInfo(AtomicRMW), MRI_ModRef);
 }
 
+class AAPassInfraTest : public testing::Test {
+protected:
+  LLVMContext &C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M;
+
+public:
+  AAPassInfraTest()
+      : C(getGlobalContext()),
+        M(parseAssemblyString("define i32 @f(i32* %x, i32* %y) {\n"
+                              "entry:\n"
+                              "  %lx = load i32, i32* %x\n"
+                              "  %ly = load i32, i32* %y\n"
+                              "  %sum = add i32 %lx, %ly\n"
+                              "  ret i32 %sum\n"
+                              "}\n",
+                              Err, C)) {
+    assert(M && "Failed to build the module!");
+  }
+};
+
+TEST_F(AAPassInfraTest, injectExternalAA) {
+  legacy::PassManager PM;
+
+  // Register our custom AA's wrapper pass manually.
+  bool IsCustomAAQueried = false;
+  PM.add(new TestCustomAAWrapperPass([&] { IsCustomAAQueried = true; }));
+
+  // Now add the external AA wrapper with a lambda which queries for the
+  // wrapper around our custom AA and adds it to the results.
+  PM.add(createExternalAAWrapperPass([](Pass &P, Function &, AAResults &AAR) {
+    if (auto *WrapperPass = P.getAnalysisIfAvailable<TestCustomAAWrapperPass>())
+      AAR.addAAResult(WrapperPass->getResult());
+  }));
+
+  // And run a pass that will make some alias queries. This will automatically
+  // trigger the rest of the alias analysis stack to be run. It is analagous to
+  // building a full pass pipeline with any of the existing pass manager
+  // builders.
+  PM.add(new AATestPass());
+  PM.run(*M);
+
+  // Finally, ensure that our custom AA was indeed queried.
+  EXPECT_TRUE(IsCustomAAQueried);
+}
+
 } // end anonymous namspace
-} // end llvm namespace