[SCEV] Apply loop guards to divisibility tests

[gilr]

Summary

Extend applyLoopGuards() to take into account conditions/assumes proving some value %v to be divisible by D by rewriting %v to (%v / D) * D. This lets the loop unroller and the loop vectorizer identify more loops as not requiring remainder loops.

[gilr]

Rebased

[fhahn]

llvm/lib/Analysis/ScalarEvolution.cpp

13271

Here we basically are matching a SCEV URem expression, right? There's a helper for that (matchURem), which should help to simplify the code a bit?

13282

Here we could just re-use the existing URem expression?

[fhahn]

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

8

could you also add version of this test that checks %cmp in the guard condition, instead of an assume?

46

is this extra branch needed?

58

the extra blocks should not be needed for unrolling.

69

can we either change %i.011 to i64 or the users to use i32 instead of the extension?

81

should not be needed for unrolling.

[gilr]

llvm/lib/Analysis/ScalarEvolution.cpp

13271

Ah, indeed. Thanks!

13282

Do you mean by computing (A / B) * B from the URem expression e.g. via A - URemExpr?

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

46

This was done in D94088 to prevent D93974 (should it ever get in) from validating the entry-block assume for GetMinTrailingZeros() w/o an assumption context. But since appyLoopGuards() uses its own assumption validation check it's indeed redundant. Will remove.

[fhahn]

llvm/lib/Analysis/ScalarEvolution.cpp

13282

Nevermind, I was thinking about something else and we probably need the div & multiply to make sure the SCEV logic can determine the right multiple.

[gilr]

Addressed comments.

[fhahn]

LGTM with the additional tests, thanks!

I did some testing and it appears this exposes a crash in matchURem. I'm taking a look at that now, I think it would be good to wait with landing the change until this is resolved.

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

8

could you add a similar negative test, e.g. where the and does not strip the lowest bit?

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

67

could you add a similar negative test, e.g. where the and does not strip the lowest bit?

Perhaps also add a test using urem?

[fhahn]

In D95521#2533414, @fhahn wrote:

I did some testing and it appears this exposes a crash in matchURem. I'm taking a look at that now, I think it would be good to wait with landing the change until this is resolved.

Should be fixed by f1e8136115ac

[gilr]

In D95521#2533644, @fhahn wrote:

In D95521#2533414, @fhahn wrote:

I did some testing and it appears this exposes a crash in matchURem. I'm taking a look at that now, I think it would be good to wait with landing the change until this is resolved.

Should be fixed by f1e8136115ac

Excellent, thanks!

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

67

could you add a similar negative test, e.g. where the and does not strip the lowest bit?

Will do.

Perhaps also add a test using urem?

Ah, good catch!
Trying VF=4, IC=3 exposed two issues:

• While applyLoopGuards() rewrites my divisible-by-12 TC correctly to (12 * ...), getURemExpr(TC, 12) doesn't fold to zero. This can perhaps be solved by using SCEVDivision instead of getURemExpr(). Will try that as an improvement over this patch.

• The test then hits the assert "VF*UF must be a power of 2 when folding tail by masking" later in LV. This isn't related to this patch (can be reproduced by forcing optsize and IC=3), so I'll upload a separate fix for it.

[fhahn]

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

67

While applyLoopGuards() rewrites my divisible-by-12 TC correctly to (12 * ...), getURemExpr(TC, 12) doesn't fold to zero. This can perhaps be solved by using SCEVDivision instead of getURemExpr(). Will try that as an improvement over this patch.

That sounds like an independent improvement to SCEV.

The test then hits the assert "VF*UF must be a power of 2 when folding tail by masking" later in LV. This isn't related to this patch (can be reproduced by forcing optsize and IC=3), so I'll upload a separate fix for it.

Sounds good, thanks!

[gilr]

Addressed comments.

llvm/include/llvm/Analysis/ScalarEvolution.h

