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

[fhahn]

Summary

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.

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.cpp

912

Is CFG.LastBB now redundant, retrievable from CurrentVectorLoop->getExitBlock()?
(Or could CurrentVectorLoop be retrieved from LastBB'd predecessor, if unique.)

llvm/lib/Transforms/Vectorize/VPlan.h

355

Worth documenting?

[Ayal]

llvm/lib/Transforms/Vectorize/VPlan.cpp

912

(State->CurrentVectorLoop itself is also caching of State->LI->getLoopFor(State->CFG.VectorHeaderBB))

[fhahn]

Address latest comments.

llvm/lib/Transforms/Vectorize/VPlan.cpp

912

Is CFG.LastBB now redundant, retrievable from CurrentVectorLoop->getExitBlock()?

Unfortunately this only works if we maintain a correct branch to the exit block throughout VPlan execution. But we change the header branch to unreachable, so the won't work I think.

(Or could CurrentVectorLoop be retrieved from LastBB'd predecessor, if unique.)
(State->CurrentVectorLoop itself is also caching of State->LI->getLoopFor(State->CFG.VectorHeaderBB))

The main motivation for this change is to be able to create new loop objects for each non-replicator region during VPRegionBlock::execute. At that point, no basic blocks for the loop exist and when creating the first block in the loop (in VPBasicBlock::execute), there's no previous block we could use to get the loop object. We could change VPRegionBlock::execute to create a dummy block on entry, but keeping track of the loop object explicitly seems cleaner to me.

llvm/lib/Transforms/Vectorize/VPlan.h

355

Thanks, I added a comment!

[Ayal]

This is fine; adding a clarifying comment.

llvm/lib/Transforms/Vectorize/VPlan.cpp

912

(This now relates to CFG.ExitBB instead of CFG.LastBB)

Ok, right; disconnecting the branch from VectorHeaderBB to middleBlock below will invalidate L->getExitBlock(), so ExitBB needs to be recorded above.

The current skeleton still creates VectorPreHeaderBB, VectorHeaderBB, and ExitBB/middleBlock, right? That's how L is retrieved above; the comment only clarifies that CurrentVectorLoop currently caches this retrieval (which should still be valid via the cumbersome
State->LI->getLoopFor(State->CFG.VectorPreHeader->getSingleSuccessor()))).
When that changes, and VectorHeaderBB is created during VPlan::execute() etc., the "Assumes a single LoopVectorBody basic-block was created" comment above should also be updated.

(In any case, no dummy blocks intended)

llvm/lib/Transforms/Vectorize/VPlan.h

struct VPTransformState {
   InnerLoopVectorizer *ILV;
 
   /// Pointer to the VPlan code is generated for.
   VPlan *Plan;
 
   /// 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
 /// predicate. Currently VPBlockUsers are used in VPBlockBase for historical
 /// reasons, but in the future the only VPUsers should either be recipes or
 /// live-outs.VPBlockBase uses.
 struct VPBlockUser : public VPUser {
   VPBlockUser() : VPUser({}, VPUserID::Block) {}

llvm/lib/Transforms/Vectorize/VPlan.cpp

if (PrevVPBB && /* A */
         PrevVPBB->getSingleHierarchicalSuccessor()) && /* B */
       !(Replica && getPredecessors().empty())) {       /* C */
     NewBB = createEmptyBasicBlock(State->CFG);
     State->Builder.SetInsertPoint(NewBB);
     // Temporarily terminate with unreachable until CFG is rewired.
     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;
   }
 
   // 2. Fill the IR basic block with IR instructions.
   LLVM_DEBUG(dbgs() << "LV: vectorizing VPBB:" << getName()
                     << " in BB:" << NewBB->getName() << '\n');
 
   State->CFG.VPBB2IRBB[this] = NewBB;
for (auto &Entry : Value2VPValue)
     State->VPValue2Value[Entry.second] = Entry.first;
 
   BasicBlock *VectorPreHeaderBB = State->CFG.PrevBB;
   State->CFG.VectorPreHeader = VectorPreHeaderBB;
   BasicBlock *VectorHeaderBB = VectorPreHeaderBB->getSingleSuccessor();
   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.
   // Temporarily terminate with unreachable until CFG is rewired.
   // Note: this asserts the generated code's assumption that
   // getFirstInsertionPt() can be dereferenced into an Instruction.
   VectorHeaderBB->getTerminator()->eraseFromParent();
   State->Builder.SetInsertPoint(VectorHeaderBB);
   UnreachableInst *Terminator = State->Builder.CreateUnreachable();
void VPReductionPHIRecipe::execute(VPTransformState &State) {
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #1: We create a new vector PHI node with no incoming edges. We'll use
   // this value when we vectorize all of the instructions that use the PHI.
   bool ScalarPHI = State.VF.isScalar() || IsInLoop;
   Type *VecTy =
       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) {
     Value *EntryPart =
         PHINode::Create(VecTy, 2, "vec.phi", &*HeaderBB->getFirstInsertionPt());
     State.set(this, EntryPart, Part);
   }