[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

@@ -19945,6 +19945,46 @@ static bool checkTreeSizes(ArrayRef<std::pair<unsigned, unsigned>> Sizes,
   return Dev * 96 / (Mean * Mean) == 0;
 }
 
+namespace {
+
+/// A reference to a store instruction and the distance of its address to a base
+/// pointer.
+struct InstAndPtrDist {
+  InstAndPtrDist(unsigned InstrIdx, int DistToBasePtr)
+      : InstrIdx(InstrIdx), DistToBasePtr(DistToBasePtr) {}
+  unsigned InstrIdx;
+  int DistToBasePtr;
+};
+struct PtrDistCompare {
+  bool operator()(const InstAndPtrDist &Op1, const InstAndPtrDist &Op2) const {
+    return Op1.DistToBasePtr < Op2.DistToBasePtr;
+  }
+};
+
+/// A group of instructions 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); }
+  void reset(unsigned NewBaseInstr) {
+    BaseInstrIdx = NewBaseInstr;
+    Instrs.clear();
+    Instrs.emplace(NewBaseInstr, 0);
+  }
+
+  // Note: PtrDistCompare ensures that there a 1:1 mapping between distance and
+  // instruction. This could be turned into a map to be more obvious.
+  using InstSet = std::set<InstAndPtrDist, PtrDistCompare>;
+
+  /// The index of the Base instruction, i.e. the one with a 0 pointer distance.
+  unsigned BaseInstrIdx;
+
+  /// Maps an instruction index to a pointer distance from \p BaseInstrIdx.
+  InstSet Instrs;
+};
+
+} // end anonymous namespace
+
 bool SLPVectorizerPass::vectorizeStores(
     ArrayRef<StoreInst *> Stores, BoUpSLP &R,
     DenseSet<std::tuple<Value *, Value *, Value *, Value *, unsigned>>
@@ -19954,31 +19994,21 @@ bool SLPVectorizerPass::vectorizeStores(
   BoUpSLP::ValueSet VectorizedStores;
   bool Changed = false;
 
-  struct StoreDistCompare {
-    bool operator()(const std::pair<unsigned, int> &Op1,
-                    const std::pair<unsigned, int> &Op2) const {
-      return Op1.second < Op2.second;
-    }
-  };
-  // 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) {
+  auto TryToVectorize = [&](const RelatedStoreInsts::InstSet &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 ||
@@ -20245,7 +20275,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<RelatedStoreInsts> 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
@@ -20279,56 +20310,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 (RelatedStoreInsts &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;
+      RelatedStoreInsts::InstSet PrevSet;
+      copy_if(StoreSeq.Instrs, std::inserter(PrevSet, PrevSet.end()),
+              [&](const InstAndPtrDist &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 InstAndPtrDist &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)) {
@@ -20338,17 +20364,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 (RelatedStoreInsts &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 (RelatedStoreInsts &StoreSeq : SortedStores)
+    TryToVectorize(StoreSeq.Instrs);
 
   return Changed;
 }