[SCEV] Support rewriting ZExt expressions with loop guard info.

[fhahn]

Summary

So far, applying loop guard information has been restricted to
SCEVUnknown. In a few cases, like PR40961 and PR52464, this leads to
SCEV failing to determine tight upper bounds for the backedge taken
count.

This patch adjusts SCEVLoopGuardRewriter and applyLoopGuards to support
re-writing ZExt expressions.

This is a first step towards fixing PR40961 and PR52464.

[nikic]

llvm/lib/Analysis/ScalarEvolution.cpp

13722

Can you please add a test where you both have a condition on %x and on zext %x? I think we'll end up losing the information from the direct condition in that case.

Something I also find odd here is the mismatch between the code collecting rewrites and the visitor. We'll collect rewrites on all SCEVs, but then only rewrite unknown and zext.

[fhahn]

Updated the code to also apply information collected to other expressions in the mas. This ensures we include information about sub-expressions when applying larger expressions. I added tests that show cases where we would otherwise miss information ( 5dfe60d171d7) and rebased on top of that.

I also restricted info collection to LHS that are either SCEVUnknown or SCEVZeroExtendExpr for now.

Compile-time wise this shows a very small increase on average: http://llvm-compile-time-tracker.com/compare.php?from=b57c22ade8673d68088f5638a1050e3f3c1f163d&to=5792b893889fbfd459b9ec062dedea6c130b9786&stat=instructions

NewPM-O3: +0.02%
NewPM-ReleaseThinLTO: +0.02%
NewPM-ReleaseLTO-g: +0.04%

[fhahn]

llvm/lib/Analysis/ScalarEvolution.cpp

13722

Can you please add a test where you both have a condition on %x and on zext %x? I think we'll end up losing the information from the direct condition in that case.

Added those tests in 5dfe60d171d7. The original version failed to apply information of sub-expressions to larger expressions, so we failed to determine the tightest bound in one the cases. I updated the patch to re-write the collected expressions with the information we collected, which should resolve that issue.

Something I also find odd here is the mismatch between the code collecting rewrites and the visitor. We'll collect rewrites on all SCEVs, but then only rewrite unknown and zext.

Good point! There are other motivating cases that require re-writing other expressions as well (for PR40961), but for now I also restricted collecting to LHSs that are either SCEVUnknown or SCEVZeroExtendExpr.

[reames]

This generally looks very reasonable.

Can I request that you split off the change which switches the map from value keyed to SCEV keyed, and land that without separate review? It seems like it should be entirely NFC. Reducing diff and isolating failures is always good.

llvm/lib/Analysis/ScalarEvolution.cpp

13923

This should follow from the empty check just above.

13932

You can fold this into the previous code by simply pushing Expr onto the replacement list last.

Hm, do we need a fixed point in the loop above? It seems like we have the possibility of one replacement enabling another. Is there something non-obvious in the above which handles that?

[fhahn]

In D113577#3127717, @reames wrote:

This generally looks very reasonable.

Can I request that you split off the change which switches the map from value keyed to SCEV keyed, and land that without separate review? It seems like it should be entirely NFC. Reducing diff and isolating failures is always good.

Thanks, I committed the SCEV->SECV map change in 03cfea68c65f and rebased the patch.

I'll respond to the inline comments in a bit.

[reames]

Thanks for the rebase, much easier to understand.

LGTM w/the previous inline style comment addressed. You can follow up on the fixed point question post-commit if desired.

[fhahn]

In D113577#3127923, @reames wrote:

Thanks for the rebase, much easier to understand.

LGTM w/the previous inline style comment addressed. You can follow up on the fixed point question post-commit if desired.

Thanks! I still need to construct a couple of test cases for the follow-up. But I thought a bit more and I think I need maintain a vector with expressions to rewrite separately, to avoid depending on the iteration order of the dense map, which may not be consistent. I'm planning to incorporate that in the commit version, independent of following up on the fixed-point iteration.

