[llvm] [VPlan] Use vputils::findCanonicalIVIncrement in more places. NFC (PR #226516)

[Luke Lau via llvm-commits]

https://github.com/lukel97 created https://github.com/llvm/llvm-project/pull/226516

In a future patch we want to enable more scenarios with early exits and tail folding. In these scenarios the terminator will no longer be a BranchOnCount/BranchOnCond but a BranchOnTwoConds, so a bunch of canonical IV matchers break. Use vputils::findCanonicalIVIncrement instead since it searches the users of the canonical IV instead of looking at the terminator.


>From 094c8507aa9077b581c863529725f466ab960c71 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Wed, 23 Sep 2026 16:58:44 +0100
Subject: [PATCH] [VPlan] Use vputils::findCanonicalIVIncrement in more places.
 NFC

In a future patch we want to enable more scenarios with early exits and tail folding. In these scenarios the terminator will no longer be a BranchOnCount/BranchOnCond but a BranchOnTwoConds, so a bunch of canonical IV matchers break. Use vputils::findCanonicalIVIncrement instead since it searches the users of the canonical IV instead of looking at the terminator.
---
 .../Transforms/Vectorize/LoopVectorize.cpp    | 11 ++--------
 .../Vectorize/VPlanConstruction.cpp           | 10 ++--------
 .../Vectorize/VPlanEVLTailFolding.cpp         | 20 +++++++------------
 3 files changed, 11 insertions(+), 30 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index d929af8afbd1da..7ea17702eae9ab 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -6615,15 +6615,8 @@ VPlanPtr LoopVectorizationPlanner::tryToBuildVPlan(VPlanPtr Plan,
   // count is >= increment and a multiple of the increment.
   VPRegionBlock *LoopRegion = Plan->getVectorLoopRegion();
   bool HasNUW = !IVUpdateMayOverflow || Style == TailFoldingStyle::None;
-  if (!HasNUW) {
-    auto *IVInc =
-        LoopRegion->getExitingBasicBlock()->getTerminator()->getOperand(0);
-    assert(match(IVInc,
-                 m_VPInstruction<Instruction::Add>(
-                     m_Specific(LoopRegion->getCanonicalIV()), m_VPValue())) &&
-           "Did not find the canonical IV increment");
-    LoopRegion->clearCanonicalIVNUW(cast<VPInstruction>(IVInc));
-  }
+  if (!HasNUW)
+    LoopRegion->clearCanonicalIVNUW(vputils::findCanonicalIVIncrement(*Plan));
 
   // ---------------------------------------------------------------------------
   // Pre-construction: record ingredients whose recipes we'll need to further
diff --git a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
index 96653c9a1c3adb..a33009d1e87fc3 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp
@@ -1353,14 +1353,8 @@ void VPlanTransforms::foldTailByMasking(VPlan &Plan) {
   Builder.createNaryOp(VPInstruction::BranchOnCond, HeaderMask);
 
   VPBasicBlock *OrigLatch = LoopRegion->getExitingBasicBlock();
-  VPValue *IVInc;
-  [[maybe_unused]] bool TermBranchOnCount =
-      match(OrigLatch->getTerminator(),
-            m_BranchOnCount(m_VPValue(IVInc),
-                            m_Specific(&Plan.getVectorTripCount())));
-  assert(TermBranchOnCount &&
-         match(IVInc, m_Add(m_Specific(LoopRegion->getCanonicalIV()),
-                            m_Specific(&Plan.getVFxUF()))) &&
+  VPValue *IVInc = vputils::findCanonicalIVIncrement(Plan);
+  assert(IVInc &&
          std::next(IVInc->getDefiningRecipe()->getIterator()) ==
              OrigLatch->getTerminator()->getIterator() &&
          "Unexpected canonical iv increment");
diff --git a/llvm/lib/Transforms/Vectorize/VPlanEVLTailFolding.cpp b/llvm/lib/Transforms/Vectorize/VPlanEVLTailFolding.cpp
index b6ede1c1edc7f4..f3572286b3834f 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanEVLTailFolding.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanEVLTailFolding.cpp
@@ -637,19 +637,13 @@ void VPlanTransforms::convertEVLExitCond(VPlan &Plan) {
   if (match(LatchBr, m_BranchOnCond(m_True())))
     return;
 
-  VPValue *CanIVInc;
-  [[maybe_unused]] bool FoundIncrement = match(
-      LatchBr,
-      m_BranchOnCond(m_SpecificCmp(CmpInst::ICMP_EQ, m_VPValue(CanIVInc),
-                                   m_Specific(&Plan.getVectorTripCount()))));
-  assert(FoundIncrement &&
-         match(CanIVInc, m_Add(m_Specific(LoopRegion->getCanonicalIV()),
-                               m_Specific(&Plan.getVFxUF()))) &&
-         "Expected BranchOnCond with ICmp comparing CanIV + VFxUF with vector "
-         "trip count");
+  VPValue *CanIVInc = vputils::findCanonicalIVIncrement(Plan);
+  auto *ExitCond = cast<VPInstruction>(findUserOf(
+      CanIVInc, m_SpecificICmp(CmpInst::ICMP_EQ, m_Specific(CanIVInc),
+                               m_Specific(&Plan.getVectorTripCount()))));
 
   Type *AVLTy = AVLNext->getScalarType();
-  VPBuilder Builder(LatchBr);
-  LatchBr->setOperand(
-      0, Builder.createICmp(CmpInst::ICMP_EQ, AVLNext, Plan.getZero(AVLTy)));
+  VPBuilder Builder(ExitCond);
+  ExitCond->replaceAllUsesWith(
+      Builder.createICmp(CmpInst::ICMP_EQ, AVLNext, Plan.getZero(AVLTy)));
 }