1 files changed, 14 insertions, 38 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
index 9e23fa9..db7a6d5 100644
--- a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
@@ -195,6 +195,9 @@ Value *SimplifyIndvar::foldIVUser(Instruction *UseInst, Instruction *IVOperand)
 
 bool SimplifyIndvar::makeIVComparisonInvariant(ICmpInst *ICmp,
                                                Instruction *IVOperand) {
+  auto *Preheader = L->getLoopPreheader();
+  if (!Preheader)
+    return false;
   unsigned IVOperIdx = 0;
   ICmpInst::Predicate Pred = ICmp->getPredicate();
   if (IVOperand != ICmp->getOperand(0)) {
@@ -209,11 +212,6 @@ bool SimplifyIndvar::makeIVComparisonInvariant(ICmpInst *ICmp,
   const Loop *ICmpLoop = LI->getLoopFor(ICmp->getParent());
   const SCEV *S = SE->getSCEVAtScope(ICmp->getOperand(IVOperIdx), ICmpLoop);
   const SCEV *X = SE->getSCEVAtScope(ICmp->getOperand(1 - IVOperIdx), ICmpLoop);
-
-  auto *PN = dyn_cast<PHINode>(IVOperand);
-  if (!PN)
-    return false;
-
   auto LIP = SE->getLoopInvariantPredicate(Pred, S, X, L, ICmp);
   if (!LIP)
     return false;
@@ -221,40 +219,18 @@ bool SimplifyIndvar::makeIVComparisonInvariant(ICmpInst *ICmp,
   const SCEV *InvariantLHS = LIP->LHS;
   const SCEV *InvariantRHS = LIP->RHS;
 
-  // Rewrite the comparison to a loop invariant comparison if it can be done
-  // cheaply, where cheaply means "we don't need to emit any new
-  // instructions".
-
-  SmallDenseMap<const SCEV*, Value*> CheapExpansions;
-  CheapExpansions[S] = ICmp->getOperand(IVOperIdx);
-  CheapExpansions[X] = ICmp->getOperand(1 - IVOperIdx);
-
-  // TODO: Support multiple entry loops?  (We currently bail out of these in
-  // the IndVarSimplify pass)
-  if (auto *BB = L->getLoopPredecessor()) {
-    const int Idx = PN->getBasicBlockIndex(BB);
-    if (Idx >= 0) {
-      Value *Incoming = PN->getIncomingValue(Idx);
-      const SCEV *IncomingS = SE->getSCEV(Incoming);
-      CheapExpansions[IncomingS] = Incoming;
-    }
-  }
-  Value *NewLHS = CheapExpansions[InvariantLHS];
-  Value *NewRHS = CheapExpansions[InvariantRHS];
-
-  if (!NewLHS)
-    if (auto *ConstLHS = dyn_cast<SCEVConstant>(InvariantLHS))
-      NewLHS = ConstLHS->getValue();
-  if (!NewRHS)
-    if (auto *ConstRHS = dyn_cast<SCEVConstant>(InvariantRHS))
-      NewRHS = ConstRHS->getValue();
-
-  if (!NewLHS || !NewRHS)
-    // We could not find an existing value to replace either LHS or RHS.
-    // Generating new instructions has subtler tradeoffs, so avoid doing that
-    // for now.
+  // Do not generate something ridiculous.
+  auto *PHTerm = Preheader->getTerminator();
+  if (Rewriter.isHighCostExpansion(InvariantLHS, L, SCEVCheapExpansionBudget,
+                                   TTI, PHTerm))
     return false;
-
+  if (Rewriter.isHighCostExpansion(InvariantRHS, L, SCEVCheapExpansionBudget,
+                                   TTI, PHTerm))
+    return false;
+  auto *NewLHS =
+      Rewriter.expandCodeFor(InvariantLHS, IVOperand->getType(), PHTerm);
+  auto *NewRHS =
+      Rewriter.expandCodeFor(InvariantRHS, IVOperand->getType(), PHTerm);
   LLVM_DEBUG(dbgs() << "INDVARS: Simplified comparison: " << *ICmp << '\n');
   ICmp->setPredicate(InvariantPredicate);
   ICmp->setOperand(0, NewLHS);