[llvm] [SCEV] Fix infinite recursion in expensive range sharpening (PR #215029)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: Kacper Doga (varev-dev)

<details>
<summary>Changes</summary>

`getRangeForAffineNoSelfWrappingAR()` collects loop guards, which queries ranges of other expressions and can end up back in the range computation for the same AddRec. Neither the guard collection depth nor `getRangeRef()` depth is carried across that boundary, so this recurses until the stack overflows.

Mark AddRecs currently being sharpened and return a full range on re-entry, following the existing Pending* idiom for recursive SCEV queries.

Fixes #<!-- -->158255 

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


3 Files Affected:

- (modified) llvm/include/llvm/Analysis/ScalarEvolution.h (+4) 
- (modified) llvm/lib/Analysis/ScalarEvolution.cpp (+9) 
- (added) llvm/test/Analysis/ScalarEvolution/range-sharpening-recursion.ll (+43) 


``````````diff
diff --git a/llvm/include/llvm/Analysis/ScalarEvolution.h b/llvm/include/llvm/Analysis/ScalarEvolution.h
index 0d7f9ae298e2a..bea83b6165e0e 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolution.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolution.h
@@ -1716,6 +1716,10 @@ class ScalarEvolution {
   // Mark SCEVUnknown Phis currently being processed by isImpliedViaMerge.
   SmallPtrSet<const PHINode *, 6> PendingMerges;
 
+  /// Mark AddRecs currently being processed by
+  /// getRangeForAffineNoSelfWrappingAR.
+  SmallPtrSet<const SCEVAddRecExpr *, 4> PendingRangeSharpening;
+
   /// Set to true by isLoopBackedgeGuardedByCond when we're walking the set of
   /// conditions dominating the backedge of a loop.
   bool WalkingBEDominatingConds = false;
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 27a1a20bcdf79..c941338aa1aaa 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -7261,6 +7261,12 @@ ConstantRange ScalarEvolution::getRangeForAffineNoSelfWrappingAR(
                                          MaxItersWithoutWrap))
     return ConstantRange::getFull(BitWidth);
 
+  // Collecting loop guards below queries ranges of other expressions, which can
+  // recursively end up back here for the same AddRec.
+  if (!PendingRangeSharpening.insert(AddRec).second)
+    return ConstantRange::getFull(BitWidth);
+  llvm::scope_exit ClearOnExit([&]() { PendingRangeSharpening.erase(AddRec); });
+
   ICmpInst::Predicate LEPred =
       IsSigned ? ICmpInst::ICMP_SLE : ICmpInst::ICMP_ULE;
   ICmpInst::Predicate GEPred =
@@ -14112,6 +14118,7 @@ ScalarEvolution::ScalarEvolution(ScalarEvolution &&Arg)
       ValueExprMap(std::move(Arg.ValueExprMap)),
       PendingLoopPredicates(std::move(Arg.PendingLoopPredicates)),
       PendingMerges(std::move(Arg.PendingMerges)),
+      PendingRangeSharpening(std::move(Arg.PendingRangeSharpening)),
       ConstantMultipleCache(std::move(Arg.ConstantMultipleCache)),
       BackedgeTakenCounts(std::move(Arg.BackedgeTakenCounts)),
       PredicatedBackedgeTakenCounts(
@@ -14155,6 +14162,8 @@ ScalarEvolution::~ScalarEvolution() {
 
   assert(PendingLoopPredicates.empty() && "isImpliedCond garbage");
   assert(PendingMerges.empty() && "isImpliedViaMerge garbage");
+  assert(PendingRangeSharpening.empty() &&
+         "getRangeForAffineNoSelfWrappingAR garbage");
   assert(!WalkingBEDominatingConds && "isLoopBackedgeGuardedByCond garbage!");
   assert(!ProvingSplitPredicate && "ProvingSplitPredicate garbage!");
 }
diff --git a/llvm/test/Analysis/ScalarEvolution/range-sharpening-recursion.ll b/llvm/test/Analysis/ScalarEvolution/range-sharpening-recursion.ll
new file mode 100644
index 0000000000000..e6a007c58a76f
--- /dev/null
+++ b/llvm/test/Analysis/ScalarEvolution/range-sharpening-recursion.ll
@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -disable-output -scalar-evolution-use-expensive-range-sharpening -passes='print<scalar-evolution>' 2>&1 | FileCheck %s
+
+; Sharpening the range of %iv collects the loop guards of %inner, which queries
+; the range of %iv again. Check we do not recurse infinitely.
+
+define void @test(i1 %c) {
+; CHECK-LABEL: 'test'
+; CHECK-NEXT:  Classifying expressions for: @test
+; CHECK-NEXT:    %p = phi i32 [ 0, %inner.exit ], [ 0, %other ], [ 0, %entry ]
+; CHECK-NEXT:    --> 0 U: [0,1) S: [0,1) Exits: 0 LoopDispositions: { %outer: Invariant, %inner: Invariant }
+; CHECK-NEXT:    %iv = phi i32 [ %iv.next, %inner ], [ 0, %outer ]
+; CHECK-NEXT:    --> {0,+,1}<nuw><nsw><%inner> U: [0,2) S: [0,2) Exits: 1 LoopDispositions: { %inner: Computable, %outer: Uniform }
+; CHECK-NEXT:    %iv.next = add i32 %iv, 1
+; CHECK-NEXT:    --> {1,+,1}<nuw><nsw><%inner> U: [1,3) S: [1,3) Exits: 2 LoopDispositions: { %inner: Computable, %outer: Uniform }
+; CHECK-NEXT:  Determining loop execution counts for: @test
+; CHECK-NEXT:  Loop %inner: backedge-taken count is i32 1
+; CHECK-NEXT:  Loop %inner: constant max backedge-taken count is i32 1
+; CHECK-NEXT:  Loop %inner: symbolic max backedge-taken count is i32 1
+; CHECK-NEXT:  Loop %inner: Trip multiple is 2
+; CHECK-NEXT:  Loop %outer: <multiple exits> Unpredictable backedge-taken count.
+; CHECK-NEXT:  Loop %outer: Unpredictable constant max backedge-taken count.
+; CHECK-NEXT:  Loop %outer: Unpredictable symbolic max backedge-taken count.
+;
+entry:
+  br i1 %c, label %other, label %outer
+
+other:
+  br label %outer
+
+outer:
+  %p = phi i32 [ 0, %inner.exit ], [ 0, %other ], [ 0, %entry ]
+  br label %inner
+
+inner:
+  %iv = phi i32 [ %iv.next, %inner ], [ 0, %outer ]
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp ult i32 %iv, 1
+  br i1 %cmp, label %inner, label %inner.exit
+
+inner.exit:
+  br label %outer
+}

``````````

</details>


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