public:
   /// to be a constant.
   Optional<APInt> computeConstantDifference(const SCEV *LHS, const SCEV *RHS);
 
   /// Update no-wrap flags of an AddRec. This may drop the cached info about
   /// this AddRec (such as range info) in case if new flags may potentially
   /// sharpen it.
   void setNoWrapFlags(SCEVAddRecExpr *AddRec, SCEV::NoWrapFlags Flags);
 
+  /// Try to apply information from loop guards for \p L to \p Expr.
+  const SCEV *applyLoopGuards(const SCEV *Expr, const Loop *L);
+
 private:
   /// A CallbackVH to arrange for ScalarEvolution to be notified whenever a
   /// Value is deleted.
   class SCEVCallbackVH final : public CallbackVH {
     ScalarEvolution *SE;
 
     void deleted() override;
     void allUsesReplacedWith(Value *New) override;
private:
   /// Find all of the loops transitively used in \p S, and update \c LoopUsers
   /// accordingly.
   void addToLoopUseLists(const SCEV *S);
 
   /// Try to match the pattern generated by getURemExpr(A, B). If successful,
   /// Assign A and B to LHS and RHS, respectively.
   bool matchURem(const SCEV *Expr, const SCEV *&LHS, const SCEV *&RHS);
 
-  /// Try to apply information from loop guards for \p L to \p Expr.
-  const SCEV *applyLoopGuards(const SCEV *Expr, const Loop *L);
-
   /// Look for a SCEV expression with type `SCEVType` and operands `Ops` in
   /// `UniqueSCEVs`.
   ///
   /// The first component of the returned tuple is the SCEV if found and null
   /// otherwise.  The second component is the `FoldingSetNodeID` that was
   /// constructed to look up the SCEV and the third component is the insertion
   /// point.
   std::tuple<SCEV *, FoldingSetNodeID, void *>

llvm/lib/Analysis/ScalarEvolution.cpp

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(applyLoopGuards(TCExpr, L)));
 
   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)
if (I == Map.end())
       return Expr;
     return I->second;
   }
 };
 
 const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
   auto CollectCondition = [&](ICmpInst::Predicate Predicate, const SCEV *LHS,
                               const SCEV *RHS, ValueToSCEVMapTy &RewriteMap) {
+    // 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)) {
+          Value *V = LHSUnknown->getValue();
+          auto Multiple =
+              getMulExpr(getUDivExpr(URemLHS, URemRHS), URemRHS,
+                         (SCEV::NoWrapFlags)(SCEV::FlagNUW | SCEV::FlagNSW));
+          RewriteMap[V] = Multiple;
+          return;
+        }
+      }
+    }
+
     if (!isa<SCEVUnknown>(LHS)) {
       std::swap(LHS, RHS);
       Predicate = CmpInst::getSwappedPredicate(Predicate);
     }
 
     // For now, limit to conditions that provide information about unknown
     // expressions.
     auto *LHSUnknown = dyn_cast<SCEVUnknown>(LHS);

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

unsigned MaxVFtimesIC =
       UserIC ? MaxVF.getFixedValue() * UserIC : MaxVF.getFixedValue();
   // Avoid tail folding if the trip count is known to be a multiple of any VF we
   // chose.
   ScalarEvolution *SE = PSE.getSE();
   const SCEV *BackedgeTakenCount = PSE.getBackedgeTakenCount();
   const SCEV *ExitCount = SE->getAddExpr(
       BackedgeTakenCount, SE->getOne(BackedgeTakenCount->getType()));
   const SCEV *Rem = SE->getURemExpr(
-      ExitCount, SE->getConstant(BackedgeTakenCount->getType(), MaxVFtimesIC));
+      SE->applyLoopGuards(ExitCount, TheLoop),
+      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;
   }
 
   // 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

