[SCEV] Use trip count information to improve shift recurrence ranges

[reames]

Summary

This patch exploits the knowledge that we may be running many fewer than bitwidth iterations of the loop, and may be able to disallow the overflow case. This patch specifically implements only the shl case, but this can be generalized to ashr and lshr without difficulty.

In my motivating example (a multiple dimension loop nest w/a shl recurrence at the outer level), this results in a massive improvement in scev-aa results.

[lebedev.ri]

llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll

1

Please precommit

[reames]

rebase over landed tests

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

5663

LI.isLoopHeader(P->getParent())?
I also wonder why is that needed. The implementation of matchSimpleRecurrence cannot recognize something which is not a header phi, because there is no other way to have BO as operand of Phi and vice versa at the same time. I guess this may be turned into assert after matchSimpleRecurrence.

5671

Looks like it's better split it into two checks, as ashr and lshr will still need BO->getOperand(0) != P.

5673

Are we missing a check that another operand is loop-invariant, or it is implied by something?

[mkazantsev]

Regarding the helper you are using: looking into matchSimpleRecurrence, I don't really think it's doing what it seems to be doing. Consider case:

loop:
  iv = phi [start][iv.next]
  x = load volatile ...  // Loop-variant load
  iv.next = shl iv, x

I didn't really run a test, but from the code I cannot figure why it wouldn't be recognized as a simple recurrence. I guess the loop-invariance check is missing inside this helper, not in your current patch. The logic there could also be simplified if we gave it LoopInfo with notion of latch (we wouldn't need 2 iterations in loop there). Please consider this refactoring.

[reames]

In D98222#2621535, @mkazantsev wrote:

Regarding the helper you are using: looking into matchSimpleRecurrence, I don't really think it's doing what it seems to be doing. Consider case:

loop:
  iv = phi [start][iv.next]
  x = load volatile ...  // Loop-variant load
  iv.next = shl iv, x

I didn't really run a test, but from the code I cannot figure why it wouldn't be recognized as a simple recurrence. I guess the loop-invariance check is missing inside this helper, not in your current patch. The logic there could also be simplified if we gave it LoopInfo with notion of latch (we wouldn't need 2 iterations in loop there). Please consider this refactoring.

The lack of an invariant check in the matching routine is deliberate. Please see usage .e.g. in computeKnownBits. Not all callers want the restriction on step being loop invariant.

In this particular case, you are correct that there is a missing invariance check. I'd originally written this for constant step (which are trivially invariant), but forgot to add back the invariance check when generalizing. Thanks for the catch!

[reames]

llvm/lib/Analysis/ScalarEvolution.cpp

5663

You're right, this can become an assert. It made since before some refactoring I did to the code before commit. Not so much anymore. Will fix.

5671

I'd prefer to leave this to when someone actually generalizes for lshr and ashr.

5673

We are, good catch.

[reames]

llvm/lib/Analysis/ScalarEvolution.cpp

5673

Actually, now that I think about it further, we don't need the check.

Consider the non-invariant case. The known bits must hold for step on each iteration individually. As a result, they must be a conservative result for the value of step on all iterations. As such, if step is loop variant, we'll get a loose bound on the total step.

Instead of adding a bailout, I'll add a test for this case.

[reames]

llvm/lib/Analysis/ScalarEvolution.cpp

5663

Hm, funny story. It turns out this needs to be a bail out after all. Not because it should be, but because asserting it exposes a bug in LoopFusion. I'm putting a careful comment in the code and have filed PR49566 to track the issue.

[reames]

address review comments

(highly suggest reading inline comments on review for context)

[mkazantsev]

I still think that trickery of definition with invariant step should be emphasised, regardless of the bug you've pointed out.

llvm/lib/Analysis/ScalarEvolution.cpp

5673

The problem is not with this implementation, the problem is with the way how this is defined. What I understood from the comments around is that a shift recurrence is something analogous to the add recurrence. However, add recurrence assumes that all components of it are loop-invariant. If you run opt -analyze -scalar-evolution on

declare i1 @cond()

define void @foo(i32* %p) {
entry:
  %x = load volatile i32, i32* %p
  br label %loop

loop:
  %iv = phi i32 [0, %entry], [%iv.next, %loop]
  %iv.next = add i32 %iv, %x
  %cond = call i1 @cond()
  br i1 %cond, label %loop, label %done

done:
  ret void
}

then you get addrec for %iv. But if you modify test to

declare i1 @cond()

define void @foo(i32* %p) {
entry:
  br label %loop

loop:
  %iv = phi i32 [0, %entry], [%iv.next, %loop]
  %x = load volatile i32, i32* %p
  %iv.next = add i32 %iv, %x
  %cond = call i1 @cond()
  br i1 %cond, label %loop, label %done

done:
  ret void
}

then IV is a SCEVUnknown. The first comment inside this function suggests that what you are trying to recognize is analogous to the AddRec, but in fact it's not.

I don't say that the way how it works is wrong, I'm saying that if the pattern recognized is semantically different from AddRec, you can't just say that you are looking for something like "a simple recurrence of the form: <start, ShiftOp, Step>". This needs a better clarification of what Step actually is.

My preference is to avoid any kind of inconsistency in way how we define different recurrencies. If you support invariant step here, it should also be supported for AddRec beforehand.

5918

nit: remove

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

5673

My preference is to avoid any kind of inconsistency in way how we define different recurrencies. If you support invariant step here, it should also be supported for AddRec beforehand.

sorry, I mean variant step off course.

[reames]

Max, on the terminology bit, I see your point. I've spent some time thinking about it, and haven't really come up with anything better naming wise. What we have is a recurrence with one extra free variable. Said another way, we have a cyclic dependence through the phi and one binop. Do you have any suggestions on naming here? The best I've got is matchSimpleRecurrenceWithOneFreeVariable but that's a bit ugly to say the least. I also considered matchNearRecurrence, but well, defining "near" is ugly too.

I would prefer to separate the naming point from this review though. While I agree we should change it, it's an issue with existing code and not specific to this use. Would you be okay with moving forward on this review under the understanding that I tackle the naming/terminology issue separately in the very near future?

[reames]

Opened a separate review to resolve the naming confusion: https://reviews.llvm.org/D98807

[mkazantsev]

We can either name it separately (and honestly, I don't think there is a better name than shift recurrence), but please make it crystal clear in the comments that this thing *looks* very much like addrec, but there is a tricky difference that it may have a variant step. Many algorithms for AddRecs rely on invariance of arguments explicitly or implicitly. And if someone ever wants to generalize them with your new recurrencies, they will keep falling into this pithole.

Please add either

• A Loop Invariant check
• A comment highlighting that the shift value might not be an invariant

Either should be fine, but the former is more reliable.
Otherwise, looks good.

[reames]

In D98222#2636487, @mkazantsev wrote:

We can either name it separately (and honestly, I don't think there is a better name than shift recurrence).

I added a clarifying description to the ValueTracking header in 9c16621c0. Will also update the patch with an explicit comment on the difference with addrec.

[uabelho]

llvm/lib/Analysis/ScalarEvolution.cpp

5678

I've seen this assertion fail when running

opt -slsr

I wrote
https://bugs.llvm.org/show_bug.cgi?id=49768
about it.

[reames]

JFYI, AFK for next few hours.  If this is causing problematic breakage,
feel free to revert.  I'll look at this later today.

Philip

[reames]

In D98222#2658825, @reames wrote:

JFYI, AFK for next few hours.  If this is causing problematic breakage,
feel free to revert.  I'll look at this later today.

Philip

JFYI, the reported bug is fixed.

[uabelho]

llvm/lib/Analysis/ScalarEvolution.cpp

5678

I reopened
https://bugs.llvm.org/show_bug.cgi?id=49768
due to the assertion triggering again for another

opt -slsr

case.

llvm/include/llvm/Analysis/ScalarEvolution.h

ConstantRange getRangeForAffineNoSelfWrappingAR(const SCEVAddRecExpr *AddRec,
                                                   RangeSignHint SignHint);
 
   /// Try to compute a range for the affine SCEVAddRecExpr {\p Start,+,\p
   /// Stop} by "factoring out" a ternary expression from the add recurrence.
   /// Helper called by \c getRange.
   ConstantRange getRangeViaFactoring(const SCEV *Start, const SCEV *Stop,
                                      const SCEV *MaxBECount, unsigned BitWidth);
 
+  /// If the unknown expression U corresponds to a simple recurrence, return
+  /// a constant range which represents the entire recurrence.  Note that
+  /// *add* recurrences with loop invariant steps aren't represented by
+  /// SCEVUnknowns and thus don't use this mechanism.
+  ConstantRange getRangeForUnknownRecurrence(const SCEVUnknown *U);
+
   /// We know that there is no SCEV for the specified value.  Analyze the
   /// expression.
   const SCEV *createSCEV(Value *V);
 
   /// Provide the special handling we need to analyze PHI SCEVs.
   const SCEV *createNodeForPHI(PHINode *PN);
 
   /// Helper function called from createNodeForPHI.

llvm/lib/Analysis/ScalarEvolution.cpp

void ScalarEvolution::setNoWrapFlags(SCEVAddRecExpr *AddRec,
                                      SCEV::NoWrapFlags Flags) {
   if (AddRec->getNoWrapFlags(Flags) != Flags) {
     AddRec->setNoWrapFlags(Flags);
     UnsignedRanges.erase(AddRec);
     SignedRanges.erase(AddRec);
   }
 }
 
+ConstantRange ScalarEvolution::
+getRangeForUnknownRecurrence(const SCEVUnknown *U) {
+  const DataLayout &DL = getDataLayout();
+
+  unsigned BitWidth = getTypeSizeInBits(U->getType());
+  ConstantRange CR(BitWidth, /*isFullSet=*/true);
+
+  // Match a simple recurrence of the form: <start, ShiftOp, Step>, and then
+  // use information about the trip count to improve our available range.  Note
+  // that the trip count independent cases are already handled by known bits.
+  // WARNING: The definition of recurrence used here is subtly different than
+  // the one used by AddRec (and thus most of this file).  Step is allowed to
+  // be arbitrarily loop varying here, where AddRec allows only loop invariant
+  // and other addrecs in the same loop (for non-affine addrecs).  The code
+  // below intentionally handles the case where step is not loop invariant.
+  auto *P = dyn_cast<PHINode>(U->getValue());
+  if (!P)
+    return CR;
+
+  BinaryOperator *BO;
+  Value *Start, *Step;
+  if (!matchSimpleRecurrence(P, BO, Start, Step))
+    return CR;
+
+  // If we found a recurrence, we must be in a loop -- unless we're
+  // in unreachable code where dominance collapses.  Note that BO might
+  // be in some subloop of L, and that's completely okay.
+  auto *L = LI.getLoopFor(P->getParent());
+  if (!L)
+    return CR;
+  assert(L->getHeader() == P->getParent());
+  if (!L->contains(BO->getParent()))
+    // NOTE: This bailout should be an assert instead.  However, asserting
+    // the condition here exposes a case where LoopFusion is querying SCEV
+    // with malformed loop information during the midst of the transform.
+    // There doesn't appear to be an obvious fix, so for the moment bailout
+    // until the caller issue can be fixed.  PR49566 tracks the bug.
+    return CR;
+
+  // TODO: Handle ashr and lshr cases to increase minimum value reported
+  if (BO->getOpcode() != Instruction::Shl || BO->getOperand(0) != P)
+    return CR;
+
+  unsigned TC = getSmallConstantMaxTripCount(L);
+  if (!TC || TC >= BitWidth)
+    return CR;
+
+  auto KnownStart = computeKnownBits(Start, DL, 0, &AC, nullptr, &DT);
+  auto KnownStep = computeKnownBits(Step, DL, 0, &AC, nullptr, &DT);
+  assert(KnownStart.getBitWidth() == BitWidth &&
+         KnownStep.getBitWidth() == BitWidth);
+
+  // Compute total shift amount, being careful of overflow and bitwidths.
+  auto MaxShiftAmt = KnownStep.getMaxValue();
+  bool Overflow = false;
+  auto TotalShift = MaxShiftAmt.umul_ov(APInt(BitWidth, TC-1, false), Overflow);
+  if (Overflow)
+    return CR;
+
+  // Iff no bits are shifted out, value increases on every shift.
+  auto KnownEnd = KnownBits::shl(KnownStart,
+                                 KnownBits::makeConstant(TotalShift));
+  if (TotalShift.ult(KnownStart.countMinLeadingZeros()))
+    CR = CR.intersectWith(ConstantRange(KnownStart.getMinValue(),
+                                        KnownEnd.getMaxValue() + 1));
+  return CR;
+}
+
+
+
 /// Determine the range for a particular SCEV.  If SignHint is
 /// HINT_RANGE_UNSIGNED (resp. HINT_RANGE_SIGNED) then getRange prefers ranges
 /// with a "cleaner" unsigned (resp. signed) representation.
 const ConstantRange &
 ScalarEvolution::getRangeRef(const SCEV *S,
                              ScalarEvolution::RangeSignHint SignHint) {
   DenseMap<const SCEV *, ConstantRange> &Cache =
       SignHint == ScalarEvolution::HINT_RANGE_UNSIGNED ? UnsignedRanges
if (AddRec->isAffine()) {
         }
       }
     }
 
     return setRange(AddRec, SignHint, std::move(ConservativeResult));
   }
 
   if (const SCEVUnknown *U = dyn_cast<SCEVUnknown>(S)) {
+
     // Check if the IR explicitly contains !range metadata.
     Optional<ConstantRange> MDRange = GetRangeFromMetadata(U->getValue());
     if (MDRange.hasValue())
       ConservativeResult = ConservativeResult.intersectWith(MDRange.getValue(),
                                                             RangeType);
 
+    // Use facts about recurrences in the underlying IR.  Note that add
+    // recurrences are AddRecExprs and thus don't hit this path.  This
+    // primarily handles shift recurrences.
+    auto CR = getRangeForUnknownRecurrence(U);
+    ConservativeResult = ConservativeResult.intersectWith(CR);
+
     // See if ValueTracking can give us a useful range.
     const DataLayout &DL = getDataLayout();
     KnownBits Known = computeKnownBits(U->getValue(), DL, 0, &AC, nullptr, &DT);
     if (Known.getBitWidth() != BitWidth)
       Known = Known.zextOrTrunc(BitWidth);
 
     // ValueTracking may be able to compute a tighter result for the number of
     // sign bits than for the value of those sign bits.

llvm/test/Analysis/ScalarEvolution/shift-recurrences.ll

 ; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
 ; RUN: opt -analyze -enable-new-pm=0 -scalar-evolution < %s | FileCheck %s
 ; RUN: opt -disable-output "-passes=print<scalar-evolution>" < %s 2>&1 | FileCheck %s
 
 define void @test_lshr() {
 ; CHECK-LABEL: 'test_lshr'
 ; CHECK-NEXT:  Classifying expressions for: @test_lshr
 ; CHECK-NEXT:    %iv.lshr = phi i64 [ 1023, %entry ], [ %iv.lshr.next, %loop ]
 ; CHECK-NEXT:    --> %iv.lshr U: [0,1024) S: [0,1024) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }

 
 ; use trip count to refine
 define void @test_shl2() {
 ; CHECK-LABEL: 'test_shl2'
 ; CHECK-NEXT:  Classifying expressions for: @test_shl2
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl = phi i64 [ 4, %entry ], [ %iv.shl.next, %loop ]
-; CHECK-NEXT:    --> %iv.shl U: [0,-3) S: [-9223372036854775808,9223372036854775801) Exits: 64 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.shl U: [4,65) S: [4,65) Exits: 64 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl.next = shl i64 %iv.shl, 1
-; CHECK-NEXT:    --> (2 * %iv.shl) U: [0,-7) S: [-9223372036854775808,9223372036854775801) Exits: 128 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> (2 * %iv.shl)<nuw><nsw> U: [8,129) S: [8,129) Exits: 128 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_shl2
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 4
 ; CHECK-NEXT:   Predicates:
 ; CHECK:       Loop %loop: Trip multiple is 5
 ;
 entry:

 define void @test_shl3(i1 %c) {
 ; CHECK-LABEL: 'test_shl3'
 ; CHECK-NEXT:  Classifying expressions for: @test_shl3
 ; CHECK-NEXT:    %shiftamt = select i1 %c, i64 1, i64 0
 ; CHECK-NEXT:    --> %shiftamt U: [0,2) S: [0,2)
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl = phi i64 [ 4, %entry ], [ %iv.shl.next, %loop ]
-; CHECK-NEXT:    --> %iv.shl U: [0,-3) S: [-9223372036854775808,9223372036854775805) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.shl U: [4,65) S: [4,65) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl.next = shl i64 %iv.shl, %shiftamt
 ; CHECK-NEXT:    --> %iv.shl.next U: [0,-3) S: [-9223372036854775808,9223372036854775805) Exits: <<Unknown>> LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_shl3
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 4
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 4

 
 ; edge case on max value not overflowing
 define void @test_shl4() {
 ; CHECK-LABEL: 'test_shl4'
 ; CHECK-NEXT:  Classifying expressions for: @test_shl4
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,61) S: [0,61) Exits: 60 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl = phi i64 [ 4, %entry ], [ %iv.shl.next, %loop ]
-; CHECK-NEXT:    --> %iv.shl U: [0,-3) S: [-9223372036854775808,9223372036854775801) Exits: 4611686018427387904 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.shl U: [4,4611686018427387905) S: [4,4611686018427387905) Exits: 4611686018427387904 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,62) S: [1,62) Exits: 61 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl.next = shl i64 %iv.shl, 1
-; CHECK-NEXT:    --> (2 * %iv.shl) U: [0,-7) S: [-9223372036854775808,9223372036854775801) Exits: -9223372036854775808 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> (2 * %iv.shl)<nuw> U: [8,-9223372036854775807) S: [-9223372036854775808,9223372036854775801) Exits: -9223372036854775808 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_shl4
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 60
 ; CHECK-NEXT:  Loop %loop: max backedge-taken count is 60
 ; CHECK-NEXT:  Loop %loop: Predicated backedge-taken count is 60
 ; CHECK-NEXT:   Predicates:
 ; CHECK:       Loop %loop: Trip multiple is 61
 ;
 entry:

 
 ; Loop varying (but tightly bounded) shift amount
 define void @test_shl6(i1 %c) {
 ; CHECK-LABEL: 'test_shl6'
 ; CHECK-NEXT:  Classifying expressions for: @test_shl6
 ; CHECK-NEXT:    %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
 ; CHECK-NEXT:    --> {0,+,1}<%loop> U: [0,5) S: [0,5) Exits: 4 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl = phi i64 [ 4, %entry ], [ %iv.shl.next, %loop ]
-; CHECK-NEXT:    --> %iv.shl U: [0,-3) S: [-9223372036854775808,9223372036854775805) Exits: 16 LoopDispositions: { %loop: Variant }
+; CHECK-NEXT:    --> %iv.shl U: [4,65) S: [4,65) Exits: 16 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:    %iv.next = add i64 %iv, 1
 ; CHECK-NEXT:    --> {1,+,1}<%loop> U: [1,6) S: [1,6) Exits: 5 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %shiftamt = and i64 %iv, 1
 ; CHECK-NEXT:    --> (zext i1 {false,+,true}<%loop> to i64) U: [0,2) S: [0,2) Exits: 0 LoopDispositions: { %loop: Computable }
 ; CHECK-NEXT:    %iv.shl.next = shl i64 %iv.shl, %shiftamt
 ; CHECK-NEXT:    --> %iv.shl.next U: [0,-3) S: [-9223372036854775808,9223372036854775805) Exits: 16 LoopDispositions: { %loop: Variant }
 ; CHECK-NEXT:  Determining loop execution counts for: @test_shl6
 ; CHECK-NEXT:  Loop %loop: backedge-taken count is 4