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

[Jack Styles via llvm-commits]

================
@@ -3072,12 +3073,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) {
+          unsigned VF = (TC - 1) / EffectiveIC;
+          if (VF >= 2 && VF <= MaxFixedVF && isPowerOf2_32(VF)) {
+            LLVM_DEBUG(dbgs() << "LV: Picking VF=" << VF
+                              << " with 1 scalar iteration remaining.\n");
+            MaxFactors.FixedVF = ElementCount::getFixed(VF);
----------------
Stylie777 wrote:

I've updated this to be `Picking MaxVF` to better describe what is being chosen here. For AArch64+sve the only situation this would happen would be MaxVF=4 and then VF ends up being VF=2. Other targets may differ, but there is an AArch64 test to cover this (`tc5_sin_f32_select_smaller_vf` in `sve-small-trip-count-vf-plus-one-cost.ll`)

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