[llvm] [VPlan] Fix isLatch to handle BranchOnTwoConds, check verifier.(NFC) (PR #192808)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/192808

>From 18c7c6c621d8ef3ce079757606f49a003e5261a7 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 28 Feb 2026 09:19:10 +0000
Subject: [PATCH] [VPlan] Fix isLatch to handle BranchOnTwoConds, add verifier
 test. (NFCI)

---
 llvm/lib/Transforms/Vectorize/VPlanUtils.cpp  |  6 +-
 .../Transforms/Vectorize/VPlanVerifier.cpp    |  7 +++
 .../Transforms/Vectorize/VPlanTestBase.h      |  6 +-
 .../Vectorize/VPlanVerifierTest.cpp           | 58 +++++++++++++++++++
 4 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
index e33a936dd9978..9f73ba77039c2 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanUtils.cpp
@@ -662,11 +662,11 @@ bool VPBlockUtils::isHeader(const VPBlockBase *VPB,
 
 bool VPBlockUtils::isLatch(const VPBlockBase *VPB,
                            const VPDominatorTree &VPDT) {
-  // A latch has a header as its second successor, with its other successor
+  // A latch has a header as its last successor, with its other successors
   // leaving the loop. A preheader OTOH has a header as its first (and only)
   // successor.
-  return VPB->getNumSuccessors() == 2 &&
-         VPBlockUtils::isHeader(VPB->getSuccessors()[1], VPDT);
+  return VPB->getNumSuccessors() >= 2 &&
+         VPBlockUtils::isHeader(VPB->getSuccessors().back(), VPDT);
 }
 
 std::optional<MemoryLocation>
diff --git a/llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp b/llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp
index 1dff1831c305c..61038bcf3f224 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanVerifier.cpp
@@ -385,6 +385,13 @@ bool VPlanVerifier::verifyBlock(const VPBlockBase *VPB) {
       return false;
     }
   }
+  // For plain CFG VPlans (no parent region), verify header/latch structure.
+  if (VPBB && !VPBB->getParent() && VPBlockUtils::isHeader(VPBB, VPDT) &&
+      !VPBlockUtils::isLatch(VPB->getPredecessors()[1], VPDT)) {
+    errs() << "Header's second predecessor must be the latch!\n";
+    return false;
+  }
+
   return !VPBB || verifyVPBasicBlock(VPBB);
 }
 
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
index 4c8e5aef2075b..2c1797a5a724e 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
+++ b/llvm/unittests/Transforms/Vectorize/VPlanTestBase.h
@@ -70,7 +70,8 @@ class VPlanTestIRBase : public testing::Test {
   /// Build the VPlan for the loop starting from \p LoopHeader.
   VPlanPtr buildVPlan(
       BasicBlock *LoopHeader,
-      UncountableExitStyle Style = UncountableExitStyle::NoUncountableExit) {
+      UncountableExitStyle Style = UncountableExitStyle::NoUncountableExit,
+      bool CreateLoopRegions = true) {
     Function &F = *LoopHeader->getParent();
     assert(!verifyFunction(F) && "input function must be valid");
     doAnalysis(F);
@@ -98,7 +99,8 @@ class VPlanTestIRBase : public testing::Test {
     VPlanTransforms::handleEarlyExits(*Plan, Style, L, PSE, *DT, AC.get());
     VPlanTransforms::addMiddleCheck(*Plan, false);
 
-    VPlanTransforms::createLoopRegions(*Plan);
+    if (CreateLoopRegions)
+      VPlanTransforms::createLoopRegions(*Plan);
     return Plan;
   }
 
diff --git a/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
index 2472e18b654e2..b22b93af8ae8e 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
@@ -444,4 +444,62 @@ TEST_F(VPIRVerifierTest, testVerifyIRPhiInExitVPIRBB) {
       ::testing::internal::GetCapturedStderr().c_str());
 #endif
 }
+
+TEST_F(VPIRVerifierTest, BranchOnTwoCondsLatchHeaderVerification) {
+  const char *ModuleString =
+      "define void @f(ptr dereferenceable(40) align 2 %pred) {\n"
+      "entry:\n"
+      "  br label %for.body\n"
+      "for.body:\n"
+      "  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.inc ]\n"
+      "  %uncountable.addr = getelementptr inbounds nuw i16, ptr %pred, i64 "
+      "%iv\n"
+      "  %uncountable.val = load i16, ptr %uncountable.addr, align 2\n"
+      "  %uncountable.cond = icmp sgt i16 %uncountable.val, 500\n"
+      "  br i1 %uncountable.cond, label %exit, label %for.inc\n"
+      "for.inc:\n"
+      "  %iv.next = add nuw nsw i64 %iv, 1\n"
+      "  %countable.cond = icmp eq i64 %iv.next, 20\n"
+      "  br i1 %countable.cond, label %exit, label %for.body\n"
+      "exit:\n"
+      "  ret void\n"
+      "}\n";
+
+  Module &M = parseModule(ModuleString);
+
+  Function *F = M.getFunction("f");
+  BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
+  // Build a plain CFG VPlan with BranchOnTwoConds as the latch terminator
+  // (3 successors), without wrapping blocks in loop regions.
+  auto Plan = buildVPlan(LoopHeader, UncountableExitStyle::ReadOnly,
+                         /*CreateLoopRegions=*/false);
+
+  auto *MiddleVPBB =
+      cast<VPBasicBlock>(Plan->getScalarPreheader()->getPredecessors()[0]);
+  auto *Latch = cast<VPBasicBlock>(MiddleVPBB->getSinglePredecessor());
+  auto *Header = Latch->getSuccessors().back();
+  ASSERT_EQ(Header->getPredecessors()[1], Latch);
+  ASSERT_EQ(Latch->getNumSuccessors(), 3u);
+  auto *Term = cast<VPInstruction>(&Latch->back());
+  EXPECT_EQ(Term->getOpcode(), VPInstruction::BranchOnTwoConds);
+
+  // Verify the plan is valid; this exercises isLatch with a 3-successor latch.
+  EXPECT_TRUE(verifyVPlanIsValid(*Plan));
+
+  // Swap the latch's first and last successors, placing the header at index 0
+  // instead of the last position. isLatch checks the last successor, so the
+  // latch is no longer recognized, triggering the header predecessor check.
+  auto &Succs = Latch->getSuccessors();
+  std::swap(Succs[0], Succs[2]);
+
+#if GTEST_HAS_STREAM_REDIRECTION
+  ::testing::internal::CaptureStderr();
+#endif
+  EXPECT_FALSE(verifyVPlanIsValid(*Plan));
+#if GTEST_HAS_STREAM_REDIRECTION
+  EXPECT_STREQ("Header's second predecessor must be the latch!\n",
+               ::testing::internal::GetCapturedStderr().c_str());
+#endif
+}
+
 } // namespace