[SLP] Special case SK_PermuteSingleSrc extract cost estimate (NFCI).

[fhahn]

Summary

Preparation for more fine-grained cost estimates for extracts that form
a SK_PermuteSingleSrc pattern.

[fhahn]

Not needed

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

if (isSplat(VL)) {
                                  0);
     }
     if (E->getOpcode() == Instruction::ExtractElement &&
         allSameType(VL) && allSameBlock(VL)) {
       SmallVector<int> Mask;
       Optional<TargetTransformInfo::ShuffleKind> ShuffleKind =
           isShuffle(VL, Mask);
       if (ShuffleKind.hasValue()) {
-        InstructionCost Cost =
-            TTI->getShuffleCost(ShuffleKind.getValue(), VecTy, Mask);
+        InstructionCost Cost = 0;
+
+        if (*ShuffleKind != TargetTransformInfo::SK_PermuteSingleSrc)
+          Cost += TTI->getShuffleCost(ShuffleKind.getValue(), VecTy, Mask);
+
+        unsigned NumOfParts = TTI->getNumberOfParts(VecTy);
+        // Compute the number of number of elements per vector register for
+        // VecTy. If that is not possible, because the number of parts for VecTy
+        // is unknown, use the maximum value for unsigned.
+        unsigned EltsPerVector =
+            NumOfParts ? VecTy->getNumElements() / NumOfParts : -1;
+        unsigned Idx = 0;
+
         for (auto *V : VL) {
           // If all users of instruction are going to be vectorized and this
           // instruction itself is not going to be vectorized, consider this
           // instruction as dead and remove its cost from the final cost of the
           // vectorized tree.
           if (areAllUsersVectorized(cast<Instruction>(V)) &&
               !ScalarToTreeEntry.count(V)) {
             auto *IO = cast<ConstantInt>(
                 cast<ExtractElementInst>(V)->getIndexOperand());
             Cost -= TTI->getVectorInstrCost(Instruction::ExtractElement, VecTy,
                                             IO->getZExtValue());
           }
+          if (*ShuffleKind == TargetTransformInfo::SK_PermuteSingleSrc) {
+            if ((Idx > 0 && (Idx + 1) % EltsPerVector == 0) ||
+                Idx + 1 == VL.size()) {
+              SmallVector<int> Mask;
+              unsigned StartIdx =
+                  Idx - std::min(EltsPerVector, VecTy->getNumElements()) + 1;
+              Optional<TargetTransformInfo::ShuffleKind> ShuffleKind =
+                  isShuffle(
+                      {&VL[StartIdx],
+                       std::min(EltsPerVector, unsigned(VL.size() - StartIdx))},
+                      Mask);
+              Cost += TTI->getShuffleCost(
+                  ShuffleKind.getValue(),
+                  FixedVectorType::get(
+                      VecTy->getElementType(),
+                      std::min(EltsPerVector, VecTy->getNumElements())),
+                  Mask);
+            }
+          }
+          ++Idx;
         }
         return ReuseShuffleCost + Cost;
       }
     }
     return ReuseShuffleCost + getGatherCost(VL);
   }
   assert((E->State == TreeEntry::Vectorize ||
           E->State == TreeEntry::ScatterVectorize) &&