[SLP] Use named structs in vectorizeStores() (NFC)

This is a mostly straightforward replacement of the previous
std::pair<int, std::set<std::pair<...>>> data structure with slightly
more readable alternatives.

Author: gbossu

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -19923,31 +19923,51 @@ bool SLPVectorizerPass::vectorizeStores(
   BoUpSLP::ValueSet VectorizedStores;
   bool Changed = false;
 
+  /// A store instruction and the distance of its address to a base pointer.
+  struct CandidateInstr {
+    CandidateInstr(unsigned InstrIdx, int DistToBasePtr)
+        : InstrIdx(InstrIdx), DistToBasePtr(DistToBasePtr) {}
+    unsigned InstrIdx;
+    int DistToBasePtr;
+  };
   struct StoreDistCompare {
-    bool operator()(const std::pair<unsigned, int> &Op1,
-                    const std::pair<unsigned, int> &Op2) const {
-      return Op1.second < Op2.second;
+    bool operator()(const CandidateInstr &Op1,
+                    const CandidateInstr &Op2) const {
+      return Op1.DistToBasePtr < Op2.DistToBasePtr;
     }
   };
-  // A set of pairs (index of store in Stores array ref, Distance of the store
-  // address relative to base store address in units).
-  using StoreIndexToDistSet =
-      std::set<std::pair<unsigned, int>, StoreDistCompare>;
-  auto TryToVectorize = [&](const StoreIndexToDistSet &Set) {
+
+  /// A group of store instructions that we'll try to bundle together.
+  /// They are ordered using their signed distance to the address of this
+  /// set's BaseInstr.
+  struct CandidateBundle {
+    CandidateBundle(unsigned BaseInstrIdx) { reset(BaseInstrIdx); }
+    void reset(unsigned NewBaseInstr) {
+      BaseInstrIdx = NewBaseInstr;
+      Instrs.clear();
+      Instrs.emplace(NewBaseInstr, 0);
+    }
+    // TODO: This should probably just be an std::map
+    using CandidateSet = std::set<CandidateInstr, StoreDistCompare>;
+    unsigned BaseInstrIdx;
+    CandidateSet Instrs;
+  };
+
+  auto TryToVectorize = [&](const CandidateBundle::CandidateSet &StoreSeq) {
     int PrevDist = -1;
     BoUpSLP::ValueList Operands;
     // Collect the chain into a list.
-    for (auto [Idx, Data] : enumerate(Set)) {
-      if (Operands.empty() || Data.second - PrevDist == 1) {
-        Operands.push_back(Stores[Data.first]);
-        PrevDist = Data.second;
-        if (Idx != Set.size() - 1)
+    for (auto [Idx, Data] : enumerate(StoreSeq)) {
+      if (Operands.empty() || Data.DistToBasePtr - PrevDist == 1) {
+        Operands.push_back(Stores[Data.InstrIdx]);
+        PrevDist = Data.DistToBasePtr;
+        if (Idx != StoreSeq.size() - 1)
           continue;
       }
       auto E = make_scope_exit([&, &DataVar = Data]() {
         Operands.clear();
-        Operands.push_back(Stores[DataVar.first]);
-        PrevDist = DataVar.second;
+        Operands.push_back(Stores[DataVar.InstrIdx]);
+        PrevDist = DataVar.DistToBasePtr;
       });
 
       if (Operands.size() <= 1 ||
@@ -20214,7 +20234,8 @@ bool SLPVectorizerPass::vectorizeStores(
   // 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.
-  SmallVector<std::pair<unsigned, StoreIndexToDistSet>> SortedStores;
+  SmallVector<CandidateBundle> SortedStores;
+
   // Inserts the specified store SI with the given index Idx to the set of the
   // stores. If the store with the same distance is found already - stop
   // insertion, try to vectorize already found stores. If some stores from this
@@ -20248,56 +20269,51 @@ 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 (std::pair<unsigned, StoreIndexToDistSet> &Set : SortedStores) {
+    for (CandidateBundle &StoreSeq : SortedStores) {
       std::optional<int> Diff = getPointersDiff(
-          Stores[Set.first]->getValueOperand()->getType(),
-          Stores[Set.first]->getPointerOperand(),
+          Stores[StoreSeq.BaseInstrIdx]->getValueOperand()->getType(),
+          Stores[StoreSeq.BaseInstrIdx]->getPointerOperand(),
           SI->getValueOperand()->getType(), SI->getPointerOperand(), *DL, *SE,
           /*StrictCheck=*/true);
       if (!Diff)
         continue;
-      auto It = Set.second.find(std::make_pair(Idx, *Diff));
-      if (It == Set.second.end()) {
-        Set.second.emplace(Idx, *Diff);
+      auto It = StoreSeq.Instrs.find({Idx, *Diff});
+      if (It == StoreSeq.Instrs.end()) {
+        StoreSeq.Instrs.emplace(Idx, *Diff);
         return;
       }
       // Try to vectorize the first found set to avoid duplicate analysis.
-      TryToVectorize(Set.second);
-      unsigned ItIdx = It->first;
-      int ItDist = It->second;
-      StoreIndexToDistSet PrevSet;
-      copy_if(Set.second, std::inserter(PrevSet, PrevSet.end()),
-              [&](const std::pair<unsigned, int> &Pair) {
-                return Pair.first > ItIdx;
-              });
-      Set.second.clear();
-      Set.first = Idx;
-      Set.second.emplace(Idx, 0);
+      TryToVectorize(StoreSeq.Instrs);
+      unsigned ItIdx = It->InstrIdx;
+      int ItDist = It->DistToBasePtr;
+      CandidateBundle::CandidateSet PrevSet;
+      copy_if(StoreSeq.Instrs, std::inserter(PrevSet, PrevSet.end()),
+              [&](const CandidateInstr &I) { return I.InstrIdx > ItIdx; });
+      StoreSeq.reset(Idx);
       // Insert stores that followed previous match to try to vectorize them
       // with this store.
       unsigned StartIdx = ItIdx + 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 (const std::pair<unsigned, int> &Pair : reverse(PrevSet)) {
+      for (const CandidateInstr &Store : reverse(PrevSet)) {
         // Do not try to vectorize sequences, we already tried.
-        if (VectorizedStores.contains(Stores[Pair.first]))
+        if (VectorizedStores.contains(Stores[Store.InstrIdx]))
           break;
-        unsigned BI = Pair.first - StartIdx;
+        unsigned BI = Store.InstrIdx - StartIdx;
         UsedStores.set(BI);
-        Dists[BI] = Pair.second - ItDist;
+        Dists[BI] = Store.DistToBasePtr - ItDist;
       }
       for (unsigned I = StartIdx; I < Idx; ++I) {
         unsigned BI = I - StartIdx;
         if (UsedStores.test(BI))
-          Set.second.emplace(I, Dists[BI]);
+          StoreSeq.Instrs.emplace(I, Dists[BI]);
       }
       return;
     }
-    auto &Res = SortedStores.emplace_back();
-    Res.first = Idx;
-    Res.second.emplace(Idx, 0);
+    // We did not find a comparable store, start a new sequence.
+    SortedStores.emplace_back(Idx);
   };
   Type *PrevValTy = nullptr;
   for (auto [I, SI] : enumerate(Stores)) {
@@ -20307,17 +20323,17 @@ bool SLPVectorizerPass::vectorizeStores(
       PrevValTy = SI->getValueOperand()->getType();
     // Check that we do not try to vectorize stores of different types.
     if (PrevValTy != SI->getValueOperand()->getType()) {
-      for (auto &Set : SortedStores)
-        TryToVectorize(Set.second);
+      for (CandidateBundle &StoreSeq : SortedStores)
+        TryToVectorize(StoreSeq.Instrs);
       SortedStores.clear();
       PrevValTy = SI->getValueOperand()->getType();
     }
     FillStoresSet(I, SI);
   }
 
   // Final vectorization attempt.
-  for (auto &Set : SortedStores)
-    TryToVectorize(Set.second);
+  for (CandidateBundle &StoreSeq : SortedStores)
+    TryToVectorize(StoreSeq.Instrs);
 
   return Changed;
 }