Clone once

Author: fhahn

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

@@ -507,7 +507,7 @@ class LoopVectorizationPlanner {
   /// returned VPlan is valid for. If no VPlan can be built for the input range,
   /// set the largest included VF to the maximum VF for which no plan could be
   /// built.
-  VPlanPtr tryToBuildVPlanWithVPRecipes(VFRange &Range);
+  VPlanPtr tryToBuildVPlanWithVPRecipes(VPlanPtr InitialPlan, VFRange &Range);
 
   /// Build VPlans for power-of-2 VF's between \p MinVF and \p MaxVF inclusive,
   /// according to the information gathered by Legal when it checked if it is

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -8786,10 +8786,13 @@ void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
                                                         ElementCount MaxVF) {
   assert(OrigLoop->isInnermost() && "Inner loop expected.");
 
+  auto VPlan0 = VPlanTransforms::buildPlainCFG(OrigLoop, *LI);
+
   auto MaxVFTimes2 = MaxVF * 2;
   for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFTimes2);) {
     VFRange SubRange = {VF, MaxVFTimes2};
-    if (auto Plan = tryToBuildVPlanWithVPRecipes(SubRange)) {
+    if (auto Plan = tryToBuildVPlanWithVPRecipes(
+            std::unique_ptr<VPlan>(VPlan0->duplicate()), SubRange)) {
       bool HasScalarVF = Plan->hasScalarVFOnly();
       // Now optimize the initial VPlan.
       if (!HasScalarVF)
@@ -9093,7 +9096,8 @@ static void addExitUsersForFirstOrderRecurrences(
 }
 
 VPlanPtr
-LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
+LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VPlanPtr Plan,
+                                                       VFRange &Range) {
 
   using namespace llvm::VPlanPatternMatch;
   SmallPtrSet<const InterleaveGroup<Instruction> *, 1> InterleaveGroups;
@@ -9115,7 +9119,6 @@ LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(VFRange &Range) {
             return !CM.requiresScalarEpilogue(VF.isVector());
           },
           Range);
-  auto Plan = VPlanTransforms::buildPlainCFG(OrigLoop, *LI);
   VPlanTransforms::introduceRegions(*Plan, Legal->getWidestInductionType(), PSE,
                                     RequiresScalarEpilogueCheck,
                                     CM.foldTailByMasking(), OrigLoop);

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -173,7 +173,6 @@ VPBasicBlock *VPBlockBase::getEntryBasicBlock() {
 }
 
 void VPBlockBase::setPlan(VPlan *ParentPlan) {
-  assert(ParentPlan->getEntry() == this && "Can only set plan on its entry.");
   Plan = ParentPlan;
 }
 
@@ -1184,11 +1183,16 @@ VPlan *VPlan::duplicate() {
   const auto &[NewEntry, __] = cloneFrom(Entry);
 
   BasicBlock *ScalarHeaderIRBB = getScalarHeader()->getIRBasicBlock();
-  VPIRBasicBlock *NewScalarHeader = cast<VPIRBasicBlock>(*find_if(
-      vp_depth_first_shallow(NewEntry), [ScalarHeaderIRBB](VPBlockBase *VPB) {
-        auto *VPIRBB = dyn_cast<VPIRBasicBlock>(VPB);
-        return VPIRBB && VPIRBB->getIRBasicBlock() == ScalarHeaderIRBB;
-      }));
+  VPIRBasicBlock *NewScalarHeader = nullptr;
+  if (getScalarHeader()->getNumPredecessors() == 0) {
+    NewScalarHeader = getScalarHeader()->clone();
+  } else {
+    NewScalarHeader = cast<VPIRBasicBlock>(*find_if(
+        vp_depth_first_shallow(NewEntry), [ScalarHeaderIRBB](VPBlockBase *VPB) {
+          auto *VPIRBB = dyn_cast<VPIRBasicBlock>(VPB);
+          return VPIRBB && VPIRBB->getIRBasicBlock() == ScalarHeaderIRBB;
+        }));
+  }
   // Create VPlan, clone live-ins and remap operands in the cloned blocks.
   auto *NewPlan = new VPlan(cast<VPBasicBlock>(NewEntry), NewScalarHeader);
   DenseMap<VPValue *, VPValue *> Old2NewVPValues;
@@ -1203,8 +1207,7 @@ VPlan *VPlan::duplicate() {
     NewPlan->BackedgeTakenCount = new VPValue();
     Old2NewVPValues[BackedgeTakenCount] = NewPlan->BackedgeTakenCount;
   }
-  assert(TripCount && "trip count must be set");
-  if (TripCount->isLiveIn())
+  if (TripCount && TripCount->isLiveIn())
     Old2NewVPValues[TripCount] =
         NewPlan->getOrAddLiveIn(TripCount->getLiveInIRValue());
   // else NewTripCount will be created and inserted into Old2NewVPValues when
@@ -1217,23 +1220,32 @@ VPlan *VPlan::duplicate() {
   NewPlan->UFs = UFs;
   // TODO: Adjust names.
   NewPlan->Name = Name;
-  assert(Old2NewVPValues.contains(TripCount) &&
-         "TripCount must have been added to Old2NewVPValues");
-  NewPlan->TripCount = Old2NewVPValues[TripCount];
+  if (TripCount) {
+    assert(Old2NewVPValues.contains(TripCount) &&
+           "TripCount must have been added to Old2NewVPValues");
+    NewPlan->TripCount = Old2NewVPValues[TripCount];
+  }
 
   // Transfer all cloned blocks (the second half of all current blocks) from
   // current to new VPlan.
   unsigned NumBlocksAfterCloning = CreatedBlocks.size();
   for (unsigned I :
-       seq<unsigned>(NumBlocksBeforeCloning, NumBlocksAfterCloning))
+       seq<unsigned>(NumBlocksBeforeCloning, NumBlocksAfterCloning)) {
+    this->CreatedBlocks[I]->setPlan(NewPlan);
     NewPlan->CreatedBlocks.push_back(this->CreatedBlocks[I]);
+  }
   CreatedBlocks.truncate(NumBlocksBeforeCloning);
 
+  for (VPBlockBase *VPB : NewPlan->CreatedBlocks) {
+    if (VPB->getNumSuccessors() == 0 && isa<VPIRBasicBlock>(VPB) &&
+        VPB != NewScalarHeader)
+      NewPlan->ExitBlocks.push_back(cast<VPIRBasicBlock>(VPB));
+  }
   return NewPlan;
 }
 
 VPIRBasicBlock *VPlan::createEmptyVPIRBasicBlock(BasicBlock *IRBB) {
-  auto *VPIRBB = new VPIRBasicBlock(IRBB);
+  auto *VPIRBB = new VPIRBasicBlock(IRBB, this);
   CreatedBlocks.push_back(VPIRBB);
   return VPIRBB;
 }

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -145,8 +145,8 @@ class VPBlockBase {
   }
 
 protected:
-  VPBlockBase(const unsigned char SC, const std::string &N)
-      : SubclassID(SC), Name(N) {}
+  VPBlockBase(const unsigned char SC, const std::string &N, VPlan *ParentPlan)
+      : SubclassID(SC), Name(N), Plan(ParentPlan) {}
 
 public:
   /// An enumeration for keeping track of the concrete subclass of VPBlockBase
@@ -959,6 +959,8 @@ class VPInstruction : public VPRecipeWithIRFlags,
   VPInstruction *clone() override {
     SmallVector<VPValue *, 2> Operands(operands());
     auto *New = new VPInstruction(Opcode, Operands, getDebugLoc(), Name);
+    if (getUnderlyingValue())
+      New->setUnderlyingValue(getUnderlyingInstr());
     New->transferFlags(*this);
     return New;
   }
@@ -1932,7 +1934,11 @@ class VPWidenPHIRecipe : public VPSingleDefRecipe {
   }
 
   VPWidenPHIRecipe *clone() override {
-    llvm_unreachable("cloning not implemented yet");
+    auto *C = new VPWidenPHIRecipe(cast<PHINode>(getUnderlyingValue()),
+                                   getOperand(0), getDebugLoc(), Name);
+    for (VPValue *Op : make_range(std::next(op_begin()), op_end()))
+      C->addOperand(Op);
+    return C;
   }
 
   ~VPWidenPHIRecipe() override = default;
@@ -3140,8 +3146,9 @@ class VPBasicBlock : public VPBlockBase {
   friend class VPlan;
 
   /// Use VPlan::createVPBasicBlock to create VPBasicBlocks.
-  VPBasicBlock(const Twine &Name = "", VPRecipeBase *Recipe = nullptr)
-      : VPBlockBase(VPBasicBlockSC, Name.str()) {
+  VPBasicBlock(VPlan *ParentPlan, const Twine &Name = "",
+               VPRecipeBase *Recipe = nullptr)
+      : VPBlockBase(VPBasicBlockSC, Name.str(), ParentPlan) {
     if (Recipe)
       appendRecipe(Recipe);
   }
@@ -3153,8 +3160,9 @@ class VPBasicBlock : public VPBlockBase {
   /// The VPRecipes held in the order of output instructions to generate.
   RecipeListTy Recipes;
 
-  VPBasicBlock(const unsigned char BlockSC, const Twine &Name = "")
-      : VPBlockBase(BlockSC, Name.str()) {}
+  VPBasicBlock(const unsigned char BlockSC, VPlan *ParentPlan,
+               const Twine &Name = "")
+      : VPBlockBase(BlockSC, Name.str(), ParentPlan) {}
 
 public:
   ~VPBasicBlock() override {
@@ -3292,8 +3300,8 @@ class VPIRBasicBlock : public VPBasicBlock {
   BasicBlock *IRBB;
 
   /// Use VPlan::createVPIRBasicBlock to create VPIRBasicBlocks.
-  VPIRBasicBlock(BasicBlock *IRBB)
-      : VPBasicBlock(VPIRBasicBlockSC,
+  VPIRBasicBlock(BasicBlock *IRBB, VPlan *ParentPlan)
+      : VPBasicBlock(VPIRBasicBlockSC, ParentPlan,
                      (Twine("ir-bb<") + IRBB->getName() + Twine(">")).str()),
         IRBB(IRBB) {}
 
@@ -3336,18 +3344,19 @@ class VPRegionBlock : public VPBlockBase {
   bool IsReplicator;
 
   /// Use VPlan::createVPRegionBlock to create VPRegionBlocks.
-  VPRegionBlock(VPBlockBase *Entry, VPBlockBase *Exiting,
+  VPRegionBlock(VPBlockBase *Entry, VPBlockBase *Exiting, VPlan *ParentPlan,
                 const std::string &Name = "", bool IsReplicator = false)
-      : VPBlockBase(VPRegionBlockSC, Name), Entry(Entry), Exiting(Exiting),
-        IsReplicator(IsReplicator) {
+      : VPBlockBase(VPRegionBlockSC, Name, ParentPlan), Entry(Entry),
+        Exiting(Exiting), IsReplicator(IsReplicator) {
     assert(Entry->getPredecessors().empty() && "Entry block has predecessors.");
     assert(Exiting->getSuccessors().empty() && "Exit block has successors.");
     Entry->setParent(this);
     Exiting->setParent(this);
   }
-  VPRegionBlock(const std::string &Name = "", bool IsReplicator = false)
-      : VPBlockBase(VPRegionBlockSC, Name), Entry(nullptr), Exiting(nullptr),
-        IsReplicator(IsReplicator) {}
+  VPRegionBlock(VPlan *ParentPlan, const std::string &Name = "",
+                bool IsReplicator = false)
+      : VPBlockBase(VPRegionBlockSC, Name, ParentPlan), Entry(nullptr),
+        Exiting(nullptr), IsReplicator(IsReplicator) {}
 
 public:
   ~VPRegionBlock() override {}
@@ -3740,7 +3749,7 @@ class VPlan {
   /// VPlan is destroyed.
   VPBasicBlock *createVPBasicBlock(const Twine &Name,
                                    VPRecipeBase *Recipe = nullptr) {
-    auto *VPB = new VPBasicBlock(Name, Recipe);
+    auto *VPB = new VPBasicBlock(this, Name, Recipe);
     CreatedBlocks.push_back(VPB);
     return VPB;
   }
@@ -3751,7 +3760,7 @@ class VPlan {
   VPRegionBlock *createVPRegionBlock(VPBlockBase *Entry, VPBlockBase *Exiting,
                                      const std::string &Name = "",
                                      bool IsReplicator = false) {
-    auto *VPB = new VPRegionBlock(Entry, Exiting, Name, IsReplicator);
+    auto *VPB = new VPRegionBlock(Entry, Exiting, this, Name, IsReplicator);
     CreatedBlocks.push_back(VPB);
     return VPB;
   }
@@ -3762,7 +3771,7 @@ class VPlan {
   /// destroyed.
   VPRegionBlock *createVPRegionBlock(const std::string &Name = "",
                                      bool IsReplicator = false) {
-    auto *VPB = new VPRegionBlock(Name, IsReplicator);
+    auto *VPB = new VPRegionBlock(this, Name, IsReplicator);
     CreatedBlocks.push_back(VPB);
     return VPB;
   }