[fhahn]

rebased.

@nikic it looks like there's a crash on the llvm-compile-time-tracker with the latest version (applying rewrites), but I was unable to reproduce this locally with various combinations of flags. What compiler/flags is used by the tracker to build the compiler used?

The log to the failure is http://llvm-compile-time-tracker.com/show_error.php?commit=d628e5e177b9a9dd4f279ce7c301fd312be0cfa7 in case you have any ideas.

[nikic]

In D113577#3131976, @fhahn wrote:

rebased.

@nikic it looks like there's a crash on the llvm-compile-time-tracker with the latest version (applying rewrites), but I was unable to reproduce this locally with various combinations of flags. What compiler/flags is used by the tracker to build the compiler used?

The log to the failure is http://llvm-compile-time-tracker.com/show_error.php?commit=d628e5e177b9a9dd4f279ce7c301fd312be0cfa7 in case you have any ideas.

Here's the (unreduced) IR, crashes for me under opt -O1: https://gist.github.com/nikic/bd117601aedea5f03b5750571d07624e

I just ran the /root/llvm-compile-time-tracker/llvm-project-build/bin/clang -DNDEBUG -fno-experimental-new-pass-manager -O3 -fomit-frame-pointer -flto -DNDEBUG -g -w -Werror=date-time -MD -MT tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o -MF tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o.d -o tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o -c /root/llvm-compile-time-tracker/llvm-test-suite/tools/fpcmp.c command after adjusting to local paths, but sometimes these things depend on external deps like the used libc.

[fhahn]

In D113577#3132006, @nikic wrote:

In D113577#3131976, @fhahn wrote:

rebased.

@nikic it looks like there's a crash on the llvm-compile-time-tracker with the latest version (applying rewrites), but I was unable to reproduce this locally with various combinations of flags. What compiler/flags is used by the tracker to build the compiler used?

The log to the failure is http://llvm-compile-time-tracker.com/show_error.php?commit=d628e5e177b9a9dd4f279ce7c301fd312be0cfa7 in case you have any ideas.

Here's the (unreduced) IR, crashes for me under opt -O1: https://gist.github.com/nikic/bd117601aedea5f03b5750571d07624e

I just ran the /root/llvm-compile-time-tracker/llvm-project-build/bin/clang -DNDEBUG -fno-experimental-new-pass-manager -O3 -fomit-frame-pointer -flto -DNDEBUG -g -w -Werror=date-time -MD -MT tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o -MF tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o.d -o tools/CMakeFiles/fpcmp-target.dir/fpcmp.c.o -c /root/llvm-compile-time-tracker/llvm-test-suite/tools/fpcmp.c command after adjusting to local paths, but sometimes these things depend on external deps like the used libc.

Thanks! I think I figured out the issue: RewriteMap[Expr] = Rewriter.visit(RewriteTo); may add an entry to RewriteMap before evaluating RewriteTo depending on the compiler used.

llvm/lib/Analysis/ScalarEvolution.cpp

if (ExitCounts.empty())
     return getCouldNotCompute();
   return getUMinFromMismatchedTypes(ExitCounts);
 }
 
 /// A rewriter to replace SCEV expressions in Map with the corresponding entry
 /// in the map. It skips AddRecExpr because we cannot guarantee that the
 /// replacement is loop invariant in the loop of the AddRec.
 ///
