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

[David Sherwood via llvm-commits]

================
@@ -390,22 +390,40 @@ BasicBlock *VPBasicBlock::createEmptyBasicBlock(VPTransformState &State) {
   return NewBB;
 }
 
+/// Determine the parent loop for \p VPBB by checking if it is an exit block
+/// or if all of its successors lead to exit blocks (following single-successor
+/// chains). Returns the innermost loop containing any of the exits if all
+/// paths lead to exits, or \p DefaultLoop otherwise.
+static Loop *getParentLoopForBlock(VPBlockBase *VPBB, VPlan &Plan, LoopInfo &LI,
+                                   Loop *DefaultLoop) {
+  Loop *ParentLoop = nullptr;
+  for (VPBlockBase *Succ : VPBB->getSuccessors()) {
+    while (Succ && !Plan.isExitBlock(Succ))
+      Succ = Succ->getSingleSuccessor();
+    if (!Succ)
+      return DefaultLoop;
+    Loop *L = LI.getLoopFor(cast<VPIRBasicBlock>(Succ)->getIRBasicBlock());
+    if (L && (!ParentLoop || ParentLoop->contains(L)))
+      ParentLoop = L;
+  }
+  return ParentLoop;
+}
+
 void VPBasicBlock::connectToPredecessors(VPTransformState &State) {
   auto &CFG = State.CFG;
   BasicBlock *NewBB = CFG.VPBB2IRBB[this];
 
   // Register NewBB in its loop. In innermost loops its the same for all
   // BB's.
-  Loop *ParentLoop = State.CurrentParentLoop;
-  // If this block has a sole successor that is an exit block or is an exit
-  // block itself then it needs adding to the same parent loop as the exit
-  // block.
-  VPBlockBase *SuccOrExitVPB = getSingleSuccessor();
-  SuccOrExitVPB = SuccOrExitVPB ? SuccOrExitVPB : this;
-  if (State.Plan->isExitBlock(SuccOrExitVPB)) {
-    ParentLoop = State.LI->getLoopFor(
-        cast<VPIRBasicBlock>(SuccOrExitVPB)->getIRBasicBlock());
-  }
+  Loop *ParentLoop = nullptr;
+  if (State.Plan->isExitBlock(this))
+    ParentLoop =
+        State.LI->getLoopFor(cast<VPIRBasicBlock>(this)->getIRBasicBlock());
+  else if (getNumSuccessors() == 0)
+    ParentLoop = State.CurrentParentLoop;
----------------
david-arm wrote:

I think it would be good to add a brief comment explaining the logic behind this. It's not obvious to me why a block with no successors can inherit the parent loop from the State, and why a block with successors cannot.

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