[VPlan] Add ParentPlan pointer to VPBlockBase.

[fhahn]

Summary

This is mostly to help the VPSlotTracker in D73078. I would slightly prefer the ParentPlan to be a reference. I think that would be fine for all current uses, except that the tests in VPlanTest.cpp need updating . They are the only place that instantiates VPBasicBlocks without a VPlan.

[fhahn]

ping

[gilr]

An alternative to adding a VPlan* to every VP block might be to reuse their existing Parent property by having (only) the topmost block's parent point to VPlan, similar to the IR hierarchy where ancestors are reached indirectly through transitivity. This would require VPlan to inherit from VPRegion, which may deserve further discussion.
For now let's go ahead with the suggested VPlan pointer in VPBlockBase, but better just call it get/Plan instead of get/ParentPlan, saving parenthood's existing semantics.
Since the only current use case for this is debugging, we can simplify the patch and maintain a future trivial switch to the above by setting only the entry block's VPlan e.g. at VPlan::setEntry() and have getPlan() traverse upwards until reaching a topmost block and backwards until reaching the entry. This also better supports potential moving blocks in and out of VPlans.
Makes sense?

[fhahn]

In D74445#1899945, @gilr wrote:

An alternative to adding a VPlan* to every VP block might be to reuse their existing Parent property by having (only) the topmost block's parent point to VPlan, similar to the IR hierarchy where ancestors are reached indirectly through transitivity. This would require VPlan to inherit from VPRegion, which may deserve further discussion.
For now let's go ahead with the suggested VPlan pointer in VPBlockBase, but better just call it get/Plan instead of get/ParentPlan, saving parenthood's existing semantics.
Since the only current use case for this is debugging, we can simplify the patch and maintain a future trivial switch to the above by setting only the entry block's VPlan e.g. at VPlan::setEntry() and have getPlan() traverse upwards until reaching a topmost block and backwards until reaching the entry. This also better supports potential moving blocks in and out of VPlans.
Makes sense?

Sounds good to me. Updated the patch accordingly. I've used getParent() in order to find the top-most entry block.

[gilr]

LGTM, with minor fix.

llvm/lib/Transforms/Vectorize/VPlan.cpp

60

nit: I'd add something like "Function is templated to support both const and non-const blocks" for clarity.

68

This would wrongfully return nullptr for multiple predecessors. Better use *getPredecessors().begin() after checking that there are any.

[fhahn]

llvm/lib/Transforms/Vectorize/VPlan.cpp

60

Will do before committing.

68

Agreed! I think currently we do not really generate such CFGs on the top-level, due to the nesting with region blocks, but I'll generalize it to explore all predecessors until we reach one without any predecessors.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
   };
 
   // ---------------------------------------------------------------------------
   // Build initial VPlan: Scan the body of the loop in a topological order to
   // visit each basic block after having visited its predecessor basic blocks.
   // ---------------------------------------------------------------------------
 
   // Create a dummy pre-entry VPBasicBlock to start building the VPlan.
+  auto Plan = std::make_unique<VPlan>();
   VPBasicBlock *VPBB = new VPBasicBlock("Pre-Entry");
-  auto Plan = std::make_unique<VPlan>(VPBB);
+  Plan->setEntry(VPBB);
 
   // Represent values that will have defs inside VPlan.
   for (Value *V : NeedDef)
     Plan->addVPValue(V);
 
   // Scan the body of the loop in a topological order to visit each basic block
   // after having visited its predecessor basic blocks.
   LoopBlocksDFS DFS(OrigLoop);

llvm/lib/Transforms/Vectorize/VPlan.h

private:
   SmallVector<VPBlockBase *, 1> Successors;
 
   /// Successor selector, null for zero or single successor blocks.
   VPValue *CondBit = nullptr;
 
   /// Current block predicate - null if the block does not need a predicate.
   VPValue *Predicate = nullptr;
 
