[VPlan] Manage scalarized values using VPValues.

[fhahn]

Summary

This patch updates codegen to use VPValues to manage the generated
scalarized instructions.

[fhahn]

Adjust vplan printing test and split off changes to use VPTransformState in a few more places to D95757. This is required to ensure we pick up the scalarized values managed directly in VPTransformState.

[a.elovikov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2929

nit: Should it be Builder.getInsertBlock() to avoid calling getParent on BasicBlock::end() iterator? It might be impossible to have the end iterator here, but the getInsertBlock version is generally safer, IMO.

2947

nit: I think commenting it out would be misleading for the reader. Can we just remove it?

[fhahn]

Use GetInsertBlock, remove misleading comment, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

2929

Good point, adjusted!

2947

Yes, the comment should be outdated now. I removed it.

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9203

Call getOperand(0)->getUnderlyingValue() directly?

9204

Can this line be brought back under the comment to minimize the diff?

9210–9211

Can the then/else logic be reversed to minimize the diff?

[fhahn]

Updated code in VPPredInstPHIRecipe::execute to reduce diff as suggested. Thanks!

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9055

No need to pass as arguments now that State itself is passed.

9217

This is porting the resetVectorValue to State operations, which I think is fine by itself (I tend to agree with passing nullptr as the underlying value to VPPredInstPHI), but something else bothers me: It seems VPPredInstPHI doesn't take over the VPReplicate as the VPValue mapped to the predicated IR value during VPlan construction, which means that recipes constructed for users of predicated instructions use the ReplicateRecipe although it doesn't dominate them. We get away with this by overriding the IR's vector value, but for correct VP def-use we should have VPPredInstPHI take over as the registered VPValue in VPlan during construction. Does all that make sense?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9217

I think the problem here is how we generate code for a scalarization region, like the one below. Uses after the region use the PHI-PREDICATED-INSTRUCTION VPValue.

But when generating code for the region itself, we execute the same blocks once for each vector element. So in the second iteration "REPLICATE ir<%div> = udiv ir<16>, ir<%counter1.0> (S->V)\l" currently needs to pack its result in the result vector from the previous iteration. Currently that is done by overwriting the state of the def "REPLICATE ir<%div> = udiv ir<16>, ir<%counter1.0> (S->V)\l" with the predicated PHI from the previous iteration, so the vector-pack logic picks it up.

It might be possible to instead change the code to just have scalar phis and do the packing as part of the phi creation. I think for transformations, the VPValue def-use chain should be correct, and it is 'just' a code-gen improvement. Do you think we need to fix this before this one goes in?

subgraph cluster_N2 {
  fontname=Courier
  label="\<xVFxUF\> pred.udiv"
  N1 [label =
    "pred.udiv.entry:\n" +
      " +
      "BRANCH-ON-MASK vp<%4>\l"\l" +
       "CondBit: vp<%4> (for.body3)\l"
  ]
  N1 -> N3 [ label="T"]
  N1 -> N4 [ label="F"]
  N3 [label =
    "pred.udiv.if:\n" +
      "REPLICATE ir<%div> = udiv ir<16>, ir<%counter1.0> (S->V)\l"
  ]
  N3 -> N4 [ label=""]
  N4 [label =
    "pred.udiv.continue:\n" +
      "PHI-PREDICATED-INSTRUCTION vp<%6> = ir<%div>\l"
  ]
}

[fhahn]

Ping :)

Rebased and added a note about why we are re-setting the value for operands of VPPRedPHI. Also added a few comments to set/reset functions and assertions.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9217

@gilr I added a note explaining why we need to also update the operand. Do you think that is sufficient to move forward with this patch?

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9055

No need to pass as arguments now that State itself is passed.

Does that make sense?

9217

Yes, thanks!

9217

Uses after the region use the PHI-PREDICATED-INSTRUCTION VPValue.

Right, I missed the part where you replace the predicated instruction's VPValue with the new PredInstRecipe.

It might be possible to instead change the code to just have scalar phis and do the packing as part of the phi creation. I think for transformations, the VPValue def-use chain should be correct, and it is 'just' a code-gen improvement. Do you think we need to fix this before this one goes in?

Not sure recipes should create values in other places than their designated location. Anyway, this shouldn't block this patch.

[fhahn]

Thanks Gil! I removed the UF/VF arguments from widenPHIInstruction, I totally missed that I did not address this comment originally.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9055

Does that make sense?

Yes, sorry I totally missed that I did not update this so far! Fixed now.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   void widenGEP(GetElementPtrInst *GEP, VPValue *VPDef, VPUser &Indices,
                 unsigned UF, ElementCount VF, bool IsPtrLoopInvariant,
                 SmallBitVector &IsIndexLoopInvariant, VPTransformState &State);
 
   /// Vectorize a single PHINode in a block. This method handles the induction
   /// variable canonicalization. It supports both VF = 1 for unrolled loops and
   /// arbitrary length vectors.
   void widenPHIInstruction(Instruction *PN, RecurrenceDescriptor *RdxDesc,
-                           Value *StartV, unsigned UF, ElementCount VF);
+                           Value *StartV, VPValue *Def,
+                           VPTransformState &State);
 
   /// A helper function to scalarize a single Instruction in the innermost loop.
   /// Generates a sequence of scalar instances for each lane between \p MinLane
   /// and \p MaxLane, times each part between \p MinPart and \p MaxPart,
   /// inclusive. Uses the VPValue operands from \p Operands instead of \p
   /// Instr's operands.
