[llvm] [ConstraintElim] Forget loop SCEV when folding an exiting condition. (PR #213945)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Uzair rehman (Uzair-90)

<details>
<summary>Changes</summary>

ConstraintElimination preserves ScalarEvolution, but when it folds the condition of a loop-exiting branch to a constant, the loop's cached backedge-taken count (computed against the original exit structure, often populated by the pass's own SCEV queries in `addInfoForInductions`) becomes stale: still sound, but less precise than fresh recomputation. Later trip-count based folds, e.g. in indvars, then fail.

For the test case from #<!-- -->213872, the cached count after constraint-elimination is
```
Loop %loop: <multiple exits> backedge-taken count is (((3 + (4 * %count))<nuw><nsw> /u 4) umin ((4 * %count) /u 4))
```
still including an exit count for the `loop.latch` exit whose branch was just folded to `br i1 true`, while fresh recomputation gives `((4 * %count) /u 4)` with the latch exit as `***COULDNOTCOMPUTE***`. The stale `umin` is what indvars then materializes in the preheader instead of folding the guard.

This patch records loops whose exiting-block terminators have their condition replaced in `ReplaceCmpWithConstant` and calls `forgetTopmostLoop` on them after the replacements. `forgetValue` on the compare would not be enough, since backedge-taken counts are expressed in terms of the compare's operands, not the compare itself.

`checkOrAndOpImpliedByOther` and `checkAndReplaceMinMax` can in principle leave similarly stale (sound but imprecise) cached SCEVs; left for a follow-up to keep this change minimal.

Fixes #<!-- -->213872.

---
Full diff: https://github.com/llvm/llvm-project/pull/213945.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/ConstraintElimination.cpp (+24-2) 
- (added) llvm/test/Transforms/ConstraintElimination/preserve-scev.ll (+32) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 945ee0351e2e3..5a5dbf5be4a28 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -1602,12 +1602,14 @@ static bool checkAndReplaceCondition(
     ConstraintInfo &Info, unsigned NumIn, unsigned NumOut,
     Instruction *ContextInst, Module *ReproducerModule,
     ArrayRef<ReproducerEntry> ReproducerCondStack, DominatorTree &DT,
+    LoopInfo &LI, ScalarEvolution &SE,
     SmallVectorImpl<Instruction *> &ToRemove) {
   auto ReplaceCmpWithConstant = [&](Instruction *CheckInst, bool IsTrue) {
     generateReproducer(CheckInst, ICmpInst::isSigned(Pred), ReproducerModule,
                        ReproducerCondStack, Info, DT);
     Constant *ConstantC = ConstantInt::getBool(
         CmpInst::makeCmpResultType(CheckInst->getType()), IsTrue);
+    SmallPtrSet<const Loop *, 2> AffectedLoops;
     bool Changed = CheckInst->replaceUsesWithIf(ConstantC, [&](Use &U) {
       auto *UserI = getContextInstForUse(U);
       auto *DTN = DT.getNode(UserI->getParent());
@@ -1620,8 +1622,27 @@ static bool checkAndReplaceCondition(
       // Conditions in an assume trivially simplify to true. Skip uses
       // in assume calls to not destroy the available information.
       auto *II = dyn_cast<IntrinsicInst>(U.getUser());
-      return !II || II->getIntrinsicID() != Intrinsic::assume;
+      if (II && II->getIntrinsicID() == Intrinsic::assume)
+        return false;
+
+      // Replacing the condition of a terminator in a loop-exiting block may
+      // change how often the loop is exited, so cached trip counts, while
+      // still correct, may be less precise than freshly computed ones. Keep
+      // track of the affected loops, so their cached info can be dropped.
+      if (auto *User = dyn_cast<Instruction>(U.getUser());
+          User && User->isTerminator()) {
+        BasicBlock *BB = User->getParent();
+        if (Loop *L = LI.getLoopFor(BB); L && L->isLoopExiting(BB))
+          AffectedLoops.insert(L);
+      }
+      return true;
     });
+
+    // The exit conditions of the affected loops (and any enclosing loops the
+    // exiting blocks may also exit) changed; drop the cached trip counts so
+    // later passes re-compute them with the now more precise exit structure.
+    for (const Loop *L : AffectedLoops)
+      SE.forgetTopmostLoop(L);
     NumCondsRemoved++;
 
     // Update the debug value records that satisfy the same condition used
@@ -2059,7 +2080,8 @@ static bool eliminateConstraints(Function &F, DominatorTree &DT, LoopInfo &LI,
       } else if (match(Inst, m_ICmpLike(Pred, m_Value(A), m_Value(B)))) {
         bool Simplified = checkAndReplaceCondition(
             Pred, A, B, Inst, Info, CB.NumIn, CB.NumOut, CB.getContextInst(),
-            ReproducerModule.get(), ReproducerCondStack, S.DT, ToRemove);
+            ReproducerModule.get(), ReproducerCondStack, S.DT, S.LI, S.SE,
+            ToRemove);
         if (!Simplified &&
             match(CB.getContextInst(), m_LogicalOp(m_Value(), m_Value()))) {
           Simplified = checkOrAndOpImpliedByOther(
diff --git a/llvm/test/Transforms/ConstraintElimination/preserve-scev.ll b/llvm/test/Transforms/ConstraintElimination/preserve-scev.ll
new file mode 100644
index 0000000000000..53f0971120dae
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/preserve-scev.ll
@@ -0,0 +1,32 @@
+; RUN: opt -passes='constraint-elimination,print<scalar-evolution>' -disable-output %s 2>&1 | FileCheck %s
+
+; Make sure ScalarEvolution's cached trip counts are dropped when
+; constraint-elimination folds the condition of a loop-exiting branch. The
+; preserved analysis would otherwise keep a stale, less precise backedge-taken
+; count (umin of both original exits), blocking later trip-count based folds,
+; e.g. in indvars.
+; See https://github.com/llvm/llvm-project/issues/213872.
+
+; CHECK-LABEL: Classifying expressions for: @multiple_pow2
+; CHECK:       Loop %loop: <multiple exits> backedge-taken count is ((4 * %count) /u 4)
+; CHECK-NEXT:    exit count for loop: ((4 * %count) /u 4)
+; CHECK-NEXT:    exit count for loop.latch: ***COULDNOTCOMPUTE***
+
+define void @multiple_pow2(i64 %count) {
+entry:
+  %end = shl i64 %count, 2
+  br label %loop
+
+loop:                                             ; preds = %loop.latch, %entry
+  %iv = phi i64 [ %iv.next, %loop.latch ], [ 0, %entry ]
+  %cmp.i.not = icmp eq i64 %iv, %end
+  br i1 %cmp.i.not, label %exit, label %loop.latch
+
+loop.latch:                                       ; preds = %loop
+  %iv.next = add i64 %iv, 4
+  %cmp2.i.i = icmp ult i64 %iv, %end
+  br i1 %cmp2.i.i, label %loop, label %exit
+
+exit:                                             ; preds = %loop.latch, %loop
+  ret void
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/213945