[SandboxVec][BottomUpVec] Separate vectorization decisions from code generation

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";