[llvm] [VPlan] Move tail folding logic out of addMiddleCheck. NFC (PR #203475)

Luke Lau via llvm-commits llvm-commits at lists.llvm.org
Sun Jun 14 19:12:01 PDT 2026


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

>From e9177e11f7aa404399aab11a80425d1830858430 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 12 Jun 2026 15:36:51 +0800
Subject: [PATCH 1/2] [VPlan] Move tail folding logic out of addMiddleCheck.
 NFC

We simplify the TripCount == VectorTrip count condition with tail
folding, but we can just do that in foldTailByMasking and keep the
logic in one place instead.
---
 .../Transforms/Vectorize/LoopVectorize.cpp    |  3 +--
 .../Vectorize/VPlanConstruction.cpp           | 21 ++++++++++++-------
 .../Transforms/Vectorize/VPlanTransforms.h    |  2 +-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index f9b325a5233e6..c3dbf0a1ee3fa 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6563,8 +6563,7 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan1() {
   // loop with an in-loop mask, then the middle check has already been
   // created to compare against the actual number of lanes executed.
   if (EEStyle != UncountableExitStyle::MaskedHandleExitInScalarLoop)
-    RUN_VPLAN_PASS(VPlanTransforms::addMiddleCheck, *VPlan0,
-                   CM.foldTailByMasking());
+    RUN_VPLAN_PASS(VPlanTransforms::addMiddleCheck, *VPlan0);
   RUN_VPLAN_PASS(VPlanTransforms::createLoopRegions, *VPlan0,
                  getDebugLocFromInstOrOperands(Legal->getPrimaryInduction()));
   if (CM.foldTailByMasking())
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index 3e41133ff7463..dd3c8bebd1cd1 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1301,7 +1301,7 @@ bool VPlanTransforms::handleEarlyExits(VPlan &Plan, UncountableExitStyle Style,
   return true;
 }
 
-void VPlanTransforms::addMiddleCheck(VPlan &Plan, bool TailFolded) {
+void VPlanTransforms::addMiddleCheck(VPlan &Plan) {
   auto *MiddleVPBB = VPBlockUtils::getPlainCFGMiddleBlock(Plan);
   // If MiddleVPBB has a single successor then the original loop does not exit
   // via the latch and the single successor must be the scalar preheader.
@@ -1332,12 +1332,9 @@ void VPlanTransforms::addMiddleCheck(VPlan &Plan, bool TailFolded) {
   auto *LatchVPBB = cast<VPBasicBlock>(MiddleVPBB->getSinglePredecessor());
   DebugLoc LatchDL = LatchVPBB->getTerminator()->getDebugLoc();
   VPBuilder Builder(MiddleVPBB);
-  VPValue *Cmp;
-  if (TailFolded)
-    Cmp = Plan.getTrue();
-  else
-    Cmp = Builder.createICmp(CmpInst::ICMP_EQ, Plan.getTripCount(),
-                             &Plan.getVectorTripCount(), LatchDL, "cmp.n");
+  VPValue *Cmp =
+      Builder.createICmp(CmpInst::ICMP_EQ, Plan.getTripCount(),
+                         &Plan.getVectorTripCount(), LatchDL, "cmp.n");
   Builder.createNaryOp(VPInstruction::BranchOnCond, {Cmp}, LatchDL);
 }
 
@@ -1449,6 +1446,16 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
         Builder.createNaryOp(VPInstruction::ExtractLane, {LastActiveLane, Op});
     R.getVPSingleValue()->replaceAllUsesWith(Ext);
   }
+
+  // VectorTripCount now equals TripCount so simplify the MiddleVPBB branch.
+  assert(match(Plan.getMiddleBlock()->getTerminator(),
+               m_BranchOnCond(m_OneUse(m_SpecificICmp(
+                   CmpInst::ICMP_EQ, m_Specific(Plan.getTripCount()),
+                   m_Specific(&Plan.getVectorTripCount()))))) &&
+         "Unexpected MiddleVPBB branch");
+  VPValue *OldCond = Plan.getMiddleBlock()->getTerminator()->getOperand(0);
+  Plan.getMiddleBlock()->getTerminator()->setOperand(0, Plan.getTrue());
+  OldCond->getDefiningRecipe()->eraseFromParent();
 }
 
 /// Insert \p CheckBlockVPBB on the edge leading to the vector preheader,
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
index b186594cd5b1b..51984b803ab43 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.h
@@ -172,7 +172,7 @@ struct VPlanTransforms {
 
   /// If a check is needed to guard executing the scalar epilogue loop, it will
   /// be added to the middle block.
-  LLVM_ABI_FOR_TEST static void addMiddleCheck(VPlan &Plan, bool TailFolded);
+  LLVM_ABI_FOR_TEST static void addMiddleCheck(VPlan &Plan);
 
   // Create a check in \p CheckBlock to see if the vector loop should be
   // executed. May create VPExpandSCEV recipes in the plan's entry block.

>From 06db36d0254d183329bdaee9dfcbbfc7744016d2 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Fri, 12 Jun 2026 18:09:56 +0800
Subject: [PATCH 2/2] Leave erasing to removeDeadRecipes, fix unit tests

---
 llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp      | 6 ++----
 llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll | 4 ++++
 llvm/unittests/Transforms/Vectorize/VPlanTestBase.h      | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index dd3c8bebd1cd1..a92ac81305c04 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1449,13 +1449,11 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
 
   // VectorTripCount now equals TripCount so simplify the MiddleVPBB branch.
   assert(match(Plan.getMiddleBlock()->getTerminator(),
-               m_BranchOnCond(m_OneUse(m_SpecificICmp(
+               m_BranchOnCond(m_SpecificICmp(
                    CmpInst::ICMP_EQ, m_Specific(Plan.getTripCount()),
-                   m_Specific(&Plan.getVectorTripCount()))))) &&
+                   m_Specific(&Plan.getVectorTripCount())))) &&
          "Unexpected MiddleVPBB branch");
-  VPValue *OldCond = Plan.getMiddleBlock()->getTerminator()->getOperand(0);
   Plan.getMiddleBlock()->getTerminator()->setOperand(0, Plan.getTrue());
-  OldCond->getDefiningRecipe()->eraseFromParent();
 }
 
 /// Insert \p CheckBlockVPBB on the edge leading to the vector preheader,
diff --git a/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll b/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
index 464bc967bc65d..c5623e3de4b86 100644
--- a/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
+++ b/llvm/test/Transforms/LoopVectorize/VPlan/tail-folding.ll
@@ -47,6 +47,7 @@ define i32 @live_out(ptr noalias %p, i32 %n) {
 ; CHECK-NEXT:    EMIT vp<[[VP10:%[0-9]+]]> = exiting-iv-value ir<%iv>
 ; CHECK-NEXT:    EMIT vp<[[VP11:%[0-9]+]]> = extract-last-part vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT vp<[[VP12:%[0-9]+]]> = extract-last-lane vp<[[VP11]]>
+; CHECK-NEXT:    EMIT vp<%cmp.n> = icmp eq ir<%n>, vp<[[VP2]]>
 ; CHECK-NEXT:    EMIT vp<[[VP13:%[0-9]+]]> = last-active-lane vp<[[VP6]]>
 ; CHECK-NEXT:    EMIT vp<[[VP14:%[0-9]+]]> = extract-lane vp<[[VP13]]>, vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT branch-on-cond ir<true>
@@ -142,6 +143,7 @@ define i32 @conditional_live_out(ptr noalias %p, i32 %n, i1 %c) {
 ; CHECK-NEXT:    EMIT vp<[[VP10:%[0-9]+]]> = exiting-iv-value ir<%iv>
 ; CHECK-NEXT:    EMIT vp<[[VP11:%[0-9]+]]> = extract-last-part vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT vp<[[VP12:%[0-9]+]]> = extract-last-lane vp<[[VP11]]>
+; CHECK-NEXT:    EMIT vp<%cmp.n> = icmp eq ir<%n>, vp<[[VP2]]>
 ; CHECK-NEXT:    EMIT vp<[[VP13:%[0-9]+]]> = last-active-lane vp<[[VP6]]>
 ; CHECK-NEXT:    EMIT vp<[[VP14:%[0-9]+]]> = extract-lane vp<[[VP13]]>, vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT branch-on-cond ir<true>
@@ -226,6 +228,7 @@ define void @header_unconditional_branch(ptr noalias %p, i32 %n) {
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  middle.block:
 ; CHECK-NEXT:    EMIT vp<[[VP9:%[0-9]+]]> = exiting-iv-value ir<%iv>
+; CHECK-NEXT:    EMIT vp<%cmp.n> = icmp eq ir<%n>, vp<[[VP2]]>
 ; CHECK-NEXT:    EMIT branch-on-cond ir<true>
 ; CHECK-NEXT:  Successor(s): ir-bb<exit>, scalar.ph
 ; CHECK-EMPTY:
@@ -305,6 +308,7 @@ define i32 @reduction(ptr noalias %p, i32 %n) {
 ; CHECK-NEXT:    EMIT vp<[[VP12:%[0-9]+]]> = extract-last-lane vp<[[VP11]]>
 ; CHECK-NEXT:    EMIT vp<[[VP13:%[0-9]+]]> = extract-last-part vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT vp<[[VP14:%[0-9]+]]> = extract-last-lane vp<[[VP13]]>
+; CHECK-NEXT:    EMIT vp<%cmp.n> = icmp eq ir<%n>, vp<[[VP2]]>
 ; CHECK-NEXT:    EMIT vp<[[VP15:%[0-9]+]]> = last-active-lane vp<[[VP6]]>
 ; CHECK-NEXT:    EMIT vp<[[VP16:%[0-9]+]]> = extract-lane vp<[[VP15]]>, vp<[[VP8]]>
 ; CHECK-NEXT:    EMIT vp<[[VP17:%[0-9]+]]> = last-active-lane vp<[[VP6]]>
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
index c38b4cc730f47..9420420c45517 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
@@ -98,7 +98,7 @@ class VPlanTestIRBase : public testing::Test {
     }
 
     VPlanTransforms::handleEarlyExits(*Plan, Style, L, PSE, *DT, AC.get());
-    VPlanTransforms::addMiddleCheck(*Plan, false);
+    VPlanTransforms::addMiddleCheck(*Plan);
 
     if (CreateLoopRegions)
       VPlanTransforms::createLoopRegions(*Plan, {});



More information about the llvm-commits mailing list