[SCEV] Assumption context for GetMinTrailingZeros

[gilr]

Summary

Let GetMinTrailingZeros() make better use of assumptions by taking an optional context instruction.
The context defaults to null, keeping the API backward-compatible. If a context is provided, the method will neither use nor update the trailing zeros cache.
The context argument is then used by the loop vectorizer and by SCEV::getSmallConstantTripMultiple() to set the context at their given loop.

[fhahn]

llvm/lib/Analysis/ScalarEvolution.cpp

6893

Why use the header's terminator here?

llvm/test/Transforms/LoopUnroll/runtime-unroll-assume-no-remainder.ll

8

I think it would be good to combine a few of those related tests in a single file :)

[gilr]

llvm/lib/Analysis/ScalarEvolution.cpp

6893

No reason, actually. Will change to first instruction in the header.

llvm/test/Transforms/LoopUnroll/runtime-unroll-assume-no-remainder.ll

8

Ah, good catch on LV (but I better do this separately to keep this patch's changes clear).
(Unroller tests running roughly the same opt command seem to have manual CHECK lines. Not sure how well it would mix with automated CHECKs that need to be updated from time to time)

[gilr]

Addressed comments.

[gilr]

Rebased on top of 7ddbe0cb905ec62d37b284a2e8daf54887a35f94 (merge ..-assume-divisible-TC.ll into ..-divisible-TC.ll)

[gilr]

Ping :)

[gilr]

ping

[lebedev.ri]

I would expect that the SCEV change could be testable with just the scev itselfs (-analyze -scalar-evolution), is it not?

llvm/lib/Analysis/ScalarEvolution.cpp

5611–5615

Doesn't this defy the point of the cache?
I think MinTrailingZerosCache's key should be changed to be a std::pair<SCEV*,CxtI*>

[fhahn]

I think the fix makes sense in isolation, but I am wondering if we could generalize this to automatically apply in more cases. Would it be possible to 'apply' the information from the loop guard directly to the SCEV expression? Like we do for some conditions already: https://github.com/llvm/llvm-project/blob/master/llvm/lib/Analysis/ScalarEvolution.cpp#L13241 ?

[gilr]

In D94088#2519311, @lebedev.ri wrote:

I would expect that the SCEV change could be testable with just the scev itselfs (-analyze -scalar-evolution), is it not?

Right, will add a unit test for this API.

In D94088#2519327, @fhahn wrote:

I think the fix makes sense in isolation, but I am wondering if we could generalize this to automatically apply in more cases. Would it be possible to 'apply' the information from the loop guard directly to the SCEV expression? Like we do for some conditions already: https://github.com/llvm/llvm-project/blob/master/llvm/lib/Analysis/ScalarEvolution.cpp#L13241 ?

Will take a look at that. Thanks!

llvm/lib/Analysis/ScalarEvolution.cpp

5611–5615

Doesn't this defy the point of the cache?

This API seems to be used extensively. Since (for now) only the unroller and vectorizer would be providing a non-null context I wasn't sure caching those calls was worth the added space/time of extending the key.

I think MinTrailingZerosCache's key should be changed to be a std::pair<SCEV*,CxtI*>

Will do.

[gilr]

Trying to apply Florian's suggestion in D95521 (only for powers of 2 as a first step). This handles the power-of-two case added by PR47904.

[lebedev.ri]

llvm/lib/Analysis/ScalarEvolution.cpp

5611–5615

Thinking about this, it's not even so much that it's pessimising the cache, it's actually straight-up upsafe to do.
E.g. given

x = y
if(!(x & 0b1))
  f(x); // x has at least one trailing bit

if we first ask GetMinTrailingZeros(X, "f(x)"), we'll get 1,
and if we then ask GetMinTrailingZeros(X, "x = y"), we'll again get 1, which is obviously not correct.

[gilr]

llvm/lib/Analysis/ScalarEvolution.cpp

5611–5615

Not sure I follow. If a context is provided the cache is neither queried nor updated, so the call GetMinTrailingZeros(X, "f(x)") shouldn't affect a later GetMinTrailingZeros(X, "x = y") call (or a later GetMinTrailingZeros(X) call). Am I missing something?

