[LV] Make VPSymbolicValue more extensible (NFCI)

[MacDue]

This refactors the current VPSymbolicValue system so that symbolic values are created lazily and adding new values is simple. A new symbolic value just requires creating a new enum case, and returning its name in VPSymbolicValue::getName().

This is intended to simplify adding an AliasMask symbolic value in #182457.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Benjamin Maxwell (MacDue)

Changes

This refactors the current VPSymbolicValue system so that symbolic values are created lazily and adding new values is simple. A new symbolic value just requires creating a new enum case, and returning its name in VPSymbolicValue::getName().

This is intended to simplify adding an AliasMask symbolic value in #182457.

---

Full diff: https://github.com/llvm/llvm-project/pull/187041.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/VPlan.cpp (+45-46)
• (modified) llvm/lib/Transforms/Vectorize/VPlan.h (+32-25)
• (modified) llvm/lib/Transforms/Vectorize/VPlanValue.h (+33-7)

diff --git a/llvm/lib/Transforms/Vectorize/VPlan.cpp b/llvm/lib/Transforms/Vectorize/VPlan.cpp
index d914e9ce290be..dc759893dbcab 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -904,7 +904,8 @@ VPlan::~VPlan() {
   }
   for (VPValue *VPV : getLiveIns())
     delete VPV;
