[LV] Avoid needless fold tail

[gilr]

Summary

When the trip-count is provably divisible by the maximal/chosen VF, folding the loop's tail during vectorization is redundant.
This commit extends the existing test for constant trip-counts to any trip-count known to be divisible by maximal/selected VF by SCEV.

[SjoerdMeijer]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5505

Was surprised to see this change, because I thought we were handling it here already.
Is this check here still relevant? Or can we "merge" this with the one that you added below?

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5505

Is this check here still relevant?

Yes, since IC may take non-power-of-2 values. Will add a test to cover that.

[gilr]

Add a test for a constant TC with IC=3.

[SjoerdMeijer]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5505

Ah yeah, I do see that now.
Thanks

[SjoerdMeijer]

LGTM

[fhahn]

LGTM, thanks

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

10 ↗

(On Diff #313091)

Is this enough? I think it might be better to check the whole vector body?

[gilr]

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

10 ↗

(On Diff #313091)

Tried to keep the checks to the minimum proving unmasked vectorization, but perhaps indeed better to check the whole context - will switch to update_test format for easy maintenance.
Thanks @SjoerdMeijer, @fhahn!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

5505

Thinking a bit more about this, I think we should be able to use ScalarEvolution::getURemExpr to check if the trip count is a multiple of any VF. That should work for both the constant and variable trip-count cases. As I missed commenting on that before the patch landed, I put up D93677

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
   }
 
   assert(!MaxVF.isScalable() &&
          "Scalable vectors do not yet support tail folding");
   assert((UserVF.isNonZero() || isPowerOf2_32(MaxVF.getFixedValue())) &&
          "MaxVF must be a power of 2");
   unsigned MaxVFtimesIC =
       UserIC ? MaxVF.getFixedValue() * UserIC : MaxVF.getFixedValue();
   if (TC > 0 && TC % 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;
   }
 
+  // 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()));
+  unsigned TCisMultipleOf = 1 << SE->GetMinTrailingZeros(ExitCount);
+  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;
     return MaxVF;
   }

llvm/test/Transforms/LoopVectorize/avoid-needless-fold-tail.ll

+; 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
+; the trip-count's lower bits being zero.
+; CHECK: vector.body:
+; CHECK: store <4 x i32>
+
+define dso_local void @alignTC(i32* noalias nocapture %A, i32 %n) optsize {
+entry:
+  %alignedTC = and i32 %n, -8
+  br label %loop
+
+loop:
+  %riv = phi i32 [ 0, %entry ], [ %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, %alignedTC
+  br i1 %cond, label %exit, label %loop
+
+exit:
+  ret void
+}