[lebedev.ri]

llvm/lib/Analysis/ScalarEvolution.cpp

5611–5615

Err, nope.

[lebedev.ri]

This review seems to be stuck/dead, consider abandoning if no longer relevant.

llvm/include/llvm/Analysis/ScalarEvolution.h

public:
   /// We don't have a way to invalidate per-loop dispositions. Clear and
   /// recompute is simpler.
   void forgetLoopDispositions(const Loop *L);
 
   /// Determine the minimum number of zero bits that S is guaranteed to end in
   /// (at every loop iteration).  It is, at the same time, the minimum number
   /// of times S is divisible by 2.  For example, given {4,+,8} it returns 2.
   /// If S is guaranteed to be 0, it returns the bitwidth of S.
-  uint32_t GetMinTrailingZeros(const SCEV *S);
+  uint32_t GetMinTrailingZeros(const SCEV *S,
+                               const Instruction *CxtI = nullptr);
 
   /// Determine the unsigned range for a particular SCEV.
   /// NOTE: This returns a copy of the reference returned by getRangeRef.
   ConstantRange getUnsignedRange(const SCEV *S) {
     return getRangeRef(S, HINT_RANGE_UNSIGNED);
   }
 
   /// Determine the min of the unsigned range for a particular SCEV.
private:
 
   /// Memoized values for the GetMinTrailingZeros
   DenseMap<const SCEV *, uint32_t> MinTrailingZerosCache;
 
   /// Return the Value set from which the SCEV expr is generated.
   SetVector<ValueOffsetPair> *getSCEVValues(const SCEV *S);
 
   /// Private helper method for the GetMinTrailingZeros method
-  uint32_t GetMinTrailingZerosImpl(const SCEV *S);
+  uint32_t GetMinTrailingZerosImpl(const SCEV *S, const Instruction *CxtI);
 
   /// Information about the number of loop iterations for which a loop exit's
   /// branch condition evaluates to the not-taken path.  This is a temporary
   /// pair of exact and max expressions that are eventually summarized in
   /// ExitNotTakenInfo and BackedgeTakenInfo.
   struct ExitLimit {
     const SCEV *ExactNotTaken; // The exit is not taken exactly this many times
     const SCEV *MaxNotTaken; // The exit is not taken at most this many times

llvm/lib/Analysis/ScalarEvolution.cpp

if (!GEP->getSourceElementType()->isSized())
     return getUnknown(GEP);
 
   SmallVector<const SCEV *, 4> IndexExprs;
   for (auto Index = GEP->idx_begin(); Index != GEP->idx_end(); ++Index)
     IndexExprs.push_back(getSCEV(*Index));
   return getGEPExpr(GEP, IndexExprs);
 }
 
