[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 83ed6e69365024969c54225fb62fde6e884410a2 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 8f6a73d0a2dd8..430dd2277c973 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
@@ -444,6 +444,16 @@ class LoopVectorizationPlanner {
   /// all profitable VFs in ProfitableVFs.
   VectorizationFactor computeBestVF();
 
+  /// 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 a28cda9fe62b3..50b7a509a7d20 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -400,6 +400,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 assumptions about SCEV
 // variables not overflowing do not hold. See `emitSCEVChecks`.
 static constexpr uint32_t SCEVCheckBypassWeights[] = {1, 127};
@@ -7365,7 +7372,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;
 
@@ -11121,6 +11139,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");