[llvm] [VPlan] Set the containing plan when creating blocks (NFC). (PR #223269)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/223269

>From e3ec1e7043ce308a5194500bf4192a26d9395f14 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Thu, 10 Sep 2026 13:49:21 +0100
Subject: [PATCH] [VPlan] Set the containing plan when creating blocks (NFC).

Now that we funnel all VP block creation through VPlan helpers, make
sure the Plan field is set for each block. This means we can remove
the CFG walk to find a Plan's entry.

Also drop an assertion in VPSlotTracker::getOrCreateName: a block
now knows its plan as soon as it is created.

This has a very small positive compile-time impact.
---
 llvm/lib/Transforms/Vectorize/VPlan.cpp | 37 ++-----------------------
 llvm/lib/Transforms/Vectorize/VPlan.h   | 16 ++++++-----
 2 files changed, 11 insertions(+), 42 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index e12bd18d495986..f3bdf910a6e7be 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -184,32 +184,6 @@ VPMultiDefValue::~VPMultiDefValue() {
   getDefiningRecipe()->removeDefinedValue(this);
 }
 
-// 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->hasPredecessors())
-      return Current;
-    auto &Predecessors = Current->getPredecessors();
-    WorkList.insert_range(Predecessors);
-  }
-
-  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;
@@ -225,11 +199,6 @@ VPBasicBlock *VPBlockBase::getEntryBasicBlock() {
   return cast<VPBasicBlock>(Block);
 }
 
-void VPBlockBase::setPlan(VPlan *ParentPlan) {
-  assert(ParentPlan->getEntry() == this && "Can only set plan on its entry.");
-  Plan = ParentPlan;
-}
-
 /// \return the VPBasicBlock that is the exit of Block, possibly indirectly.
 const VPBasicBlock *VPBlockBase::getExitingBasicBlock() const {
   const VPBlockBase *Block = this;
@@ -1318,6 +1287,7 @@ VPlan *VPlan::duplicate() {
   unsigned NumBlocksAfterCloning = CreatedBlocks.size();
   for (unsigned I :
        seq<unsigned>(NumBlocksBeforeCloning, NumBlocksAfterCloning)) {
+    this->CreatedBlocks[I]->setPlan(NewPlan);
     this->CreatedBlocks[I]->setNumber(NewPlan->CreatedBlocks.size());
     NewPlan->CreatedBlocks.push_back(this->CreatedBlocks[I]);
   }
@@ -1335,6 +1305,7 @@ VPlan *VPlan::duplicate() {
 
 VPIRBasicBlock *VPlan::createEmptyVPIRBasicBlock(BasicBlock *IRBB) {
   auto *VPIRBB = new VPIRBasicBlock(IRBB);
+  VPIRBB->setPlan(this);
   VPIRBB->setNumber(CreatedBlocks.size());
   CreatedBlocks.push_back(VPIRBB);
   return VPIRBB;
@@ -1653,10 +1624,6 @@ std::string VPSlotTracker::getOrCreateName(const VPValue *V) const {
   // TODO: Update VPSlotTracker constructor to assign names to recipes &
   // VPValues not associated with a VPlan, instead of constructing names ad-hoc
   // here.
-  const VPRecipeBase *DefR = V->getDefiningRecipe();
-  (void)DefR;
-  assert((!DefR || !DefR->getParent() || !DefR->getParent()->getPlan()) &&
-         "VPValue defined by a recipe in a VPlan?");
 
   // Use the underlying value's name, if there is one.
   if (auto *UV = V->getUnderlyingValue()) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 2393403b3a8370..a4b2596a025d06 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -120,8 +120,8 @@ class LLVM_ABI_FOR_TEST VPBlockBase {
   /// List of successor blocks.
   SmallVector<VPBlockBase *, 1> Successors;
 
-  /// VPlan containing the block. Can only be set on the entry block of the
-  /// plan.
+  /// VPlan containing the block. Set when the block is created via VPlan
+  /// helpers.
   VPlan *Plan = nullptr;
 
   /// Subclass identifier (for isa/dyn_cast).
@@ -194,12 +194,11 @@ class LLVM_ABI_FOR_TEST VPBlockBase {
   const VPRegionBlock *getParent() const { return Parent; }
 
   /// \return A pointer to the plan containing the current block.
-  VPlan *getPlan();
-  const VPlan *getPlan() const;
+  VPlan *getPlan() { return Plan; }
+  const VPlan *getPlan() const { return Plan; }
 
-  /// Sets the pointer of the plan containing the block. The block must be the
-  /// entry block into the VPlan.
-  void setPlan(VPlan *ParentPlan);
+  /// Sets the pointer of the plan containing the block.
+  void setPlan(VPlan *ParentPlan) { Plan = ParentPlan; }
 
   void setParent(VPRegionBlock *P) { Parent = P; }
 
@@ -5192,6 +5191,7 @@ class VPlan {
   VPBasicBlock *createVPBasicBlock(const Twine &Name,
                                    VPRecipeBase *Recipe = nullptr) {
     auto *VPB = new VPBasicBlock(Name, Recipe);
+    VPB->setPlan(this);
     VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;
@@ -5206,6 +5206,7 @@ class VPlan {
                                   VPBlockBase *Entry = nullptr,
                                   VPBlockBase *Exiting = nullptr) {
     auto *VPB = new VPRegionBlock(CanIVTy, DL, Entry, Exiting, Name);
+    VPB->setPlan(this);
     VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;
@@ -5217,6 +5218,7 @@ class VPlan {
   VPRegionBlock *createReplicateRegion(VPBlockBase *Entry, VPBlockBase *Exiting,
                                        const std::string &Name = "") {
     auto *VPB = new VPRegionBlock(Entry, Exiting, Name);
+    VPB->setPlan(this);
     VPB->setNumber(CreatedBlocks.size());
     CreatedBlocks.push_back(VPB);
     return VPB;