-uint32_t ScalarEvolution::GetMinTrailingZerosImpl(const SCEV *S) {
+uint32_t ScalarEvolution::GetMinTrailingZerosImpl(const SCEV *S,
+                                                  const Instruction *CxtI) {
   if (const SCEVConstant *C = dyn_cast<SCEVConstant>(S))
     return C->getAPInt().countTrailingZeros();
 
   if (const SCEVPtrToIntExpr *I = dyn_cast<SCEVPtrToIntExpr>(S))
-    return GetMinTrailingZeros(I->getOperand());
+    return GetMinTrailingZeros(I->getOperand(), CxtI);
 
   if (const SCEVTruncateExpr *T = dyn_cast<SCEVTruncateExpr>(S))
-    return std::min(GetMinTrailingZeros(T->getOperand()),
+    return std::min(GetMinTrailingZeros(T->getOperand(), CxtI),
                     (uint32_t)getTypeSizeInBits(T->getType()));
 
   if (const SCEVZeroExtendExpr *E = dyn_cast<SCEVZeroExtendExpr>(S)) {
-    uint32_t OpRes = GetMinTrailingZeros(E->getOperand());
+    uint32_t OpRes = GetMinTrailingZeros(E->getOperand(), CxtI);
     return OpRes == getTypeSizeInBits(E->getOperand()->getType())
                ? getTypeSizeInBits(E->getType())
                : OpRes;
   }
 
   if (const SCEVSignExtendExpr *E = dyn_cast<SCEVSignExtendExpr>(S)) {
-    uint32_t OpRes = GetMinTrailingZeros(E->getOperand());
+    uint32_t OpRes = GetMinTrailingZeros(E->getOperand(), CxtI);
     return OpRes == getTypeSizeInBits(E->getOperand()->getType())
                ? getTypeSizeInBits(E->getType())
                : OpRes;
   }
 
   if (const SCEVAddExpr *A = dyn_cast<SCEVAddExpr>(S)) {
     // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
+    uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0), CxtI);
     for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
+      MinOpRes =
+          std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i), CxtI));
     return MinOpRes;
   }
 
   if (const SCEVMulExpr *M = dyn_cast<SCEVMulExpr>(S)) {
     // The result is the sum of all operands results.
     uint32_t SumOpRes = GetMinTrailingZeros(M->getOperand(0));
     uint32_t BitWidth = getTypeSizeInBits(M->getType());
     for (unsigned i = 1, e = M->getNumOperands();
          SumOpRes != BitWidth && i != e; ++i)
-      SumOpRes =
-          std::min(SumOpRes + GetMinTrailingZeros(M->getOperand(i)), BitWidth);
+      SumOpRes = std::min(
+          SumOpRes + GetMinTrailingZeros(M->getOperand(i), CxtI), BitWidth);
     return SumOpRes;
   }
 
   if (const SCEVAddRecExpr *A = dyn_cast<SCEVAddRecExpr>(S)) {
     // The result is the min of all operands results.
     uint32_t MinOpRes = GetMinTrailingZeros(A->getOperand(0));
     for (unsigned i = 1, e = A->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i)));
+      MinOpRes =
+          std::min(MinOpRes, GetMinTrailingZeros(A->getOperand(i), CxtI));
     return MinOpRes;
   }
 
   if (const SCEVSMaxExpr *M = dyn_cast<SCEVSMaxExpr>(S)) {
     // The result is the min of all operands results.
-    uint32_t MinOpRes = GetMinTrailingZeros(M->getOperand(0));
+    uint32_t MinOpRes = GetMinTrailingZeros(M->getOperand(0), CxtI);
     for (unsigned i = 1, e = M->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i)));
+      MinOpRes =
+          std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i), CxtI));
     return MinOpRes;
   }
 
   if (const SCEVUMaxExpr *M = dyn_cast<SCEVUMaxExpr>(S)) {
     // The result is the min of all operands results.
     uint32_t MinOpRes = GetMinTrailingZeros(M->getOperand(0));
     for (unsigned i = 1, e = M->getNumOperands(); MinOpRes && i != e; ++i)
-      MinOpRes = std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i)));
+      MinOpRes =
+          std::min(MinOpRes, GetMinTrailingZeros(M->getOperand(i), CxtI));
     return MinOpRes;
   }
 
   if (const SCEVUnknown *U = dyn_cast<SCEVUnknown>(S)) {
     // For a SCEVUnknown, ask ValueTracking.
-    KnownBits Known = computeKnownBits(U->getValue(), getDataLayout(), 0, &AC, nullptr, &DT);
+    KnownBits Known =
+        computeKnownBits(U->getValue(), getDataLayout(), 0, &AC, CxtI, &DT);
     return Known.countMinTrailingZeros();
   }
 
   // SCEVUDivExpr
   return 0;
 }
 
