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

[David Sherwood via llvm-commits]

================
@@ -1696,3 +1702,79 @@ void VPlanTransforms::createInterleaveGroups(
       }
   }
 }
+
+void VPlanTransforms::convertToMultiCond(VPlan &Plan, ScalarEvolution &SE,
+                                         Loop *OrigLoop,
+                                         VPRecipeBuilder &RecipeBuilder) {
+  auto *LatchVPBB =
+      cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getExiting());
+  VPBuilder Builder(LatchVPBB->getTerminator());
+  auto *MiddleVPBB =
+      cast<VPBasicBlock>(Plan.getVectorLoopRegion()->getSingleSuccessor());
+
+  VPRegionBlock *LoopRegion = Plan.getVectorLoopRegion();
+
+  const SCEV *BackedgeTakenCount =
+      SE.getExitCount(OrigLoop, OrigLoop->getLoopLatch());
+  const SCEV *TripCount = SE.getTripCountFromExitCount(
+      BackedgeTakenCount, Plan.getCanonicalIV()->getScalarType(), OrigLoop);
+  VPValue *NewTC = vputils::getOrCreateVPValueForSCEVExpr(Plan, TripCount, SE);
+  Plan.getTripCount()->replaceAllUsesWith(NewTC);
+  Plan.resetTripCount(NewTC);
+
+  VPValue *EarlyExitTaken = nullptr;
+  SmallVector<BasicBlock *> ExitingBBs;
+  OrigLoop->getExitingBlocks(ExitingBBs);
+  for (BasicBlock *Exiting : ExitingBBs) {
+    auto *ExitingTerm = cast<BranchInst>(Exiting->getTerminator());
+    BasicBlock *TrueSucc = ExitingTerm->getSuccessor(0);
+    BasicBlock *FalseSucc = ExitingTerm->getSuccessor(1);
+    VPIRBasicBlock *VPExitBlock;
+    if (OrigLoop->getUniqueExitBlock())
----------------
david-arm wrote:

Again, this doesn't seem to interact well with loops that have countable early exits, which we already support vectorising. I will be adding support for loops with a mixture of countable and uncountable early exits (https://github.com/llvm/llvm-project/pull/88385), which will require a scalar epilogue for the final iteration (assuming we didn't leave via the uncountable exit). I imagine that's more efficient than testing the early countable exit in each vector iteration.

For the purposes of adding initial support for multiple loop region successors, you could also potentially simplify the algorithm by requiring `canVectorizeEarlyExit` to only permit one early exit? Or if you prefer to add support for multiple early exits then it makes sense to add tests for them. I think at the moment the tests only have a single, countable early exit.

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