llvm/test/Analysis/ScalarEvolution/trip-multiple-guard-info.ll

 ; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
 ; 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
 
 ; Tests for PR47904.
 
 define void @test_trip_multiple_4(i32 %num) {
 ; CHECK-LABEL: @test_trip_multiple_4
 ; CHECK:       Loop %for.body: backedge-taken count is (-1 + %num)
 ; CHECK-NEXT:  Loop %for.body: max backedge-taken count is -2
 ; CHECK-NEXT:  Loop %for.body: Predicated backedge-taken count is (-1 + %num)
-; CHECK:       Loop %for.body: Trip multiple is 1
+; CHECK:       Loop %for.body: Trip multiple is 4
 ;
 entry:
   %u = urem i32 %num, 4
   %cmp = icmp eq i32 %u, 0
   tail call void @llvm.assume(i1 %cmp)
   %cmp.1 = icmp uge i32 %num, 4
   tail call void @llvm.assume(i1 %cmp.1)
   br label %for.body

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 %p, i32 %q) local_unnamed_addr {
+; CHECK-LABEL: @assumeDivisibleTC(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[P:%.*]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[GUARDED:%.*]], label [[EXIT:%.*]]
+; CHECK:       guarded:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 [[Q:%.*]], 2
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[REM]], 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:    [[CMP110:%.*]] = icmp sgt i32 [[N]], 0
+; CHECK-NEXT:    br i1 [[CMP110]], label [[FOR_BODY_PREHEADER:%.*]], label [[EXIT]]
+; CHECK:       for.body.preheader:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_011:%.*]] = phi i32 [ 0, [[FOR_BODY_PREHEADER]] ], [ [[INC_1:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i8, i8* [[B:%.*]], i32 [[I_011]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i8, i8* [[ARRAYIDX]], align 1
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[TMP0]], 3
+; CHECK-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds i8, i8* [[A:%.*]], i32 [[I_011]]
+; CHECK-NEXT:    store i8 [[ADD]], i8* [[ARRAYIDX4]], align 1
+; CHECK-NEXT:    [[INC:%.*]] = add nuw nsw i32 [[I_011]], 1
+; CHECK-NEXT:    [[ARRAYIDX_1:%.*]] = getelementptr inbounds i8, i8* [[B]], i32 [[INC]]
+; 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]], i32 [[INC]]
+; 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 [[EXIT_LOOPEXIT:%.*]], [[LOOP0:!llvm.loop !.*]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %and = and i32 %p, 1
+  %cmp = icmp eq i32 %and, 0
+  br i1 %cmp, label %guarded, label %exit
+
+guarded:
+  %rem = urem i32 %q, 2
+  %cmp2 = icmp eq i32 %rem, 0
+  tail call void @llvm.assume(i1 %cmp2)
+  %gt = icmp sgt i32 %p, %q
+  %n = select i1 %gt, i32 %p, i32 %q
+  %cmp110 = icmp sgt i32 %n, 0
+  br i1 %cmp110, label %for.body, label %exit
+
+for.body:
+  %i.011 = phi i32 [ %inc, %for.body ], [ 0, %guarded ]
+  %arrayidx = getelementptr inbounds i8, i8* %b, i32 %i.011
+  %0 = load i8, i8* %arrayidx, align 1
+  %add = add i8 %0, 3
+  %arrayidx4 = getelementptr inbounds i8, i8* %a, i32 %i.011
+  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 %exit
+
+exit:
+  ret void
+}
+
+; Make sure the loop is unrolled with a remainder loop when the trip-count
+; is not provably divisible by the unroll factor.
+
+define dso_local void @cannotProveDivisibleTC(i8* noalias nocapture %a, i8* noalias nocapture readonly %b, i32 %p, i32 %q) local_unnamed_addr {
+; CHECK-LABEL: @cannotProveDivisibleTC(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[P:%.*]], 6
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[AND]], 0
+; CHECK-NEXT:    br i1 [[CMP]], label [[GUARDED:%.*]], label [[EXIT:%.*]]
+; CHECK:       guarded:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 [[Q:%.*]], 2
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[REM]], 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:    [[CMP110:%.*]] = icmp sgt i32 [[N]], 0
+; CHECK-NEXT:    br i1 [[CMP110]], label [[FOR_BODY_PREHEADER:%.*]], label [[EXIT]]
+; CHECK:       for.body.preheader:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[N]], -1
+; CHECK-NEXT:    [[XTRAITER:%.*]] = and i32 [[N]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ult i32 [[TMP0]], 1
+; CHECK-NEXT:    br i1 [[TMP1]], label [[EXIT_LOOPEXIT_UNR_LCSSA:%.*]], label [[FOR_BODY_PREHEADER_NEW:%.*]]
+; CHECK:       for.body.preheader.new:
+; CHECK-NEXT:    [[UNROLL_ITER:%.*]] = sub i32 [[N]], [[XTRAITER]]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[I_011:%.*]] = phi i32 [ 0, [[FOR_BODY_PREHEADER_NEW]] ], [ [[INC_1:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[NITER:%.*]] = phi i32 [ [[UNROLL_ITER]], [[FOR_BODY_PREHEADER_NEW]] ], [ [[NITER_NSUB_1:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i8, i8* [[B:%.*]], i32 [[I_011]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i8, i8* [[ARRAYIDX]], align 1
+; CHECK-NEXT:    [[ADD:%.*]] = add i8 [[TMP2]], 3
+; CHECK-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds i8, i8* [[A:%.*]], i32 [[I_011]]
+; CHECK-NEXT:    store i8 [[ADD]], i8* [[ARRAYIDX4]], align 1
+; CHECK-NEXT:    [[INC:%.*]] = add nuw nsw i32 [[I_011]], 1
+; CHECK-NEXT:    [[NITER_NSUB:%.*]] = sub i32 [[NITER]], 1
+; CHECK-NEXT:    [[ARRAYIDX_1:%.*]] = getelementptr inbounds i8, i8* [[B]], i32 [[INC]]
+; CHECK-NEXT:    [[TMP3:%.*]] = load i8, i8* [[ARRAYIDX_1]], align 1
+; CHECK-NEXT:    [[ADD_1:%.*]] = add i8 [[TMP3]], 3
+; CHECK-NEXT:    [[ARRAYIDX4_1:%.*]] = getelementptr inbounds i8, i8* [[A]], i32 [[INC]]
+; CHECK-NEXT:    store i8 [[ADD_1]], i8* [[ARRAYIDX4_1]], align 1
+; CHECK-NEXT:    [[INC_1]] = add nuw nsw i32 [[INC]], 1
+; CHECK-NEXT:    [[NITER_NSUB_1]] = sub i32 [[NITER_NSUB]], 1
+; CHECK-NEXT:    [[NITER_NCMP_1:%.*]] = icmp ne i32 [[NITER_NSUB_1]], 0
+; CHECK-NEXT:    br i1 [[NITER_NCMP_1]], label [[FOR_BODY]], label [[EXIT_LOOPEXIT_UNR_LCSSA_LOOPEXIT:%.*]], [[LOOP2:!llvm.loop !.*]]
+; CHECK:       exit.loopexit.unr-lcssa.loopexit:
+; CHECK-NEXT:    [[I_011_UNR_PH:%.*]] = phi i32 [ [[INC_1]], [[FOR_BODY]] ]
+; CHECK-NEXT:    br label [[EXIT_LOOPEXIT_UNR_LCSSA]]
+; CHECK:       exit.loopexit.unr-lcssa:
+; CHECK-NEXT:    [[I_011_UNR:%.*]] = phi i32 [ 0, [[FOR_BODY_PREHEADER]] ], [ [[I_011_UNR_PH]], [[EXIT_LOOPEXIT_UNR_LCSSA_LOOPEXIT]] ]
+; CHECK-NEXT:    [[LCMP_MOD:%.*]] = icmp ne i32 [[XTRAITER]], 0
+; CHECK-NEXT:    br i1 [[LCMP_MOD]], label [[FOR_BODY_EPIL_PREHEADER:%.*]], label [[EXIT_LOOPEXIT:%.*]]
+; CHECK:       for.body.epil.preheader:
+; CHECK-NEXT:    br label [[FOR_BODY_EPIL:%.*]]
+; CHECK:       for.body.epil:
+; CHECK-NEXT:    [[I_011_EPIL:%.*]] = phi i32 [ [[I_011_UNR]], [[FOR_BODY_EPIL_PREHEADER]] ]
+; CHECK-NEXT:    [[ARRAYIDX_EPIL:%.*]] = getelementptr inbounds i8, i8* [[B]], i32 [[I_011_EPIL]]
+; CHECK-NEXT:    [[TMP4:%.*]] = load i8, i8* [[ARRAYIDX_EPIL]], align 1
+; CHECK-NEXT:    [[ADD_EPIL:%.*]] = add i8 [[TMP4]], 3
+; CHECK-NEXT:    [[ARRAYIDX4_EPIL:%.*]] = getelementptr inbounds i8, i8* [[A]], i32 [[I_011_EPIL]]
+; CHECK-NEXT:    store i8 [[ADD_EPIL]], i8* [[ARRAYIDX4_EPIL]], align 1
+; CHECK-NEXT:    [[INC_EPIL:%.*]] = add nuw nsw i32 [[I_011_EPIL]], 1
+; CHECK-NEXT:    [[CMP1_EPIL:%.*]] = icmp slt i32 [[INC_EPIL]], [[N]]
+; CHECK-NEXT:    br label [[EXIT_LOOPEXIT_EPILOG_LCSSA:%.*]]
+; CHECK:       exit.loopexit.epilog-lcssa:
+; CHECK-NEXT:    br label [[EXIT_LOOPEXIT]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %and = and i32 %p, 6
+  %cmp = icmp eq i32 %and, 0
+  br i1 %cmp, label %guarded, label %exit
+
+guarded:
+  %rem = urem i32 %q, 2
+  %cmp2 = icmp eq i32 %rem, 0
+  tail call void @llvm.assume(i1 %cmp2)
+  %gt = icmp sgt i32 %p, %q
+  %n = select i1 %gt, i32 %p, i32 %q
+  %cmp110 = icmp sgt i32 %n, 0
+  br i1 %cmp110, label %for.body, label %exit
+
+for.body:
+  %i.011 = phi i32 [ %inc, %for.body ], [ 0, %guarded ]
+  %arrayidx = getelementptr inbounds i8, i8* %b, i32 %i.011
+  %0 = load i8, i8* %arrayidx, align 1
+  %add = add i8 %0, 3
+  %arrayidx4 = getelementptr inbounds i8, i8* %a, i32 %i.011
+  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 %exit
+
+exit:
+  ret void
+}
+
+declare void @llvm.assume(i1 noundef) nofree nosync nounwind willreturn

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

 
 exit:
   ret void
 }
 
 ; 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) 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:    [[AND1:%.*]] = and i32 [[P:%.*]], 3
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[AND1]], 0
+; CHECK-NEXT:    br i1 [[CMP1]], label [[GUARDED:%.*]], label [[EXIT:%.*]]
+; CHECK:       guarded:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 [[Q:%.*]], 8
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[REM]], 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:    [[CMP110:%.*]] = icmp sgt i32 [[N]], 0
+; CHECK-NEXT:    br i1 [[CMP110]], label [[LOOP_PREHEADER:%.*]], label [[EXIT]]
+; CHECK:       loop.preheader:
 ; 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-NEXT:    br i1 [[TMP4]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP4:!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-NEXT:    br i1 [[CMP_N]], label [[EXIT_LOOPEXIT:%.*]], 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, [[LOOP_PREHEADER]] ]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[RIV:%.*]] = phi i32 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[RIVPLUS1:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[RIV:%.*]] = phi i32 [ [[RIVPLUS1:%.*]], [[LOOP]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]
 ; 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-NEXT:    br i1 [[COND]], label [[EXIT_LOOPEXIT]], label [[LOOP]], [[LOOP5:!llvm.loop !.*]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    br label [[EXIT]]
 ; 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 label %loop
+  %and1 = and i32 %p, 3
+  %cmp1 = icmp eq i32 %and1, 0
+  br i1 %cmp1, label %guarded, label %exit
+
+guarded:
+  %rem = urem i32 %q, 8
+  %cmp2 = icmp eq i32 %rem, 0
+  tail call void @llvm.assume(i1 %cmp2)
+  %gt = icmp sgt i32 %p, %q
+  %n = select i1 %gt, i32 %p, i32 %q
+  %cmp110 = icmp sgt i32 %n, 0
+  br i1 %cmp110, label %loop, label %exit
 
 loop:
-  %riv = phi i32 [ 0, %entry ], [ %rivPlus1, %loop ]
+  %riv = phi i32 [ 0, %guarded ], [ %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
+}
+
+; Make sure the loop's tail is folded when vectorized under -Os based on its trip-count's
+; not being provably divisible by chosen VF.
+
+define dso_local void @cannotProveAlignedTC(i32* noalias nocapture %A, i32 %p, i32 %q) optsize {
+; CHECK-LABEL: @cannotProveAlignedTC(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[P:%.*]], 3
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq i32 [[AND1]], 0
+; CHECK-NEXT:    br i1 [[CMP1]], label [[GUARDED:%.*]], label [[EXIT:%.*]]
+; CHECK:       guarded:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 [[Q:%.*]], 3
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq i32 [[REM]], 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:    [[CMP110:%.*]] = icmp sgt i32 [[N]], 0
+; CHECK-NEXT:    br i1 [[CMP110]], label [[LOOP_PREHEADER:%.*]], label [[EXIT]]
+; CHECK:       loop.preheader:
+; CHECK-NEXT:    br i1 false, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[N_RND_UP:%.*]] = add i32 [[N]], 3
+; CHECK-NEXT:    [[N_MOD_VF:%.*]] = urem i32 [[N_RND_UP]], 4
+; CHECK-NEXT:    [[N_VEC:%.*]] = sub i32 [[N_RND_UP]], [[N_MOD_VF]]
+; CHECK-NEXT:    [[TRIP_COUNT_MINUS_1:%.*]] = sub i32 [[N]], 1
+; CHECK-NEXT:    [[BROADCAST_SPLATINSERT:%.*]] = insertelement <4 x i32> poison, i32 [[TRIP_COUNT_MINUS_1]], i32 0
+; CHECK-NEXT:    [[BROADCAST_SPLAT:%.*]] = shufflevector <4 x i32> [[BROADCAST_SPLATINSERT]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i32 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[PRED_STORE_CONTINUE6:%.*]] ]
+; CHECK-NEXT:    [[VEC_IND:%.*]] = phi <4 x i32> [ <i32 0, i32 1, i32 2, i32 3>, [[VECTOR_PH]] ], [ [[VEC_IND_NEXT:%.*]], [[PRED_STORE_CONTINUE6]] ]
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ule <4 x i32> [[VEC_IND]], [[BROADCAST_SPLAT]]
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <4 x i1> [[TMP0]], i32 0
+; CHECK-NEXT:    br i1 [[TMP1]], label [[PRED_STORE_IF:%.*]], label [[PRED_STORE_CONTINUE:%.*]]
+; CHECK:       pred.store.if:
+; CHECK-NEXT:    [[TMP2:%.*]] = add i32 [[INDEX]], 0
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, i32* [[A:%.*]], i32 [[TMP2]]
+; CHECK-NEXT:    store i32 13, i32* [[TMP3]], align 1
+; CHECK-NEXT:    br label [[PRED_STORE_CONTINUE]]
+; CHECK:       pred.store.continue:
+; CHECK-NEXT:    [[TMP4:%.*]] = extractelement <4 x i1> [[TMP0]], i32 1
+; CHECK-NEXT:    br i1 [[TMP4]], label [[PRED_STORE_IF1:%.*]], label [[PRED_STORE_CONTINUE2:%.*]]
+; CHECK:       pred.store.if1:
+; CHECK-NEXT:    [[TMP5:%.*]] = add i32 [[INDEX]], 1
+; CHECK-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[A]], i32 [[TMP5]]
+; CHECK-NEXT:    store i32 13, i32* [[TMP6]], align 1
+; CHECK-NEXT:    br label [[PRED_STORE_CONTINUE2]]
+; CHECK:       pred.store.continue2:
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <4 x i1> [[TMP0]], i32 2
+; CHECK-NEXT:    br i1 [[TMP7]], label [[PRED_STORE_IF3:%.*]], label [[PRED_STORE_CONTINUE4:%.*]]
+; CHECK:       pred.store.if3:
+; CHECK-NEXT:    [[TMP8:%.*]] = add i32 [[INDEX]], 2
+; CHECK-NEXT:    [[TMP9:%.*]] = getelementptr inbounds i32, i32* [[A]], i32 [[TMP8]]
+; CHECK-NEXT:    store i32 13, i32* [[TMP9]], align 1
+; CHECK-NEXT:    br label [[PRED_STORE_CONTINUE4]]
+; CHECK:       pred.store.continue4:
+; CHECK-NEXT:    [[TMP10:%.*]] = extractelement <4 x i1> [[TMP0]], i32 3
+; CHECK-NEXT:    br i1 [[TMP10]], label [[PRED_STORE_IF5:%.*]], label [[PRED_STORE_CONTINUE6]]
+; CHECK:       pred.store.if5:
+; CHECK-NEXT:    [[TMP11:%.*]] = add i32 [[INDEX]], 3
+; CHECK-NEXT:    [[TMP12:%.*]] = getelementptr inbounds i32, i32* [[A]], i32 [[TMP11]]
+; CHECK-NEXT:    store i32 13, i32* [[TMP12]], align 1
+; CHECK-NEXT:    br label [[PRED_STORE_CONTINUE6]]
+; CHECK:       pred.store.continue6:
+; CHECK-NEXT:    [[INDEX_NEXT]] = add i32 [[INDEX]], 4
+; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <4 x i32> [[VEC_IND]], <i32 4, i32 4, i32 4, i32 4>
+; CHECK-NEXT:    [[TMP13:%.*]] = icmp eq i32 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP13]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP6:!llvm.loop !.*]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    br i1 true, label [[EXIT_LOOPEXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i32 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[LOOP_PREHEADER]] ]
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[RIV:%.*]] = phi i32 [ [[RIVPLUS1:%.*]], [[LOOP]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]
+; 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_LOOPEXIT]], label [[LOOP]], [[LOOP7:!llvm.loop !.*]]
+; CHECK:       exit.loopexit:
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %and1 = and i32 %p, 3
+  %cmp1 = icmp eq i32 %and1, 0
+  br i1 %cmp1, label %guarded, label %exit
+
+guarded:
+  %rem = urem i32 %q, 3
+  %cmp2 = icmp eq i32 %rem, 0
+  tail call void @llvm.assume(i1 %cmp2)
+  %gt = icmp sgt i32 %p, %q
+  %n = select i1 %gt, i32 %p, i32 %q
+  %cmp110 = icmp sgt i32 %n, 0
+  br i1 %cmp110, label %loop, label %exit
+
+loop:
+  %riv = phi i32 [ 0, %guarded ], [ %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