Reapply "[VPlan] Support cloning initial VPlan (NFC)."

This reverts commit 204252e.

Recommit with a fix for the leak in a unit test.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -1168,11 +1168,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 = createVPIRBasicBlock(ScalarHeaderIRBB);
+  } 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;
@@ -1187,8 +1192,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
@@ -1201,9 +1205,11 @@ 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.

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

@@ -904,6 +904,49 @@ No successors
     EXPECT_EQ(ExpectedStr, FullDump);
   }
 }
+
+TEST_F(VPBasicBlockTest, cloneAndPrint) {
+  VPlan &Plan = getPlan(nullptr);
+  VPBasicBlock *VPBB0 = Plan.getEntry();
+
+  VPInstruction *I1 = new VPInstruction(Instruction::Add, {});
+  VPInstruction *I2 = new VPInstruction(Instruction::Sub, {I1});
+  VPInstruction *I3 = new VPInstruction(Instruction::Br, {I1, I2});
+
+  VPBasicBlock *VPBB1 = Plan.createVPBasicBlock("");
+  VPBB1->appendRecipe(I1);
+  VPBB1->appendRecipe(I2);
+  VPBB1->appendRecipe(I3);
+  VPBB1->setName("bb1");
+  VPBlockUtils::connectBlocks(VPBB0, VPBB1);
+
+  const char *ExpectedStr = R"(digraph VPlan {
+graph [labelloc=t, fontsize=30; label="Vectorization Plan\n for UF\>=1\n"]
+node [shape=rect, fontname=Courier, fontsize=30]
+edge [fontname=Courier, fontsize=30]
+compound=true
+  N0 [label =
+    "preheader:\l" +
+    "Successor(s): bb1\l"
+  ]
+  N0 -> N1 [ label=""]
+  N1 [label =
+    "bb1:\l" +
+    "  EMIT vp\<%1\> = add\l" +
+    "  EMIT vp\<%2\> = sub vp\<%1\>\l" +
+    "  EMIT br vp\<%1\>, vp\<%2\>\l" +
+    "No successors\l"
+  ]
+}
+)";
+  // Check that printing a cloned plan produces the same output.
+  std::string FullDump;
+  raw_string_ostream OS(FullDump);
+  VPlan *Clone = Plan.duplicate();
+  Clone->printDOT(OS);
+  EXPECT_EQ(ExpectedStr, FullDump);
+  delete Clone;
+}
 #endif
 
 using VPRecipeTest = VPlanTestBase;