[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)
----------------
david-arm wrote:

If we bail out early here should we assert that ParentLoop is nullptr? What if we've already processed an exit block before getting here - won't this return the incorrect loop?

Also, do we have tests for all cases where `getSingleSuccessor` would return nullptr? For example, if there is no successor or if there is more than one successor?

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