[Utils] Make -assume-builder/-assume-simplify actually work on Old-PM

clang w/ old-pm currently would simply crash
when -mllvm  -enable-knowledge-retention=true is specified.

Clearly, these two passes had no Old-PM test coverage,
which would have shown the problem - not requiring AssumptionCacheTracker,
but then trying to always get it.

Also, why try to get domtree only if it's cached,
but at the same time marking it as required?

5 files changed, 97 insertions, 28 deletions

diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index 3b8bc9f..f0d5acc 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -72,6 +72,7 @@ void initializeAliasSetPrinterPass(PassRegistry&);
 void initializeAlignmentFromAssumptionsPass(PassRegistry&);
 void initializeAlwaysInlinerLegacyPassPass(PassRegistry&);
 void initializeAssumeSimplifyPassLegacyPassPass(PassRegistry &);
+void initializeAssumeBuilderPassLegacyPassPass(PassRegistry &);
 void initializeOpenMPOptLegacyPassPass(PassRegistry &);
 void initializeArgPromotionPass(PassRegistry&);
 void initializeAssumptionCacheTrackerPass(PassRegistry&);
diff --git a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
index ad45744..f2208ed 100644
--- a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
+++ b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
@@ -525,14 +525,16 @@ public:
   bool runOnFunction(Function &F) override {
     if (skipFunction(F) || !EnableKnowledgeRetention)
       return false;
-    DominatorTreeWrapperPass *DT =
-        getAnalysisIfAvailable<DominatorTreeWrapperPass>();
     AssumptionCache &AC =
         getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
-    return simplifyAssumes(F, &AC, DT ? &DT->getDomTree() : nullptr);
+    DominatorTreeWrapperPass *DTWP =
+        getAnalysisIfAvailable<DominatorTreeWrapperPass>();
+    return simplifyAssumes(F, &AC, DTWP ? &DTWP->getDomTree() : nullptr);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
+
     AU.setPreservesAll();
   }
 };
@@ -542,7 +544,6 @@ char AssumeSimplifyPassLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AssumeSimplifyPassLegacyPass, "assume-simplify",
                       "Assume Simplify", false, false)
-INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_END(AssumeSimplifyPassLegacyPass, "assume-simplify",
                     "Assume Simplify", false, false)
@@ -553,9 +554,43 @@ FunctionPass *llvm::createAssumeSimplifyPass() {
 
 PreservedAnalyses AssumeBuilderPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
-  AssumptionCache* AC = AM.getCachedResult<AssumptionAnalysis>(F);
+  AssumptionCache *AC = &AM.getResult<AssumptionAnalysis>(F);
   DominatorTree* DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
   for (Instruction &I : instructions(F))
     salvageKnowledge(&I, AC, DT);
   return PreservedAnalyses::all();
 }
