[LoopVectorize] Optimise away the icmp when tail-folding for some low trip counts

[david-arm]

Summary

For low trip counts the vectoriser will attempt to create a single
predicated loop that folds the scalar tail into the vector body. For
some combinations of the trip count and the VF it is possible to
determine at compile time if there will only be a single vector
iteration. If so, we can avoid creating the comparison at the end of
the loop and just always branch to the loop exit. This improves the
code quality for smaller loops with low trip counts because the
compare + branch add a relatively high cost to the loop.

This optimisation may also apply for unpredicated vector loops with
low trip counts too, hence the change in test X86/pr42674.ll.

[sdesmalen]

Thanks for this patch Dave, I've left a few comments.

Perhaps it's worth highlighting in the commit message that this is more of a problem for scalable vectors, where the branch/compare isn't constant folded or instcombined away that easily as it is for fixed-width vectors.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8809

Initially I wanted to suggest creating a new (unconditional) branch instruction here, but at this point you don't know which VF will be chosen for the given VPlan, so we have to defer this decision until codegen.

8812

There's nothing wrong with it, but perhaps it's a bit unfortunate this has to be passed into the VPInstruction as an operand. It seems more like a bit of useful knowledge of the loop that the VPInstruction could use to optimise the branch, rather than an actual operand for the conceptual BranchOnCount intrinsic. If the trip-count of the loop is constant, maybe we can store that information as a piece of state somewhere.

@fhahn what would be the desired place to add such information about the loop? My understanding was that the uses of the InnerLoopVectorizer in VPlan were being phased out, is that right?

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

nit: How about

Value *ConstCmp = nullptr;
if (auto *C = dyn_cast<ConstantInt>(TC))
  if (C->getZExtValue() <= State.UF * State.VF.getKnownMinValue())
    ConstCmp = Builder.getInt1(true);
Value *Cond = ConstCmp ? ConstCmp : Builder.CreateICmpEQ(IV, VTC);

978

I noticed some of this code here has been removed and I'm not sure if this change is then still required. In any case this patch needs a rebase.

[david-arm]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8812

Yeah I'm not terribly happy about adding this as an operand, but I wasn't sure how else to pass this. I could possibly add a constant trip count to the VPTransformState, which is available when generating the instruction?

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

It would have to be something like this I think:

Value *ConstCmp = nullptr;
if (auto *C = dyn_cast<ConstantInt>(TC)) {
  uint64_t TCVal = C->getZExtValue();
  if (TCVal && TCVal  <= State.UF * State.VF.getKnownMinValue())
    ConstCmp = Builder.getInt1(true);
}
Value *Cond = ConstCmp ? ConstCmp : Builder.CreateICmpEQ(IV, VTC);

because I'm pretty sure I found a test where the trip count was actually defined as a zero constant. I'm happy to change it to the code above, but not sure it looks much better to be honest. :)

978

Yeah I noticed that too. Will fix in a new patch!

[sdesmalen]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

If the trip-count is zero, wouldn't true be the correct value for the condition?

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

Hmm, I am not sure if doing this late in codegen is the right place. If we simplify the condition and effectively remove the loop, it would be good to remove the branch-on-count recipe directly in VPlan. I *think* we should have almost everything needed to do that already in place. Let me check.

[david-arm]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

Hi @fhahn, have you had a chance to look into this at all?

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

Unfortunately not yet, I've still got a backlog of other things to work through :(. Hopefully I'll be able to check this week, otherwise I think we should proceed with this patch next week.

[david-arm]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

So the problem with treating zero-trip counts the same as non-zero trip counts is that the loop vectoriser considers loops like this as having a zero trip count:

entry:
  br label %loop.body

loop.body:
  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.body ]
  %a = extractvalue { i64, i64 } %sv, 0
  %b = extractvalue { i64, i64 } %sv, 1
  %addr = getelementptr i64, i64* %dst, i32 %iv
  %add = add i64 %a, %b
  store i64 %add, i64* %addr
  %iv.next = add nsw i32 %iv, 1
  %cond = icmp ne i32 %iv.next, 0
  br i1 %cond, label %loop.body, label %exit

exit:
  ret void

even though this loop is actually going to execute UINT32_MAX times before finally the IV overflows back to 0! This came from a real test:

Transforms/LoopVectorize/AArch64/extractvalue-no-scalarization-required.ll

If I fold away the comparison and always branch out of the loop then I've changed the behaviour of the original scalar loop. So I think we either have to:

1. Fix the vectoriser to report the trip count as UINT32_MAX for loops like this, or
2. Only apply my optimisation for low trip counts > 0

[david-arm]

• I've removed the need for a separate TripCount VPValue in the VPlan class because we're always going to need the original scalar trip count, and the value is the same for each Part anyway. Now we just store a copy in the VPTransformState so that the execute() functions can access it.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

