From 2b896fc0d5bf5091c1933bd7f5bbeeccca031d46 Mon Sep 17 00:00:00 2001
From: Alexey Bataev <a.bataev@outlook.com>
Date: Fri, 23 Apr 2021 12:33:41 -0700
Subject: [PATCH] [SLP]Improved isGatherShuffledEntry, NFC.

Reworked isGatherShuffledEntry function, simplified and moved
common code to the lambda (it shall go away when non-power-2 patch will
be landed).
---
 lib/Transforms/Vectorize/SLPVectorizer.cpp | 98 +++++++++++-----------
 1 file changed, 48 insertions(+), 50 deletions(-)

diff --git a/lib/Transforms/Vectorize/SLPVectorizer.cpp b/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 2eb5b8d0a17..900df48c80d 100644
--- a/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -4296,70 +4296,68 @@ InstructionCost BoUpSLP::getTreeCost() {
 Optional<TargetTransformInfo::ShuffleKind>
 BoUpSLP::isGatherShuffledEntry(const TreeEntry *TE, SmallVectorImpl<int> &Mask,
                                SmallVectorImpl<const TreeEntry *> &Entries) {
-  auto *VLIt = find_if(VectorizableTree,
-                       [TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
-                         return EntryPtr.get() == TE;
-                       });
-  assert(VLIt != VectorizableTree.end() &&
-         "Gathered values should be in the tree.");
   Mask.assign(TE->Scalars.size(), UndefMaskElem);
   Entries.clear();
-  DenseMap<const TreeEntry *, int> Used;
-  int NumShuffles = 0;
+  DenseMap<Value *, const TreeEntry *> UsedValuesEntry;
+  unsigned VF = 0;
+  // FIXME: Shall be replaced by GetVF function once non-power-2 patch is
+  // landed.
+  auto &&GetVF = [](const TreeEntry *TE) {
+    if (!TE->ReuseShuffleIndices.empty())
+      return TE->ReuseShuffleIndices.size();
+    return TE->Scalars.size();
+  };
   for (int I = 0, E = TE->Scalars.size(); I < E; ++I) {
     Value *V = TE->Scalars[I];
     if (isa<UndefValue>(V))
       continue;
-    const TreeEntry *VTE = getTreeEntry(V);
+    const TreeEntry *VTE = UsedValuesEntry.lookup(V);
     if (!VTE) {
-      // Check if it is used in one of the gathered entries.
-      const auto *It =
-          find_if(make_range(VectorizableTree.begin(), VLIt),
-                  [V](const std::unique_ptr<TreeEntry> &EntryPtr) {
-                    return EntryPtr->State == TreeEntry::NeedToGather &&
-                           is_contained(EntryPtr->Scalars, V);
-                  });
-      if (It != VLIt)
-        VTE = It->get();
-    }
-    if (VTE) {
-      auto Res = Used.try_emplace(VTE, NumShuffles);
-      if (Res.second) {
-        Entries.push_back(VTE);
-        ++NumShuffles;
-        if (NumShuffles > 2)
-          return None;
-        if (NumShuffles == 2) {
-          unsigned FirstSz = Entries.front()->Scalars.size();
-          if (!Entries.front()->ReuseShuffleIndices.empty())
-            FirstSz = Entries.front()->ReuseShuffleIndices.size();
-          unsigned SecondSz = Entries.back()->Scalars.size();
-          if (!Entries.back()->ReuseShuffleIndices.empty())
-            SecondSz = Entries.back()->ReuseShuffleIndices.size();
-          if (FirstSz != SecondSz)
-            return None;
-        }
-      }
-      int FoundLane =
-          findLaneForValue(VTE->Scalars, VTE->ReuseShuffleIndices, V);
-      unsigned Sz = VTE->Scalars.size();
-      if (!VTE->ReuseShuffleIndices.empty())
-        Sz = VTE->ReuseShuffleIndices.size();
-      Mask[I] = Res.first->second * Sz + FoundLane;
-      // Extra check required by isSingleSourceMaskImpl function (called by
-      // ShuffleVectorInst::isSingleSourceMask).
-      if (Mask[I] >= 2 * E)
+      if (Entries.size() == 2)
         return None;
-      continue;
+      VTE = getTreeEntry(V);
+      if (!VTE || find_if(
+                      VectorizableTree,
+                      [VTE, TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
+                        return EntryPtr.get() == VTE || EntryPtr.get() == TE;
+                      })->get() == TE) {
+        // Check if it is used in one of the gathered entries.
+        const auto *It =
+            find_if(VectorizableTree,
+                    [V, TE](const std::unique_ptr<TreeEntry> &EntryPtr) {
+                      return EntryPtr.get() == TE ||
+                             (EntryPtr->State == TreeEntry::NeedToGather &&
+                              is_contained(EntryPtr->Scalars, V));
+                    });
+        // The vector factor of shuffled entries must be the same.
+        if (It->get() == TE)
+          return None;
+        VTE = It->get();
+      }
+      Entries.push_back(VTE);
+      if (Entries.size() == 1) {
+        VF = GetVF(VTE);
+      } else if (VF != GetVF(VTE)) {
+        assert(Entries.size() == 2 && "Expected shuffle of 1 or 2 entries.");
+        assert(VF > 0 && "Expected non-zero vector factor.");
+        return None;
+      }
+      for (Value *SV : VTE->Scalars)
+        UsedValuesEntry.try_emplace(SV, VTE);
     }
-    return None;
+    int FoundLane = findLaneForValue(VTE->Scalars, VTE->ReuseShuffleIndices, V);
+    Mask[I] = (Entries.front() == VTE ? 0 : VF) + FoundLane;
+    // Extra check required by isSingleSourceMaskImpl function (called by
+    // ShuffleVectorInst::isSingleSourceMask).
+    if (Mask[I] >= 2 * E)
+      return None;
   }
-  if (NumShuffles == 1) {
+  if (Entries.size() == 1) {
     if (ShuffleVectorInst::isReverseMask(Mask))
       return TargetTransformInfo::SK_Reverse;
     return TargetTransformInfo::SK_PermuteSingleSrc;
   }
-  if (NumShuffles == 2) {
+  if (Entries.size() == 2) {
     if (ShuffleVectorInst::isSelectMask(Mask))
       return TargetTransformInfo::SK_Select;
     if (ShuffleVectorInst::isTransposeMask(Mask))