[llvm] [LoopVectorize] Improve Vectorization of Low Trip Count Loops (PR #195823)

[Jack Styles via llvm-commits]

================
@@ -3063,12 +3064,35 @@ LoopVectorizationCostModel::computeMaxVF(ElementCount UserVF, unsigned UserIC) {
       // the trip count but the scalable factor does not, use the fixed-width
       // factor in preference to allow the generation of a non-predicated loop.
       if (EpilogueLoweringStatus == CM_EpilogueNotAllowedLowTripLoop &&
-          NoScalarEpilogueNeeded(MaxFactors.FixedVF.getFixedValue())) {
+          NoScalarEpilogueNeeded(MaxFactors.FixedVF.getFixedValue(),
+                                 EffectiveIC)) {
         LLVM_DEBUG(dbgs() << "LV: Picking a fixed-width so that no tail will "
                              "remain for any chosen VF.\n");
         MaxFactors.ScalableVF = ElementCount::getScalable(0);
         return MaxFactors;
       }
+      // Allow cases where the ExactTC == VF + 1. VF can be any power of
+      // 2 between 2 and MaxVF.
+      //
+      // This produces 1 vector iteration, and 1 scalar iteration with
+      // no remainder. Later passes will eliminate the loop and leave
+      // straight-line code as the both iteration counts are statically known.
+      ElementCount ExactTC = getSmallConstantTripCount(PSE.getSE(), TheLoop);
+      if (EpilogueLoweringStatus == CM_EpilogueNotAllowedLowTripLoop &&
+          ExactTC.getFixedValue() > 1) {
+        unsigned TC = ExactTC.getFixedValue();
+        unsigned MaxFixedVF = MaxFactors.FixedVF.getFixedValue();
+        if ((TC - 1) % EffectiveIC == 0) {
----------------
Stylie777 wrote:

I've refactored this for the Lambda function to return the remainder and it can be checked from there. Its also allowed me to simplify the calculation of the VF for the TC == (VF * IC) + 1.

https://github.com/llvm/llvm-project/pull/195823