-  void scalarizeInstruction(Instruction *Instr, VPUser &Operands,
+  void scalarizeInstruction(Instruction *Instr, VPValue *Def, VPUser &Operands,
                             const VPIteration &Instance, bool IfPredicateInstr,
                             VPTransformState &State);
 
   /// Widen an integer or floating-point induction variable \p IV. If \p Trunc
   /// is provided, the integer induction variable will first be truncated to
   /// the corresponding type.
   void widenIntOrFpInduction(PHINode *IV, Value *Start, TruncInst *Trunc,
                              VPValue *Def, VPValue *CastDef,
if (CreateGatherScatter) {
       if (Reverse)
         NewLI = reverseVector(NewLI);
     }
 
     State.set(Def, Instr, NewLI, Part);
   }
 }
 
-void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr, VPUser &User,
+void InnerLoopVectorizer::scalarizeInstruction(Instruction *Instr, VPValue *Def,
+                                               VPUser &User,
                                                const VPIteration &Instance,
                                                bool IfPredicateInstr,
                                                VPTransformState &State) {
   assert(!Instr->getType()->isAggregateType() && "Can't handle vectors");
 
   // llvm.experimental.noalias.scope.decl intrinsics must only be duplicated for
   // the first lane and part.
   if (isa<NoAliasScopeDeclInst>(Instr))
     if (!Instance.isFirstIteration())
       return;
 
   setDebugLocFromInst(Builder, Instr);
 
   // Does this instruction return a value ?
   bool IsVoidRetTy = Instr->getType()->isVoidTy();
 
   Instruction *Cloned = Instr->clone();
   if (!IsVoidRetTy)
     Cloned->setName(Instr->getName() + ".cloned");
 
+  State.Builder.SetInsertPoint(Builder.GetInsertBlock(),
+                               Builder.GetInsertPoint());
   // Replace the operands of the cloned instructions with their scalar
   // equivalents in the new loop.
   for (unsigned op = 0, e = User.getNumOperands(); op != e; ++op) {
     auto *Operand = dyn_cast<Instruction>(Instr->getOperand(op));
     auto InputInstance = Instance;
     if (!Operand || !OrigLoop->contains(Operand) ||
         (Cost->isUniformAfterVectorization(Operand, State.VF)))
       InputInstance.Lane = 0;
     auto *NewOp = State.get(User.getOperand(op), InputInstance);
     Cloned->setOperand(op, NewOp);
   }
   addNewMetadata(Cloned, Instr);
 
   // Place the cloned scalar in the new loop.
   Builder.Insert(Cloned);
 
-  // TODO: Set result for VPValue of VPReciplicateRecipe. This requires
-  // representing scalar values in VPTransformState. Add the cloned scalar to
-  // the scalar map entry.
-  VectorLoopValueMap.setScalarValue(Instr, Instance, Cloned);
+  State.set(Def, Instr, Cloned, Instance);
 
   // If we just cloned a new assumption, add it the assumption cache.
   if (auto *II = dyn_cast<IntrinsicInst>(Cloned))
     if (II->getIntrinsicID() == Intrinsic::assume)
       AC->registerAssumption(II);
 
   // End if-block.
   if (IfPredicateInstr)
for (unsigned Part = 0; Part < UF; ++Part) {
       State.set(VPDef, GEP, NewGEP, Part);
       addMetadata(NewGEP, GEP);
     }
   }
 }
 
 void InnerLoopVectorizer::widenPHIInstruction(Instruction *PN,
                                               RecurrenceDescriptor *RdxDesc,
-                                              Value *StartV, unsigned UF,
-                                              ElementCount VF) {
-  assert(!VF.isScalable() && "scalable vectors not yet supported.");
+                                              Value *StartV, VPValue *Def,
+                                              VPTransformState &State) {
+  assert(!State.VF.isScalable() && "scalable vectors not yet supported.");
   PHINode *P = cast<PHINode>(PN);
   if (EnableVPlanNativePath) {
     // Currently we enter here in the VPlan-native path for non-induction
     // PHIs where all control flow is uniform. We simply widen these PHIs.
     // Create a vector phi with no operands - the vector phi operands will be
     // set at the end of vector code generation.
-    Type *VecTy =
-        (VF.isScalar()) ? PN->getType() : VectorType::get(PN->getType(), VF);
+    Type *VecTy = (State.VF.isScalar())
+                      ? PN->getType()
+                      : VectorType::get(PN->getType(), State.VF);
     Value *VecPhi = Builder.CreatePHI(VecTy, PN->getNumOperands(), "vec.phi");
-    VectorLoopValueMap.setVectorValue(P, 0, VecPhi);
+    State.set(Def, P, VecPhi, 0);
     OrigPHIsToFix.push_back(P);
 
     return;
   }
 
   assert(PN->getParent() == OrigLoop->getHeader() &&
          "Non-header phis should have been handled elsewhere");
 
   // In order to support recurrences we need to be able to vectorize Phi nodes.
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #1: We create a new vector PHI node with no incoming edges. We'll use
   // this value when we vectorize all of the instructions that use the PHI.
   if (RdxDesc || Legal->isFirstOrderRecurrence(P)) {
     Value *Iden = nullptr;
     bool ScalarPHI =
-        (VF.isScalar()) || Cost->isInLoopReduction(cast<PHINode>(PN));
+        (State.VF.isScalar()) || Cost->isInLoopReduction(cast<PHINode>(PN));
     Type *VecTy =
-        ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), VF);
+        ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), State.VF);
 
     if (RdxDesc) {
       assert(Legal->isReductionVariable(P) && StartV &&
              "RdxDesc should only be set for reduction variables; in that case "
              "a StartV is also required");
       RecurKind RK = RdxDesc->getRecurrenceKind();
       if (RecurrenceDescriptor::isMinMaxRecurrenceKind(RK)) {
         // MinMax reduction have the start value as their identify.
         if (ScalarPHI) {
           Iden = StartV;
         } else {
           IRBuilderBase::InsertPointGuard IPBuilder(Builder);
           Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator());
-          StartV = Iden = Builder.CreateVectorSplat(VF, StartV, "minmax.ident");
+          StartV = Iden =
+              Builder.CreateVectorSplat(State.VF, StartV, "minmax.ident");
         }
       } else {
         Constant *IdenC = RecurrenceDescriptor::getRecurrenceIdentity(
             RK, VecTy->getScalarType());
         Iden = IdenC;
 
         if (!ScalarPHI) {
-          Iden = ConstantVector::getSplat(VF, IdenC);
+          Iden = ConstantVector::getSplat(State.VF, IdenC);
           IRBuilderBase::InsertPointGuard IPBuilder(Builder);
           Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator());
           Constant *Zero = Builder.getInt32(0);
           StartV = Builder.CreateInsertElement(Iden, StartV, Zero);
         }
       }
     }
 
