[llvm] [VPlan] Add BranchOnMultiCond, use for early exit plans. (PR #172750)

[Florian Hahn via llvm-commits]

================
@@ -3694,6 +3694,60 @@ void VPlanTransforms::dissolveLoopRegions(VPlan &Plan) {
     R->dissolveToCFGLoop();
 }
 
+void VPlanTransforms::expandBranchOnMultiCond(VPlan &Plan) {
+  // Expand BranchOnMultiCond instructions into explicit CFG with branching.
+  // This must run after dissolveLoopRegions, when successors have been
+  // restored.
+  SmallVector<VPBasicBlock *> WorkList;
+  for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(
+           vp_depth_first_shallow(Plan.getEntry()))) {
+    if (VPBB->empty())
+      continue;
+    if (match(&VPBB->back(), m_BranchOnMultiCond()))
+      WorkList.push_back(VPBB);
+  }
+
+  for (VPBasicBlock *VPBB : WorkList) {
+    auto *BMC = cast<VPInstruction>(&VPBB->back());
+
+    unsigned NumConds = BMC->getNumOperands();
+    assert(NumConds > 0 &&
+           "BranchOnMultiCond must have at least one condition");
+    DebugLoc DL = BMC->getDebugLoc();
+    const auto Successors = to_vector(VPBB->getSuccessors());
+    assert(Successors.size() == NumConds + 1 &&
+           "BranchOnMultiCond must have N operands and N+1 successors");
+
+    // Disconnect all successors.
+    for (VPBlockBase *Succ : Successors)
+      VPBlockUtils::disconnectBlocks(VPBB, Succ);
+
+    // Create chain of conditional branches. Each condition branches to its
+    // corresponding successor on true, and to the next condition block (or the
+    // final successor) on false.
+    VPBasicBlock *CurrentBlock = VPBB;
+    for (const auto &[I, Cond] : enumerate(BMC->operands())) {
+      VPBlockBase *TrueSucc = Successors[I];
+      VPBlockBase *FalseSucc = Successors[NumConds];
+      if (I + 1 != NumConds) {
+        auto *FalseBlock = Plan.createVPBasicBlock(
+            (Twine(VPBB->getName()) + ".multi.cond." + Twine(I + 1)).str());
+        FalseBlock->setParent(VPBB->getParent());
----------------
fhahn wrote:

I've removed the loop for now.

https://github.com/llvm/llvm-project/pull/172750