-  delete BackedgeTakenCount;
+  for (VPSymbolicValue *SV : getSymbolicValues())
+    delete SV;
 }
 
 VPIRBasicBlock *VPlan::getExitBlock(BasicBlock *IRBB) const {
@@ -1076,38 +1077,25 @@ const VPRegionBlock *VPlan::getVectorLoopRegion() const {
   return nullptr;
 }
 
+static SmallVector<VPSymbolicValue *> sortedSymbolicValues(const VPlan &Plan) {
+  // Sort symbolic values by their ID (enum order in VPSymbolicValue::ID).
+  SmallVector<VPSymbolicValue *> SortedValues(Plan.getSymbolicValues());
+  sort(SortedValues, [](VPSymbolicValue *A, VPSymbolicValue *B) {
+    return A->getID() < B->getID();
+  });
+  return SortedValues;
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPlan::printLiveIns(raw_ostream &O) const {
   VPSlotTracker SlotTracker(this);
 
-  if (VF.getNumUsers() > 0) {
-    O << "\nLive-in ";
-    VF.printAsOperand(O, SlotTracker);
-    O << " = VF";
-  }
-
-  if (UF.getNumUsers() > 0) {
-    O << "\nLive-in ";
-    UF.printAsOperand(O, SlotTracker);
-    O << " = UF";
-  }
-
-  if (VFxUF.getNumUsers() > 0) {
-    O << "\nLive-in ";
-    VFxUF.printAsOperand(O, SlotTracker);
-    O << " = VF * UF";
-  }
-
-  if (VectorTripCount.getNumUsers() > 0) {
-    O << "\nLive-in ";
-    VectorTripCount.printAsOperand(O, SlotTracker);
-    O << " = vector-trip-count";
-  }
-
-  if (BackedgeTakenCount && BackedgeTakenCount->getNumUsers()) {
-    O << "\nLive-in ";
-    BackedgeTakenCount->printAsOperand(O, SlotTracker);
-    O << " = backedge-taken count";
+  for (VPSymbolicValue *SV : sortedSymbolicValues(*this)) {
+    if (SV->getNumUsers() > 0) {
+      O << "\nLive-in ";
+      SV->printAsOperand(O, SlotTracker);
+      O << " = " << SV->getName();
+    }
   }
 
   O << "\n";
@@ -1230,14 +1218,8 @@ VPlan *VPlan::duplicate() {
   DenseMap<VPValue *, VPValue *> Old2NewVPValues;
   for (VPIRValue *OldLiveIn : getLiveIns())
     Old2NewVPValues[OldLiveIn] = NewPlan->getOrAddLiveIn(OldLiveIn);
-  Old2NewVPValues[&VectorTripCount] = &NewPlan->VectorTripCount;
-  Old2NewVPValues[&VF] = &NewPlan->VF;
-  Old2NewVPValues[&UF] = &NewPlan->UF;
-  Old2NewVPValues[&VFxUF] = &NewPlan->VFxUF;
-  if (BackedgeTakenCount) {
-    NewPlan->BackedgeTakenCount = new VPSymbolicValue();
-    Old2NewVPValues[BackedgeTakenCount] = NewPlan->BackedgeTakenCount;
-  }
+  for (VPSymbolicValue *SV : getSymbolicValues())
+    Old2NewVPValues[SV] = &NewPlan->getOrCreateSymbolicValue(SV->getID());
   if (auto *TripCountIRV = dyn_cast_or_null<VPIRValue>(TripCount))
     Old2NewVPValues[TripCountIRV] = NewPlan->getOrAddLiveIn(TripCountIRV);
   // else NewTripCount will be created and inserted into Old2NewVPValues when
@@ -1522,16 +1504,33 @@ void VPSlotTracker::assignName(const VPValue *V) {
   }
 }
 
+StringRef VPSymbolicValue::getName() const {
+  switch (ValueID) {
+  case VPSymbolicValue::VF:
+    return "VF";
+  case VPSymbolicValue::UF:
+    return "UF";
+  case VPSymbolicValue::VFxUF:
+    return "VF * UF";
+  case VPSymbolicValue::VectorTripCount:
+    return "vector-trip-count";
+  case VPSymbolicValue::BackedgeTakenCount:
+    return "backedge-taken count";
+  default:
+    llvm_unreachable("unamed symbolic value");
+  }
+}
+
 void VPSlotTracker::assignNames(const VPlan &Plan) {
-  if (Plan.VF.getNumUsers() > 0)
-    assignName(&Plan.VF);
-  if (Plan.UF.getNumUsers() > 0)
-    assignName(&Plan.UF);
-  if (Plan.VFxUF.getNumUsers() > 0)
-    assignName(&Plan.VFxUF);
-  assignName(&Plan.VectorTripCount);
-  if (Plan.BackedgeTakenCount)
-    assignName(Plan.BackedgeTakenCount);
+  for (VPSymbolicValue *SV : sortedSymbolicValues(Plan))
+    if (SV->getNumUsers() > 0)
+      assignName(SV);
+
+  // FIXME: Remove this and update the test expectations in
+  // llvm/unittests/Transforms/Vectorize. There is a lot of hardcoded
+  // expectations that assume the first slot starts at 1.
+  NextSlot = std::max(NextSlot, 1U);
+
   for (VPValue *LI : Plan.getLiveIns())
     assignName(LI);
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 80df058dfcf66..61af3c7c7bfdc 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -4587,22 +4587,6 @@ class VPlan {
   /// the tail.
   VPValue *TripCount = nullptr;
 
-  /// Represents the backedge taken count of the original loop, for folding
-  /// the tail. It equals TripCount - 1.
-  VPSymbolicValue *BackedgeTakenCount = nullptr;
-
-  /// Represents the vector trip count.
-  VPSymbolicValue VectorTripCount;
-
-  /// Represents the vectorization factor of the loop.
-  VPSymbolicValue VF;
-
-  /// Represents the unroll factor of the loop.
-  VPSymbolicValue UF;
-
-  /// Represents the loop-invariant VF * UF of the vector loop region.
-  VPSymbolicValue VFxUF;
-
   /// Contains all the external definitions created for this VPlan, as a mapping
   /// from IR Values to VPIRValues.
   SmallMapVector<Value *, VPIRValue *, 16> LiveIns;
@@ -4611,6 +4595,10 @@ class VPlan {
   /// VPlan is destroyed.
   SmallVector<VPBlockBase *> CreatedBlocks;
 
+  /// Contains all the symbolic values created for this plan. Mapping from the
+  /// Value ID to the symbolic value.
+  SmallDenseMap<VPSymbolicValue::ID, VPSymbolicValue *, 4> SymbolicValues;
+
   /// Construct a VPlan with \p Entry to the plan and with \p ScalarHeader
   /// wrapping the original header of the scalar loop.
   VPlan(VPBasicBlock *Entry, VPIRBasicBlock *ScalarHeader)
@@ -4724,26 +4712,43 @@ class VPlan {
     TripCount = NewTripCount;
   }
 
+  VPSymbolicValue &getOrCreateSymbolicValue(VPSymbolicValue::ID ID) {
+    auto [It, Inserted] = SymbolicValues.try_emplace(ID);
+    if (Inserted)
+      It->second = new VPSymbolicValue(ID);
+    return *It->second;
+  }
+
   /// The backedge taken count of the original loop.
   VPValue *getOrCreateBackedgeTakenCount() {
-    if (!BackedgeTakenCount)
-      BackedgeTakenCount = new VPSymbolicValue();
-    return BackedgeTakenCount;
+    return &getOrCreateSymbolicValue(VPSymbolicValue::BackedgeTakenCount);
+  }
+  VPValue *getBackedgeTakenCount() const {
+    return SymbolicValues.lookup(VPSymbolicValue::BackedgeTakenCount);
   }
-  VPValue *getBackedgeTakenCount() const { return BackedgeTakenCount; }
 
   /// The vector trip count.
-  VPSymbolicValue &getVectorTripCount() { return VectorTripCount; }
+  VPSymbolicValue &getVectorTripCount() {
+    return getOrCreateSymbolicValue(VPSymbolicValue::VectorTripCount);
+  }
 
   /// Returns the VF of the vector loop region.
-  VPSymbolicValue &getVF() { return VF; };
-  const VPSymbolicValue &getVF() const { return VF; };
+  VPSymbolicValue &getVF() {
+    return getOrCreateSymbolicValue(VPSymbolicValue::VF);
+  }
+  const VPSymbolicValue &getVF() const {
+    return const_cast<VPlan &>(*this).getVF();
+  }
 
   /// Returns the UF of the vector loop region.
-  VPSymbolicValue &getUF() { return UF; };
+  VPSymbolicValue &getUF() {
+    return getOrCreateSymbolicValue(VPSymbolicValue::UF);
+  }
 
   /// Returns VF * UF of the vector loop region.
-  VPSymbolicValue &getVFxUF() { return VFxUF; }
+  VPSymbolicValue &getVFxUF() {
+    return getOrCreateSymbolicValue(VPSymbolicValue::VFxUF);
+  }
 
   LLVMContext &getContext() const {
     return getScalarHeader()->getIRBasicBlock()->getContext();
@@ -4753,6 +4758,8 @@ class VPlan {
     return getScalarHeader()->getIRBasicBlock()->getDataLayout();
   }
 
+  auto getSymbolicValues() const { return SymbolicValues.values(); }
+
   void addVF(ElementCount VF) { VFs.insert(VF); }
 
   void setVF(ElementCount VF) {
diff --git a/llvm/lib/Transforms/Vectorize/VPlanValue.h b/llvm/lib/Transforms/Vectorize/VPlanValue.h
index 56e512849cb08..1228164460ae7 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanValue.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanValue.h
@@ -246,25 +246,51 @@ struct VPConstantInt : public VPIRValue {
 /// A symbolic live-in VPValue, used for values like vector trip count, VF, and
 /// VFxUF.
 struct VPSymbolicValue : public VPValue {
-  VPSymbolicValue() : VPValue(VPVSymbolicSC, nullptr) {}
+  enum ID : uint8_t {
+    /// Represents the vectorization factor of the loop.
+    VF,
+    /// Represents the unroll factor of the loop.
+    UF,
+    /// Represents the loop-invariant VF * UF of the vector loop region.
+    VFxUF,
+    /// Represents the vector trip count.
+    VectorTripCount,
+    /// Represents the backedge taken count of the original loop, for folding
+    /// the tail. It equals TripCount - 1.
+    BackedgeTakenCount,
+    /// Represents a placeholder value.
+    Placeholder,
+    /// Marks a VPSymbolicValue as materialized.
+    Materialized,
+  };
+
+  VPSymbolicValue() : VPSymbolicValue(Placeholder) {}
+  VPSymbolicValue(ID ID) : VPValue(VPVSymbolicSC, nullptr), ValueID{ID} {}
 
   static bool classof(const VPValue *V) {
     return V->getVPValueID() == VPVSymbolicSC;
   }
 
   /// Returns true if this symbolic value has been materialized.
-  bool isMaterialized() const { return Materialized; }
+  bool isMaterialized() const { return ValueID == Materialized; }
 
   /// Mark this symbolic value as materialized.
   void markMaterialized() {
-    assert(!Materialized && "VPSymbolicValue already materialized");
-    Materialized = true;
+    assert(ValueID != Materialized && "VPSymbolicValue already materialized");
+    ValueID = Materialized;
   }
 
+  /// Returns the name for this symbolic value.
+  StringRef getName() const;
+
+  /// Returns the value ID for this symbolic value.
+  VPSymbolicValue::ID getID() const { return VPSymbolicValue::ID(ValueID); }
+
 private:
-  /// Track whether this symbolic value has been materialized (replaced).
-  /// After materialization, accessing users should trigger an assertion.
-  bool Materialized = false;
+  /// The value ID for this symbolic value. This is set to "Materialized" after
+  /// the value has been materialized (replaced). After materialization,
+  /// accessing users should trigger an assertion.
+  VPSymbolicValue::ID ValueID;
 };
 
 /// A VPValue defined by a recipe that produces one or more values.