750–760

Took a bit longer, but D126680 is what simplification directly on VPlan would look like. It depends on a few upcoming improvements though.

[david-arm]

Already fixed by an alternative implementation - D126680

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

static void addCanonicalIVRecipes(VPlan &Plan, Type *IdxTy, DebugLoc DL,
 
   VPBasicBlock *EB = TopRegion->getExitBasicBlock();
   if (IsVPlanNative) {
     EB = cast<VPBasicBlock>(EB->getSinglePredecessor());
     EB->setCondBit(nullptr);
   }
   EB->appendRecipe(CanonicalIVIncrement);
 
   auto *BranchOnCount =
       new VPInstruction(VPInstruction::BranchOnCount,
-                        {CanonicalIVIncrement, &Plan.getVectorTripCount()}, DL);
+                        {CanonicalIVIncrement, &Plan.getVectorTripCount(),
+                         Plan.getOrCreateTripCount()},
+                        DL);
   EB->appendRecipe(BranchOnCount);
 }
 
 VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
     VFRange &Range, SmallPtrSetImpl<Instruction *> &DeadInstructions,
     const MapVector<Instruction *, Instruction *> &SinkAfter) {
 
   SmallPtrSet<const InterleaveGroup<Instruction> *, 1> InterleaveGroups;

llvm/lib/Transforms/Vectorize/VPlan.cpp

case VPInstruction::CanonicalIVIncrementNUW: {
     }
 
     State.set(this, Next, Part);
     break;
   }
   case VPInstruction::BranchOnCount: {
     if (Part != 0)
       break;
-    // First create the compare.
+    // First create the compare if necessary.
     Value *IV = State.get(getOperand(0), Part);
-    Value *TC = State.get(getOperand(1), Part);
-    Value *Cond = Builder.CreateICmpEQ(IV, TC);
+    Value *VTC = State.get(getOperand(1), Part);
+    Value *TC = State.get(getOperand(2), Part);
+
+    uint64_t TCVal = 0;
+    if (auto *TC2 = dyn_cast<ConstantInt>(TC))
+      TCVal = TC2->getZExtValue();
+
+    Value *Cond;
+    // When we know there will only be one vector iteration there is no need to
+    // create the comparison, since we already know the answer.
+    if (TCVal && TCVal <= (State.UF * State.VF.getKnownMinValue()))
+      Cond = Builder.getInt1(true);
+    else
+      Cond = Builder.CreateICmpEQ(IV, VTC);
 
     // Now create the branch.
     auto *Plan = getParent()->getPlan();
     VPRegionBlock *TopRegion = Plan->getVectorLoopRegion();
     VPBasicBlock *Header = TopRegion->getEntry()->getEntryBasicBlock();
     if (Header->empty()) {
       assert(EnableVPlanNativePath &&
              "empty entry block only expected in VPlanNativePath");
void VPlan::execute(VPTransformState *State) {
   // Move both the branch and check from LastBB to VectorLatchBB.
   auto *LastBranch = cast<BranchInst>(LastBB->getTerminator());
   LastBranch->moveBefore(VectorLatchBB->getTerminator());
   VectorLatchBB->getTerminator()->eraseFromParent();
   // Move condition so it is guaranteed to be next to branch. This is only done
   // to avoid excessive test updates.
   // TODO: Remove special handling once the increments for all inductions are
   // modeled explicitly in VPlan.
+  if (isa<Instruction>(LastBranch->getCondition()))
     cast<Instruction>(LastBranch->getCondition())->moveBefore(LastBranch);
   // Connect LastBB to VectorLatchBB to facilitate their merge.
   BranchInst::Create(VectorLatchBB, LastBB);
 
   // Merge LastBB with Latch.
   bool Merged = MergeBlockIntoPredecessor(VectorLatchBB, nullptr, State->LI);
   (void)Merged;
   assert(Merged && "Could not merge last basic block with latch.");
   VectorLatchBB = LastBB;

llvm/test/Transforms/LoopVectorize/AArch64/sve-low-trip-count.ll

 define void @trip5_i8(i8* noalias nocapture noundef %dst, i8* noalias nocapture noundef readonly %src) #0 {
 ; CHECK-LABEL: @trip5_i8(
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %vector.body ]
 ; CHECK:         [[ACTIVE_LANE_MASK:%.*]] = call <vscale x 16 x i1> @llvm.get.active.lane.mask.nxv16i1.i64(i64 {{%.*}}, i64 5)
 ; CHECK:         {{%.*}} = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]], <vscale x 16 x i8> poison)
 ; CHECK:         {{%.*}} = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]], <vscale x 16 x i8> poison)
 ; CHECK:         call void @llvm.masked.store.nxv16i8.p0nxv16i8(<vscale x 16 x i8> {{%.*}}, <vscale x 16 x i8>* {{%.*}}, i32 1, <vscale x 16 x i1> [[ACTIVE_LANE_MASK]])
