[llvm] [VPlan] Dispatch to multiple exit blocks via middle blocks. (PR #112138)

[via llvm-commits]

================
@@ -1820,6 +1820,57 @@ void VPlanTransforms::createInterleaveGroups(
   }
 }
 
+void VPlanTransforms::handleUncountableEarlyExit(
+    VPlan &Plan, ScalarEvolution &SE, Loop *OrigLoop,
+    BasicBlock *UncountableExitingBlock, VPRecipeBuilder &RecipeBuilder) {
+  auto *LatchVPBB =
+      cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getExiting());
+  VPBuilder Builder(LatchVPBB->getTerminator());
+  auto *MiddleVPBB = Plan.getMiddleBlock();
+  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+  VPValue *EarlyExitTaken = nullptr;
+
+  // Process the uncountable exiting block. Update EarlyExitTaken, which tracks
+  // if any uncountable early exit has been taken. Also split the middle block
+  // and branch to the exit block for the early exit if it has been taken.
+  auto *ExitingTerm =
+      cast<BranchInst>(UncountableExitingBlock->getTerminator());
+  BasicBlock *TrueSucc = ExitingTerm->getSuccessor(0);
+  BasicBlock *FalseSucc = ExitingTerm->getSuccessor(1);
+  VPIRBasicBlock *VPExitBlock;
+  if (OrigLoop->getUniqueExitBlock()) {
+    VPExitBlock = cast<VPIRBasicBlock>(MiddleVPBB->getSuccessors()[0]);
+  } else {
+    VPExitBlock = VPIRBasicBlock::fromBasicBlock(
+        !OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
+  }
+
+  VPValue *M = RecipeBuilder.getBlockInMask(
+      OrigLoop->contains(TrueSucc) ? TrueSucc : FalseSucc);
----------------
ayalz wrote:

The successor being wrapped is currently the latch.

The mask we seek is that of the edge from the early exiting block to its successor-outside-the-loop (whose VPBB may only have been constructed above, so its edge mask may not have been computed), which is the negation of the edge mask to its successor-in-the-loop (assuming block in mask of early exiting block is full), which is this block in mask.

Can the early exiting block be conditional or must it currently post-dominate the header? I.e., is M currently (the widening of the possible negation of) the early exit condition?

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