[llvm] [LV] codegen for tail-folded epilogue loop (PR #208764)

[Florian Hahn via llvm-commits]

================
@@ -5580,12 +5581,73 @@ void LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
     CM.collectNonVectorizedAndSetWideningDecisions(VF);
   }
 
-  buildVPlans(*VPlan1, ElementCount::getFixed(1), MaxFactors.FixedVF);
-  buildVPlans(*VPlan1, ElementCount::getScalable(1), MaxFactors.ScalableVF);
+  buildVPlans(*VPlan1, ElementCount::getFixed(1), MaxFactors.FixedVF, CM);
+  buildVPlans(*VPlan1, ElementCount::getScalable(1), MaxFactors.ScalableVF, CM);
 
   LLVM_DEBUG(printPlans(dbgs()));
 }
 
+bool LoopVectorizationPlanner::planForEpilogueTF(
+    ElementCount UserVF, unsigned UserIC, ElementCount EpilogueUserVF,
+    LoopVectorizationCostModel &EpilogueCM) {
+  if (VPlans.empty())
+    return false;
+  if (!OrigLoop->isInnermost())
+    return false;
+
+  if (!EpilogueUserVF.isVector() ||
+      ElementCount::isKnownGE(EpilogueUserVF, UserVF))
+    return false;
+
+  EpilogueCM.ValuesToIgnore.insert_range(CM.ValuesToIgnore);
+  EpilogueCM.VecValuesToIgnore.insert_range(CM.VecValuesToIgnore);
+
+  FixedScalableVFPair MaxFactors =
+      EpilogueCM.computeMaxVF(EpilogueUserVF, UserIC);
+  if (!MaxFactors ||
+      !EpilogueCM.foldTailByMasking()) // Cases that should not to be vectorized
+                                       // nor tail-folded.
+    return false;
+
+  auto VPlan1 = tryToBuildVPlan1(EpilogueCM);
+  if (!VPlan1)
+    return false;
+
+  // Invalidate interleave groups if all blocks of loop will be predicated.
+  if (EpilogueCM.blockNeedsPredicationForAnyReason(OrigLoop->getHeader()) &&
+      !useMaskedInterleavedAccesses(TTI)) {
+    LLVM_DEBUG(
+        dbgs() << "LV: [EpilogueTF] Invalidate all interleaved groups due to "
+                  "fold-tail "
+                  "by masking which requires masked-interleaved support.\n");
+    if (EpilogueCM.InterleaveInfo.invalidateGroups())
+      // Invalidating interleave groups also requires invalidating all decisions
+      // based on them, which includes widening decisions and uniform and scalar
+      // values.
+      EpilogueCM.invalidateCostModelingDecisions();
+  }
+
+  if (EpilogueCM.foldTailByMasking())
+    Legal->prepareToFoldTailByMasking();
+
+  // Collect the instructions (and their associated costs) that will be more
+  // profitable to scalarize.
+  EpilogueCM.collectNonVectorizedAndSetWideningDecisions(EpilogueUserVF);
+
+  assert(VPlans.size() == 2 &&
+         "For tail-folded epilogue, VPlans size is expected to be 2");
+  // remove last vplan which should be the epilogue plan to replace it by our
+  // tail-folded vplan:
+  assert(VPlans.back()->getSingleVF() == EpilogueUserVF &&
+         "For tail-folded epilogue, first vplan is expected to have "
+         "EpilogueUserVF");
+  VPlans.pop_back();
----------------
fhahn wrote:

Would it not be possible to call this in `::plan` in case epilogue VF is forced? The function should construct a tail-folded epilogue plan if the user requested it and in that case we are done. otherwise we construct the non-tail-folded plan 

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