-uint32_t ScalarEvolution::GetMinTrailingZeros(const SCEV *S) {
+uint32_t ScalarEvolution::GetMinTrailingZeros(const SCEV *S,
+                                              const Instruction *CxtI) {
+  if (!CxtI) {
     auto I = MinTrailingZerosCache.find(S);
     if (I != MinTrailingZerosCache.end())
       return I->second;
+  }
+
+  uint32_t Result = GetMinTrailingZerosImpl(S, CxtI);
+  if (CxtI)
+    return Result;
 
-  uint32_t Result = GetMinTrailingZerosImpl(S);
   auto InsertPair = MinTrailingZerosCache.insert({S, Result});
   assert(InsertPair.second && "Should insert a new key");
   return InsertPair.first->second;
 }
 
 /// Helper method to assign a range to V from metadata present in the IR.
 static Optional<ConstantRange> GetRangeFromMetadata(Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
ScalarEvolution::getSmallConstantTripMultiple(const Loop *L,
   // Get the trip count from the BE count by adding 1.
   const SCEV *TCExpr = getAddExpr(ExitCount, getOne(ExitCount->getType()));
 
   const SCEVConstant *TC = dyn_cast<SCEVConstant>(TCExpr);
   if (!TC)
     // Attempt to factor more general cases. Returns the greatest power of
     // two divisor. If overflow happens, the trip count expression is still
     // divisible by the greatest power of 2 divisor returned.
-    return 1U << std::min((uint32_t)31, GetMinTrailingZeros(TCExpr));
+    return 1U << std::min((uint32_t)31, GetMinTrailingZeros(
+                                            TCExpr, &L->getHeader()->front()));
 
   ConstantInt *Result = TC->getValue();
 
   // Guard against huge trip counts (this requires checking
   // for zero to handle the case where the trip count == -1 and the
   // addition wraps).
   if (!Result || Result->getValue().getActiveBits() > 32 ||
       Result->getValue().getActiveBits() == 0)
     return 1;

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
   const SCEV *ExitCount = SE->getAddExpr(
       BackedgeTakenCount, SE->getOne(BackedgeTakenCount->getType()));
   const SCEV *Rem = SE->getURemExpr(
       ExitCount, SE->getConstant(BackedgeTakenCount->getType(), MaxVFtimesIC));
   if (Rem->isZero()) {
     // Accept MaxVF if we do not have a tail.
     LLVM_DEBUG(dbgs() << "LV: No tail will remain for any chosen VF.\n");
     return MaxVF;
+  } else if (isPowerOf2_32(MaxVFtimesIC)) {
+    // Since getURemExpr() doesn't take an assumption context it may have
+    // ignored relevant assumptions. For powers of two also try proving
+    // divisibility using GetMinTrailingZeros() with the header's terminator as
+    // context.
+    unsigned TCisMultipleOf =
+        1 << SE->GetMinTrailingZeros(ExitCount, &TheLoop->getHeader()->front());
+    if (TCisMultipleOf % MaxVFtimesIC == 0) {
+      // Accept MaxVF if we do not have a tail.
+      LLVM_DEBUG(dbgs() << "LV: No tail will remain for any chosen VF.\n");
+      return MaxVF;
+    }
   }
 
   // If we don't know the precise trip count, or if the trip count that we
   // found modulo the vectorization factor is not zero, try to fold the tail
   // by masking.
   // FIXME: look for a smaller MaxVF that does divide TC rather than masking.
   if (Legal->prepareToFoldTailByMasking()) {
     FoldTailByMasking = true;

llvm/test/Transforms/LoopUnroll/runtime-unroll-assume-no-remainder.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -loop-unroll -unroll-runtime=true -unroll-runtime-epilog=true -unroll-count=2 | FileCheck %s
+
+; Make sure the loop is unrolled without a remainder loop based on an assumption
+; that the least significant bit is known to be zero.
+
+define dso_local void @assumeDivisibleTC(i8* noalias nocapture %a, i8* noalias nocapture readonly %b, i32 %n, i1 %c) local_unnamed_addr {
+; CHECK-LABEL: @assumeDivisibleTC(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[BAIL_OUT:%.*]], label [[DO_WORK:%.*]]
+; CHECK:       bail.out:
+; CHECK-NEXT:    ret void
+; CHECK:       do.work:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[N:%.*]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[CMP110:%.*]] = icmp sgt i32 [[N]], 0
+; CHECK-NEXT:    br i1 [[CMP110]], label [[FOR_BODY_PREHEADER:%.*]], label [[FOR_COND_CLEANUP:%.*]]
+; CHECK:       for.body.preheader:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.cond.cleanup.loopexit:
+; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    ret void
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_011:%.*]] = phi i32 [ 0, [[FOR_BODY_PREHEADER]] ], [ [[INC_1:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[IDXPROM:%.*]] = zext i32 [[I_011]] to i64
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i8, i8* [[B:%.*]], i64 [[IDXPROM]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, i8* [[ARRAYIDX]], align 1
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[TMP0]], 3
+; CHECK-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds i8, i8* [[A:%.*]], i64 [[IDXPROM]]
+; CHECK-NEXT:    store i8 [[ADD]], i8* [[ARRAYIDX4]], align 1
+; CHECK-NEXT:    [[INC:%.*]] = add nuw nsw i32 [[I_011]], 1
+; CHECK-NEXT:    [[IDXPROM_1:%.*]] = zext i32 [[INC]] to i64
+; CHECK-NEXT:    [[ARRAYIDX_1:%.*]] = getelementptr inbounds i8, i8* [[B]], i64 [[IDXPROM_1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load i8, i8* [[ARRAYIDX_1]], align 1
+; CHECK-NEXT:    [[ADD_1:%.*]] = add i8 [[TMP1]], 3
+; CHECK-NEXT:    [[ARRAYIDX4_1:%.*]] = getelementptr inbounds i8, i8* [[A]], i64 [[IDXPROM_1]]
+; CHECK-NEXT:    store i8 [[ADD_1]], i8* [[ARRAYIDX4_1]], align 1
+; CHECK-NEXT:    [[INC_1]] = add nuw nsw i32 [[INC]], 1
+; CHECK-NEXT:    [[CMP1_1:%.*]] = icmp slt i32 [[INC_1]], [[N]]
+; CHECK-NEXT:    br i1 [[CMP1_1]], label [[FOR_BODY]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], [[LOOP0:!llvm.loop !.*]]
+;
+entry:
+  br i1 %c, label %bail.out, label %do.work
+
+bail.out:
+  ret void
+
+do.work:
+  %and = and i32 %n, 1
+  %cmp = icmp eq i32 %and, 0
+  tail call void @llvm.assume(i1 %cmp)
+  %cmp110 = icmp sgt i32 %n, 0
+  br i1 %cmp110, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %do.work
+  br label %for.body
+
+for.cond.cleanup.loopexit:                        ; preds = %for.body
+  br label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.cond.cleanup.loopexit, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %i.011 = phi i32 [ %inc, %for.body ], [ 0, %for.body.preheader ]
+  %idxprom = zext i32 %i.011 to i64
+  %arrayidx = getelementptr inbounds i8, i8* %b, i64 %idxprom
+  %0 = load i8, i8* %arrayidx, align 1
+  %add = add i8 %0, 3
+  %arrayidx4 = getelementptr inbounds i8, i8* %a, i64 %idxprom
+  store i8 %add, i8* %arrayidx4, align 1
+  %inc = add nuw nsw i32 %i.011, 1
+  %cmp1 = icmp slt i32 %inc, %n
+  br i1 %cmp1, label %for.body, label %for.cond.cleanup.loopexit, !llvm.loop !0
+}
+
+declare void @llvm.assume(i1 noundef) nofree nosync nounwind willreturn
+!0 = distinct !{!0, !1, !2}
+!1 = !{!"llvm.loop.mustprogress"}
+!2 = !{!"llvm.loop.unroll.count", i32 4}

llvm/test/Transforms/LoopVectorize/dont-fold-tail-for-assumed-divisible-TC.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -loop-vectorize -force-vector-width=4 -S | FileCheck %s
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 
 ; Make sure the loop is vectorized under -Os without folding its tail based on
 ; its trip-count's lower bits assumed to be zero.
 
-define dso_local void @assumeAlignedTC(i32* noalias nocapture %A, i32* %p) optsize {
+define dso_local void @assumeAlignedTC(i32* noalias nocapture %A, i32 %p, i32 %q, i1 %c) optsize {
 ; CHECK-LABEL: @assumeAlignedTC(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[N:%.*]] = load i32, i32* [[P:%.*]], align 4
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[N]], 3
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
-; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[BAIL_OUT:%.*]], label [[DO_WORK:%.*]]
+; CHECK:       bail.out:
+; CHECK-NEXT:    ret void
+; CHECK:       do.work:
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[P:%.*]], 3
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[AND1]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP1]])
+; CHECK-NEXT:    [[AND2:%.*]] = and i32 [[Q:%.*]], 7
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[AND2]], 0
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[CMP2]])
+; CHECK-NEXT:    [[GT:%.*]] = icmp sgt i32 [[P]], [[Q]]
+; CHECK-NEXT:    [[N:%.*]] = select i1 [[GT]], i32 [[P]], i32 [[Q]]
 ; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ult i32 [[N]], 4
 ; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
 ; CHECK:       vector.ph:
 ; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i32 [[N]], 4
 ; CHECK-NEXT:    [[N_VEC:%.*]] = sub i32 [[N]], [[N_MOD_VF]]
 ; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i32 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
 ; CHECK-NEXT:    [[BROADCAST_SPLATINSERT:%.*]] = insertelement <4 x i32> poison, i32 [[INDEX]], i32 0
 ; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT]], <4 x i32> poison, <4 x i32> zeroinitializer
 ; CHECK-NEXT:    [[INDUCTION:%.*]] = add <4 x i32> [[BROADCAST_SPLAT]], <i32 0, i32 1, i32 2, i32 3>
 ; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[INDEX]], 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i32 [[TMP0]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i32, i32* [[TMP1]], i32 0
 ; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i32* [[TMP2]] to <4 x i32>*
 ; CHECK-NEXT:    store <4 x i32> <i32 13, i32 13, i32 13, i32 13>, <4 x i32>* [[TMP3]], align 1
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add i32 [[INDEX]], 4
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp eq i32 [[INDEX_NEXT]], [[N_VEC]]
 ; CHECK-NEXT:    br i1 [[TMP4]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP0:!llvm.loop !.*]]
 ; CHECK:       middle.block:
 ; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i32 [[N]], [[N_VEC]]
 ; CHECK-NEXT:    br i1 [[CMP_N]], label [[EXIT:%.*]], label [[SCALAR_PH]]
 ; CHECK:       scalar.ph:
-; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i32 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i32 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[DO_WORK]] ]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[RIV:%.*]] = phi i32 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[RIVPLUS1:%.*]], [[LOOP]] ]
 ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* [[A]], i32 [[RIV]]
 ; CHECK-NEXT:    store i32 13, i32* [[ARRAYIDX]], align 1
 ; CHECK-NEXT:    [[RIVPLUS1]] = add nuw nsw i32 [[RIV]], 1
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[RIVPLUS1]], [[N]]
 ; CHECK-NEXT:    br i1 [[COND]], label [[EXIT]], label [[LOOP]], [[LOOP2:!llvm.loop !.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret void
 ;
 entry:
-  %n = load i32, i32* %p
-  %and = and i32 %n, 3
-  %cmp = icmp eq i32 %and, 0
-  tail call void @llvm.assume(i1 %cmp)
+  br i1 %c, label %bail.out, label %do.work
+
+bail.out:
+  ret void
+
+do.work:
+  %and1 = and i32 %p, 3
+  %cmp1 = icmp eq i32 %and1, 0
+  tail call void @llvm.assume(i1 %cmp1)
+  %and2 = and i32 %q, 7
+  %cmp2 = icmp eq i32 %and2, 0
+  tail call void @llvm.assume(i1 %cmp2)
+  %gt = icmp sgt i32 %p, %q
+  %n = select i1 %gt, i32 %p, i32 %q
   br label %loop
 
 loop:
-  %riv = phi i32 [ 0, %entry ], [ %rivPlus1, %loop ]
+  %riv = phi i32 [ 0, %do.work ], [ %rivPlus1, %loop ]
   %arrayidx = getelementptr inbounds i32, i32* %A, i32 %riv
   store i32 13, i32* %arrayidx, align 1
   %rivPlus1 = add nuw nsw i32 %riv, 1
   %cond = icmp eq i32 %rivPlus1, %n
   br i1 %cond, label %exit, label %loop
 
 exit:
   ret void
 }
 
 declare void @llvm.assume(i1 noundef) nofree nosync nounwind willreturn