+  /// VPlan containing the block. Can only be set on the entry block of the
+  /// plan.
+  VPlan *Plan = nullptr;
+
   /// Add \p Successor as the last successor to this block.
   void appendSuccessor(VPBlockBase *Successor) {
     assert(Successor && "Cannot add nullptr successor!");
     Successors.push_back(Successor);
   }
 
   /// Add \p Predecessor as the last predecessor to this block.
   void appendPredecessor(VPBlockBase *Predecessor) {
public:
   /// \return an ID for the concrete type of this object.
   /// This is used to implement the classof checks. This should not be used
   /// for any other purpose, as the values may change as LLVM evolves.
   unsigned getVPBlockID() const { return SubclassID; }
 
   VPRegionBlock *getParent() { return Parent; }
   const VPRegionBlock *getParent() const { return Parent; }
 
+  /// \return A pointer to the plan containing the current block.
+  VPlan *getPlan();
+  const VPlan *getPlan() const;
+
+  /// Sets the pointer of the plan containing the block. The block must be the
+  /// entry block into the VPlan.
+  void setPlan(VPlan *ParentPlan);
+
   void setParent(VPRegionBlock *P) { Parent = P; }
 
   /// \return the VPBasicBlock that is the entry of this VPBlockBase,
   /// recursively, if the latter is a VPRegionBlock. Otherwise, if this
   /// VPBlockBase is a VPBasicBlock, it is returned.
   const VPBasicBlock *getEntryBasicBlock() const;
   VPBasicBlock *getEntryBasicBlock();
 
public:
   }
 
   /// Generate the IR code for this VPlan.
   void execute(struct VPTransformState *State);
 
   VPBlockBase *getEntry() { return Entry; }
   const VPBlockBase *getEntry() const { return Entry; }
 
-  VPBlockBase *setEntry(VPBlockBase *Block) { return Entry = Block; }
+  VPBlockBase *setEntry(VPBlockBase *Block) {
+    Entry = Block;
+    Block->setPlan(this);
+    return Entry;
+  }
 
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {
     if (!BackedgeTakenCount)
       BackedgeTakenCount = new VPValue();
     return BackedgeTakenCount;
   }
 

llvm/lib/Transforms/Vectorize/VPlan.cpp

 raw_ostream &llvm::operator<<(raw_ostream &OS, const VPValue &V) {
   if (const VPInstruction *Instr = dyn_cast<VPInstruction>(&V))
     Instr->print(OS);
   else
     V.printAsOperand(OS);
   return OS;
 }
 
+// Get the top-most entry block of \p Start. This is the entry block of the
+// containing VPlan. This function is templated to support both const and non-const blocks
+template <typename T> static T *getPlanEntry(T *Start) {
+  T *Next = Start;
+  T *Current = Start;
+  while ((Next = Next->getParent()))
+    Current = Next;
+
+  SmallSetVector<T *, 8> WorkList;
+  WorkList.insert(Current);
+
+  for (unsigned i = 0; i < WorkList.size(); i++) {
+    T *Current = WorkList[i];
+    if (Current->getNumPredecessors() == 0)
+      return Current;
+    auto &Predecessors = Current->getPredecessors();
+    WorkList.insert(Predecessors.begin(), Predecessors.end());
+  }
+
+  llvm_unreachable("VPlan without any entry node without predecessors");
+}
+
+VPlan *VPBlockBase::getPlan() { return getPlanEntry(this)->Plan; }
+
+const VPlan *VPBlockBase::getPlan() const { return getPlanEntry(this)->Plan; }
+
 /// \return the VPBasicBlock that is the entry of Block, possibly indirectly.
 const VPBasicBlock *VPBlockBase::getEntryBasicBlock() const {
   const VPBlockBase *Block = this;
   while (const VPRegionBlock *Region = dyn_cast<VPRegionBlock>(Block))
     Block = Region->getEntry();
   return cast<VPBasicBlock>(Block);
 }
 
 VPBasicBlock *VPBlockBase::getEntryBasicBlock() {
   VPBlockBase *Block = this;
   while (VPRegionBlock *Region = dyn_cast<VPRegionBlock>(Block))
     Block = Region->getEntry();
   return cast<VPBasicBlock>(Block);
 }
 