-; CHECK:         [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    [[VF:%.*]] = mul i64 [[VSCALE]], 16
-; CHECK-NEXT:    [[INDEX_NEXT]] = add i64 [[INDEX]], [[VF]]
-; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[INDEX_NEXT]], {{%.*}}
-; CHECK-NEXT:    br i1 [[COND]], label %middle.block, label %vector.body
+; CHECK:         br i1 true, label %middle.block, label %vector.body
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body
   %i.08 = phi i64 [ 0, %entry ], [ %inc, %for.body ]
   %arrayidx = getelementptr inbounds i8, i8* %src, i64 %i.08
   %0 = load i8, i8* %arrayidx, align 1

llvm/test/Transforms/LoopVectorize/X86/constant-fold.ll

 ; CHECK-NEXT:    [[OFFSET_IDX:%.*]] = trunc i32 [[INDEX]] to i16
 ; CHECK-NEXT:    [[TMP0:%.*]] = add i16 [[OFFSET_IDX]], 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = sext i16 [[TMP0]] to i64
 ; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr [2 x i16*], [2 x i16*]* @b, i16 0, i64 [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i16*, i16** [[TMP2]], i32 0
 ; CHECK-NEXT:    [[TMP4:%.*]] = bitcast i16** [[TMP3]] to <2 x i16*>*
 ; CHECK-NEXT:    store <2 x i16*> <i16* getelementptr inbounds ([1 x %rec8], [1 x %rec8]* @a, i32 0, i32 0, i32 0), i16* getelementptr inbounds ([1 x %rec8], [1 x %rec8]* @a, i32 0, i32 0, i32 0)>, <2 x i16*>* [[TMP4]], align 8
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i32 [[INDEX]], 2
-; CHECK-NEXT:    [[TMP5:%.*]] = icmp eq i32 [[INDEX_NEXT]], 2
-; CHECK-NEXT:    br i1 [[TMP5]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP0:!llvm.loop !.*]]
+; CHECK-NEXT:    br i1 true, label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP0:!llvm.loop !.*]]
 ; CHECK:       middle.block:
 ; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i32 2, 2
 ; CHECK-NEXT:    br i1 [[CMP_N]], label [[BB3:%.*]], label [[SCALAR_PH]]
 ; CHECK:       scalar.ph:
 ; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i16 [ 2, [[MIDDLE_BLOCK]] ], [ 0, [[BB1:%.*]] ]
 ; CHECK-NEXT:    br label [[BB2:%.*]]
 ; CHECK:       bb2:
 ; CHECK-NEXT:    [[C_1_0:%.*]] = phi i16 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[_TMP9:%.*]], [[BB2]] ]

llvm/test/Transforms/LoopVectorize/X86/outer_loop_test1_no_explicit_vect_width.ll

 ; AVX: %[[InnerPhiNext]] = add nuw nsw <8 x i64> %[[InnerPhi]], <i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1, i64 1>
 ; AVX: %[[VecCond:.*]] = icmp eq <8 x i64> %[[InnerPhiNext]], <i64 8, i64 8, i64 8, i64 8, i64 8, i64 8, i64 8, i64 8>
 ; AVX: %[[InnerCond:.*]] = extractelement <8 x i1> %[[VecCond]], i32 0
 ; AVX: br i1 %[[InnerCond]], label %[[ForInc]], label %[[InnerLoop]]
 
 ; AVX: [[ForInc]]:
 ; AVX: %[[IndNext]] = add nuw i64 %[[Ind]], 8
 ; AVX: %[[VecIndNext]] = add <8 x i64> %[[VecInd]], <i64 8, i64 8, i64 8, i64 8, i64 8, i64 8, i64 8, i64 8>
