[VPlan] Track current vector loop in VPTransformState (NFC).

Instead of looking up the vector loop using the header, keep track of
the current vector loop in VPTransformState. This removes the
requirement for the vector header block being part of the loop up front.

A follow-up patch will move the code to generate the Loop object for the
vector loop to VPRegionBlock.

Depends on D121619.

Reviewed By: Ayal

Differential Revision: https://reviews.llvm.org/D121621

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -319,8 +319,7 @@ void VPBasicBlock::execute(VPTransformState *State) {
     UnreachableInst *Terminator = State->Builder.CreateUnreachable();
     State->Builder.SetInsertPoint(Terminator);
     // Register NewBB in its loop. In innermost loops its the same for all BB's.
-    Loop *L = State->LI->getLoopFor(State->CFG.PrevBB);
-    L->addBasicBlockToLoop(NewBB, *State->LI);
+    State->CurrentVectorLoop->addBasicBlockToLoop(NewBB, *State->LI);
     State->CFG.PrevBB = NewBB;
   }
 
@@ -909,6 +908,7 @@ void VPlan::execute(VPTransformState *State) {
   assert(VectorHeaderBB && "Loop preheader does not have a single successor.");
 
   Loop *L = State->LI->getLoopFor(VectorHeaderBB);
+  State->CurrentVectorLoop = L;
   State->CFG.ExitBB = L->getExitBlock();
 
   // Remove the edge between Header and Latch to allow other connections.
@@ -1543,7 +1543,7 @@ void VPReductionPHIRecipe::execute(VPTransformState &State) {
       ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), State.VF);
 
   BasicBlock *HeaderBB = State.CFG.PrevBB;
-  assert(State.LI->getLoopFor(HeaderBB)->getHeader() == HeaderBB &&
+  assert(State.CurrentVectorLoop->getHeader() == HeaderBB &&
          "recipe must be in the vector loop header");
   unsigned LastPartForNewPhi = isOrdered() ? 1 : State.UF;
   for (unsigned Part = 0; Part < LastPartForNewPhi; ++Part) {

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -352,6 +352,9 @@ struct VPTransformState {
   /// Holds recipes that may generate a poison value that is used after
   /// vectorization, even when their operands are not poison.
   SmallPtrSet<VPRecipeBase *, 16> MayGeneratePoisonRecipes;
+
+  /// The loop object for the current parent region, or nullptr.
+  Loop *CurrentVectorLoop = nullptr;
 };
 
 /// VPUsers instance used by VPBlockBase to manage CondBit and the block