[VPlan] Introduce code to limit querying VPValues using IR references.

[fhahn]

Summary

After applying VPlan-to-VPlan transformations, using IR references to
query VPlan values may be incorrect, as the IR is not in sync with the
VPlan any longer.

To better detect such mis-matches, this patch introduces a new flag to
VPlans to indicate whether it is safe to query VPValues using IR values.

getVPValue is updated to assert if it is called when the flag indicates
it is not safe any longer.

There is an escape hatch via an extra argument, because there are 3
places that need to be fixed first. Those are

1. truncateToMinimalBitwidths
2. clearReductionWrapFlags
3. fixLCSSAPHIs

As a first step, this flag will help preventing new code from violating
this property.

Any suggestions with respect to naming very welcome!

[Ayal]

This follows-up on https://reviews.llvm.org/rGd995d63767624a60a5d3276f9f16d7b995435af1#inline-6235

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9489

Perhaps refactor to have buildVPlanWithVPRecipes() take care of building the initial unoptimized VPlan only, ending with marking IRToVPlan disallowed. I.e., invoke sinkScalarOperands() and mergeReplicateRegion() optimizations after buildVPlanWithVPRecipes() returns?

llvm/lib/Transforms/Vectorize/VPlan.h

2150

disableGetVPValue()?

Some related thoughts: should addVPValue() and removeVPValue() also be disallowed, i.e., disable all accesses to Value2VPValue? Conceptually, Value2VPValue can be destroyed, and potentially another type of VPlan which does not support Value2VPValue created and provided instead.

[fhahn]

Thanks for taking a look! I changed the name to Value2VPValueOutdated & added assertions to addVPValueFor and removeVPValue.

[fhahn]

Adjust naming again to`Value2VPValueEnabled`.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9489

I'm not sure, introduction of FirstOrderRecurrenceSplice could also be moved after disallowing Value2VPValue access, but it is not an optimization. Should we do the potential split separately, to keep the diff smaller?

llvm/lib/Transforms/Vectorize/VPlan.h

2150

Some related thoughts: should addVPValue() and removeVPValue() also be disallowed, i.e., disable all accesses to Value2VPValue

Sounds good! I added assertions there and renamed the code to use disableValue2VPValue/Value2VPValueEnabled.

Conceptually, Value2VPValue can be destroyed, and potentially another type of VPlan which does not support Value2VPValue created and provided instead.

Agreed, once the offending pieces of code have been removed, we can clear the map instead.

[Ayal]

This looks good to me, thanks for following-up!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9489

Sure, refactoring can be done separately. And indeed may include non-optimization stages that are pure VPlan2VPlan transformations.

llvm/lib/Transforms/Vectorize/VPlan.h

2193