-; AVX: %[[Cmp:.*]] = icmp eq i64 %[[IndNext]], 8
-; AVX: br i1 %[[Cmp]], label %middle.block, label %vector.body
+; AVX: br i1 true, label %middle.block, label %vector.body
 
 @arr2 = external global [8 x i32], align 16
 @arr = external global [8 x [8 x i32]], align 16
 
 ; Function Attrs: norecurse nounwind uwtable
 define void @foo(i32 %n) {
 entry:
   br label %for.body

llvm/test/Transforms/LoopVectorize/X86/pr34438.ll

 ; CHECK-NEXT:    [[TMP4:%.*]] = getelementptr inbounds float, float* [[A:%.*]], i64 [[TMP0]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds float, float* [[TMP4]], i32 0
 ; CHECK-NEXT:    [[TMP6:%.*]] = bitcast float* [[TMP5]] to <8 x float>*
 ; CHECK-NEXT:    [[WIDE_LOAD1:%.*]] = load <8 x float>, <8 x float>* [[TMP6]], align 4, !llvm.access.group !0
 ; CHECK-NEXT:    [[TMP7:%.*]] = fadd fast <8 x float> [[WIDE_LOAD]], [[WIDE_LOAD1]]
 ; CHECK-NEXT:    [[TMP8:%.*]] = bitcast float* [[TMP5]] to <8 x float>*
 ; CHECK-NEXT:    store <8 x float> [[TMP7]], <8 x float>* [[TMP8]], align 4, !llvm.access.group !0
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 8
-; CHECK-NEXT:    [[TMP9:%.*]] = icmp eq i64 [[INDEX_NEXT]], 8
-; CHECK-NEXT:    br i1 [[TMP9]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP1:!llvm.loop !.*]]
+; CHECK-NEXT:    br i1 true, label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], [[LOOP1:!llvm.loop !.*]]
 ; CHECK:       middle.block:
 ; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 8, 8
 ; CHECK-NEXT:    br i1 [[CMP_N]], label [[FOR_END:%.*]], label [[SCALAR_PH]]
 ; CHECK:       scalar.ph:
 ; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ 8, [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[INDVARS_IV_NEXT:%.*]], [[FOR_BODY]] ]

llvm/test/Transforms/LoopVectorize/X86/pr42674.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt %s -loop-vectorize -instcombine -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -mtriple=x86_64-unknown-linux-gnu -mattr=avx512vl,avx512dq,avx512bw -S | FileCheck %s
 
 @bytes = global [128 x i8] zeroinitializer, align 16
 
 ; Make sure we end up with vector code for this loop. We used to try to create
 ; a VF=64,UF=4 loop, but the scalar trip count is only 128 so
 ; the vector loop was dead code leaving only a scalar remainder.
 define zeroext i8 @sum() {
 ; CHECK-LABEL: @sum(
 ; CHECK-NEXT:  iter.check:
-; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
-; CHECK:       vector.body:
-; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[VEC_PHI:%.*]] = phi <64 x i8> [ zeroinitializer, [[ENTRY]] ], [ [[TMP4:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[VEC_PHI1:%.*]] = phi <64 x i8> [ zeroinitializer, [[ENTRY]] ], [ [[TMP5:%.*]], [[VECTOR_BODY]] ]
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds [128 x i8], [128 x i8]* @bytes, i64 0, i64 [[INDEX]]
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds [128 x i8], [128 x i8]* @bytes, i64 0, i64 0
 ; CHECK-NEXT:    [[TMP1:%.*]] = bitcast i8* [[TMP0]] to <64 x i8>*
 ; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <64 x i8>, <64 x i8>* [[TMP1]], align 16
 ; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr inbounds i8, i8* [[TMP0]], i64 64
 ; CHECK-NEXT:    [[TMP3:%.*]] = bitcast i8* [[TMP2]] to <64 x i8>*
 ; CHECK-NEXT:    [[WIDE_LOAD2:%.*]] = load <64 x i8>, <64 x i8>* [[TMP3]], align 16
-; CHECK-NEXT:    [[TMP4]] = add <64 x i8> [[WIDE_LOAD]], [[VEC_PHI]]
-; CHECK-NEXT:    [[TMP5]] = add <64 x i8> [[WIDE_LOAD2]], [[VEC_PHI1]]
-; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 128
-; CHECK-NEXT:    [[TMP6:%.*]] = icmp eq i64 [[INDEX]], 0
-; CHECK-NEXT:    br i1 [[TMP6]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop !0
-; CHECK:       middle.block:
+; CHECK-NEXT:    [[TMP4:%.*]] = add <64 x i8> [[WIDE_LOAD]], zeroinitializer
+; CHECK-NEXT:    [[TMP5:%.*]] = add <64 x i8> [[WIDE_LOAD2]], zeroinitializer
+; CHECK-NEXT:    [[INDEX_NEXT:%.*]] = add nuw i64 0, 128
 ; CHECK-NEXT:    [[BIN_RDX:%.*]] = add <64 x i8> [[TMP5]], [[TMP4]]
-; CHECK-NEXT:    [[TMP7:%.*]] = call i8 @llvm.vector.reduce.add.v64i8(<64 x i8> [[BIN_RDX]])
-; CHECK-NEXT:    ret i8 [[TMP7]]
+; CHECK-NEXT:    [[TMP6:%.*]] = call i8 @llvm.vector.reduce.add.v64i8(<64 x i8> [[BIN_RDX]])
+; CHECK-NEXT:    ret i8 [[TMP6]]
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %for.body, %entry
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %r.010 = phi i8 [ 0, %entry ], [ %add, %for.body ]
   %arrayidx = getelementptr inbounds [128 x i8], [128 x i8]* @bytes, i64 0, i64 %indvars.iv