1 files changed, 12 insertions, 9 deletions
diff --git a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
index 31503a00379..c1df3173c0f 100644
--- a/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
+++ b/llvm/lib/Transforms/Scalar/AlignmentFromAssumptions.cpp
@@ -24,6 +24,7 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -53,6 +54,7 @@ struct AlignmentFromAssumptions : public FunctionPass {
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<DominatorTreeWrapperPass>();
 
@@ -72,6 +74,7 @@ char AlignmentFromAssumptions::ID = 0;
 static const char aip_name[] = "Alignment from assumptions";
 INITIALIZE_PASS_BEGIN(AlignmentFromAssumptions, AA_NAME,
                       aip_name, false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
 INITIALIZE_PASS_END(AlignmentFromAssumptions, AA_NAME,
@@ -405,13 +408,15 @@ bool AlignmentFromAssumptions::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
+  auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
-  return Impl.runImpl(F, SE, DT);
+  return Impl.runImpl(F, AC, SE, DT);
 }
 
-bool AlignmentFromAssumptionsPass::runImpl(Function &F, ScalarEvolution *SE_,
+bool AlignmentFromAssumptionsPass::runImpl(Function &F, AssumptionCache &AC,
+                                           ScalarEvolution *SE_,
                                            DominatorTree *DT_) {
   SE = SE_;
   DT = DT_;
@@ -420,12 +425,9 @@ bool AlignmentFromAssumptionsPass::runImpl(Function &F, ScalarEvolution *SE_,
   NewSrcAlignments.clear();
 
   bool Changed = false;
-
-  for (auto &B : F)
-    for (auto &I : B)
-      if (auto *II = dyn_cast<IntrinsicInst>(&I))
-        if (II->getIntrinsicID() == Intrinsic::assume)
-          Changed |= processAssumption(II);
+  for (auto &AssumeVH : AC.assumptions())
+    if (AssumeVH)
+      Changed |= processAssumption(cast<CallInst>(AssumeVH));
 
   return Changed;
 }
@@ -433,9 +435,10 @@ bool AlignmentFromAssumptionsPass::runImpl(Function &F, ScalarEvolution *SE_,
 PreservedAnalyses
 AlignmentFromAssumptionsPass::run(Function &F, FunctionAnalysisManager &AM) {
 
+  AssumptionCache &AC = AM.getResult<AssumptionAnalysis>(F);
   ScalarEvolution &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
   DominatorTree &DT = AM.getResult<DominatorTreeAnalysis>(F);
-  bool Changed = runImpl(F, &SE, &DT);
+  bool Changed = runImpl(F, AC, &SE, &DT);
 
   // FIXME: We need to invalidate this to avoid PR28400. Is there a better
   // solution?