[SandboxVec][BottomUpVec] Separate vectorization decisions from code generation (#127727)

Up until now the generation of vector instructions was taking place
during the top-down post-order traversal of vectorizeRec(). The issue
with this approach is that the vector instructions emitted during the
traversal can be reordered by the scheduler, making it challenging to
place them without breaking the def-before-uses rule.

With this patch we separate the vectorization decisions (done in
`vectorizeRec()`) from the code generation phase (`emitVectors()`). The
vectorization decisions are stored in the `Actions` vector and are used
by `emitVectors()` to drive code generation.

Author: vporpo

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/InstrMaps.h

@@ -23,57 +23,54 @@
 
 namespace llvm::sandboxir {
 
+class LegalityResult;
+
+struct Action {
+  unsigned Idx = 0;
+  const LegalityResult *LegalityRes = nullptr;
+  SmallVector<Value *, 4> Bndl;
+  SmallVector<Value *> UserBndl;
+  unsigned Depth;
+  SmallVector<Action *> Operands;
+  Value *Vec = nullptr;
+  Action(const LegalityResult *LR, ArrayRef<Value *> B, ArrayRef<Value *> UB,
+         unsigned Depth)
+      : LegalityRes(LR), Bndl(B), UserBndl(UB), Depth(Depth) {}
+#ifndef NDEBUG
+  void print(raw_ostream &OS) const;
+  void dump() const;
+  friend raw_ostream &operator<<(raw_ostream &OS, const Action &A) {
+    A.print(OS);
+    return OS;
+  }
+#endif // NDEBUG
+};
+
 /// Maps the original instructions to the vectorized instrs and the reverse.
 /// For now an original instr can only map to a single vector.
 class InstrMaps {
   /// A map from the original values that got combined into vectors, to the
-  /// vector value(s).
-  DenseMap<Value *, Value *> OrigToVectorMap;
-  /// A map from the vector value to a map of the original value to its lane.
+  /// vectorization Action.
+  DenseMap<Value *, Action *> OrigToVectorMap;
+  /// A map from the vec Action to a map of the original value to its lane.
   /// Please note that for constant vectors, there may multiple original values
   /// with the same lane, as they may be coming from vectorizing different
   /// original values.
-  DenseMap<Value *, DenseMap<Value *, unsigned>> VectorToOrigLaneMap;
-  Context &Ctx;
+  DenseMap<Action *, DenseMap<Value *, unsigned>> VectorToOrigLaneMap;
   std::optional<Context::CallbackID> EraseInstrCB;
 
-private:
-  void notifyEraseInstr(Value *V) {
-    // We don't know if V is an original or a vector value.
-    auto It = OrigToVectorMap.find(V);
-    if (It != OrigToVectorMap.end()) {
-      // V is an original value.
-      // Remove it from VectorToOrigLaneMap.
-      Value *Vec = It->second;
-      VectorToOrigLaneMap[Vec].erase(V);
-      // Now erase V from OrigToVectorMap.
-      OrigToVectorMap.erase(It);
-    } else {
-      // V is a vector value.
-      // Go over the original values it came from and remove them from
-      // OrigToVectorMap.
-      for (auto [Orig, Lane] : VectorToOrigLaneMap[V])
-        OrigToVectorMap.erase(Orig);
-      // Now erase V from VectorToOrigLaneMap.
-      VectorToOrigLaneMap.erase(V);
-    }
-  }
-
 public:
-  InstrMaps(Context &Ctx) : Ctx(Ctx) {
-    EraseInstrCB = Ctx.registerEraseInstrCallback(
-        [this](Instruction *I) { notifyEraseInstr(I); });
-  }
-  ~InstrMaps() { Ctx.unregisterEraseInstrCallback(*EraseInstrCB); }
+  InstrMaps() = default;
+  ~InstrMaps() = default;
   /// \Returns the vector value that we got from vectorizing \p Orig, or
   /// nullptr if not found.
-  Value *getVectorForOrig(Value *Orig) const {
+  Action *getVectorForOrig(Value *Orig) const {
     auto It = OrigToVectorMap.find(Orig);
     return It != OrigToVectorMap.end() ? It->second : nullptr;
   }
   /// \Returns the lane of \p Orig before it got vectorized into \p Vec, or
   /// nullopt if not found.
-  std::optional<unsigned> getOrigLane(Value *Vec, Value *Orig) const {
+  std::optional<unsigned> getOrigLane(Action *Vec, Value *Orig) const {
     auto It1 = VectorToOrigLaneMap.find(Vec);
     if (It1 == VectorToOrigLaneMap.end())
       return std::nullopt;
@@ -84,7 +81,7 @@ class InstrMaps {
     return It2->second;
   }
   /// Update the map to reflect that \p Origs got vectorized into \p Vec.
-  void registerVector(ArrayRef<Value *> Origs, Value *Vec) {
+  void registerVector(ArrayRef<Value *> Origs, Action *Vec) {
     auto &OrigToLaneMap = VectorToOrigLaneMap[Vec];
     unsigned Lane = 0;
     for (Value *Orig : Origs) {

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h

@@ -17,6 +17,7 @@
 #include "llvm/IR/DataLayout.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrMaps.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
 
 namespace llvm::sandboxir {
@@ -206,30 +207,30 @@ class Widen final : public LegalityResult {
 
 class DiamondReuse final : public LegalityResult {
   friend class LegalityAnalysis;
-  Value *Vec;
-  DiamondReuse(Value *Vec)
+  Action *Vec;
+  DiamondReuse(Action *Vec)
       : LegalityResult(LegalityResultID::DiamondReuse), Vec(Vec) {}
 
 public:
   static bool classof(const LegalityResult *From) {
     return From->getSubclassID() == LegalityResultID::DiamondReuse;
   }
-  Value *getVector() const { return Vec; }
+  Action *getVector() const { return Vec; }
 };
 
 class DiamondReuseWithShuffle final : public LegalityResult {
   friend class LegalityAnalysis;
-  Value *Vec;
+  Action *Vec;
   ShuffleMask Mask;
-  DiamondReuseWithShuffle(Value *Vec, const ShuffleMask &Mask)
+  DiamondReuseWithShuffle(Action *Vec, const ShuffleMask &Mask)
       : LegalityResult(LegalityResultID::DiamondReuseWithShuffle), Vec(Vec),
         Mask(Mask) {}
 
 public:
   static bool classof(const LegalityResult *From) {
     return From->getSubclassID() == LegalityResultID::DiamondReuseWithShuffle;
   }
-  Value *getVector() const { return Vec; }
+  Action *getVector() const { return Vec; }
   const ShuffleMask &getMask() const { return Mask; }
 };
 
@@ -250,18 +251,18 @@ class CollectDescr {
   /// Describes how to get a value element. If the value is a vector then it
   /// also provides the index to extract it from.
   class ExtractElementDescr {
-    Value *V;
+    PointerUnion<Action *, Value *> V = nullptr;
     /// The index in `V` that the value can be extracted from.
-    /// This is nullopt if we need to use `V` as a whole.
-    std::optional<int> ExtractIdx;
+    int ExtractIdx = 0;
 
   public:
-    ExtractElementDescr(Value *V, int ExtractIdx)
+    ExtractElementDescr(Action *V, int ExtractIdx)
         : V(V), ExtractIdx(ExtractIdx) {}
-    ExtractElementDescr(Value *V) : V(V), ExtractIdx(std::nullopt) {}
-    Value *getValue() const { return V; }
-    bool needsExtract() const { return ExtractIdx.has_value(); }
-    int getExtractIdx() const { return *ExtractIdx; }
+    ExtractElementDescr(Value *V) : V(V) {}
+    Action *getValue() const { return cast<Action *>(V); }
+    Value *getScalar() const { return cast<Value *>(V); }
+    bool needsExtract() const { return isa<Action *>(V); }
+    int getExtractIdx() const { return ExtractIdx; }
   };
 
   using DescrVecT = SmallVector<ExtractElementDescr, 4>;
@@ -272,11 +273,11 @@ class CollectDescr {
       : Descrs(std::move(Descrs)) {}
   /// If all elements come from a single vector input, then return that vector
   /// and also the shuffle mask required to get them in order.
-  std::optional<std::pair<Value *, ShuffleMask>> getSingleInput() const {
+  std::optional<std::pair<Action *, ShuffleMask>> getSingleInput() const {
     const auto &Descr0 = *Descrs.begin();
-    Value *V0 = Descr0.getValue();
     if (!Descr0.needsExtract())
       return std::nullopt;
+    auto *V0 = Descr0.getValue();
     ShuffleMask::IndicesVecT MaskIndices;
     MaskIndices.push_back(Descr0.getExtractIdx());
     for (const auto &Descr : drop_begin(Descrs)) {

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h

@@ -58,9 +58,33 @@ class BottomUpVec final : public RegionPass {
   /// function helps collect these instructions (along with the pointer operands
   /// for loads/stores) so that they can be cleaned up later.
   void collectPotentiallyDeadInstrs(ArrayRef<Value *> Bndl);
-  /// Recursively try to vectorize \p Bndl and its operands.
-  Value *vectorizeRec(ArrayRef<Value *> Bndl, ArrayRef<Value *> UserBndl,
-                      unsigned Depth);
+
+  /// Helper class describing how(if) to vectorize the code.
+  class ActionsVector {
+  private:
+    SmallVector<std::unique_ptr<Action>, 16> Actions;
+
+  public:
+    auto begin() const { return Actions.begin(); }
+    auto end() const { return Actions.end(); }
+    void push_back(std::unique_ptr<Action> &&ActPtr) {
+      ActPtr->Idx = Actions.size();
+      Actions.push_back(std::move(ActPtr));
+    }
+    void clear() { Actions.clear(); }
+#ifndef NDEBUG
+    void print(raw_ostream &OS) const;
+    void dump() const;
+#endif // NDEBUG
+  };
+  ActionsVector Actions;
+  /// Recursively try to vectorize \p Bndl and its operands. This populates the
+  /// `Actions` vector.
+  Action *vectorizeRec(ArrayRef<Value *> Bndl, ArrayRef<Value *> UserBndl,
+                       unsigned Depth);
+  /// Generate vector instructions based on `Actions` and return the last vector
+  /// created.
+  Value *emitVectors();
   /// Entry point for vectorization starting from \p Seeds.
   bool tryVectorize(ArrayRef<Value *> Seeds);
 

llvm/lib/Transforms/Vectorize/SandboxVectorizer/InstrMaps.cpp

@@ -8,10 +8,23 @@
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/InstrMaps.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h"
 
 namespace llvm::sandboxir {
 
 #ifndef NDEBUG
+void Action::print(raw_ostream &OS) const {
+  OS << Idx << ". " << *LegalityRes << " Depth:" << Depth << "\n";
+  OS.indent(2) << "Bndl:\n";
+  for (Value *V : Bndl)
+    OS.indent(4) << *V << "\n";
+  OS.indent(2) << "UserBndl:\n";
+  for (Value *V : UserBndl)
+    OS.indent(4) << *V << "\n";
+}
+
+void Action::dump() const { print(dbgs()); }
+
 void InstrMaps::dump() const {
   print(dbgs());
   dbgs() << "\n";