[llvm] [VPlan] Account for early-exit dispatch blocks when updating LI. (PR #185618)

[Luke Lau via llvm-commits]

================
@@ -965,6 +965,32 @@ void VPlan::execute(VPTransformState *State) {
   for (VPBlockBase *Block : RPOT)
     Block->execute(State);
 
+  if (hasEarlyExit()) {
+    // Fix up LoopInfo for extra dispatch blocks when vectorizing loops with
+    // early exits. For dispatch blocks, we need to find the smallest common
+    // loop of all successors. Note: we only need to update loop info for blocks
+    // after the middle block, but there is no easy way to get those at this
+    // point.
+    for (VPBlockBase *VPB : reverse(RPOT)) {
+      auto *VPBB = dyn_cast<VPBasicBlock>(VPB);
+      if (!VPBB || isa<VPIRBasicBlock>(VPBB))
+        continue;
+      BasicBlock *BB = State->CFG.VPBB2IRBB[VPBB];
+      Loop *L = State->LI->getLoopFor(BB);
+      if (!L || any_of(successors(BB),
+                       [L](BasicBlock *Succ) { return L->contains(Succ); }))
+        continue;
+      // Find the innermost loop containing all successors.
+      Loop *Target = State->LI->getLoopFor(*succ_begin(BB));
----------------
lukel97 wrote:

Just confirming if a block doesn't have successors it can't be part of a loop to begin with, so we would exit with the !L check above

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