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

[Jack Styles via llvm-commits]

================
@@ -5790,14 +5814,78 @@ LoopVectorizationPlanner::computeBestVF() {
     return {VectorizationFactor::Disabled(), nullptr};
   // If there is a single VPlan with a single VF, return it directly.
   VPlan &FirstPlan = *VPlans[0];
+  auto IsUnprofitableOneScalarTail =
+      [&](const VectorizationFactor &CurrentFactor, bool HasTail,
+          bool ForceVectorization, const ElementCount &ExactTC,
+          const VectorizationFactor &ScalarFactor,
+          const InstructionCost &ScalarCost, unsigned int UserIC) {
+        if (ForceVectorization || !HasTail || !ExactTC.isFixed() ||
+            CurrentFactor.Width.isScalable())
+          return false;
+
+        unsigned TC = ExactTC.getFixedValue();
+        if (TC == 0 || TC > TTI.getMinTripCountTailFoldingThreshold())
+          return false;
+
+        unsigned EstimatedWidth = estimateElementCount(
+            CurrentFactor.Width, Config.getVScaleForTuning());
+        if (TC != (EstimatedWidth * UserIC) + 1)
+          return false;
+
+        InstructionCost VectorCost =
+            getCostForKnownTripCount(CurrentFactor, TC, HasTail);
+        InstructionCost ScalarCostForTC =
+            getCostForKnownTripCount(ScalarFactor, TC, /*HasTail=*/false);
+        // Be conservative for the one-scalar-tail shape. It introduces extra
+        // control flow and a scalar epilogue for a single element, so require
+        // the vectorized form to save at least one scalar iteration.
+        InstructionCost AdjustedVectorCost = VectorCost + ScalarCost;
----------------
Stylie777 wrote:

I don't think this does. It accounts for it if there is a tail, but my understanding is in this case we have stopped tail folding, so HasTail will be false, and we need to account for the Scalar Cost.

The other option is to pass `true` to `getCostForKnownTripCount` for this case, however when I tested this it has issues. It does not consider the cost of the extra control flow needed for the single iteration epilogue. It is this that means cases such as `tc3_smin_i8_reject` should be rejected.Because of this I feel we need to keep `AdjustedVectorCost` to ensure we are making a clear saving to account for the setup costs. If there is a way you feel this can be done better I am open to suggesstions here.

Adding `AdjustedVectorCost` stopped the examples of loops that were better as scalar, these are covered by tests such as `tc3_udiv_i8_user_vf2` in `sve-small-trip-count-vf-plus-one-cost.ll` to ensure they stay scalar. I have added supporting IR tests to ensure the scalar loop is still there in `sve-small-trip-count-vf-plus-one.ll`.

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