2 files changed, 47 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index c69bd601aff..e872e81a5e7 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -4756,11 +4756,26 @@ const SCEV *ScalarEvolution::createAddRecFromPHI(PHINode *PN) {
           Ops.push_back(Add->getOperand(i));
       const SCEV *Accum = getAddExpr(Ops);
 
+      bool InvariantF = isLoopInvariant(Accum, L);
+
+      if (!InvariantF && Accum->getSCEVType() == scZeroExtend) {
+        const SCEV *Op = dyn_cast<SCEVZeroExtendExpr>(Accum)->getOperand();
+        const SCEVUnknown *Un = dyn_cast<SCEVUnknown>(Op);
+        if (Un && Un->getValue() && isa<Instruction>(Un->getValue()) &&
+            dyn_cast<Instruction>(Un->getValue())->getOpcode() ==
+                Instruction::ICmp) {
+          const SCEV *ICmpSC = evaluateForICmp(cast<ICmpInst>(Un->getValue()));
+          bool IsConstSC = ICmpSC->getSCEVType() == scConstant;
+          Accum =
+              IsConstSC ? getZeroExtendExpr(ICmpSC, Accum->getType()) : Accum;
+          InvariantF = IsConstSC ? true : false;
+        }
+      }
+
       // This is not a valid addrec if the step amount is varying each
       // loop iteration, but is not itself an addrec in this loop.
-      if (isLoopInvariant(Accum, L) ||
-          (isa<SCEVAddRecExpr>(Accum) &&
-           cast<SCEVAddRecExpr>(Accum)->getLoop() == L)) {
+      if (InvariantF || (isa<SCEVAddRecExpr>(Accum) &&
+                         cast<SCEVAddRecExpr>(Accum)->getLoop() == L)) {
         SCEV::NoWrapFlags Flags = SCEV::FlagAnyWrap;
 
         if (auto BO = MatchBinaryOp(BEValueV, DT)) {
@@ -6443,6 +6458,30 @@ void ScalarEvolution::forgetLoop(const Loop *L) {
   }
 }
 
+
+const SCEV *ScalarEvolution::evaluateForICmp(ICmpInst *IC) {
+  BasicBlock *Latch = nullptr;
+  const Loop *L = LI.getLoopFor(IC->getParent());
+
+  // If compare instruction is same or inverse of the compare in the
+  // branch of the loop latch, then return a constant evolution
+  // node. This shall facilitate computations of loop exit counts
+  // in cases where compare appears in the evolution chain of induction
+  // variables.
+  if (L && (Latch = L->getLoopLatch())) {
+    BranchInst *BI = dyn_cast<BranchInst>(Latch->getTerminator());
+    if (BI && BI->isConditional() && BI->getCondition() == IC) {
+      if (BI->getSuccessor(0) != L->getHeader())
+        return getZero(Type::getInt1Ty(getContext()));
+      else
+        return getOne(Type::getInt1Ty(getContext()));
+    }
+  }
+
+  return getUnknown(IC);
+}
+
+
 void ScalarEvolution::forgetValue(Value *V) {
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I) return;
diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 6462e3fb856..8551392baa8 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -2973,8 +2973,11 @@ void LSRInstance::CollectChains() {
       // Ignore users that are part of a SCEV expression. This way we only
       // consider leaf IV Users. This effectively rediscovers a portion of
       // IVUsers analysis but in program order this time.
-      if (SE.isSCEVable(I.getType()) && !isa<SCEVUnknown>(SE.getSCEV(&I)))
-        continue;
+      if (SE.isSCEVable(I.getType())) {
+        const SCEV *SI = SE.getSCEV(&I);
+        if (!isa<SCEVUnknown>(SI) && !isa<SCEVConstant>(SI))
+          continue;
+      }
 
       // Remove this instruction from any NearUsers set it may be in.
       for (unsigned ChainIdx = 0, NChains = IVChainVec.size();