[SLP]Improve isGatherShuffledEntry to handles shuffles between different length vectors.

[ABataev]

Summary

Originally isGatherShuffledEntry supports shuffles between 2 input
vectors of the same size. Actually, we can support the shuffles between
vectors of 2 different sizes. One of them (or both, for variable length
vector) just must be resized to match the sizes. Improves cost and the
vectorization itself.

[RKSimon]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

7377–7379

min index? Aren't you finding the max_element here?

7408

Duplicate code - worth pulling this out as a helper?

7914

We do this often enough - I'm wondering if we just add a helper to IRBuilder that creates an identity shuffle that pads/extracts a source Value to requested #elements size?

7946

Worth putting all of this inside ShuffleInstructionBuilder ?

[ABataev]

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

7946

Currently it does not support shuffles between 2 input vectors. There are some other patches, where the lambda is added

[ABataev]

Rebase

[ABataev]

All changes were committed already in separate patches

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

BoUpSLP::isGatherShuffledEntry(const TreeEntry *TE, SmallVectorImpl<int> &Mask,
 
   if (UsedTEs.empty()) {
     assert(all_of(TE->Scalars, UndefValue::classof) &&
            "Expected vector of undefs only.");
     return None;
   }
 
   unsigned VF = 0;
+  auto &&FindMaxIdxTreeNode = [](SmallPtrSet<const TreeEntry *, 4> &TEs) {
+    return *std::max_element(TEs.begin(), TEs.end(),
+                             [](const TreeEntry *TE1, const TreeEntry *TE2) {
+                               return TE1->Idx < TE2->Idx;
+                             });
+  };
   if (UsedTEs.size() == 1) {
     // Try to find the perfect match in another gather node at first.
     auto It = find_if(UsedTEs.front(), [TE](const TreeEntry *EntryPtr) {
       return EntryPtr->isSame(TE->Scalars);
     });
     if (It != UsedTEs.front().end()) {
       Entries.push_back(*It);
       std::iota(Mask.begin(), Mask.end(), 0);
       return TargetTransformInfo::SK_PermuteSingleSrc;
     }
-    // No perfect match, just shuffle, so choose the first tree node.
-    Entries.push_back(*UsedTEs.front().begin());
+    // No perfect match, just shuffle, so choose the tree node with max index.
+    Entries.push_back(FindMaxIdxTreeNode(UsedTEs.front()));
+    VF = Entries.back()->getVectorFactor();
   } else {
     // Try to find nodes with the same vector factor.
     assert(UsedTEs.size() == 2 && "Expected at max 2 permuted entries.");
     DenseMap<int, const TreeEntry *> VFToTE;
-    for (const TreeEntry *TE : UsedTEs.front())
-      VFToTE.try_emplace(TE->getVectorFactor(), TE);
-    for (const TreeEntry *TE : UsedTEs.back()) {
+    for (const TreeEntry *TE : UsedTEs.front()) {
+      auto Res = VFToTE.try_emplace(TE->getVectorFactor(), TE);
+      if (!Res.second && TE->Idx > Res.first->second->Idx)
+        Res.first->getSecond() = TE;
+    }
+    SmallVector<const TreeEntry *> UsedTEsBack(UsedTEs.back().begin(),
+                                               UsedTEs.back().end());
+    // Sort the elements of UsedTEsBack by the idx to keep the determinism.
+    sort(UsedTEsBack, [](const TreeEntry *TE1, const TreeEntry *TE2) {
+      return TE1->Idx > TE2->Idx;
+    });
+    for (const TreeEntry *TE : UsedTEsBack) {
       auto It = VFToTE.find(TE->getVectorFactor());
       if (It != VFToTE.end()) {
         VF = It->first;
         Entries.push_back(It->second);
         Entries.push_back(TE);
         break;
       }
     }
-    // No 2 source vectors with the same vector factor - give up and do regular
-    // gather.
-    if (Entries.empty())
-      return None;
+    // No 2 source vectors with the same vector factor - just choose 2 with max
+    // index.
+    if (Entries.empty()) {
+      Entries.push_back(FindMaxIdxTreeNode(UsedTEs.front()));
+      Entries.push_back(FindMaxIdxTreeNode(UsedTEs.back()));
+      VF = std::max(Entries.front()->getVectorFactor(),
+                    Entries.back()->getVectorFactor());
+    }
   }
 
   // Build a shuffle mask for better cost estimation and vector emission.
   for (int I = 0, E = TE->Scalars.size(); I < E; ++I) {
     Value *V = TE->Scalars[I];
     if (isa<UndefValue>(V))
       continue;
     unsigned Idx = UsedValuesEntry.lookup(V);
if (E->State == TreeEntry::NeedToGather) {
     Value *Vec;
     SmallVector<int> Mask;
     SmallVector<const TreeEntry *> Entries;
     Optional<TargetTransformInfo::ShuffleKind> Shuffle =
         isGatherShuffledEntry(E, Mask, Entries);
     if (Shuffle) {
       assert((Entries.size() == 1 || Entries.size() == 2) &&
              "Expected shuffle of 1 or 2 entries.");
-      Vec = Builder.CreateShuffleVector(Entries.front()->VectorizedValue,
-                                        Entries.back()->VectorizedValue, Mask);
+      Value *Vec1 = Entries.front()->VectorizedValue;
+      Value *Vec2 = Entries.back()->VectorizedValue;
+      // Resize smaller vector if the types are not equal.
+      if (Vec1->getType() != Vec2->getType()) {
+        auto *FixedTy1 = dyn_cast<FixedVectorType>(Vec1->getType());
+        auto *FixedTy2 = dyn_cast<FixedVectorType>(Vec2->getType());
+        if (!FixedTy1 || !FixedTy2) {
+          // Not a fixed vector - resize both vectors to the size of the
+          // resulting vector.
+          SmallVector<int> ResizeMask(Mask.size(), UndefMaskElem);
+          std::iota(ResizeMask.begin(), ResizeMask.end(), 0);
+          Vec1 = Builder.CreateShuffleVector(Vec1, ResizeMask);
+          if (auto *I = dyn_cast<Instruction>(Vec1)) {
+            GatherShuffleSeq.insert(I);
+            CSEBlocks.insert(I->getParent());
+          }
+          Vec2 = Builder.CreateShuffleVector(Vec2, ResizeMask);
+          if (auto *I = dyn_cast<Instruction>(Vec1)) {
+            GatherShuffleSeq.insert(I);
+            CSEBlocks.insert(I->getParent());
+          }
+        } else {
+          // Resize smallest vector to the size of the greatest one.
+          unsigned VF1 = FixedTy1->getNumElements();
+          unsigned VF2 = FixedTy2->getNumElements();
+          SmallVector<int> ResizeMask(std::max(VF1, VF2), UndefMaskElem);
+          std::iota(ResizeMask.begin(),
+                    std::next(ResizeMask.begin(), std::min(VF1, VF2)), 0);
+          if (cast<FixedVectorType>(Vec1->getType())->getNumElements() <
+              cast<FixedVectorType>(Vec2->getType())->getNumElements()) {
+            Vec1 = Builder.CreateShuffleVector(Vec1, ResizeMask);
+            if (auto *I = dyn_cast<Instruction>(Vec1)) {
+              GatherShuffleSeq.insert(I);
+              CSEBlocks.insert(I->getParent());
+            }
+          } else {
+            Vec2 = Builder.CreateShuffleVector(Vec2, ResizeMask);
+            if (auto *I = dyn_cast<Instruction>(Vec2)) {
+              GatherShuffleSeq.insert(I);
+              CSEBlocks.insert(I->getParent());
+            }
+          }
+        }
+      }
+      Vec = Builder.CreateShuffleVector(Vec1, Vec2, Mask);
       if (auto *I = dyn_cast<Instruction>(Vec)) {
         GatherShuffleSeq.insert(I);
         CSEBlocks.insert(I->getParent());
       }
     } else {
       Vec = gather(E->Scalars);
     }
     if (NeedToShuffleReuses) {

llvm/test/Transforms/SLPVectorizer/X86/shuffled-gathers-diff-size.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -slp-vectorizer -S -mtriple=x86_64-unknown-linux -slp-threshold=-2 | FileCheck %s
 
 define void @foo(i32* noalias nocapture writeonly %B, i32* noalias nocapture readonly %A, i32* noalias nocapture readonly %C, i32 %n, i32 %m) {
 ; CHECK-LABEL: @foo(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[A:%.*]], align 4
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[TMP0]], [[N:%.*]]
-; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* [[C:%.*]], align 4
-; CHECK-NEXT:    [[MUL2:%.*]] = mul nsw i32 [[TMP1]], [[M:%.*]]
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[MUL2]], [[MUL]]
-; CHECK-NEXT:    [[ARRAYIDX3:%.*]] = getelementptr inbounds i32, i32* [[A]], i64 1
-; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX3]], align 4
-; CHECK-NEXT:    [[MUL4:%.*]] = mul nsw i32 [[ADD]], [[TMP2]]
-; CHECK-NEXT:    store i32 [[MUL4]], i32* [[B:%.*]], align 4
-; CHECK-NEXT:    [[ARRAYIDX8:%.*]] = getelementptr inbounds i32, i32* [[C]], i64 1
-; CHECK-NEXT:    [[TMP3:%.*]] = load i32, i32* [[ARRAYIDX8]], align 4
-; CHECK-NEXT:    [[MUL9:%.*]] = mul nsw i32 [[TMP3]], [[M]]
-; CHECK-NEXT:    [[ADD10:%.*]] = add nsw i32 [[MUL9]], [[MUL]]
-; CHECK-NEXT:    [[ARRAYIDX11:%.*]] = getelementptr inbounds i32, i32* [[C]], i64 2
-; CHECK-NEXT:    [[TMP4:%.*]] = load i32, i32* [[ARRAYIDX11]], align 4
-; CHECK-NEXT:    [[MUL12:%.*]] = mul nsw i32 [[ADD10]], [[TMP4]]
-; CHECK-NEXT:    [[ARRAYIDX13:%.*]] = getelementptr inbounds i32, i32* [[B]], i64 1
-; CHECK-NEXT:    store i32 [[MUL12]], i32* [[ARRAYIDX13]], align 4
-; CHECK-NEXT:    [[MUL15:%.*]] = mul nsw i32 [[TMP2]], [[N]]
-; CHECK-NEXT:    [[MUL17:%.*]] = mul nsw i32 [[TMP4]], [[M]]
-; CHECK-NEXT:    [[ADD18:%.*]] = add nsw i32 [[MUL17]], [[MUL15]]
-; CHECK-NEXT:    [[MUL20:%.*]] = mul nsw i32 [[ADD18]], [[TMP0]]
-; CHECK-NEXT:    [[ARRAYIDX21:%.*]] = getelementptr inbounds i32, i32* [[B]], i64 2
-; CHECK-NEXT:    store i32 [[MUL20]], i32* [[ARRAYIDX21]], align 4
-; CHECK-NEXT:    [[ARRAYIDX24:%.*]] = getelementptr inbounds i32, i32* [[C]], i64 3
-; CHECK-NEXT:    [[TMP5:%.*]] = load i32, i32* [[ARRAYIDX24]], align 4
-; CHECK-NEXT:    [[MUL25:%.*]] = mul nsw i32 [[TMP5]], [[M]]
-; CHECK-NEXT:    [[ADD26:%.*]] = add nsw i32 [[MUL25]], [[MUL15]]
-; CHECK-NEXT:    [[MUL28:%.*]] = mul nsw i32 [[ADD26]], [[TMP1]]
-; CHECK-NEXT:    [[ARRAYIDX29:%.*]] = getelementptr inbounds i32, i32* [[B]], i64 3
-; CHECK-NEXT:    store i32 [[MUL28]], i32* [[ARRAYIDX29]], align 4
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[A:%.*]] to <2 x i32>*
+; CHECK-NEXT:    [[TMP1:%.*]] = load <2 x i32>, <2 x i32>* [[TMP0]], align 4
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x i32> poison, i32 [[N:%.*]], i32 0
+; CHECK-NEXT:    [[TMP3:%.*]] = insertelement <2 x i32> [[TMP2]], i32 [[N]], i32 1
+; CHECK-NEXT:    [[TMP4:%.*]] = mul nsw <2 x i32> [[TMP1]], [[TMP3]]
+; CHECK-NEXT:    [[SHUFFLE1:%.*]] = shufflevector <2 x i32> [[TMP4]], <2 x i32> poison, <4 x i32> <i32 0, i32 0, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP5:%.*]] = bitcast i32* [[C:%.*]] to <4 x i32>*
+; CHECK-NEXT:    [[TMP6:%.*]] = load <4 x i32>, <4 x i32>* [[TMP5]], align 4
+; CHECK-NEXT:    [[TMP7:%.*]] = insertelement <4 x i32> poison, i32 [[M:%.*]], i32 0
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i32> [[TMP7]], <4 x i32> poison, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP8:%.*]] = mul nsw <4 x i32> [[TMP6]], [[SHUFFLE]]
+; CHECK-NEXT:    [[TMP9:%.*]] = add nsw <4 x i32> [[TMP8]], [[SHUFFLE1]]
+; CHECK-NEXT:    [[TMP10:%.*]] = shufflevector <2 x i32> [[TMP1]], <2 x i32> poison, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+; CHECK-NEXT:    [[TMP11:%.*]] = shufflevector <4 x i32> [[TMP10]], <4 x i32> [[TMP6]], <4 x i32> <i32 1, i32 6, i32 0, i32 4>
+; CHECK-NEXT:    [[TMP12:%.*]] = mul nsw <4 x i32> [[TMP9]], [[TMP11]]
+; CHECK-NEXT:    [[TMP13:%.*]] = bitcast i32* [[B:%.*]] to <4 x i32>*
+; CHECK-NEXT:    store <4 x i32> [[TMP12]], <4 x i32>* [[TMP13]], align 4
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load i32, i32* %A, align 4
   %mul = mul nsw i32 %0, %n
   %1 = load i32, i32* %C, align 4
   %mul2 = mul nsw i32 %1, %m
   %add = add nsw i32 %mul2, %mul