[SCEV] Also use info from assumes in applyLoopGuards.

[fhahn]

Summary

Similar to collecting information from branches guarding a loop, we can
also collect information from assumes dominating the loop header.

Fixes PR47247.

[jdoerfert]

commit message?

llvm/lib/Analysis/ScalarEvolution.cpp

12600

Nit: perform the dyn_cast first, dominance check is probably more costly.

[fhahn]

In D87854#2280261, @jdoerfert wrote:

commit message?

It was not quite ready yet, but now everything should be in place and this patch only contains the important change. I'll update the title/description shortly after updating the diff.

[jdoerfert]

LGTM

llvm/lib/Analysis/ScalarEvolution.cpp

for (BasicBlock *ExitingBB : ExitingBlocks) {
     }
   }
   if (ExitCounts.empty())
     return getCouldNotCompute();
   return getUMinFromMismatchedTypes(ExitCounts);
 }
 
 const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
+  auto ApplyCondition = [&](ICmpInst *Cmp) {
+    Value *V = Cmp->getOperand(0);
+    const SCEV *Op0 = getSCEV(Cmp->getOperand(0));
+    const SCEV *Op1 = getSCEV(Cmp->getOperand(1));
+    auto Predicate = Cmp->getPredicate();
+    if (isa<SCEVConstant>(Op0)) {
+      std::swap(Op0, Op1);
+      Predicate = CmpInst::getSwappedPredicate(Predicate);
+      V = Cmp->getOperand(1);
+    }
+    if (!isa<SCEVUnknown>(Op0) || !isa<SCEVConstant>(Op1))
+      return;
+
+    // TODO: use information from more predicates.
+    switch (Predicate) {
+    case CmpInst::ICMP_ULT: {
+      ValueToSCEVMapTy RewriteMap;
+      RewriteMap[V] =
+          getUMinExpr(Op0, getMinusSCEV(Op1, getOne(Op1->getType())));
+      Expr = SCEVParameterRewriter::rewrite(Expr, *this, RewriteMap);
+      break;
+    }
+
+    case CmpInst::ICMP_EQ: {
+      ValueToSCEVMapTy RewriteMap;
+      RewriteMap[V] = getUMaxExpr(Op0, getAddExpr(Op1, getOne(Op1->getType())));
+      Expr = SCEVParameterRewriter::rewrite(Expr, *this, RewriteMap);
+      break;
+    }
+    default:
+      break;
+    }
+  };
   // Starting at the loop predecessor, climb up the predecessor chain, as long
   // as there are predecessors that can be found that have unique successors
   // leading to the original header.
   // TODO: share this logic with isLoopEntryGuardedByCond.
   for (std::pair<const BasicBlock *, const BasicBlock *> Pair(
            L->getLoopPredecessor(), L->getHeader());
        Pair.first; Pair = getPredecessorWithUniqueSuccessorForBB(Pair.first)) {
 
     const BranchInst *LoopEntryPredicate =
         dyn_cast<BranchInst>(Pair.first->getTerminator());
     if (!LoopEntryPredicate || LoopEntryPredicate->isUnconditional())
       continue;
 
     // TODO: use information from more complex conditions, e.g. AND expressions.
     auto *Cmp = dyn_cast<ICmpInst>(LoopEntryPredicate->getCondition());
     if (!Cmp)
       continue;
 
-    // TODO: use information from more predicates.
-    switch (Cmp->getPredicate()) {
-    case CmpInst::ICMP_ULT: {
-      const SCEV *LHS = getSCEV(Cmp->getOperand(0));
-      const SCEV *RHS = getSCEV(Cmp->getOperand(1));
-      if (isa<SCEVUnknown>(LHS)) {
-        ValueToSCEVMapTy RewriteMap;
-        RewriteMap[Cmp->getOperand(0)] =
-            getUMinExpr(LHS, getMinusSCEV(RHS, getOne(RHS->getType())));
-        Expr = SCEVParameterRewriter::rewrite(Expr, *this, RewriteMap);
-      }
+    ApplyCondition(Cmp);
 
-      break;
-    }
-    default:
-      break;
-    }
+  for (auto &AssumeVH : AC.assumptions()) {
+    if (!AssumeVH)
+      continue;
+    auto *CI = cast<CallInst>(AssumeVH);
+    if (!DT.dominates(CI, L->getHeader()))
+      continue;
+    auto *Cmp = dyn_cast<ICmpInst>(CI->getOperand(0));
+    if (!Cmp)
+      continue;
+
+    ApplyCondition(Cmp);
   }
 
   return Expr;
 }

llvm/test/Analysis/ScalarEvolution/max-backedge-taken-count-guard-info.ll

 }
 
 ; Test case for PR47247. Both the guard condition and the assume limit the
 ; max backedge-taken count.
 
 define void @test_guard_and_assume(i32* nocapture readonly %data, i64 %count) {
 ; CHECK-LABEL: Determining loop execution counts for: @test_guard_and_assume
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is (-1 + %count)
-; CHECK-NEXT:  Loop %loop: max backedge-taken count is -2
+; CHECK-NEXT:  Loop %loop: max backedge-taken count is 3
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is (-1 + %count)
 ;
 entry:
   %cmp = icmp ult i64 %count, 5
   tail call void @llvm.assume(i1 %cmp)
   %cmp18.not = icmp eq i64 %count, 0
   br i1 %cmp18.not, label %exit, label %loop