+
+namespace {
+class AssumeBuilderPassLegacyPass : public FunctionPass {
+public:
+  static char ID;
+
+  AssumeBuilderPassLegacyPass() : FunctionPass(ID) {
+    initializeAssumeBuilderPassLegacyPassPass(*PassRegistry::getPassRegistry());
+  }
+  bool runOnFunction(Function &F) override {
+    AssumptionCache &AC =
+        getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+    DominatorTreeWrapperPass *DTWP =
+        getAnalysisIfAvailable<DominatorTreeWrapperPass>();
+    for (Instruction &I : instructions(F))
+      salvageKnowledge(&I, &AC, DTWP ? &DTWP->getDomTree() : nullptr);
+    return true;
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
+
+    AU.setPreservesAll();
+  }
+};
+} // namespace
+
+char AssumeBuilderPassLegacyPass::ID = 0;
+
+INITIALIZE_PASS_BEGIN(AssumeBuilderPassLegacyPass, "assume-builder",
+                      "Assume Builder", false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
+INITIALIZE_PASS_END(AssumeBuilderPassLegacyPass, "assume-builder",
+                    "Assume Builder", false, false)
diff --git a/llvm/lib/Transforms/Utils/Utils.cpp b/llvm/lib/Transforms/Utils/Utils.cpp
index fc1e824..ce98a73 100644
--- a/llvm/lib/Transforms/Utils/Utils.cpp
+++ b/llvm/lib/Transforms/Utils/Utils.cpp
@@ -25,6 +25,7 @@ using namespace llvm;
 void llvm::initializeTransformUtils(PassRegistry &Registry) {
   initializeAddDiscriminatorsLegacyPassPass(Registry);
   initializeAssumeSimplifyPassLegacyPassPass(Registry);
+  initializeAssumeBuilderPassLegacyPassPass(Registry);
   initializeBreakCriticalEdgesPass(Registry);
   initializeCanonicalizeAliasesLegacyPassPass(Registry);
   initializeCanonicalizeFreezeInLoopsPass(Registry);
diff --git a/llvm/test/Transforms/Util/assume-builder.ll b/llvm/test/Transforms/Util/assume-builder.ll
index 33cc3ff..9917686 100644
--- a/llvm/test/Transforms/Util/assume-builder.ll
+++ b/llvm/test/Transforms/Util/assume-builder.ll
@@ -1,8 +1,14 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; RUN: opt -assume-builder -verify --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=BASIC
+; RUN: opt -assume-builder -verify --enable-knowledge-retention --assume-preserve-all -S %s | FileCheck %s --check-prefixes=ALL
+; RUN: opt -assumption-cache-tracker -assume-builder -verify --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=WITH-AC
+; RUN: opt -domtree -assumption-cache-tracker -assume-builder -verify --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=CROSS-BLOCK,CROSS-BLOCK-OLDPM
+; RUN: opt -assume-builder -domtree -assumption-cache-tracker -assume-simplify -verify --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=FULL-SIMPLIFY
+
 ; RUN: opt -passes='assume-builder,verify' --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=BASIC
 ; RUN: opt -passes='assume-builder,verify' --enable-knowledge-retention --assume-preserve-all -S %s | FileCheck %s --check-prefixes=ALL
 ; RUN: opt -passes='require<assumptions>,assume-builder,verify' --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=WITH-AC
-; RUN: opt -passes='require<domtree>,require<assumptions>,assume-builder,verify' --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=CROSS-BLOCK
+; RUN: opt -passes='require<domtree>,require<assumptions>,assume-builder,verify' --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=CROSS-BLOCK,CROSS-BLOCK-NEWMP
 ; RUN: opt -passes='assume-builder,require<domtree>,require<assumptions>,assume-simplify,verify' --enable-knowledge-retention -S %s | FileCheck %s --check-prefixes=FULL-SIMPLIFY
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
@@ -641,28 +647,30 @@ define dso_local i32 @_Z6squarePi(i32* %P, i32* %P1, i1 %cond) {
 ; WITH-AC-NEXT:    store i32 0, i32* [[P1]], align 4
 ; WITH-AC-NEXT:    ret i32 0
 ;
-; CROSS-BLOCK-LABEL: define {{[^@]+}}@_Z6squarePi
-; CROSS-BLOCK-SAME: (i32* [[P:%.*]], i32* [[P1:%.*]], i1 [[COND:%.*]])
-; CROSS-BLOCK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]), "align"(i32* [[P]], i64 4) ]
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P]], align 4
-; CROSS-BLOCK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]), "align"(i32* [[P1]], i64 8) ]
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P1]], align 8
-; CROSS-BLOCK-NEXT:    br i1 [[COND]], label [[A:%.*]], label [[B:%.*]]
-; CROSS-BLOCK:       A:
-; CROSS-BLOCK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 8) ]
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P]], align 8
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P1]], align 4
-; CROSS-BLOCK-NEXT:    br i1 [[COND]], label [[C:%.*]], label [[B]]
-; CROSS-BLOCK:       B:
-; CROSS-BLOCK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 8) ]
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P]], align 8
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P1]], align 8
-; CROSS-BLOCK-NEXT:    br label [[C]]
-; CROSS-BLOCK:       C:
-; CROSS-BLOCK-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 32) ]
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P]], align 32
-; CROSS-BLOCK-NEXT:    store i32 0, i32* [[P1]], align 4
-; CROSS-BLOCK-NEXT:    ret i32 0
+; CROSS-BLOCK-OLDPM-LABEL: define {{[^@]+}}@_Z6squarePi
+; CROSS-BLOCK-OLDPM-SAME: (i32* [[P:%.*]], i32* [[P1:%.*]], i1 [[COND:%.*]])
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]), "align"(i32* [[P]], i64 4) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P]], align 4
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]), "align"(i32* [[P1]], i64 8) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P1]], align 8
+; CROSS-BLOCK-OLDPM-NEXT:    br i1 [[COND]], label [[A:%.*]], label [[B:%.*]]
+; CROSS-BLOCK-OLDPM:       A:
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 8) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P]], align 8
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P1]], align 4
+; CROSS-BLOCK-OLDPM-NEXT:    br i1 [[COND]], label [[C:%.*]], label [[B]]
+; CROSS-BLOCK-OLDPM:       B:
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]), "align"(i32* [[P]], i64 8) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P]], align 8
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]), "align"(i32* [[P1]], i64 8) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P1]], align 8
+; CROSS-BLOCK-OLDPM-NEXT:    br label [[C]]
+; CROSS-BLOCK-OLDPM:       C:
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]), "align"(i32* [[P]], i64 32) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P]], align 32
+; CROSS-BLOCK-OLDPM-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]), "align"(i32* [[P1]], i64 4) ]
+; CROSS-BLOCK-OLDPM-NEXT:    store i32 0, i32* [[P1]], align 4
+; CROSS-BLOCK-OLDPM-NEXT:    ret i32 0
 ;
 ; FULL-SIMPLIFY-LABEL: define {{[^@]+}}@_Z6squarePi
 ; FULL-SIMPLIFY-SAME: (i32* nonnull align 4 dereferenceable(4) [[P:%.*]], i32* nonnull align 8 dereferenceable(4) [[P1:%.*]], i1 [[COND:%.*]])