-    for (unsigned Part = 0; Part < UF; ++Part) {
+    for (unsigned Part = 0; Part < State.UF; ++Part) {
       // This is phase one of vectorizing PHIs.
       Value *EntryPart = PHINode::Create(
           VecTy, 2, "vec.phi", &*LoopVectorBody->getFirstInsertionPt());
-      VectorLoopValueMap.setVectorValue(P, Part, EntryPart);
+      State.set(Def, P, EntryPart, Part);
       if (StartV) {
         // Make sure to add the reduction start value only to the
         // first unroll part.
         Value *StartVal = (Part == 0) ? StartV : Iden;
         cast<PHINode>(EntryPart)->addIncoming(StartVal, LoopVectorPreHeader);
       }
     }
     return;
case InductionDescriptor::IK_NoInduction:
     llvm_unreachable("Unknown induction");
   case InductionDescriptor::IK_IntInduction:
   case InductionDescriptor::IK_FpInduction:
     llvm_unreachable("Integer/fp induction is handled elsewhere.");
   case InductionDescriptor::IK_PtrInduction: {
     // Handle the pointer induction variable case.
     assert(P->getType()->isPointerTy() && "Unexpected type.");
 
-    if (Cost->isScalarAfterVectorization(P, VF)) {
+    if (Cost->isScalarAfterVectorization(P, State.VF)) {
       // This is the normalized GEP that starts counting at zero.
       Value *PtrInd =
           Builder.CreateSExtOrTrunc(Induction, II.getStep()->getType());
       // Determine the number of scalars we need to generate for each unroll
       // iteration. If the instruction is uniform, we only need to generate the
       // first lane. Otherwise, we generate all VF values.
-      unsigned Lanes =
-          Cost->isUniformAfterVectorization(P, VF) ? 1 : VF.getKnownMinValue();
+      unsigned Lanes = Cost->isUniformAfterVectorization(P, State.VF)
+                           ? 1
+                           : State.VF.getKnownMinValue();
       for (unsigned Part = 0; Part < UF; ++Part) {
         for (unsigned Lane = 0; Lane < Lanes; ++Lane) {
-          Constant *Idx = ConstantInt::get(PtrInd->getType(),
-                                           Lane + Part * VF.getKnownMinValue());
+          Constant *Idx = ConstantInt::get(
+              PtrInd->getType(), Lane + Part * State.VF.getKnownMinValue());
           Value *GlobalIdx = Builder.CreateAdd(PtrInd, Idx);
           Value *SclrGep =
               emitTransformedIndex(Builder, GlobalIdx, PSE.getSE(), DL, II);
           SclrGep->setName("next.gep");
-          VectorLoopValueMap.setScalarValue(P, VPIteration(Part, Lane),
-                                            SclrGep);
+          State.set(Def, P, SclrGep, VPIteration(Part, Lane));
         }
       }
       return;
     }
     assert(isa<SCEVConstant>(II.getStep()) &&
            "Induction step not a SCEV constant!");
     Type *PhiType = II.getStep()->getType();
 
case InductionDescriptor::IK_PtrInduction: {
     const SCEV *ScalarStep = II.getStep();
     SCEVExpander Exp(*PSE.getSE(), DL, "induction");
     Value *ScalarStepValue =
         Exp.expandCodeFor(ScalarStep, PhiType, InductionLoc);
     Value *InductionGEP = GetElementPtrInst::Create(
         ScStValueType->getPointerElementType(), NewPointerPhi,
         Builder.CreateMul(
             ScalarStepValue,
-            ConstantInt::get(PhiType, VF.getKnownMinValue() * UF)),
+            ConstantInt::get(PhiType, State.VF.getKnownMinValue() * State.UF)),
         "ptr.ind", InductionLoc);
     NewPointerPhi->addIncoming(InductionGEP, LoopLatch);
 
     // Create UF many actual address geps that use the pointer
     // phi as base and a vectorized version of the step value
     // (<step*0, ..., step*N>) as offset.
-    for (unsigned Part = 0; Part < UF; ++Part) {
+    for (unsigned Part = 0; Part < State.UF; ++Part) {
       SmallVector<Constant *, 8> Indices;
       // Create a vector of consecutive numbers from zero to VF.
-      for (unsigned i = 0; i < VF.getKnownMinValue(); ++i)
+      for (unsigned i = 0; i < State.VF.getKnownMinValue(); ++i)
         Indices.push_back(
-            ConstantInt::get(PhiType, i + Part * VF.getKnownMinValue()));
+            ConstantInt::get(PhiType, i + Part * State.VF.getKnownMinValue()));
       Constant *StartOffset = ConstantVector::get(Indices);
 
       Value *GEP = Builder.CreateGEP(
           ScStValueType->getPointerElementType(), NewPointerPhi,
-          Builder.CreateMul(
-              StartOffset,
-              Builder.CreateVectorSplat(VF.getKnownMinValue(), ScalarStepValue),
+          Builder.CreateMul(StartOffset,
+                            Builder.CreateVectorSplat(
+                                State.VF.getKnownMinValue(), ScalarStepValue),
                             "vector.gep"));
-      VectorLoopValueMap.setVectorValue(P, Part, GEP);
+      State.set(Def, P, GEP, Part);
     }
   }
   }
 }
 
 /// A helper function for checking whether an integer division-related
 /// instruction may divide by zero (in which case it must be predicated if
 /// executed conditionally in the scalar code).
VPRegionBlock *VPRecipeBuilder::createReplicateRegion(Instruction *Instr,
   // Build the triangular if-then region.
   std::string RegionName = (Twine("pred.") + Instr->getOpcodeName()).str();
   assert(Instr->getParent() && "Predicated instruction not in any basic block");
   auto *BOMRecipe = new VPBranchOnMaskRecipe(BlockInMask);
   auto *Entry = new VPBasicBlock(Twine(RegionName) + ".entry", BOMRecipe);
   auto *PHIRecipe = Instr->getType()->isVoidTy()
                         ? nullptr
                         : new VPPredInstPHIRecipe(Plan->getOrAddVPValue(Instr));
+  if (PHIRecipe) {
+    Plan->removeVPValueFor(Instr);
+    Plan->addVPValue(Instr, PHIRecipe);
+  }
   auto *Exit = new VPBasicBlock(Twine(RegionName) + ".continue", PHIRecipe);
   auto *Pred = new VPBasicBlock(Twine(RegionName) + ".if", PredRecipe);
   VPRegionBlock *Region = new VPRegionBlock(Entry, Exit, RegionName, true);
 
   // Note: first set Entry as region entry and then connect successors starting
   // from it in order, to propagate the "parent" of each VPBasicBlock.
   VPBlockUtils::insertTwoBlocksAfter(Pred, Exit, BlockInMask, Entry);
   VPBlockUtils::connectBlocks(Pred, Exit);
void VPWidenIntOrFpInductionRecipe::execute(VPTransformState &State) {
   State.ILV->widenIntOrFpInduction(IV, getStartValue()->getLiveInIRValue(),
                                    getTruncInst(), getVPValue(0),
                                    getCastValue(), State);
 }
 
 void VPWidenPHIRecipe::execute(VPTransformState &State) {
   Value *StartV =
       getStartValue() ? getStartValue()->getLiveInIRValue() : nullptr;
-  State.ILV->widenPHIInstruction(Phi, RdxDesc, StartV, State.UF, State.VF);
+  State.ILV->widenPHIInstruction(cast<PHINode>(getUnderlyingValue()), RdxDesc,
+                                 StartV, this, State);
 }
 
 void VPBlendRecipe::execute(VPTransformState &State) {
   State.ILV->setDebugLocFromInst(State.Builder, Phi);
   // We know that all PHIs in non-header blocks are converted into
   // selects, so we don't have to worry about the insertion order and we
   // can just use the builder.
   // At this point we generate the predication tree. There may be
for (unsigned Part = 0; Part < State.UF; ++Part) {
         // based on the incoming mask.
         Value *Cond = State.get(getMask(In), Part);
         Entry[Part] =
             State.Builder.CreateSelect(Cond, In0, Entry[Part], "predphi");
       }
     }
   }
   for (unsigned Part = 0; Part < State.UF; ++Part)
-    State.ValueMap.setVectorValue(Phi, Part, Entry[Part]);
+    State.set(this, Phi, Entry[Part], Part);
 }
 
 void VPInterleaveRecipe::execute(VPTransformState &State) {
   assert(!State.Instance && "Interleave group being replicated.");
   State.ILV->vectorizeInterleaveGroup(IG, definedValues(), State, getAddr(),
                                       getStoredValues(), getMask());
 }
 