+void VPBlockBase::setPlan(VPlan *ParentPlan) {
+  assert(ParentPlan->getEntry() == this &&
+         "Can only set plan on its entry block.");
+  Plan = ParentPlan;
+}
+
 /// \return the VPBasicBlock that is the exit of Block, possibly indirectly.
 const VPBasicBlock *VPBlockBase::getExitBasicBlock() const {
   const VPBlockBase *Block = this;
   while (const VPRegionBlock *Region = dyn_cast<VPRegionBlock>(Block))
     Block = Region->getExit();
   return cast<VPBasicBlock>(Block);
 }
 

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

TEST_F(VPlanHCFGTest, testBuildHCFGInnerLoop) {
   EXPECT_EQ(0u, Entry->getNumPredecessors());
   EXPECT_EQ(1u, Entry->getNumSuccessors());
   EXPECT_EQ(nullptr, Entry->getCondBit());
 
   VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
   EXPECT_EQ(7u, VecBB->size());
   EXPECT_EQ(2u, VecBB->getNumPredecessors());
   EXPECT_EQ(2u, VecBB->getNumSuccessors());
+  EXPECT_EQ(&*Plan, VecBB->getPlan());
 
   auto Iter = VecBB->begin();
   VPInstruction *Phi = dyn_cast<VPInstruction>(&*Iter++);
   EXPECT_EQ(Instruction::PHI, Phi->getOpcode());
 
   VPInstruction *Idx = dyn_cast<VPInstruction>(&*Iter++);
   EXPECT_EQ(Instruction::GetElementPtr, Idx->getOpcode());
   EXPECT_EQ(2u, Idx->getNumOperands());

llvm/unittests/Transforms/Vectorize/VPlanPredicatorTest.cpp

TEST_F(VPlanPredicatorTest, BasicPredicatorTest) {
   EXPECT_EQ(InnerAnd->getOperand(0), CBV1);
   EXPECT_EQ(InnerAnd->getOperand(1), CBV2);
   EXPECT_EQ(InnerIfPred, InnerAnd);
 
   // Test Linearization
   EXPECT_EQ(InnerLoopLinSucc, OuterIf);
   EXPECT_EQ(OuterIfLinSucc, InnerIf);
   EXPECT_EQ(InnerIfLinSucc, InnerLoopLatch);
+
+  // Check that the containing VPlan is set correctly.
+  EXPECT_EQ(&*Plan, InnerLoopLinSucc->getPlan());
+  EXPECT_EQ(&*Plan, OuterIfLinSucc->getPlan());
+  EXPECT_EQ(&*Plan, InnerIfLinSucc->getPlan());
+  EXPECT_EQ(&*Plan, InnerIf->getPlan());
+  EXPECT_EQ(&*Plan, InnerLoopLatch->getPlan());
 }
 
 // Test generation of Not and Or during predication.
 TEST_F(VPlanPredicatorTest, PredicatorNegOrTest) {
   const char *ModuleString =
       "@arr = common global [100 x [100 x i32]] zeroinitializer, align 16\n"
       "@arr2 = common global [100 x [100 x i32]] zeroinitializer, align 16\n"
       "@arr3 = common global [100 x [100 x i32]] zeroinitializer, align 16\n"

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

TEST(VPInstructionTest, moveAfter) {
 
   I3->moveAfter(I4);
 
   CHECK_ITERATOR(VPBB1, I2, I1);
   CHECK_ITERATOR(VPBB2, I4, I3, I5);
   EXPECT_EQ(I3->getParent(), I4->getParent());
 }
 
+TEST(VPBasicBlockTest, getPlan) {
+  {
+    VPBasicBlock *VPBB1 = new VPBasicBlock();
+    VPBasicBlock *VPBB2 = new VPBasicBlock();
+    VPBasicBlock *VPBB3 = new VPBasicBlock();
+    VPBasicBlock *VPBB4 = new VPBasicBlock();
+
+    //     VPBB1
+    //     /   \
+    // VPBB2  VPBB3
+    //    \    /
+    //    VPBB4
+    VPBlockUtils::connectBlocks(VPBB1, VPBB2);
+    VPBlockUtils::connectBlocks(VPBB1, VPBB3);
+    VPBlockUtils::connectBlocks(VPBB2, VPBB4);
+    VPBlockUtils::connectBlocks(VPBB3, VPBB4);
+
+    VPlan Plan;
+    Plan.setEntry(VPBB1);
+
+    EXPECT_EQ(&Plan, VPBB1->getPlan());
+    EXPECT_EQ(&Plan, VPBB2->getPlan());
+    EXPECT_EQ(&Plan, VPBB3->getPlan());
+    EXPECT_EQ(&Plan, VPBB4->getPlan());
+  }
+
+  {
+    // Region block is entry into VPlan.
+    VPBasicBlock *R1BB1 = new VPBasicBlock();
+    VPBasicBlock *R1BB2 = new VPBasicBlock();
+    VPRegionBlock *R1 = new VPRegionBlock(R1BB1, R1BB2, "R1");
+    VPBlockUtils::connectBlocks(R1BB1, R1BB2);
+
+    VPlan Plan;
+    Plan.setEntry(R1);
+    EXPECT_EQ(&Plan, R1->getPlan());
+    EXPECT_EQ(&Plan, R1BB1->getPlan());
+    EXPECT_EQ(&Plan, R1BB2->getPlan());
+  }
+
+  {
+    // VPBasicBlock is the entry into the VPlan, followed by a region.
+    VPBasicBlock *R1BB1 = new VPBasicBlock();
+    VPBasicBlock *R1BB2 = new VPBasicBlock();
+    VPRegionBlock *R1 = new VPRegionBlock(R1BB1, R1BB2, "R1");
+    VPBlockUtils::connectBlocks(R1BB1, R1BB2);
+
+    VPBasicBlock *VPBB1 = new VPBasicBlock();
+    VPBlockUtils::connectBlocks(VPBB1, R1);
+
+    VPlan Plan;
+    Plan.setEntry(VPBB1);
+    EXPECT_EQ(&Plan, VPBB1->getPlan());
+    EXPECT_EQ(&Plan, R1->getPlan());
+    EXPECT_EQ(&Plan, R1BB1->getPlan());
+    EXPECT_EQ(&Plan, R1BB2->getPlan());
+  }
+
+  {
+    VPBasicBlock *R1BB1 = new VPBasicBlock();
+    VPBasicBlock *R1BB2 = new VPBasicBlock();
+    VPRegionBlock *R1 = new VPRegionBlock(R1BB1, R1BB2, "R1");
+    VPBlockUtils::connectBlocks(R1BB1, R1BB2);
+
+    VPBasicBlock *R2BB1 = new VPBasicBlock();
+    VPBasicBlock *R2BB2 = new VPBasicBlock();
+    VPRegionBlock *R2 = new VPRegionBlock(R2BB1, R2BB2, "R2");
+    VPBlockUtils::connectBlocks(R2BB1, R2BB2);
+
+    VPBasicBlock *VPBB1 = new VPBasicBlock();
+    VPBlockUtils::connectBlocks(VPBB1, R1);
+    VPBlockUtils::connectBlocks(VPBB1, R2);
+
+    VPBasicBlock *VPBB2 = new VPBasicBlock();
+    VPBlockUtils::connectBlocks(R1, VPBB2);
+    VPBlockUtils::connectBlocks(R2, VPBB2);
+
+    VPlan Plan;
+    Plan.setEntry(VPBB1);
+    EXPECT_EQ(&Plan, VPBB1->getPlan());
+    EXPECT_EQ(&Plan, R1->getPlan());
+    EXPECT_EQ(&Plan, R1BB1->getPlan());
+    EXPECT_EQ(&Plan, R1BB2->getPlan());
+    EXPECT_EQ(&Plan, R2->getPlan());
+    EXPECT_EQ(&Plan, R2BB1->getPlan());
+    EXPECT_EQ(&Plan, R2BB2->getPlan());
+    EXPECT_EQ(&Plan, VPBB2->getPlan());
+  }
+}
+
 } // namespace
 } // namespace llvm