@@ -685,6 +693,29 @@ define dso_local i32 @_Z6squarePi(i32* %P, i32* %P1, i1 %cond) {
 ; FULL-SIMPLIFY-NEXT:    store i32 0, i32* [[P1]], align 4
 ; FULL-SIMPLIFY-NEXT:    ret i32 0
 ;
+; CROSS-BLOCK-NEWMP-LABEL: define {{[^@]+}}@_Z6squarePi
+; CROSS-BLOCK-NEWMP-SAME: (i32* [[P:%.*]], i32* [[P1:%.*]], i1 [[COND:%.*]])
+; CROSS-BLOCK-NEWMP-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P]], i64 4), "nonnull"(i32* [[P]]), "align"(i32* [[P]], i64 4) ]
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P]], align 4
+; CROSS-BLOCK-NEWMP-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(i32* [[P1]], i64 4), "nonnull"(i32* [[P1]]), "align"(i32* [[P1]], i64 8) ]
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P1]], align 8
+; CROSS-BLOCK-NEWMP-NEXT:    br i1 [[COND]], label [[A:%.*]], label [[B:%.*]]
+; CROSS-BLOCK-NEWMP:       A:
+; CROSS-BLOCK-NEWMP-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 8) ]
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P]], align 8
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P1]], align 4
+; CROSS-BLOCK-NEWMP-NEXT:    br i1 [[COND]], label [[C:%.*]], label [[B]]
+; CROSS-BLOCK-NEWMP:       B:
+; CROSS-BLOCK-NEWMP-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 8) ]
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P]], align 8
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P1]], align 8
+; CROSS-BLOCK-NEWMP-NEXT:    br label [[C]]
+; CROSS-BLOCK-NEWMP:       C:
+; CROSS-BLOCK-NEWMP-NEXT:    call void @llvm.assume(i1 true) [ "align"(i32* [[P]], i64 32) ]
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P]], align 32
+; CROSS-BLOCK-NEWMP-NEXT:    store i32 0, i32* [[P1]], align 4
+; CROSS-BLOCK-NEWMP-NEXT:    ret i32 0
+;
   store i32 0, i32* %P, align 4
   store i32 0, i32* %P1, align 8
   br i1 %cond, label %A, label %B
diff --git a/llvm/test/Transforms/Util/assume-simplify.ll b/llvm/test/Transforms/Util/assume-simplify.ll
index 3005de3..5f5eece 100644
--- a/llvm/test/Transforms/Util/assume-simplify.ll
+++ b/llvm/test/Transforms/Util/assume-simplify.ll
@@ -1,4 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; RUN: opt -domtree -assumption-cache-tracker -assume-simplify -verify --enable-knowledge-retention -S %s | FileCheck %s
 ; RUN: opt -passes='require<domtree>,require<assumptions>,assume-simplify,verify' --enable-knowledge-retention -S %s | FileCheck %s
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"