for (unsigned Part = 0; Part < State.UF; ++Part) {
     }
     State.set(this, getUnderlyingInstr(), NextInChain, Part);
   }
 }
 
 void VPReplicateRecipe::execute(VPTransformState &State) {
   if (State.Instance) { // Generate a single instance.
     assert(!State.VF.isScalable() && "Can't scalarize a scalable vector");
-    State.ILV->scalarizeInstruction(getUnderlyingInstr(), *this,
+    State.ILV->scalarizeInstruction(getUnderlyingInstr(), this, *this,
                                     *State.Instance, IsPredicated, State);
     // Insert scalar instance packing it into a vector.
     if (AlsoPack && State.VF.isVector()) {
       // If we're constructing lane 0, initialize to start from poison.
       if (State.Instance->Lane == 0) {
         assert(!State.VF.isScalable() && "VF is assumed to be non scalable.");
         Value *Poison = PoisonValue::get(
             VectorType::get(getUnderlyingValue()->getType(), State.VF));
-        State.ValueMap.setVectorValue(getUnderlyingInstr(),
-                                      State.Instance->Part, Poison);
+        State.set(this, getUnderlyingInstr(), Poison, State.Instance->Part);
       }
-      State.ILV->packScalarIntoVectorValue(getUnderlyingInstr(),
-                                           *State.Instance);
+      State.ILV->packScalarIntoVectorValue(this, *State.Instance, State);
     }
     return;
   }
 
   // Generate scalar instances for all VF lanes of all UF parts, unless the
   // instruction is uniform inwhich case generate only the first lane for each
   // of the UF parts.
   unsigned EndLane = IsUniform ? 1 : State.VF.getKnownMinValue();
   assert((!State.VF.isScalable() || IsUniform) &&
          "Can't scalarize a scalable vector");
   for (unsigned Part = 0; Part < State.UF; ++Part)
     for (unsigned Lane = 0; Lane < EndLane; ++Lane)
-      State.ILV->scalarizeInstruction(getUnderlyingInstr(), *this,
+      State.ILV->scalarizeInstruction(getUnderlyingInstr(), this, *this,
                                       VPIteration(Part, Lane), IsPredicated,
                                       State);
 }
 
 void VPBranchOnMaskRecipe::execute(VPTransformState &State) {
   assert(State.Instance && "Branch on Mask works only on single instance.");
 
   unsigned Part = State.Instance->Part;

 
 void VPPredInstPHIRecipe::execute(VPTransformState &State) {
   assert(State.Instance && "Predicated instruction PHI works per instance.");
   Instruction *ScalarPredInst =
       cast<Instruction>(State.get(getOperand(0), *State.Instance));
   BasicBlock *PredicatedBB = ScalarPredInst->getParent();
   BasicBlock *PredicatingBB = PredicatedBB->getSinglePredecessor();
   assert(PredicatingBB && "Predicated block has no single predecessor.");
+  assert(isa<VPReplicateRecipe>(getOperand(0)) &&
+         "operand must be VPReplicateRecipe");
 
   // By current pack/unpack logic we need to generate only a single phi node: if
   // a vector value for the predicated instruction exists at this point it means
   // the instruction has vector users only, and a phi for the vector value is
   // needed. In this case the recipe of the predicated instruction is marked to
   // also do that packing, thereby "hoisting" the insert-element sequence.
   // Otherwise, a phi node for the scalar value is needed.
   unsigned Part = State.Instance->Part;
-  Instruction *PredInst =
-      cast<Instruction>(getOperand(0)->getUnderlyingValue());
-  if (State.ValueMap.hasVectorValue(PredInst, Part)) {
-    Value *VectorValue = State.ValueMap.getVectorValue(PredInst, Part);
+  if (State.hasVectorValue(getOperand(0), Part)) {
+    Value *VectorValue = State.get(getOperand(0), Part);
     InsertElementInst *IEI = cast<InsertElementInst>(VectorValue);
     PHINode *VPhi = State.Builder.CreatePHI(IEI->getType(), 2);
     VPhi->addIncoming(IEI->getOperand(0), PredicatingBB); // Unmodified vector.
     VPhi->addIncoming(IEI, PredicatedBB); // New vector with inserted element.
-    State.ValueMap.resetVectorValue(PredInst, Part, VPhi); // Update cache.
+    if (State.hasVectorValue(this, Part))
+      State.reset(this, VPhi, Part);
+    else
+      State.set(this, VPhi, Part);
+    // NOTE: Currently we need to update the value of the operand, so the next
+    // predicated iteration inserts its generated value in the correct vector.
+    State.reset(getOperand(0), VPhi, Part);
   } else {
-    Type *PredInstType = PredInst->getType();
+    Type *PredInstType = getOperand(0)->getUnderlyingValue()->getType();
     PHINode *Phi = State.Builder.CreatePHI(PredInstType, 2);
-    Phi->addIncoming(PoisonValue::get(ScalarPredInst->getType()), PredicatingBB);
+    Phi->addIncoming(PoisonValue::get(ScalarPredInst->getType()),
+                     PredicatingBB);
     Phi->addIncoming(ScalarPredInst, PredicatedBB);
-    State.ValueMap.resetScalarValue(PredInst, *State.Instance, Phi);
+    if (State.hasScalarValue(this, *State.Instance))
+      State.reset(this, Phi, *State.Instance);
+    else
+      State.set(this, Phi, *State.Instance);
+    // NOTE: Currently we need to update the value of the operand, so the next
+    // predicated iteration inserts its generated value in the correct vector.
+    State.reset(getOperand(0), Phi, *State.Instance);
   }
 }
 
 void VPWidenMemoryInstructionRecipe::execute(VPTransformState &State) {
   VPValue *StoredValue = isStore() ? getStoredValue() : nullptr;
   State.ILV->vectorizeMemoryInstruction(&Ingredient, State,
                                         StoredValue ? nullptr : getVPValue(),
                                         getAddr(), StoredValue, getMask());
void VPTransformState::reset(VPValue *Def, Value *IRDef, Value *V,
   ILV->resetVectorValue(IRDef, Part, V);
 }
 
 Value *VPTransformState::get(VPValue *Def, unsigned Part) {
   // If Values have been set for this Def return the one relevant for \p Part.
   if (hasVectorValue(Def, Part))
     return Data.PerPartOutput[Def][Part];
 
-  // TODO: Remove the callback once all scalar recipes are managed using
-  // VPValues.
   if (!hasScalarValue(Def, {Part, 0}))
     return Callback.getOrCreateVectorValues(VPValue2Value[Def], Part);
 
   Value *ScalarValue = get(Def, {Part, 0});
   // If we aren't vectorizing, we can just copy the scalar map values over
   // to the vector map.
   if (VF.isScalar()) {
     set(Def, ScalarValue, Part);
Value *VPTransformState::get(VPValue *Def, unsigned Part) {
   // VectorLoopValueMap, we will only generate the insertelements once.
   Value *VectorValue = nullptr;
   if (IsUniform) {
     VectorValue = ILV->getBroadcastInstrs(ScalarValue);
     set(Def, VectorValue, Part);
   } else {
     // Initialize packing with insertelements to start from undef.
     assert(!VF.isScalable() && "VF is assumed to be non scalable.");
-    Value *Undef = UndefValue::get(VectorType::get(LastInst->getType(), VF));
+    Value *Undef = PoisonValue::get(VectorType::get(LastInst->getType(), VF));
     set(Def, Undef, Part);
     for (unsigned Lane = 0; Lane < VF.getKnownMinValue(); ++Lane)
       ILV->packScalarIntoVectorValue(Def, {Part, Lane}, *this);
     VectorValue = get(Def, Part);
   }
   Builder.restoreIP(OldIP);
   return VectorValue;
 }

llvm/lib/Transforms/Vectorize/VPlan.h

struct VPTransformState {
   /// Set the generated Value for a given VPValue and a given Part.
   void set(VPValue *Def, Value *V, unsigned Part) {
     if (!Data.PerPartOutput.count(Def)) {
       DataState::PerPartValuesTy Entry(UF);
       Data.PerPartOutput[Def] = Entry;
     }
     Data.PerPartOutput[Def][Part] = V;
   }
+  /// Reset an existing vector value for \p Def and a given \p Part.
+  void reset(VPValue *Def, Value *V, unsigned Part) {
+    auto Iter = Data.PerPartOutput.find(Def);
+    assert(Iter != Data.PerPartOutput.end() &&
+           "need to overwrite existing value");
+    Iter->second[Part] = V;
+  }
+
   void set(VPValue *Def, Value *IRDef, Value *V, unsigned Part);
   void reset(VPValue *Def, Value *IRDef, Value *V, unsigned Part);
-  void set(VPValue *Def, Value *IRDef, Value *V, const VPIteration &Instance);
 
+  /// Set the generated scalar \p V for \p Def and \p IRDef and the given \p
+  /// Instance.
+  void set(VPValue *Def, Value *IRDef, Value *V, const VPIteration &Instance);
+  /// Set the generated scalar \p V for \p Def and the given \p Instance.
   void set(VPValue *Def, Value *V, const VPIteration &Instance) {
     auto Iter = Data.PerPartScalars.insert({Def, {}});
     auto &PerPartVec = Iter.first->second;
     while (PerPartVec.size() <= Instance.Part)
       PerPartVec.emplace_back();
     auto &Scalars = PerPartVec[Instance.Part];
     while (Scalars.size() <= Instance.Lane)
       Scalars.push_back(nullptr);
+    assert(!Scalars[Instance.Lane] && "should overwrite existing value");
     Scalars[Instance.Lane] = V;
   }
 
+  /// Reset an existing scalar value for \p Def and a given \p Instance.
+  void reset(VPValue *Def, Value *V, const VPIteration &Instance) {
+    auto Iter = Data.PerPartScalars.find(Def);
+    assert(Iter != Data.PerPartScalars.end() &&
+           "need to overwrite existing value");
+    assert(Instance.Part < Iter->second.size() &&
+           "need to overwrite existing value");
+    assert(Instance.Lane < Iter->second[Instance.Part].size() &&
+           "need to overwrite existing value");
+    Iter->second[Instance.Part][Instance.Lane] = V;
+  }
+
   /// Hold state information used when constructing the CFG of the output IR,
   /// traversing the VPBasicBlocks and generating corresponding IR BasicBlocks.
   struct CFGState {
     /// The previous VPBasicBlock visited. Initially set to null.
     VPBasicBlock *PrevVPBB = nullptr;
 
     /// The previous IR BasicBlock created or used. Initially set to the new
     /// header BasicBlock.
public:
   const TruncInst *getTruncInst() const {
     return dyn_cast_or_null<TruncInst>(getVPValue(0)->getUnderlyingValue());
   }
 };
 
 /// A recipe for handling all phi nodes except for integer and FP inductions.
 /// For reduction PHIs, RdxDesc must point to the corresponding recurrence
 /// descriptor and the start value is the first operand of the recipe.
-class VPWidenPHIRecipe : public VPRecipeBase {
+class VPWidenPHIRecipe : public VPRecipeBase, public VPValue {
   PHINode *Phi;
 
   /// Descriptor for a reduction PHI.
   RecurrenceDescriptor *RdxDesc = nullptr;
 
 public:
   /// Create a new VPWidenPHIRecipe for the reduction \p Phi described by \p
   /// RdxDesc.
   VPWidenPHIRecipe(PHINode *Phi, RecurrenceDescriptor &RdxDesc, VPValue &Start)
       : VPWidenPHIRecipe(Phi) {
     this->RdxDesc = &RdxDesc;
     addOperand(&Start);
   }
 
   /// Create a VPWidenPHIRecipe for \p Phi
-  VPWidenPHIRecipe(PHINode *Phi) : VPRecipeBase(VPWidenPHISC, {}), Phi(Phi) {
-    new VPValue(Phi, this);
-  }
-
+  VPWidenPHIRecipe(PHINode *Phi)
+      : VPRecipeBase(VPWidenPHISC, {}),
+        VPValue(VPValue::VPVWidenPHISC, Phi, this), Phi(Phi) {}
   ~VPWidenPHIRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPWidenPHISC;
   }
 
   /// Generate the phi/select nodes.
   void execute(VPTransformState &State) override;
 
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 
   /// Returns the start value of the phi, if it is a reduction.
   VPValue *getStartValue() {
     return getNumOperands() == 0 ? nullptr : getOperand(0);
   }
 };
 
 /// A recipe for vectorizing a phi-node as a sequence of mask-based select
 /// instructions.
-class VPBlendRecipe : public VPRecipeBase {
+class VPBlendRecipe : public VPRecipeBase, public VPValue {
   PHINode *Phi;
 
 public:
   /// The blend operation is a User of the incoming values and of their
   /// respective masks, ordered [I0, M0, I1, M1, ...]. Note that a single value
   /// might be incoming with a full mask for which there is no VPValue.
   VPBlendRecipe(PHINode *Phi, ArrayRef<VPValue *> Operands)
-      : VPRecipeBase(VPBlendSC, Operands), Phi(Phi) {
-    new VPValue(Phi, this);
+      : VPRecipeBase(VPBlendSC, Operands),
+        VPValue(VPValue::VPVBlendSC, Phi, this), Phi(Phi) {
     assert(Operands.size() > 0 &&
            ((Operands.size() == 1) || (Operands.size() % 2 == 0)) &&
            "Expected either a single incoming value or a positive even number "
            "of operands");
   }
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
public:
 
   void setAlsoPack(bool Pack) { AlsoPack = Pack; }
 
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 
   bool isUniform() const { return IsUniform; }
+
+  bool isPacked() const { return AlsoPack; }
 };
 
 /// A recipe for generating conditional branches on the bits of a mask.
 class VPBranchOnMaskRecipe : public VPRecipeBase {
 public:
   VPBranchOnMaskRecipe(VPValue *BlockInMask)
       : VPRecipeBase(VPBranchOnMaskSC, {}) {
     if (BlockInMask) // nullptr means all-one mask.
public:
   }
 };
 
 /// VPPredInstPHIRecipe is a recipe for generating the phi nodes needed when
 /// control converges back from a Branch-on-Mask. The phi nodes are needed in
 /// order to merge values that are set under such a branch and feed their uses.
 /// The phi nodes can be scalar or vector depending on the users of the value.
 /// This recipe works in concert with VPBranchOnMaskRecipe.
-class VPPredInstPHIRecipe : public VPRecipeBase {
-
+class VPPredInstPHIRecipe : public VPRecipeBase, public VPValue {
 public:
   /// Construct a VPPredInstPHIRecipe given \p PredInst whose value needs a phi
   /// nodes after merging back from a Branch-on-Mask.
-  VPPredInstPHIRecipe(VPValue *PredV) : VPRecipeBase(VPPredInstPHISC, PredV) {
-    new VPValue(PredV->getUnderlyingValue(), this);
-  }
+  VPPredInstPHIRecipe(VPValue *PredV)
+      : VPRecipeBase(VPPredInstPHISC, PredV),
+        VPValue(VPValue::VPVPredInstPHI, nullptr, this) {}
   ~VPPredInstPHIRecipe() override = default;
 
   /// Method to support type inquiry through isa, cast, and dyn_cast.
   static inline bool classof(const VPDef *D) {
     return D->getVPDefID() == VPRecipeBase::VPPredInstPHISC;
   }
 
   /// Generates phi nodes for live-outs as needed to retain SSA form.

llvm/lib/Transforms/Vectorize/VPlan.cpp

 
 Value *VPTransformState::get(VPValue *Def, const VPIteration &Instance) {
   if (!Def->getDef())
     return Def->getLiveInIRValue();
 
   if (hasScalarValue(Def, Instance))
     return Data.PerPartScalars[Def][Instance.Part][Instance.Lane];
 
-  if (hasVectorValue(Def, Instance.Part)) {
-    assert(Data.PerPartOutput.count(Def));
-    auto *VecPart = Data.PerPartOutput[Def][Instance.Part];
-    if (!VecPart->getType()->isVectorTy()) {
-      assert(Instance.Lane == 0 && "cannot get lane > 0 for scalar");
-      return VecPart;
-    }
-    // TODO: Cache created scalar values.
-    return Builder.CreateExtractElement(VecPart,
-                                        Builder.getInt32(Instance.Lane));
-  }
-  return Callback.getOrCreateScalarValue(VPValue2Value[Def], Instance);
+  assert(hasVectorValue(Def, Instance.Part));
+  auto *Extract =
+      Builder.CreateExtractElement(VecPart, Builder.getInt32(Instance.Lane));
+  // set(Def, Extract, Instance);
+  return Extract;
 }
 
 BasicBlock *
 VPBasicBlock::createEmptyBasicBlock(VPTransformState::CFGState &CFG) {
   // BB stands for IR BasicBlocks. VPBB stands for VPlan VPBasicBlocks.
   // Pred stands for Predessor. Prev stands for Previous - last visited/created.
   BasicBlock *PrevBB = CFG.PrevBB;
   BasicBlock *NewBB = BasicBlock::Create(PrevBB->getContext(), getName(),
void VPWidenGEPRecipe::print(raw_ostream &O, const Twine &Indent,
   O << " ";
   printAsOperand(O, SlotTracker);
   O << " = getelementptr ";
   printOperands(O, SlotTracker);
 }
 
 void VPWidenPHIRecipe::print(raw_ostream &O, const Twine &Indent,
                              VPSlotTracker &SlotTracker) const {
-  O << "WIDEN-PHI " << VPlanIngredient(Phi);
+  O << "WIDEN-PHI " << VPlanIngredient(getUnderlyingValue());
 }
 
 void VPBlendRecipe::print(raw_ostream &O, const Twine &Indent,
                           VPSlotTracker &SlotTracker) const {
   O << "BLEND ";
   Phi->printAsOperand(O, false);
   O << " =";
   if (getNumIncomingValues() == 1) {
void VPReplicateRecipe::print(raw_ostream &O, const Twine &Indent,
 
   if (AlsoPack)
     O << " (S->V)";
 }
 
 void VPPredInstPHIRecipe::print(raw_ostream &O, const Twine &Indent,
                                 VPSlotTracker &SlotTracker) const {
   O << "PHI-PREDICATED-INSTRUCTION ";
+  printAsOperand(O, SlotTracker);
+  O << " = ";
   printOperands(O, SlotTracker);
 }
 
 void VPWidenMemoryInstructionRecipe::print(raw_ostream &O, const Twine &Indent,
                                            VPSlotTracker &SlotTracker) const {
   O << "WIDEN ";
 
   if (!isStore()) {

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (auto I = VPBB->begin(), E = VPBB->end(); I != E;) {
 
       NewRecipe->insertBefore(Ingredient);
       if (NewRecipe->getNumDefinedValues() == 1)
         VPInst->replaceAllUsesWith(NewRecipe->getVPValue());
       else
         assert(NewRecipe->getNumDefinedValues() == 0 &&
                "Only recpies with zero or one defined values expected");
       Ingredient->eraseFromParent();
+      Plan->removeVPValueFor(Inst);
+      for (auto *Def : NewRecipe->definedValues()) {
+        Plan->addVPValue(Inst, Def);
+      }
     }
   }
 }

llvm/lib/Transforms/Vectorize/VPlanValue.h

public:
   const Value *getUnderlyingValue() const { return UnderlyingVal; }
 
   /// An enumeration for keeping track of the concrete subclass of VPValue that
   /// are actually instantiated. Values of this enumeration are kept in the
   /// SubclassID field of the VPValue objects. They are used for concrete
   /// type identification.
   enum {
     VPValueSC,
+    VPVBlendSC,
     VPVInstructionSC,
     VPVMemoryInstructionSC,
+    VPVPredInstPHI,
     VPVReductionSC,
     VPVReplicateSC,
     VPVWidenSC,
     VPVWidenCallSC,
     VPVWidenGEPSC,
+    VPVWidenIntOrFpIndcutionSC,
+    VPVWidenPHISC,
     VPVWidenSelectSC,
   };
 
   VPValue(Value *UV = nullptr, VPDef *Def = nullptr)
       : VPValue(VPValueSC, UV, Def) {}
   VPValue(const VPValue &) = delete;
   VPValue &operator=(const VPValue &) = delete;
 

llvm/test/Transforms/LoopVectorize/vplan-printing.ll

 ;
 ; CHECK:       N4 [label =
 ; CHECK-NEXT:    "pred.udiv.if:\n" +
 ; CHECK-NEXT:      "REPLICATE ir<%tmp4> = udiv ir<%n>, ir<%i> (S->V)\l"
 ; CHECK-NEXT:  ]
 ;
 ; CHECK:       N5 [label =
 ; CHECK-NEXT:    "pred.udiv.continue:\n" +
-; CHECK-NEXT:      "PHI-PREDICATED-INSTRUCTION ir<%tmp4>\l"
+; CHECK-NEXT:      "PHI-PREDICATED-INSTRUCTION vp<%3> = ir<%tmp4>\l"
 ; CHECK-NEXT:  ]
 ;
 ; CHECK:       N7 [label =
 ; CHECK-NEXT:    "for.inc:\n" +
 ; CHECK-NEXT:      "EMIT vp<%4> = not ir<%cmp>\l" +
-; CHECK-NEXT:      "BLEND %d = ir<0>/vp<%4> ir<%tmp4>/ir<%cmp>\l" +
+; CHECK-NEXT:      "BLEND %d = ir<0>/vp<%4> vp<%3>/ir<%cmp>\l" +
 ; CHECK-NEXT:      "CLONE ir<%idx> = getelementptr ir<%x>, ir<%i>\l" +
 ; CHECK-NEXT:      "WIDEN store ir<%idx>, ir<%d>\l"
 ; CHECK-NEXT:  ]
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %for.inc, %entry