[llvm] [VPlan] Compute UncountableExitStyle in VPlan. NFC (PR #206020)

Luke Lau via llvm-commits llvm-commits at lists.llvm.org
Tue Jun 30 03:18:35 PDT 2026


https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/206020

>From ff58cbe84d6d6d98060fd0d1f41bf91532435b62 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 26 Jun 2026 17:39:53 +0800
Subject: [PATCH 1/3] [VPlan] Split handleEarlyExits into countable and
 uncountable passes. NFC

This allows us to remove UncountableExitStyle::NoUncountableExit, and isolate it to just the uncountable exit path. This should make it easier to choose a style from within VPlan alone later on.

This also allows us to plug UncountableExitStyle into a TTI hook or CLI flag eventually, as I don't think we want to expose NoUncoutableExit.
---
 .../Transforms/Vectorize/LoopVectorize.cpp    | 21 ++++++-----
 llvm/lib/Transforms/Vectorize/VPlan.h         |  1 -
 .../Vectorize/VPlanConstruction.cpp           | 35 ++++---------------
 .../Transforms/Vectorize/VPlanTransforms.cpp  | 14 ++++++--
 .../Transforms/Vectorize/VPlanTransforms.h    | 21 ++++++++---
 .../VPlan/vplan-print-before-after-all.ll     |  2 +-
 6 files changed, 48 insertions(+), 46 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index f1ca4061bfd9e..fa4000494c5d2 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6579,15 +6579,18 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan1() {
   // TODO: Remove this once we can properly check the VPlan itself for both
   //       the presence of an uncountable exit and the presence of stores in
   //       the loop inside handleEarlyExits itself.
-  UncountableExitStyle EEStyle = UncountableExitStyle::NoUncountableExit;
-  if (Legal->hasUncountableEarlyExit())
-    EEStyle = Legal->hasUncountableExitWithSideEffects()
-                  ? UncountableExitStyle::MaskedHandleExitInScalarLoop
-                  : UncountableExitStyle::ReadOnly;
-
-  if (!RUN_VPLAN_PASS(VPlanTransforms::handleEarlyExits, *VPlan0, EEStyle,
-                      OrigLoop, PSE, *DT, Legal->getAssumptionCache())) {
-    return nullptr;
+  if (Legal->hasUncountableEarlyExit()) {
+    // TODO: Check target preference for style.
+    UncountableExitStyle EEStyle =
+        Legal->hasUncountableExitWithSideEffects()
+            ? UncountableExitStyle::MaskedHandleExitInScalarLoop
+            : UncountableExitStyle::ReadOnly;
+    if (!RUN_VPLAN_PASS(VPlanTransforms::handleUncountableEarlyExits, *VPlan0,
+                        OrigLoop, PSE, *DT, Legal->getAssumptionCache(),
+                        EEStyle))
+      return nullptr;
+  } else {
+    RUN_VPLAN_PASS(VPlanTransforms::handleCountableEarlyExits, *VPlan0);
   }
 
   RUN_VPLAN_PASS(VPlanTransforms::createLoopRegions, *VPlan0,
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index eaf9d1433aff7..869cbe1ac49e9 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -77,7 +77,6 @@ using VPlanPtr = std::unique_ptr<VPlan>;
 /// Different methods of handling early exits.
 ///
 enum class UncountableExitStyle {
-  NoUncountableExit = 0,
   /// No side effects to worry about, so we can process any uncountable exits
   /// in the loop and branch either to the middle block if the trip count was
   /// reached, or an early exitblock to determine which exit was taken.
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index 619fea8c10b4d..1cd14df7629aa 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1216,12 +1216,11 @@ void VPlanTransforms::createInLoopReductionRecipes(VPlan &Plan,
     R->eraseFromParent();
 }
 
-/// Check if all loads in the loop are dereferenceable. Iterates over the
-/// loop body blocks reachable from \p HeaderVPBB. Returns false if any
-/// non-dereferenceable load is found.
-static bool areAllLoadsDereferenceable(VPBasicBlock *HeaderVPBB, Loop *TheLoop,
-                                       PredicatedScalarEvolution &PSE,
-                                       DominatorTree &DT, AssumptionCache *AC) {
+bool VPlanTransforms::areAllLoadsDereferenceable(VPBasicBlock *HeaderVPBB,
+                                                 Loop *TheLoop,
+                                                 PredicatedScalarEvolution &PSE,
+                                                 DominatorTree &DT,
+                                                 AssumptionCache *AC) {
   ScalarEvolution &SE = *PSE.getSE();
   const DataLayout &DL = TheLoop->getHeader()->getDataLayout();
   for (VPBasicBlock *VPBB : vp_rpo_plain_cfg_loop_body(HeaderVPBB)) {
@@ -1260,29 +1259,8 @@ static bool areAllLoadsDereferenceable(VPBasicBlock *HeaderVPBB, Loop *TheLoop,
   return true;
 }
 
-bool VPlanTransforms::handleEarlyExits(VPlan &Plan, UncountableExitStyle Style,
-                                       Loop *TheLoop,
-                                       PredicatedScalarEvolution &PSE,
-                                       DominatorTree &DT, AssumptionCache *AC) {
+void VPlanTransforms::handleCountableEarlyExits(VPlan &Plan) {
   auto *MiddleVPBB = VPBlockUtils::getPlainCFGMiddleBlock(Plan);
-  auto [HeaderVPBB, LatchVPBB] = VPBlockUtils::getPlainCFGHeaderAndLatch(Plan);
-
-  // TODO: We would like to detect uncountable exits and stores within loops
-  //       with such exits from the VPlan alone. Exit detection can be moved
-  //       here from handleUncountableEarlyExits, but we need to improve
-  //       detection of recipes which may write to memory.
-  if (Style != UncountableExitStyle::NoUncountableExit) {
-    // Dereferenceability is checked separately for uncountable exit loops with
-    // stores, as only the loads contributing to the exit condition need to
-    // be checked.
-    if (Style == UncountableExitStyle::ReadOnly &&
-        !areAllLoadsDereferenceable(HeaderVPBB, TheLoop, PSE, DT, AC))
-      return false;
-    // TODO: Check target preference for style.
-    return handleUncountableEarlyExits(Plan, HeaderVPBB, LatchVPBB, MiddleVPBB,
-                                       TheLoop, PSE, DT, AC, Style);
-  }
-
   // Disconnect countable early exits from the loop, leaving it with a single
   // exit from the latch. Countable early exits are left for a scalar epilog.
   for (VPIRBasicBlock *EB : Plan.getExitBlocks()) {
@@ -1298,7 +1276,6 @@ bool VPlanTransforms::handleEarlyExits(VPlan &Plan, UncountableExitStyle Style,
       VPBlockUtils::disconnectBlocks(Pred, EB);
     }
   }
-  return true;
 }
 
 void VPlanTransforms::addMiddleCheck(VPlan &Plan) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 5834f38e96b84..7b413ed84450c 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -4605,12 +4605,22 @@ static bool handleUncountableExitsWithSideEffects(
 }
 
 bool VPlanTransforms::handleUncountableEarlyExits(
-    VPlan &Plan, VPBasicBlock *HeaderVPBB, VPBasicBlock *LatchVPBB,
-    VPBasicBlock *MiddleVPBB, Loop *TheLoop, PredicatedScalarEvolution &PSE,
+    VPlan &Plan, Loop *TheLoop, PredicatedScalarEvolution &PSE,
     DominatorTree &DT, AssumptionCache *AC, UncountableExitStyle Style) {
 #ifndef NDEBUG
   VPDominatorTree VPDT(Plan);
 #endif
+
+  auto *MiddleVPBB = VPBlockUtils::getPlainCFGMiddleBlock(Plan);
+  auto [HeaderVPBB, LatchVPBB] = VPBlockUtils::getPlainCFGHeaderAndLatch(Plan);
+
+  // Dereferenceability is checked separately for uncountable exit loops with
+  // stores, as only the loads contributing to the exit condition need to
+  // be checked.
+  if (Style == UncountableExitStyle::ReadOnly &&
+      !areAllLoadsDereferenceable(HeaderVPBB, TheLoop, PSE, DT, AC))
+    return false;
+
   VPBuilder LatchBuilder(LatchVPBB->getTerminator());
   SmallVector<EarlyExitInfo> Exits;
   for (VPIRBasicBlock *ExitBlock : Plan.getExitBlocks()) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index 3260526552281..f873aa1215fab 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -357,14 +357,27 @@ struct VPlanTransforms {
   /// Remove dead recipes from \p Plan.
   static void removeDeadRecipes(VPlan &Plan);
 
+  /// Check if all loads in the loop are dereferenceable. Iterates over the
+  /// loop body blocks reachable from \p HeaderVPBB. Returns false if any
+  /// non-dereferenceable load is found.
+  static bool areAllLoadsDereferenceable(VPBasicBlock *HeaderVPBB,
+                                         Loop *TheLoop,
+                                         PredicatedScalarEvolution &PSE,
+                                         DominatorTree &DT,
+                                         AssumptionCache *AC);
+
   /// Update \p Plan to account for uncountable early exits by introducing
   /// appropriate branching logic in the latch that handles early exits and the
   /// latch exit condition. Multiple exits are handled with a dispatch block
   /// that determines which exit to take based on lane-by-lane semantics.
-  static bool handleUncountableEarlyExits(
-      VPlan &Plan, VPBasicBlock *HeaderVPBB, VPBasicBlock *LatchVPBB,
-      VPBasicBlock *MiddleVPBB, Loop *TheLoop, PredicatedScalarEvolution &PSE,
-      DominatorTree &DT, AssumptionCache *AC, UncountableExitStyle Style);
+  static bool handleUncountableEarlyExits(VPlan &Plan, Loop *TheLoop,
+                                          PredicatedScalarEvolution &PSE,
+                                          DominatorTree &DT,
+                                          AssumptionCache *AC,
+                                          UncountableExitStyle Style);
+
+  /// Disconnect countable early exits from the loop.
+  static void handleCountableEarlyExits(VPlan &Plan);
 
   /// Replaces the exit condition from
   ///   (branch-on-cond eq CanonicalIVInc, VectorTripCount)
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
index 2a250d4c896ff..3175f5c89b6d9 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/vplan-print-before-after-all.ll
@@ -13,7 +13,7 @@
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::replaceSymbolicStrides
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::finalizeSCEVPredicates
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::addMiddleCheck
-; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::handleEarlyExits
+; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::handleCountableEarlyExits
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::createLoopRegions
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::introduceMasksAndLinearize
 ; CHECK: VPlan for loop in 'foo' [[BEFORE_OR_AFTER]] VPlanTransforms::createInLoopReductionRecipes

>From e2710052fb354b9cc933a298fa3047607f26d688 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 26 Jun 2026 18:08:42 +0800
Subject: [PATCH 2/3] Fix unit tests

---
 .../Transforms/Vectorize/VPlanTestBase.h        | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
index 3bcbd3ca6c937..a135fa3a232f2 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
@@ -68,10 +68,9 @@ class VPlanTestIRBase : public testing::Test {
   }
 
   /// Build the VPlan for the loop starting from \p LoopHeader.
-  VPlanPtr buildVPlan(
-      BasicBlock *LoopHeader,
-      UncountableExitStyle Style = UncountableExitStyle::NoUncountableExit,
-      bool CreateLoopRegions = true) {
+  VPlanPtr buildVPlan(BasicBlock *LoopHeader,
+                      std::optional<UncountableExitStyle> Style = std::nullopt,
+                      bool CreateLoopRegions = true) {
     Function &F = *LoopHeader->getParent();
     assert(!verifyFunction(F) && "input function must be valid");
     doAnalysis(F);
@@ -81,9 +80,9 @@ class VPlanTestIRBase : public testing::Test {
     auto Plan =
         VPlanTransforms::buildVPlan0(L, *LI, IntegerType::get(*Ctx, 64), PSE);
 
-    if (Style != UncountableExitStyle::NoUncountableExit) {
+    if (Style) {
       Inductions.clear();
-      // handleEarlyExits requires induction phi recipes.
+      // handleUncountableEarlyExits requires induction phi recipes.
       for (PHINode &Phi : LoopHeader->phis()) {
         InductionDescriptor ID;
         if (InductionDescriptor::isInductionPHI(&Phi, L, PSE, ID))
@@ -97,7 +96,11 @@ class VPlanTestIRBase : public testing::Test {
           /*AllowReordering=*/false);
     }
 
-    VPlanTransforms::handleEarlyExits(*Plan, Style, L, PSE, *DT, AC.get());
+    if (Style)
+      VPlanTransforms::handleUncountableEarlyExits(*Plan, L, PSE, *DT, AC.get(),
+                                                   *Style);
+    else
+      VPlanTransforms::handleCountableEarlyExits(*Plan);
     VPlanTransforms::addMiddleCheck(*Plan);
 
     if (CreateLoopRegions)

>From 647872e7c490fa22e78f203a9698e4eadd60d785 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 26 Jun 2026 17:59:59 +0800
Subject: [PATCH 3/3] [VPlan] Compute UncountableExitStyle in VPlan. NFC

Stacked on #206017

Currently we choose MaskedHandleExitInScalarLoop basde on a simple Instruction::mayWriteToMemory check in LoopVectorizationLegality. We can just compute this in VPlan from a fairly simple loop over the loop body's recipes. This makes it easier to choose additional uncountable exit styles in future.
---
 llvm/lib/Transforms/Vectorize/LoopVectorize.cpp | 13 +++----------
 .../Transforms/Vectorize/VPlanTransforms.cpp    | 17 ++++++++++++++++-
 llvm/lib/Transforms/Vectorize/VPlanTransforms.h |  3 +--
 .../Transforms/Vectorize/VPlanTestBase.h        |  4 ++--
 4 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index fa4000494c5d2..024dd34541cce 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6576,18 +6576,11 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan1() {
 
   // If we're vectorizing a loop with an uncountable exit, make sure that the
   // recipes are safe to handle.
-  // TODO: Remove this once we can properly check the VPlan itself for both
-  //       the presence of an uncountable exit and the presence of stores in
-  //       the loop inside handleEarlyExits itself.
+  // TODO: Remove this once we can properly check the VPlan itself for the
+  //       presence of an uncountable exit.
   if (Legal->hasUncountableEarlyExit()) {
-    // TODO: Check target preference for style.
-    UncountableExitStyle EEStyle =
-        Legal->hasUncountableExitWithSideEffects()
-            ? UncountableExitStyle::MaskedHandleExitInScalarLoop
-            : UncountableExitStyle::ReadOnly;
     if (!RUN_VPLAN_PASS(VPlanTransforms::handleUncountableEarlyExits, *VPlan0,
-                        OrigLoop, PSE, *DT, Legal->getAssumptionCache(),
-                        EEStyle))
+                        OrigLoop, PSE, *DT, Legal->getAssumptionCache()))
       return nullptr;
   } else {
     RUN_VPLAN_PASS(VPlanTransforms::handleCountableEarlyExits, *VPlan0);
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 7b413ed84450c..a6aaad929b60b 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -4604,9 +4604,18 @@ static bool handleUncountableExitsWithSideEffects(
   return true;
 }
 
+/// Returns true if any non-branch recipe in the loop may have side effects.
+static bool loopHasSideEffects(VPBasicBlock *HeaderVPBB) {
+  for (VPBasicBlock *VPBB : vp_rpo_plain_cfg_loop_body(HeaderVPBB))
+    for (VPRecipeBase &R : *VPBB)
+      if (R.mayHaveSideEffects() && &R != VPBB->getTerminator())
+        return true;
+  return false;
+}
+
 bool VPlanTransforms::handleUncountableEarlyExits(
     VPlan &Plan, Loop *TheLoop, PredicatedScalarEvolution &PSE,
-    DominatorTree &DT, AssumptionCache *AC, UncountableExitStyle Style) {
+    DominatorTree &DT, AssumptionCache *AC) {
 #ifndef NDEBUG
   VPDominatorTree VPDT(Plan);
 #endif
@@ -4614,6 +4623,12 @@ bool VPlanTransforms::handleUncountableEarlyExits(
   auto *MiddleVPBB = VPBlockUtils::getPlainCFGMiddleBlock(Plan);
   auto [HeaderVPBB, LatchVPBB] = VPBlockUtils::getPlainCFGHeaderAndLatch(Plan);
 
+  // TODO: Check target preference for style.
+  UncountableExitStyle Style =
+      loopHasSideEffects(HeaderVPBB)
+          ? UncountableExitStyle::MaskedHandleExitInScalarLoop
+          : UncountableExitStyle::ReadOnly;
+
   // Dereferenceability is checked separately for uncountable exit loops with
   // stores, as only the loads contributing to the exit condition need to
   // be checked.
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index f873aa1215fab..7cd540ac19423 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -373,8 +373,7 @@ struct VPlanTransforms {
   static bool handleUncountableEarlyExits(VPlan &Plan, Loop *TheLoop,
                                           PredicatedScalarEvolution &PSE,
                                           DominatorTree &DT,
-                                          AssumptionCache *AC,
-                                          UncountableExitStyle Style);
+                                          AssumptionCache *AC);
 
   /// Disconnect countable early exits from the loop.
   static void handleCountableEarlyExits(VPlan &Plan);
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
index a135fa3a232f2..799e6d664200d 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
@@ -97,8 +97,8 @@ class VPlanTestIRBase : public testing::Test {
     }
 
     if (Style)
-      VPlanTransforms::handleUncountableEarlyExits(*Plan, L, PSE, *DT, AC.get(),
-                                                   *Style);
+      VPlanTransforms::handleUncountableEarlyExits(*Plan, L, PSE, *DT,
+                                                   AC.get());
     else
       VPlanTransforms::handleCountableEarlyExits(*Plan);
     VPlanTransforms::addMiddleCheck(*Plan);



More information about the llvm-commits mailing list