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

[Florian Hahn via llvm-commits]

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

Update `isLatch` to handle BranchOnTwoConds,
in preparation to enable running the verifier on earlier VPlans.

>From 6cb487f1d6692ac780b507df4ac14318c27ea6ec 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 ++
 .../Vectorize/VPlanVerifierTest.cpp           | 70 +++++++++++++++++++
 3 files changed, 80 insertions(+), 3 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/VPlanVerifierTest.cpp b/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
index 2472e18b654e2..8f5c3280d3f22 100644
--- a/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/VPlanVerifierTest.cpp
@@ -444,4 +444,74 @@ TEST_F(VPIRVerifierTest, testVerifyIRPhiInExitVPIRBB) {
       ::testing::internal::GetCapturedStderr().c_str());
 #endif
 }
+
+TEST_F(VPIRVerifierTest, BranchOnTwoCondsLatchHeaderVerification) {
+  const char *ModuleString =
+      "define void @f(ptr %array, ptr %pred) {\n"
+      "entry:\n"
+      "  br label %for.body\n"
+      "for.body:\n"
+      "  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.inc ]\n"
+      "  %st.addr = getelementptr inbounds i16, ptr %array, i64 %iv\n"
+      "  %data = load i16, ptr %st.addr, align 2\n"
+      "  %inc = add nsw i16 %data, 1\n"
+      "  store i16 %inc, ptr %st.addr, align 2\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();
+  doAnalysis(*F);
+
+  Loop *L = LI->getLoopFor(LoopHeader);
+  PredicatedScalarEvolution PSE(*SE, *L);
+
+  // Build a plain CFG VPlan with BranchOnTwoConds as the latch terminator
+  // (3 successors), without wrapping blocks in loop regions.
+  auto Plan =
+      VPlanTransforms::buildVPlan0(L, *LI, IntegerType::get(*Ctx, 64), {}, PSE);
+  VPlanTransforms::handleEarlyExits(*Plan, /*HasUncountableExit=*/true);
+  VPlanTransforms::addMiddleCheck(*Plan, /*RequiresScalarEpilogueCheck=*/true,
+                                  /*TailFolded=*/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