[SLP] More OOP to simplify vectorizeStores() (NFC)

[gbossu]

This moves more code into the RelatedStoreInsts helper class. The FillStoresSet lambda is now only a couple of lines and is easier to read.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Gaëtan Bossu (gbossu)

Changes

This moves more code into the RelatedStoreInsts helper class. The FillStoresSet lambda is now only a couple of lines and is easier to read.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+80-57)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 31c684e16f051..1746a703a0e26 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20341,9 +20341,15 @@ namespace {
 /// A group of stores that we'll try to bundle together using vector ops.
 /// They are ordered using the signed distance of their address operand to the
 /// address of this group's BaseInstr.
-struct RelatedStoreInsts {
-  RelatedStoreInsts(unsigned BaseInstrIdx) { reset(BaseInstrIdx); }
+class RelatedStoreInsts {
+public:
+  RelatedStoreInsts(unsigned BaseInstrIdx, ArrayRef<StoreInst *> AllStores)
+      : AllStores(AllStores) {
+    reset(BaseInstrIdx);
+  }
+
   void reset(unsigned NewBaseInstr) {
+    assert(NewBaseInstr < AllStores.size());
     BaseInstrIdx = NewBaseInstr;
     Instrs.clear();
     insertOrLookup(NewBaseInstr, 0);
@@ -20358,12 +20364,54 @@ struct RelatedStoreInsts {
     return Inserted ? std::nullopt : std::optional<unsigned>(It->second);
   }
 
+  StoreInst &getBaseStore() const { return *AllStores[BaseInstrIdx]; }
+  using DistToInstMap = std::map<int, unsigned>;
+  const DistToInstMap &getStores() const { return Instrs; }
+
+  /// Recompute the pointer distances to be based on \p NewBaseInstIdx.
+  /// Stores whose index is less than \p MinSafeIdx will be dropped.
+  void rebase(unsigned MinSafeIdx, unsigned NewBaseInstIdx,
+              int DistFromCurBase) {
+    DistToInstMap PrevSet = std::move(Instrs);
+    reset(NewBaseInstIdx);
+
+    // Re-insert stores that come after MinSafeIdx to try and vectorize them
+    // again. Their distance will be "rebased" to use NewBaseInstIdx as
+    // reference.
+    for (auto [Dist, InstIdx] : PrevSet) {
+      if (InstIdx >= MinSafeIdx) {
+        insertOrLookup(InstIdx, Dist - DistFromCurBase);
+      }
+    }
+  }
+
+  /// Remove all stores that have been vectorized from this group.
+  void clearVectorizedStores(const BoUpSLP::ValueSet &VectorizedStores) {
+    const auto Begin = Instrs.begin();
+    auto NonVectorizedStore = Instrs.end();
+
+    while (NonVectorizedStore != Begin) {
+      const auto Prev = std::prev(NonVectorizedStore);
+      unsigned InstrIdx = Prev->second;
+      if (VectorizedStores.contains(AllStores[InstrIdx])) {
+        // NonVectorizedStore is the last scalar instruction.
+        // Erase all stores before it so we don't try to vectorize them again.
+        Instrs.erase(Begin, NonVectorizedStore);
+        return;
+      }
+      NonVectorizedStore = Prev;
+    }
+  }
+
+private:
   /// The index of the Base instruction, i.e. the one with a 0 pointer distance.
   unsigned BaseInstrIdx;
 
   /// Maps a pointer distance from \p BaseInstrIdx to an instruction index.
-  using DistToInstMap = std::map<int, unsigned>;
   DistToInstMap Instrs;
+
+  /// Reference to all the stores in the BB being analyzed.
+  ArrayRef<StoreInst *> AllStores;
 };
 
 } // end anonymous namespace
@@ -20651,14 +20699,7 @@ bool SLPVectorizerPass::vectorizeStores(
     }
   };
 
-  // Stores pair (first: index of the store into Stores array ref, address of
-  // which taken as base, second: sorted set of pairs {index, dist}, which are
-  // indices of stores in the set and their store location distances relative to
-  // the base address).
-
-  // Need to store the index of the very first store separately, since the set
-  // may be reordered after the insertion and the first store may be moved. This
-  // container allows to reduce number of calls of getPointersDiff() function.
+  /// Groups of stores to vectorize
   SmallVector<RelatedStoreInsts> SortedStores;
 
   // Inserts the specified store SI with the given index Idx to the set of the
@@ -20694,52 +20735,34 @@ bool SLPVectorizerPass::vectorizeStores(
   // dependencies and no need to waste compile time to try to vectorize them.
   // - Try to vectorize the sequence {1, {1, 0}, {3, 2}}.
   auto FillStoresSet = [&](unsigned Idx, StoreInst *SI) {
-    for (RelatedStoreInsts &StoreSeq : SortedStores) {
-      std::optional<int> Diff = getPointersDiff(
-          Stores[StoreSeq.BaseInstrIdx]->getValueOperand()->getType(),
-          Stores[StoreSeq.BaseInstrIdx]->getPointerOperand(),
-          SI->getValueOperand()->getType(), SI->getPointerOperand(), *DL, *SE,
-          /*StrictCheck=*/true);
-      if (!Diff)
-        continue;
-      std::optional<unsigned> PrevInst =
-          StoreSeq.insertOrLookup(/*InstrIdx=*/Idx, /*PtrDist=*/*Diff);
-      if (!PrevInst) {
-        // No store was associated to that distance. Keep collecting.
-        return;
-      }
-      // Try to vectorize the first found set to avoid duplicate analysis.
-      TryToVectorize(StoreSeq.Instrs);
-      RelatedStoreInsts::DistToInstMap PrevSet;
-      copy_if(StoreSeq.Instrs, std::inserter(PrevSet, PrevSet.end()),
-              [&](const std::pair<int, unsigned> &DistAndIdx) {
-                return DistAndIdx.second > *PrevInst;
-              });
-      StoreSeq.reset(Idx);
-      // Insert stores that followed previous match to try to vectorize them
-      // with this store.
-      unsigned StartIdx = *PrevInst + 1;
-      SmallBitVector UsedStores(Idx - StartIdx);
-      // Distances to previously found dup store (or this store, since they
-      // store to the same addresses).
-      SmallVector<int> Dists(Idx - StartIdx, 0);
-      for (auto [PtrDist, InstIdx] : reverse(PrevSet)) {
-        // Do not try to vectorize sequences, we already tried.
-        if (VectorizedStores.contains(Stores[InstIdx]))
-          break;
-        unsigned BI = InstIdx - StartIdx;
-        UsedStores.set(BI);
-        Dists[BI] = PtrDist - *Diff;
-      }
-      for (unsigned I = StartIdx; I < Idx; ++I) {
-        unsigned BI = I - StartIdx;
-        if (UsedStores.test(BI))
-          StoreSeq.insertOrLookup(I, Dists[BI]);
-      }
+    std::optional<int> Diff;
+    auto *RelatedStores =
+        find_if(SortedStores, [&](const RelatedStoreInsts &StoreSeq) {
+          StoreInst &BaseStore = StoreSeq.getBaseStore();
+          Diff = getPointersDiff(BaseStore.getValueOperand()->getType(),
+                                 BaseStore.getPointerOperand(),
+                                 SI->getValueOperand()->getType(),
+                                 SI->getPointerOperand(), *DL, *SE,
+                                 /*StrictCheck=*/true);
+          return Diff.has_value();
+        });
+
+    // We did not find a comparable store, start a new group.
+    if (RelatedStores == SortedStores.end()) {
+      SortedStores.emplace_back(Idx, Stores);
       return;
     }
-    // We did not find a comparable store, start a new sequence.
-    SortedStores.emplace_back(Idx);
+
+    // If there is already a store in the group with the same PtrDiff, try to
+    // vectorize the existing instructions before adding the current store.
+    if (std::optional<unsigned> PrevInst =
+            RelatedStores->insertOrLookup(Idx, *Diff)) {
+      TryToVectorize(RelatedStores->getStores());
+      RelatedStores->clearVectorizedStores(VectorizedStores);
+      RelatedStores->rebase(/*MinSafeIdx=*/*PrevInst + 1,
+                            /*NewBaseInstIdx=*/Idx,
+                            /*DistFromCurBase=*/*Diff);
+    }
   };
   Type *PrevValTy = nullptr;
   for (auto [I, SI] : enumerate(Stores)) {
@@ -20750,7 +20773,7 @@ bool SLPVectorizerPass::vectorizeStores(
     // Check that we do not try to vectorize stores of different types.
     if (PrevValTy != SI->getValueOperand()->getType()) {
       for (RelatedStoreInsts &StoreSeq : SortedStores)
-        TryToVectorize(StoreSeq.Instrs);
+        TryToVectorize(StoreSeq.getStores());
       SortedStores.clear();
       PrevValTy = SI->getValueOperand()->getType();
     }
@@ -20759,7 +20782,7 @@ bool SLPVectorizerPass::vectorizeStores(
 
   // Final vectorization attempt.
   for (RelatedStoreInsts &StoreSeq : SortedStores)
-    TryToVectorize(StoreSeq.Instrs);
+    TryToVectorize(StoreSeq.getStores());
 
   return Changed;
 }

[llvmbot]

@llvm/pr-subscribers-vectorizers

Author: Gaëtan Bossu (gbossu)

Changes

This moves more code into the RelatedStoreInsts helper class. The FillStoresSet lambda is now only a couple of lines and is easier to read.

---

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

1 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp (+80-57)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 31c684e16f051..1746a703a0e26 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -20341,9 +20341,15 @@ namespace {
 /// A group of stores that we'll try to bundle together using vector ops.
 /// They are ordered using the signed distance of their address operand to the
 /// address of this group's BaseInstr.
-struct RelatedStoreInsts {
-  RelatedStoreInsts(unsigned BaseInstrIdx) { reset(BaseInstrIdx); }
+class RelatedStoreInsts {
+public:
+  RelatedStoreInsts(unsigned BaseInstrIdx, ArrayRef<StoreInst *> AllStores)
+      : AllStores(AllStores) {
+    reset(BaseInstrIdx);
+  }
+
   void reset(unsigned NewBaseInstr) {
+    assert(NewBaseInstr < AllStores.size());
     BaseInstrIdx = NewBaseInstr;
     Instrs.clear();
     insertOrLookup(NewBaseInstr, 0);
@@ -20358,12 +20364,54 @@ struct RelatedStoreInsts {
     return Inserted ? std::nullopt : std::optional<unsigned>(It->second);
   }
 
+  StoreInst &getBaseStore() const { return *AllStores[BaseInstrIdx]; }
+  using DistToInstMap = std::map<int, unsigned>;
+  const DistToInstMap &getStores() const { return Instrs; }
+
+  /// Recompute the pointer distances to be based on \p NewBaseInstIdx.
+  /// Stores whose index is less than \p MinSafeIdx will be dropped.
+  void rebase(unsigned MinSafeIdx, unsigned NewBaseInstIdx,
+              int DistFromCurBase) {
+    DistToInstMap PrevSet = std::move(Instrs);
+    reset(NewBaseInstIdx);
+
+    // Re-insert stores that come after MinSafeIdx to try and vectorize them
+    // again. Their distance will be "rebased" to use NewBaseInstIdx as
+    // reference.
+    for (auto [Dist, InstIdx] : PrevSet) {
+      if (InstIdx >= MinSafeIdx) {
+        insertOrLookup(InstIdx, Dist - DistFromCurBase);
+      }
+    }
+  }
+
+  /// Remove all stores that have been vectorized from this group.
+  void clearVectorizedStores(const BoUpSLP::ValueSet &VectorizedStores) {
+    const auto Begin = Instrs.begin();
+    auto NonVectorizedStore = Instrs.end();
+
+    while (NonVectorizedStore != Begin) {
+      const auto Prev = std::prev(NonVectorizedStore);
+      unsigned InstrIdx = Prev->second;
+      if (VectorizedStores.contains(AllStores[InstrIdx])) {
+        // NonVectorizedStore is the last scalar instruction.
+        // Erase all stores before it so we don't try to vectorize them again.
+        Instrs.erase(Begin, NonVectorizedStore);
+        return;
+      }
+      NonVectorizedStore = Prev;
+    }
+  }
+
+private:
   /// The index of the Base instruction, i.e. the one with a 0 pointer distance.
   unsigned BaseInstrIdx;
 
   /// Maps a pointer distance from \p BaseInstrIdx to an instruction index.
-  using DistToInstMap = std::map<int, unsigned>;
   DistToInstMap Instrs;
+
+  /// Reference to all the stores in the BB being analyzed.
+  ArrayRef<StoreInst *> AllStores;
 };
 
 } // end anonymous namespace
@@ -20651,14 +20699,7 @@ bool SLPVectorizerPass::vectorizeStores(
     }
   };
 
-  // Stores pair (first: index of the store into Stores array ref, address of
-  // which taken as base, second: sorted set of pairs {index, dist}, which are
-  // indices of stores in the set and their store location distances relative to
-  // the base address).
-
-  // Need to store the index of the very first store separately, since the set
-  // may be reordered after the insertion and the first store may be moved. This
-  // container allows to reduce number of calls of getPointersDiff() function.
+  /// Groups of stores to vectorize
   SmallVector<RelatedStoreInsts> SortedStores;
 
   // Inserts the specified store SI with the given index Idx to the set of the
@@ -20694,52 +20735,34 @@ bool SLPVectorizerPass::vectorizeStores(
   // dependencies and no need to waste compile time to try to vectorize them.
   // - Try to vectorize the sequence {1, {1, 0}, {3, 2}}.
   auto FillStoresSet = [&](unsigned Idx, StoreInst *SI) {
-    for (RelatedStoreInsts &StoreSeq : SortedStores) {
-      std::optional<int> Diff = getPointersDiff(
-          Stores[StoreSeq.BaseInstrIdx]->getValueOperand()->getType(),
-          Stores[StoreSeq.BaseInstrIdx]->getPointerOperand(),
-          SI->getValueOperand()->getType(), SI->getPointerOperand(), *DL, *SE,
-          /*StrictCheck=*/true);
-      if (!Diff)
-        continue;
-      std::optional<unsigned> PrevInst =
-          StoreSeq.insertOrLookup(/*InstrIdx=*/Idx, /*PtrDist=*/*Diff);
-      if (!PrevInst) {
-        // No store was associated to that distance. Keep collecting.
-        return;
-      }
-      // Try to vectorize the first found set to avoid duplicate analysis.
-      TryToVectorize(StoreSeq.Instrs);
-      RelatedStoreInsts::DistToInstMap PrevSet;
-      copy_if(StoreSeq.Instrs, std::inserter(PrevSet, PrevSet.end()),
-              [&](const std::pair<int, unsigned> &DistAndIdx) {
-                return DistAndIdx.second > *PrevInst;
-              });
-      StoreSeq.reset(Idx);
-      // Insert stores that followed previous match to try to vectorize them
-      // with this store.
-      unsigned StartIdx = *PrevInst + 1;
-      SmallBitVector UsedStores(Idx - StartIdx);
-      // Distances to previously found dup store (or this store, since they
-      // store to the same addresses).
-      SmallVector<int> Dists(Idx - StartIdx, 0);
-      for (auto [PtrDist, InstIdx] : reverse(PrevSet)) {
-        // Do not try to vectorize sequences, we already tried.
-        if (VectorizedStores.contains(Stores[InstIdx]))
-          break;
-        unsigned BI = InstIdx - StartIdx;
-        UsedStores.set(BI);
-        Dists[BI] = PtrDist - *Diff;
-      }
-      for (unsigned I = StartIdx; I < Idx; ++I) {
-        unsigned BI = I - StartIdx;
-        if (UsedStores.test(BI))
-          StoreSeq.insertOrLookup(I, Dists[BI]);
-      }
+    std::optional<int> Diff;
+    auto *RelatedStores =
+        find_if(SortedStores, [&](const RelatedStoreInsts &StoreSeq) {
+          StoreInst &BaseStore = StoreSeq.getBaseStore();
+          Diff = getPointersDiff(BaseStore.getValueOperand()->getType(),
+                                 BaseStore.getPointerOperand(),
+                                 SI->getValueOperand()->getType(),
+                                 SI->getPointerOperand(), *DL, *SE,
+                                 /*StrictCheck=*/true);
+          return Diff.has_value();
+        });
+
+    // We did not find a comparable store, start a new group.
+    if (RelatedStores == SortedStores.end()) {
+      SortedStores.emplace_back(Idx, Stores);
       return;
     }
-    // We did not find a comparable store, start a new sequence.
-    SortedStores.emplace_back(Idx);
+
+    // If there is already a store in the group with the same PtrDiff, try to
+    // vectorize the existing instructions before adding the current store.
+    if (std::optional<unsigned> PrevInst =
+            RelatedStores->insertOrLookup(Idx, *Diff)) {
+      TryToVectorize(RelatedStores->getStores());
+      RelatedStores->clearVectorizedStores(VectorizedStores);
+      RelatedStores->rebase(/*MinSafeIdx=*/*PrevInst + 1,
+                            /*NewBaseInstIdx=*/Idx,
+                            /*DistFromCurBase=*/*Diff);
+    }
   };
   Type *PrevValTy = nullptr;
   for (auto [I, SI] : enumerate(Stores)) {
@@ -20750,7 +20773,7 @@ bool SLPVectorizerPass::vectorizeStores(
     // Check that we do not try to vectorize stores of different types.
     if (PrevValTy != SI->getValueOperand()->getType()) {
       for (RelatedStoreInsts &StoreSeq : SortedStores)
-        TryToVectorize(StoreSeq.Instrs);
+        TryToVectorize(StoreSeq.getStores());
       SortedStores.clear();
       PrevValTy = SI->getValueOperand()->getType();
     }
@@ -20759,7 +20782,7 @@ bool SLPVectorizerPass::vectorizeStores(
 
   // Final vectorization attempt.
   for (RelatedStoreInsts &StoreSeq : SortedStores)
-    TryToVectorize(StoreSeq.Instrs);
+    TryToVectorize(StoreSeq.getStores());
 
   return Changed;
 }

[gbossu]

@alexey-bataev As discussed in #132781, this is a follow-up to refactor more code into RelatedStoreInsts

[gbossu]

Now that FillStoresSet() is more concise and has simpler control flow, I'm wondering if this long comment on top of it could be shortened. What do you think @alexey-bataev?

[gbossu]

Hi @alexey-bataev, will you have some time for a quick review? 🙂

[gbossu]

Thanks for the initial review @alexey-bataev. Let me know if the refactoring itself sounds good to you, otherwise I'll fix the style issues and update the PR 🙂

[gbossu]

Any further comments @alexey-bataev? I've addressed everything so far, I'd be happy to get that committed 🙂

[alexey-bataev]

I don't like this approach. Prefer lambda instead of find_if with side effects.

[gbossu]

What would you suggest exactly? Wouldn't the lambda have side effects as well, just better hidden? I also don't really think that having a lambda in a lambda is very clear, especially if it implicitly captures.

[gbossu]

One compromise I would be happy with is:

    auto *RelatedStores =
        find_if(SortedStores, [&](const RelatedStoreInsts &StoreSeq) {
          return StoreSeq.getPointerDiff(*SI, *DL, *SE).has_value();
        });

    ...

    int PtrDist = RelatedStores->getPointerDiff(*SI, *DL, *SE).value();

This is much clearer and has no side effects. The only downside is that we would compute one extra pointer distance. I think this is negligeable though, as I don't believe that this part of the code is what dominates the runtime of the SLPVectorizer, and it's just one extra time anyway.

Let me know if you're happy with this, I don't see how a lambda would be clearer.

[alexey-bataev]

I mean, replace find_if completely by a loop in the lambda. Better to avoid calculating the distance twice.

[gbossu]

I do not want to wrap the code around a loop, as it makes it seem we will process the same store instruction multiple times, while we are just searching for a group of stores to insert it in. The code in this PR isn't perfect (and I'm not content with it), but at least it does communicate the intent better: we are just searching for a RelatedStoreInsts group to add this store to.

As I said, I'd be happy with the compromise in #134605 (comment), it is even clearer and there will be no impact to compile time, this isn't the hot part of the code, and anyway getPointerDiff relies on SCEV, which uses memoization.

Let me know what you think, maybe we'll come to an even better solution.

[sdesmalen-arm]

I don't see the issue with using find_if as long as the lambda is trivial, which it is. To avoid concerns with side-effects, maybe you could capture some of the values explicitly?

Similar to your suggestion above, I think adding a (simpler) pointer-diff method in RelatedStoreInsts would be useful, which tells whether a given Store is related (and if so, returns an offset). I was thinking something like bool isRelatedStore(const StoreInst *SI, int &PtrDiff) (+ args for SE/DL), but if you prefer returning it directly through std::optional I'm happy with that too.

[gbossu]

Rebased and added one fixup to explicitly capture the PtrDist side effect. I think this is a good compromise and I'd grateful if this change could be committed now. It's not really clear to me why there is this much push back against making the code ever so slightly more maintainable.

[sdesmalen-arm]

I don't see the issue with using find_if as long as the lambda is trivial, which it is. To avoid concerns with side-effects, maybe you could capture some of the values explicitly?

Similar to your suggestion above, I think adding a (simpler) pointer-diff method in RelatedStoreInsts would be useful, which tells whether a given Store is related (and if so, returns an offset). I was thinking something like bool isRelatedStore(const StoreInst *SI, int &PtrDiff) (+ args for SE/DL), but if you prefer returning it directly through std::optional I'm happy with that too.

[sdesmalen-arm]

Thanks for moving this functionality to separate methods in RelatedStores. That definitely helps clarify the code to me.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.