-/// At the moment only rewriting SCEVUnknown is supported.
+/// At the moment only rewriting SCEVUnknown and SCEVZeroExtendExpr is
+/// supported.
 class SCEVLoopGuardRewriter : public SCEVRewriteVisitor<SCEVLoopGuardRewriter> {
   const DenseMap<const SCEV *, const SCEV *> &Map;
 
 public:
   SCEVLoopGuardRewriter(ScalarEvolution &SE,
                         DenseMap<const SCEV *, const SCEV *> &M)
       : SCEVRewriteVisitor(SE), Map(M) {}
 
   const SCEV *visitAddRecExpr(const SCEVAddRecExpr *Expr) { return Expr; }
 
   const SCEV *visitUnknown(const SCEVUnknown *Expr) {
     auto I = Map.find(Expr);
     if (I == Map.end())
       return Expr;
     return I->second;
   }
+
+  const SCEV *visitZeroExtendExpr(const SCEVZeroExtendExpr *Expr) {
+    auto I = Map.find(Expr);
+    if (I == Map.end())
+      return SCEVRewriteVisitor<SCEVLoopGuardRewriter>::visitZeroExtendExpr(
+          Expr);
+    return I->second;
+  }
 };
 
 const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
+  SmallVector<const SCEV *> ExprsToRewrite;
   auto CollectCondition = [&](ICmpInst::Predicate Predicate, const SCEV *LHS,
                               const SCEV *RHS,
                               DenseMap<const SCEV *, const SCEV *>
                                   &RewriteMap) {
     // WARNING: It is generally unsound to apply any wrap flags to the proposed
     // replacement SCEV which isn't directly implied by the structure of that
     // SCEV.  In particular, using contextual facts to imply flags is *NOT*
     // legal.  See the scoping rules for flags in the header to understand why.
 
     // If we have LHS == 0, check if LHS is computing a property of some unknown
     // SCEV %v which we can rewrite %v to express explicitly.
     const SCEVConstant *RHSC = dyn_cast<SCEVConstant>(RHS);
     if (Predicate == CmpInst::ICMP_EQ && RHSC &&
         RHSC->getValue()->isNullValue()) {
       // If LHS is A % B, i.e. A % B == 0, rewrite A to (A /u B) * B to
       // explicitly express that.
       const SCEV *URemLHS = nullptr;
       const SCEV *URemRHS = nullptr;
       if (matchURem(LHS, URemLHS, URemRHS)) {
         if (const SCEVUnknown *LHSUnknown = dyn_cast<SCEVUnknown>(URemLHS)) {
           auto Multiple = getMulExpr(getUDivExpr(URemLHS, URemRHS), URemRHS);
           RewriteMap[LHSUnknown] = Multiple;
+          ExprsToRewrite.push_back(LHSUnknown);
           return;
         }
       }
     }
 
     if (!isa<SCEVUnknown>(LHS) && isa<SCEVUnknown>(RHS)) {
       std::swap(LHS, RHS);
       Predicate = CmpInst::getSwappedPredicate(Predicate);
     }
 
     // Check for a condition of the form (-C1 + X < C2).  InstCombine will
     // create this form when combining two checks of the form (X u< C2 + C1) and
     // (X >=u C1).
-    auto MatchRangeCheckIdiom = [this, Predicate, LHS, RHS, &RewriteMap]() {
+    auto MatchRangeCheckIdiom = [this, Predicate, LHS, RHS, &RewriteMap,
+                                 &ExprsToRewrite]() {
       auto *AddExpr = dyn_cast<SCEVAddExpr>(LHS);
       if (!AddExpr || AddExpr->getNumOperands() != 2)
         return false;
 
       auto *C1 = dyn_cast<SCEVConstant>(AddExpr->getOperand(0));
       auto *LHSUnknown = dyn_cast<SCEVUnknown>(AddExpr->getOperand(1));
       auto *C2 = dyn_cast<SCEVConstant>(RHS);
       if (!C1 || !C2 || !LHSUnknown)
         return false;
 
       auto ExactRegion =
           ConstantRange::makeExactICmpRegion(Predicate, C2->getAPInt())
               .sub(C1->getAPInt());
 
       // Bail out, unless we have a non-wrapping, monotonic range.
       if (ExactRegion.isWrappedSet() || ExactRegion.isFullSet())
         return false;
       auto I = RewriteMap.find(LHSUnknown);
       const SCEV *RewrittenLHS = I != RewriteMap.end() ? I->second : LHSUnknown;
       RewriteMap[LHSUnknown] = getUMaxExpr(
           getConstant(ExactRegion.getUnsignedMin()),
           getUMinExpr(RewrittenLHS, getConstant(ExactRegion.getUnsignedMax())));
+      ExprsToRewrite.push_back(LHSUnknown);
       return true;
     };
     if (MatchRangeCheckIdiom())
       return;
 
-    // For now, limit to conditions that provide information about unknown
-    // expressions. RHS also cannot contain add recurrences.
-    auto *LHSUnknown = dyn_cast<SCEVUnknown>(LHS);
-    if (!LHSUnknown || containsAddRecurrence(RHS))
+    // If RHS is SCEVUnknown, make sure the information is applied to it.
+    if (isa<SCEVUnknown>(RHS)) {
+      std::swap(LHS, RHS);
+      Predicate = CmpInst::getSwappedPredicate(Predicate);
+    }
+    // If LHS is a constant, apply information to the other expression.
+    if (isa<SCEVConstant>(LHS)) {
+      std::swap(LHS, RHS);
+      Predicate = CmpInst::getSwappedPredicate(Predicate);
+    }
+    // Do not apply information for constants or if RHS contains an AddRec.
+    if (isa<SCEVConstant>(LHS) || containsAddRecurrence(RHS))
+      return;
+
+    // Limit to expressions that can be rewritten.
+    if (!isa<SCEVUnknown>(LHS) && !isa<SCEVZeroExtendExpr>(LHS))
       return;
 
     // Check whether LHS has already been rewritten. In that case we want to
     // chain further rewrites onto the already rewritten value.
     auto I = RewriteMap.find(LHS);
     const SCEV *RewrittenLHS = I != RewriteMap.end() ? I->second : LHS;
+
     const SCEV *RewrittenRHS = nullptr;
     switch (Predicate) {
     case CmpInst::ICMP_ULT:
       RewrittenRHS =
           getUMinExpr(RewrittenLHS, getMinusSCEV(RHS, getOne(RHS->getType())));
       break;
     case CmpInst::ICMP_SLT:
       RewrittenRHS =
case CmpInst::ICMP_NE:
       if (isa<SCEVConstant>(RHS) &&
           cast<SCEVConstant>(RHS)->getValue()->isNullValue())
         RewrittenRHS = getUMaxExpr(RewrittenLHS, getOne(RHS->getType()));
       break;
     default:
       break;
     }
 
-    if (RewrittenRHS)
+    if (RewrittenRHS) {
       RewriteMap[LHS] = RewrittenRHS;
+      if (LHS == RewrittenLHS)
+        ExprsToRewrite.push_back(LHS);
+    }
   };
   // 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.
   DenseMap<const SCEV *, const SCEV *> RewriteMap;
   for (std::pair<const BasicBlock *, const BasicBlock *> Pair(
            L->getLoopPredecessor(), L->getHeader());
for (auto &AssumeVH : AC.assumptions()) {
     if (!Cmp || !DT.dominates(AssumeI, L->getHeader()))
       continue;
     CollectCondition(Cmp->getPredicate(), getSCEV(Cmp->getOperand(0)),
                      getSCEV(Cmp->getOperand(1)), RewriteMap);
   }
 
   if (RewriteMap.empty())
     return Expr;
+
+  // Now that all rewrite information is collect, rewrite the collected
+  // expressions with the information in the map. This applies information to
+  // sub-expressions.
+  if (ExprsToRewrite.size() > 1) {
+    for (const SCEV *Expr : ExprsToRewrite) {
+      const SCEV *RewriteTo = RewriteMap[Expr];
+      RewriteMap.erase(Expr);
+      SCEVLoopGuardRewriter Rewriter(*this, RewriteMap);
+      RewriteMap.insert({Expr, Rewriter.visit(RewriteTo)});
+    }
+  }
+
   SCEVLoopGuardRewriter Rewriter(*this, RewriteMap);
   return Rewriter.visit(Expr);
 }

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

 ; RUN: opt -analyze -scalar-evolution %s -enable-new-pm=0 | FileCheck %s
 ; RUN: opt -passes='print<scalar-evolution>' -disable-output %s 2>&1 | FileCheck %s
 
 ; Test cases that require rewriting zext SCEV expression with infomration from
 ; the loop guards.
 
 define void @rewrite_zext(i32 %n) {
 ; CHECK-LABEL: Determining loop execution counts for: @rewrite_zext
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
-; CHECK-NEXT:  Loop %loop: max backedge-taken count is 2305843009213693951
+; CHECK-NEXT:  Loop %loop: max backedge-taken count is 2
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
 ; CHECK-NEXT:  Predicates:
 ; CHECK:        Loop %loop: Trip multiple is 1
 ;
 entry:
   %ext = zext i32 %n to i64
   %cmp5 = icmp ule i64 %ext, 24
   br i1 %cmp5, label %check, label %exit

 exit:
   ret void
 }
 
 ; Test case from PR40961.
 define i32 @rewrite_zext_min_max(i32 %N, i32* %arr) {
 ; CHECK-LABEL:  Determining loop execution counts for: @rewrite_zext_min_max
 ; CHECK-NEXT:   Loop %loop: backedge-taken count is ((-4 + (4 * ((zext i32 (16 umin %N) to i64) /u 4))<nuw><nsw>)<nsw> /u 4)
-; CHECK-NEXT:   Loop %loop: max backedge-taken count is 4611686018427387903
+; CHECK-NEXT:   Loop %loop: max backedge-taken count is 3
 ; CHECK-NEXT:   Loop %loop: Predicated backedge-taken count is ((-4 + (4 * ((zext i32 (16 umin %N) to i64) /u 4))<nuw><nsw>)<nsw> /u 4)
 ; CHECK-NEXT:   Predicates:
 ; CHECK:         Loop %loop: Trip multiple is 1
 ;
 entry:
   %umin = call i32 @llvm.umin.i32(i32 %N, i32 16)
   %ext = zext i32 %umin to i64
   %min.iters.check = icmp ult i64 %ext, 4

 exit:
   ret i32 0
 }
 
 ; Make sure no information is lost for conditions on both %n and (zext %n).
 define void @rewrite_zext_and_base_1(i32 %n) {
 ; CHECK-LABEL: Determining loop execution counts for: @rewrite_zext_and_base
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
-; CHECK-NEXT:  Loop %loop: max backedge-taken count is 2305843009213693951
+; CHECK-NEXT:  Loop %loop: max backedge-taken count is 3
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
 ; CHECK-NEXT:  Predicates:
 ; CHECK:        Loop %loop: Trip multiple is 1
 ;
 entry:
   %ext = zext i32 %n to i64
   %cmp5 = icmp ule i64 %ext, 48
   br i1 %cmp5, label %check.1, label %exit

 exit:
   ret void
 }
 
 ; Make sure no information is lost for conditions on both %n and (zext %n).
 define void @rewrite_zext_and_base_2(i32 %n) {
 ; CHECK-LABEL: Determining loop execution counts for: @rewrite_zext_and_base
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
-; CHECK-NEXT:  Loop %loop: max backedge-taken count is 2305843009213693951
+; CHECK-NEXT:  Loop %loop: max backedge-taken count is 3
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is ((-8 + (8 * ((zext i32 %n to i64) /u 8))<nuw><nsw>)<nsw> /u 8)
 ; CHECK-NEXT:  Predicates:
 ; CHECK:        Loop %loop: Trip multiple is 1
 ;
 entry:
   %ext = zext i32 %n to i64
   %cmp5 = icmp ule i64 %ext, 32
   br i1 %cmp5, label %check.1, label %exit