Note that calling getOrAddVPValue() with OverrideAllowed allows to get a VPValue but not add one. This is fine, as the override is temporary anyhow.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::truncateToMinimalBitwidths(VPTransformState &State) {
   // For every instruction `I` in MinBWs, truncate the operands, create a
   // truncated version of `I` and reextend its result. InstCombine runs
   // later and will remove any ext/trunc pairs.
   SmallPtrSet<Value *, 4> Erased;
   for (const auto &KV : Cost->getMinimalBitwidths()) {
     // If the value wasn't vectorized, we must maintain the original scalar
     // type. The absence of the value from State indicates that it
     // wasn't vectorized.
-    VPValue *Def = State.Plan->getVPValue(KV.first);
+    // FIXME: Should not rely on getVPValue at this point.
+    VPValue *Def = State.Plan->getVPValue(KV.first, true);
     if (!State.hasAnyVectorValue(Def))
       continue;
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *I = State.get(Def, Part);
       if (Erased.count(I) || I->use_empty() || !isa<Instruction>(I))
         continue;
       Type *OriginalTy = I->getType();
       Type *ScalarTruncatedTy =
for (const auto &KV : Cost->getMinimalBitwidths()) {
     }
   }
 
   // We'll have created a bunch of ZExts that are now parentless. Clean up.
   for (const auto &KV : Cost->getMinimalBitwidths()) {
     // If the value wasn't vectorized, we must maintain the original scalar
     // type. The absence of the value from State indicates that it
     // wasn't vectorized.
-    VPValue *Def = State.Plan->getVPValue(KV.first);
+    // FIXME: Should not rely on getVPValue at this point.
+    VPValue *Def = State.Plan->getVPValue(KV.first, true);
     if (!State.hasAnyVectorValue(Def))
       continue;
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *I = State.get(Def, Part);
       ZExtInst *Inst = dyn_cast<ZExtInst>(I);
       if (Inst && Inst->use_empty()) {
         Value *NewI = Inst->getOperand(0);
         Inst->eraseFromParent();
void InnerLoopVectorizer::clearReductionWrapFlags(const RecurrenceDescriptor &RdxDesc,
   SmallPtrSet<Instruction *, 8> Visited;
   Worklist.push_back(LoopExitInstr);
   Visited.insert(LoopExitInstr);
 
   while (!Worklist.empty()) {
     Instruction *Cur = Worklist.pop_back_val();
     if (isa<OverflowingBinaryOperator>(Cur))
       for (unsigned Part = 0; Part < UF; ++Part) {
-        Value *V = State.get(State.Plan->getVPValue(Cur), Part);
+        // FIXME: Should not rely on getVPValue at this point.
+        Value *V = State.get(State.Plan->getVPValue(Cur, true), Part);
         cast<Instruction>(V)->dropPoisonGeneratingFlags();
       }
 
     for (User *U : Cur->users()) {
       Instruction *UI = cast<Instruction>(U);
       if ((Cur != LoopExitInstr || OrigLoop->contains(UI->getParent())) &&
           Visited.insert(UI).second)
         Worklist.push_back(UI);
for (PHINode &LCSSAPhi : LoopExitBlock->phis()) {
     VPLane Lane = VPLane::getFirstLane();
     if (isa<Instruction>(IncomingValue) &&
         !Cost->isUniformAfterVectorization(cast<Instruction>(IncomingValue),
                                            VF))
       Lane = VPLane::getLastLaneForVF(VF);
 
     // Can be a loop invariant incoming value or the last scalar value to be
     // extracted from the vectorized loop.
+    // FIXME: Should not rely on getVPValue at this point.
     Builder.SetInsertPoint(LoopMiddleBlock->getTerminator());
     Value *lastIncomingValue =
         OrigLoop->isLoopInvariant(IncomingValue)
             ? IncomingValue
-            : State.get(State.Plan->getVPValue(IncomingValue),
+            : State.get(State.Plan->getVPValue(IncomingValue, true),
                         VPIteration(UF - 1, Lane));
     LCSSAPhi.addIncoming(lastIncomingValue, LoopMiddleBlock);
   }
 }
 
 void InnerLoopVectorizer::sinkScalarOperands(Instruction *PredInst) {
   // The basic block and loop containing the predicated instruction.
   auto *PredBB = PredInst->getParent();
for (unsigned i = 0; i < IG->getFactor(); ++i)
           Plan->addVPValue(Member, VPIG->getVPValue(J));
           OriginalV->replaceAllUsesWith(VPIG->getVPValue(J));
           J++;
         }
         RecipeBuilder.getRecipe(Member)->eraseFromParent();
       }
   }
 
+  // From this point onwards, VPlan-to-VPlan transformations may change the plan
+  // in ways that accessing values using original IR values is incorrect.
+  Plan->disableValue2VPValue();
+
   VPlanTransforms::sinkScalarOperands(*Plan);
   VPlanTransforms::mergeReplicateRegions(*Plan);
 
   std::string PlanName;
   raw_string_ostream RSO(PlanName);
   ElementCount VF = Range.Start;
   Plan->addVF(VF);
   RSO << "Initial VPlan for VF={" << VF;

llvm/lib/Transforms/Vectorize/VPlan.h

class VPlan {
 
   /// Contains all VPValues that been allocated by addVPValue directly and need
   /// to be free when the plan's destructor is called.
   SmallVector<VPValue *, 16> VPValuesToFree;
 
   /// Holds the VPLoopInfo analysis for this VPlan.
   VPLoopInfo VPLInfo;
 
+  /// Indicates whether it is safe use the Value2VPValue mapping or if the
+  /// mapping cannot be used any longer, because it is stale.
+  bool Value2VPValueEnabled = true;
+
 public:
   VPlan(VPBlockBase *Entry = nullptr) : Entry(Entry) {
     if (Entry)
       Entry->setPlan(this);
   }
 
   ~VPlan() {
     if (Entry) {
public:
 
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {
     if (!BackedgeTakenCount)
       BackedgeTakenCount = new VPValue();
     return BackedgeTakenCount;
   }
 
+  /// Mark the plan to indicate that using Value2VPValue is not safe any
+  /// longer, because it may be stale.
+  void disableValue2VPValue() { Value2VPValueEnabled = false; }
+
   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); }
 
   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;
     VPValuesToFree.push_back(VPV);
   }
 
   void addVPValue(Value *V, VPValue *VPV) {
+    assert(Value2VPValueEnabled && "Value2VPValue 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");
     Value2VPValue[V] = VPV;
   }
 
-  VPValue *getVPValue(Value *V) {
+  /// Returns the VPValue for \p V. \p OverrideAllowed can be used to disable
+  /// checking whether it is safe to query VPValues using IR Values.
+  VPValue *getVPValue(Value *V, bool OverrideAllowed = false) {
+    assert((OverrideAllowed || isa<Constant>(V) || Value2VPValueEnabled) &&
+           "Value2VPValue mapping may be out of date!");
     assert(V && "Trying to get the VPValue of a null Value");
     assert(Value2VPValue.count(V) && "Value does not exist in VPlan");
     return Value2VPValue[V];
   }
 
-  VPValue *getOrAddVPValue(Value *V) {
+  /// Gets the VPValue or adds a new one (if none exists yet) for \p V. \p
+  /// OverrideAllowed can be used to disable checking whether it is safe to
+  /// query VPValues using IR Values.
+  VPValue *getOrAddVPValue(Value *V, bool OverrideAllowed = false) {
+    assert((OverrideAllowed || isa<Constant>(V) || Value2VPValueEnabled) &&
+           "Value2VPValue mapping may be out of date!");
     assert(V && "Trying to get or add the VPValue of a null Value");
     if (!Value2VPValue.count(V))
       addVPValue(V);
     return getVPValue(V);
   }
 
-  void removeVPValueFor(Value *V) { Value2VPValue.erase(V); }
+  void removeVPValueFor(Value *V) {
+    assert(Value2VPValueEnabled &&
+           "IR value to VPValue mapping may be out of date!");
+    Value2VPValue.erase(V);
+  }
 
   /// Return the VPLoopInfo analysis for this VPlan.
   VPLoopInfo &getVPLoopInfo() { return VPLInfo; }
   const VPLoopInfo &getVPLoopInfo() const { return VPLInfo; }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print this VPlan to \p O.
   void print(raw_ostream &O) const;