[VPlan] Turn external defs in Value -> VPValue mapping.

[fhahn]

Summary

This addresses an existing TODO by keeping a mapping of external IR
Value * definitions wrapped in VPValues for use in a VPlan.

[Ayal]

Ship it!
minor nits.

llvm/lib/Transforms/Vectorize/VPlan.h

2467–2468

Can keep the explanation that external defs are immutable wrappers of underlying values.

2566–2567

Comment deserves updating.

2567–2573

nit: clearer to start code block on next line.

[fhahn]

In D123700#3449883, @Ayal wrote:

Ship it!
minor nits.

Thanks, they should be addressed in the committed version!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

if (!VectorizeLoop) {
         // Ensure that the start values for any VPReductionPHIRecipes are
         // updated before vectorising the epilogue loop.
         VPBasicBlock *Header =
             BestEpiPlan.getVectorLoopRegion()->getEntryBasicBlock();
         for (VPRecipeBase &R : Header->phis()) {
           if (auto *ReductionPhi = dyn_cast<VPReductionPHIRecipe>(&R)) {
             if (auto *Resume = MainILV.getReductionResumeValue(
                     ReductionPhi->getRecurrenceDescriptor())) {
-              VPValue *StartVal = new VPValue(Resume);
-              BestEpiPlan.addExternalDef(StartVal);
+              VPValue *StartVal = BestEpiPlan.getOrAddExternalDef(Resume);
               ReductionPhi->setOperand(0, StartVal);
             }
           }
         }
 
         LVP.executePlan(EPI.EpilogueVF, EPI.EpilogueUF, BestEpiPlan, EpilogILV,
                         DT);
         ++LoopsEpilogueVectorized;

llvm/lib/Transforms/Vectorize/VPlan.h

class VPlan {
   /// Hold the single entry to the Hierarchical CFG of the VPlan.
   VPBlockBase *Entry;
 
   /// Holds the VFs applicable to this VPlan.
   SmallSetVector<ElementCount, 2> VFs;
 
   /// Holds the name of the VPlan, for printing.
   std::string Name;
 
-  /// Holds all the external definitions created for this VPlan.
-  // TODO: Introduce a specific representation for external definitions in
-  // VPlan. External definitions must be immutable and hold a pointer to its
-  // underlying IR that will be used to implement its structural comparison
-  // (operators '==' and '<').
-  SetVector<VPValue *> VPExternalDefs;
+  /// Holds all the external definitions created for this VPlan. External
+  /// definitions must be immutable and hold a pointer to their underlying IR.
+  DenseMap<Value *, VPValue *> VPExternalDefs;
 
   /// Represents the trip count of the original loop, for folding
   /// the tail.
   VPValue *TripCount = nullptr;
 
   /// Represents the backedge taken count of the original loop, for folding
   /// the tail. It equals TripCount - 1.
   VPValue *BackedgeTakenCount = nullptr;
if (Entry) {
       VPBlockBase::deleteCFG(Entry);
     }
     for (VPValue *VPV : VPValuesToFree)
       delete VPV;
     if (TripCount)
       delete TripCount;
     if (BackedgeTakenCount)
       delete BackedgeTakenCount;
-    for (VPValue *Def : VPExternalDefs)
-      delete Def;
+    for (auto &P : VPExternalDefs)
+      delete P.second;
   }
 
   /// Prepare the plan for execution, setting up the required live-in values.
   void prepareToExecute(Value *TripCount, Value *VectorTripCount,
                         Value *CanonicalIVStartValue, VPTransformState &State);
 
   /// Generate the IR code for this VPlan.
   void execute(struct VPTransformState *State);
public:
 
   void addVF(ElementCount VF) { VFs.insert(VF); }
 
   bool hasVF(ElementCount VF) { return VFs.count(VF); }
 
   const std::string &getName() const { return Name; }
 
   void setName(const Twine &newName) { Name = newName.str(); }
 
-  /// Add \p VPVal to the pool of external definitions if it's not already
-  /// in the pool.
-  void addExternalDef(VPValue *VPVal) { VPExternalDefs.insert(VPVal); }
+  /// Get the existing or add a new external definition for \p V.
+  VPValue *getOrAddExternalDef(Value *V) {
+    auto I = VPExternalDefs.insert({V, nullptr});
+    if (I.second)
+      I.first->second = new VPValue(V);
+    return I.first->second;
+  }
 
   void addVPValue(Value *V) {
     assert(Value2VPValueEnabled &&
            "IR value to VPValue mapping may be out of date!");
     assert(V && "Trying to add a null Value to VPlan");
     assert(!Value2VPValue.count(V) && "Value already exists in VPlan");
     VPValue *VPV = new VPValue(V);
     Value2VPValue[V] = VPV;

llvm/lib/Transforms/Vectorize/VPlan.cpp

void VPlan::prepareToExecute(Value *TripCountV, Value *VectorTripCountV,
   }
 
   for (unsigned Part = 0, UF = State.UF; Part < UF; ++Part)
     State.set(&VectorTripCount, VectorTripCountV, Part);
 
   // When vectorizing the epilogue loop, the canonical induction start value
   // needs to be changed from zero to the value after the main vector loop.
   if (CanonicalIVStartValue) {
-    VPValue *VPV = new VPValue(CanonicalIVStartValue);
-    addExternalDef(VPV);
+    VPValue *VPV = getOrAddExternalDef(CanonicalIVStartValue);
     auto *IV = getCanonicalIV();
     assert(all_of(IV->users(),
                   [](const VPUser *U) {
                     if (isa<VPScalarIVStepsRecipe>(U))
                       return true;
                     auto *VPI = cast<VPInstruction>(U);
                     return VPI->getOpcode() ==
                                VPInstruction::CanonicalIVIncrement ||

 
 void VPSlotTracker::assignSlot(const VPValue *V) {
   assert(Slots.find(V) == Slots.end() && "VPValue already has a slot!");
   Slots[V] = NextSlot++;
 }
 
 void VPSlotTracker::assignSlots(const VPlan &Plan) {
 
-  for (const VPValue *V : Plan.VPExternalDefs)
-    assignSlot(V);
+  for (const auto &P : Plan.VPExternalDefs)
+    assignSlot(P.second);
 
   assignSlot(&Plan.VectorTripCount);
   if (Plan.BackedgeTakenCount)
     assignSlot(Plan.BackedgeTakenCount);
 
   ReversePostOrderTraversal<
       VPBlockRecursiveTraversalWrapper<const VPBlockBase *>>
       RPOT(VPBlockRecursiveTraversalWrapper<const VPBlockBase *>(

llvm/lib/Transforms/Vectorize/VPlanHCFGBuilder.cpp

VPValue *PlainCFGBuilder::getOrCreateVPOperand(Value *IRVal) {
   //   B) any other Value without specific representation in VPlan.
   // For now, we use VPValue to represent A and B and classify both as external
   // definitions. We may introduce specific VPValue subclasses for them in the
   // future.
   assert(isExternalDef(IRVal) && "Expected external definition as operand.");
 
   // A and B: Create VPValue and add it to the pool of external definitions and
   // to the Value->VPValue map.
-  VPValue *NewVPVal = new VPValue(IRVal);
-  Plan.addExternalDef(NewVPVal);
+  VPValue *NewVPVal = Plan.getOrAddExternalDef(IRVal);
   IRDef2VPValue[IRVal] = NewVPVal;
   return NewVPVal;
 }
 
 // Create new VPInstructions in a VPBasicBlock, given its BasicBlock
 // counterpart. This function must be invoked in RPO so that the operands of a
 // VPInstruction in \p BB have been visited before (except for Phi nodes).
 void PlainCFGBuilder::createVPInstructionsForVPBB(VPBasicBlock *VPBB,
VPRegionBlock *PlainCFGBuilder::buildPlainCFG() {
   // LoopBlocksDFS.
   BasicBlock *PreheaderBB = TheLoop->getLoopPreheader();
   assert((PreheaderBB->getTerminator()->getNumSuccessors() == 1) &&
          "Unexpected loop preheader");
   VPBasicBlock *PreheaderVPBB = getOrCreateVPBB(PreheaderBB);
   for (auto &I : *PreheaderBB) {
     if (I.getType()->isVoidTy())
       continue;
-    VPValue *VPV = new VPValue(&I);
-    Plan.addExternalDef(VPV);
-    IRDef2VPValue[&I] = VPV;
+    IRDef2VPValue[&I] = Plan.getOrAddExternalDef(&I);
   }
   // Create empty VPBB for Loop H so that we can link PH->H.
   VPBlockBase *HeaderVPBB = getOrCreateVPBB(TheLoop->getHeader());
   HeaderVPBB->setName("vector.body");
   // Preheader's predecessors will be set during the loop RPO traversal below.
   PreheaderVPBB->setOneSuccessor(HeaderVPBB);
 
   LoopBlocksRPO RPO(TheLoop);

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (VPRecipeBase &Phi : HeaderVPBB->phis()) {
     auto *IV = dyn_cast<VPWidenIntOrFpInductionRecipe>(&Phi);
     if (!IV || !IV->needsScalarIV())
       continue;
 
     const InductionDescriptor &ID = IV->getInductionDescriptor();
     const SCEV *StepSCEV = ID.getStep();
     VPValue *Step = nullptr;
     if (auto *E = dyn_cast<SCEVConstant>(StepSCEV)) {
-      Step = new VPValue(E->getValue());
-      Plan.addExternalDef(Step);
+      Step = Plan.getOrAddExternalDef(E->getValue());
     } else if (auto *E = dyn_cast<SCEVUnknown>(StepSCEV)) {
-      Step = new VPValue(E->getValue());
-      Plan.addExternalDef(Step);
+      Step = Plan.getOrAddExternalDef(E->getValue());
     } else {
       Step = new VPExpandSCEVRecipe(StepSCEV, SE);
     }
 
     Instruction *TruncI = IV->getTruncInst();
     VPScalarIVStepsRecipe *Steps = new VPScalarIVStepsRecipe(
         IV->getPHINode()->getType(), ID, Plan.getCanonicalIV(),
         IV->getStartValue(), Step, TruncI ? TruncI->getType() : nullptr);

llvm/unittests/Transforms/Vectorize/VPlanTest.cpp

TEST(VPRecipeTest, dump) {
 
   LLVMContext C;
 
   IntegerType *Int32 = IntegerType::get(C, 32);
   auto *AI =
       BinaryOperator::CreateAdd(UndefValue::get(Int32), UndefValue::get(Int32));
   AI->setName("a");
   SmallVector<VPValue *, 2> Args;
-  VPValue *ExtVPV1 = new VPValue();
-  VPValue *ExtVPV2 = new VPValue();
-  Plan.addExternalDef(ExtVPV1);
-  Plan.addExternalDef(ExtVPV2);
+  VPValue *ExtVPV1 = Plan.getOrAddExternalDef(ConstantInt::get(Int32, 1));
+  VPValue *ExtVPV2 = Plan.getOrAddExternalDef(ConstantInt::get(Int32, 2));
   Args.push_back(ExtVPV1);
   Args.push_back(ExtVPV2);
   VPWidenRecipe *WidenR =
       new VPWidenRecipe(*AI, make_range(Args.begin(), Args.end()));
   VPBB1->appendRecipe(WidenR);
 
   {
     // Use EXPECT_EXIT to capture stderr and compare against expected output.
     //
     // Test VPValue::dump().
     VPValue *VPV = WidenR;
     EXPECT_EXIT(
         {
           VPV->dump();
           exit(0);
         },
-        testing::ExitedWithCode(0), "WIDEN ir<%a> = add vp<%0>, vp<%1>");
+        testing::ExitedWithCode(0), "WIDEN ir<%a> = add ir<1>, ir<2>");
 
     // Test VPRecipeBase::dump().
     VPRecipeBase *R = WidenR;
     EXPECT_EXIT(
         {
           R->dump();
           exit(0);
         },
-        testing::ExitedWithCode(0), "WIDEN ir<%a> = add vp<%0>, vp<%1>");
+        testing::ExitedWithCode(0), "WIDEN ir<%a> = add ir<1>, ir<2>");
 
     // Test VPDef::dump().
     VPDef *D = WidenR;
     EXPECT_EXIT(
         {
           D->dump();
           exit(0);
         },
-        testing::ExitedWithCode(0), "WIDEN ir<%a> = add vp<%0>, vp<%1>");
+        testing::ExitedWithCode(0), "WIDEN ir<%a> = add ir<1>, ir<2>");
   }
 
   delete AI;
 }
 #endif
 
 TEST(VPRecipeTest, CastVPReductionRecipeToVPUser) {
   LLVMContext C;