[llvm] [vectorization] More flexibility for VFxIC (PR #138709)

[Serval MARTINOT-LAGARDE via llvm-commits]

https://github.com/Serval6 updated https://github.com/llvm/llvm-project/pull/138709

>From 5e4bd3267adc26fcd16a95e3ccffc98c2de38d93 Mon Sep 17 00:00:00 2001
From: Serval Martinot-Lagarde <serval.martinot-lagarde at sipearl.com>
Date: Tue, 29 Apr 2025 10:45:10 +0200
Subject: [PATCH] [vectorization] Change vectorized to custom coefficients in
 LoopVectorizer

---
 .../Vectorize/LoopVectorizationPlanner.h      | 10 +++++
 .../Transforms/Vectorize/LoopVectorize.cpp    | 40 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
index 838476dcae661..c5f280ea8e604 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
@@ -516,6 +516,16 @@ class LoopVectorizationPlanner {
   unsigned selectInterleaveCount(VPlan &Plan, ElementCount VF,
                                  InstructionCost LoopCost);
 
+  /// Search in \p ProfitableVFs if the selected \p VF exists.
+  std::optional<VectorizationFactor> getProfitableVF(ElementCount VF) const {
+    if (!hasPlanWithVF(VF))
+      return std::nullopt;
+    for (const auto &P : ProfitableVFs)
+      if (P.Width == VF)
+        return P;
+    return std::nullopt;
+  }
+
   /// Generate the IR code for the vectorized loop captured in VPlan \p BestPlan
   /// according to the best selected \p VF and  \p UF.
   ///
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index 70f884016d08c..d5e8d97233ff1 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -402,6 +402,13 @@ static cl::opt<bool> EnableEarlyExitVectorization(
     cl::desc(
         "Enable vectorization of early exit loops with uncountable exits."));
 
+static cl::opt<unsigned>
+    ChangeVectToCustomVF("change-vectorize-to-custom-VF", cl::init(0),
+                         cl::Hidden, cl::desc("Change vectorize to custom VF"));
+static cl::opt<unsigned>
+    ChangeVectToCustomIC("change-vectorize-to-custom-IC", cl::init(0),
+                         cl::Hidden, cl::desc("Change vectorize to custom IC"));
+
 // Likelyhood of bypassing the vectorized loop because there are zero trips left
 // after prolog. See `emitIterationCountCheck`.
 static constexpr uint32_t MinItersBypassWeights[] = {1, 127};
@@ -6692,7 +6699,18 @@ void LoopVectorizationPlanner::plan(ElementCount UserVF, unsigned UserIC) {
   CM.collectValuesToIgnore();
   CM.collectElementTypesForWidening();
 
-  FixedScalableVFPair MaxFactors = CM.computeMaxVF(UserVF, UserIC);
+  // Change User(VF/IC) to ChangeVect(VF/IC) if defined and User(VF/IC) is 0,
+  // used in computeMaxVF
+  auto CV2CVF = UserVF;
+  auto CV2CIC = UserIC;
+  if (ChangeVectToCustomVF > 0 && UserVF.isZero())
+    CV2CVF =
+        ElementCount::get(ChangeVectToCustomVF < 2 ? 2 : ChangeVectToCustomVF,
+                          UserVF.isScalable());
+  if (ChangeVectToCustomIC > 0 && UserIC == 0)
+    CV2CIC = ChangeVectToCustomIC;
+
+  FixedScalableVFPair MaxFactors = CM.computeMaxVF(CV2CVF, CV2CIC);
   if (!MaxFactors) // Cases that should not to be vectorized nor interleaved.
     return;
 
@@ -10254,6 +10272,26 @@ bool LoopVectorizePass::processLoop(Loop *L) {
   bool DisableRuntimeUnroll = false;
   MDNode *OrigLoopID = L->getLoopID();
   {
+    if (ChangeVectToCustomIC != 0 || ChangeVectToCustomVF != 0) {
+      LLVM_DEBUG(dbgs() << "LV: ChangeVectorizedToCustom (CustomVF:"
+                        << ChangeVectToCustomVF
+                        << ", CustomIC:" << ChangeVectToCustomIC
+                        << ", UserVF:" << UserVF << ", UserIC:" << UserIC
+                        << ") on (VF:" << VF.Width << ", IC:" << IC << "): ");
+      auto CVF = ElementCount::get(ChangeVectToCustomVF, VF.Width.isScalable());
+      if (ChangeVectToCustomVF == 0)
+        CVF = VF.Width;
+      std::optional<VectorizationFactor> MaybeVF = LVP.getProfitableVF(CVF);
+      if (MaybeVF) {
+        VF = *MaybeVF;
+        if (ChangeVectToCustomIC > 0)
+          IC = ChangeVectToCustomIC;
+        LLVM_DEBUG(dbgs() << "APPLIED (VF:" << VF.Width << ", IC:" << IC
+                          << ")\n");
+      } else
+        LLVM_DEBUG(dbgs() << "INVALID\n");
+    }
+
     using namespace ore;
     if (!VectorizeLoop) {
       assert(IC > 1 && "interleave count should not be 1 or 0");