[SLP]Improve findReusedOrderedScalars processing, NFCI.

Tries to simplify structural complexity of the findReusedOrderedScalars function.

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -3758,61 +3758,87 @@ BoUpSLP::findReusedOrderedScalars(const BoUpSLP::TreeEntry &TE) {
   OrdersType CurrentOrder(NumScalars, NumScalars);
   SmallVector<int> Positions;
   SmallBitVector UsedPositions(NumScalars);
-  const TreeEntry *STE = nullptr;
+  DenseMap<const TreeEntry *, unsigned> UsedEntries;
+  DenseMap<Value *, std::pair<const TreeEntry *, unsigned>> ValueToEntryPos;
+  for (Value *V : TE.Scalars) {
+    if (!isa<LoadInst, ExtractElementInst, ExtractValueInst>(V))
+      continue;
+    const auto *LocalSTE = getTreeEntry(V);
+    if (!LocalSTE)
+      continue;
+    unsigned Lane =
+        std::distance(LocalSTE->Scalars.begin(), find(LocalSTE->Scalars, V));
+    if (Lane >= NumScalars)
+      continue;
+    ++UsedEntries.try_emplace(LocalSTE, 0).first->getSecond();
+    ValueToEntryPos.try_emplace(V, LocalSTE, Lane);
+  }
+  if (UsedEntries.empty())
+    return std::nullopt;
+  const TreeEntry &BestSTE =
+      *std::max_element(UsedEntries.begin(), UsedEntries.end(),
+                        [](const std::pair<const TreeEntry *, unsigned> &P1,
+                           const std::pair<const TreeEntry *, unsigned> &P2) {
+                          return P1.second < P2.second;
+                        })
+           ->first;
+  UsedEntries.erase(&BestSTE);
+  const TreeEntry *SecondBestSTE = nullptr;
+  if (!UsedEntries.empty())
+    SecondBestSTE =
+        std::max_element(UsedEntries.begin(), UsedEntries.end(),
+                         [](const std::pair<const TreeEntry *, unsigned> &P1,
+                            const std::pair<const TreeEntry *, unsigned> &P2) {
+                           return P1.second < P2.second;
+                         })
+            ->first;
   // Try to find all gathered scalars that are gets vectorized in other
   // vectorize node. Here we can have only one single tree vector node to
   // correctly identify order of the gathered scalars.
   for (unsigned I = 0; I < NumScalars; ++I) {
     Value *V = TE.Scalars[I];
     if (!isa<LoadInst, ExtractElementInst, ExtractValueInst>(V))
       continue;
-    if (const auto *LocalSTE = getTreeEntry(V)) {
-      if (!STE)
-        STE = LocalSTE;
-      else if (STE != LocalSTE)
-        // Take the order only from the single vector node.
-        return std::nullopt;
-      unsigned Lane =
-          std::distance(STE->Scalars.begin(), find(STE->Scalars, V));
-      if (Lane >= NumScalars)
-        return std::nullopt;
-      if (CurrentOrder[Lane] != NumScalars) {
-        if (Lane != I)
-          continue;
-        UsedPositions.reset(CurrentOrder[Lane]);
-      }
-      // The partial identity (where only some elements of the gather node are
-      // in the identity order) is good.
-      CurrentOrder[Lane] = I;
-      UsedPositions.set(I);
+    const auto [LocalSTE, Lane] = ValueToEntryPos.lookup(V);
+    if (!LocalSTE || (LocalSTE != &BestSTE && LocalSTE != SecondBestSTE))
+      continue;
+    if (CurrentOrder[Lane] != NumScalars) {
+      if ((CurrentOrder[Lane] >= BestSTE.Scalars.size() ||
+           BestSTE.Scalars[CurrentOrder[Lane]] == V) &&
+          (Lane != I || LocalSTE == SecondBestSTE))
+        continue;
+      UsedPositions.reset(CurrentOrder[Lane]);
     }
+    // The partial identity (where only some elements of the gather node are
+    // in the identity order) is good.
+    CurrentOrder[Lane] = I;
+    UsedPositions.set(I);
   }
   // Need to keep the order if we have a vector entry and at least 2 scalars or
   // the vectorized entry has just 2 scalars.
-  if (STE && (UsedPositions.count() > 1 || STE->Scalars.size() == 2)) {
-    auto &&IsIdentityOrder = [NumScalars](ArrayRef<unsigned> CurrentOrder) {
-      for (unsigned I = 0; I < NumScalars; ++I)
-        if (CurrentOrder[I] != I && CurrentOrder[I] != NumScalars)
-          return false;
-      return true;
-    };
-    if (IsIdentityOrder(CurrentOrder))
-      return OrdersType();
-    auto *It = CurrentOrder.begin();
-    for (unsigned I = 0; I < NumScalars;) {
-      if (UsedPositions.test(I)) {
-        ++I;
-        continue;
-      }
-      if (*It == NumScalars) {
-        *It = I;
-        ++I;
-      }
-      ++It;
+  if (BestSTE.Scalars.size() != 2 && UsedPositions.count() <= 1)
+    return std::nullopt;
+  auto IsIdentityOrder = [&](ArrayRef<unsigned> CurrentOrder) {
+    for (unsigned I = 0; I < NumScalars; ++I)
+      if (CurrentOrder[I] != I && CurrentOrder[I] != NumScalars)
+        return false;
+    return true;
+  };
+  if (IsIdentityOrder(CurrentOrder))
+    return OrdersType();
+  auto *It = CurrentOrder.begin();
+  for (unsigned I = 0; I < NumScalars;) {
+    if (UsedPositions.test(I)) {
+      ++I;
+      continue;
+    }
+    if (*It == NumScalars) {
+      *It = I;
+      ++I;
     }
-    return std::move(CurrentOrder);
+    ++It;
   }
-  return std::nullopt;
+  return std::move(CurrentOrder);
 }
 
 namespace {