[RISCV][TTI] Use processShuffleMask for shuffle legalization estimate (#136191)

We had some code which tried to estimate legalization costs for
illegally typed shuffles, but it only handled the case of a widening
shuffle, and used a somewhat adhoc heuristic. We can reuse the
processShuffleMask utility (which we already use for individual vector
register splitting when exact VLEN is known) to perform the same
splitting given the legal vector type as the unit of split instead. This
makes the costing both simpler and more robust.

Note that this swings costs for illegal shuffles pretty wildly as we
were previously sometimes hitting the adhoc code, and sometimes falling
through into generic scalarization costing. I don't know that any of the
costs for the individual tests in tree are significant, but the test
which which triggered me finding this was reported to me by Alexey
reduced from something triggering a bad choice in SLP for x264. So this
has the potential to be somewhat high impact.

Author: preames

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

@@ -385,6 +385,62 @@ static VectorType *getVRGatherIndexType(MVT DataVT, const RISCVSubtarget &ST,
   return cast<VectorType>(EVT(IndexVT).getTypeForEVT(C));
 }
 
+/// Attempt to approximate the cost of a shuffle which will require splitting
+/// during legalization.  Note that processShuffleMasks is not an exact proxy
+/// for the algorithm used in LegalizeVectorTypes, but hopefully it's a
+/// reasonably close upperbound.
+static InstructionCost costShuffleViaSplitting(RISCVTTIImpl &TTI, MVT LegalVT,
+                                               VectorType *Tp,
+                                               ArrayRef<int> Mask,
+                                               TTI::TargetCostKind CostKind) {
+  assert(LegalVT.isFixedLengthVector() && !Mask.empty() &&
+         "Expected fixed vector type and non-empty mask");
+  unsigned LegalNumElts = LegalVT.getVectorNumElements();
+  // Number of destination vectors after legalization:
+  unsigned NumOfDests = divideCeil(Mask.size(), LegalNumElts);
+  // We are going to permute multiple sources and the result will be in
+  // multiple destinations. Providing an accurate cost only for splits where
+  // the element type remains the same.
+  if (NumOfDests <= 1 ||
+      LegalVT.getVectorElementType().getSizeInBits() !=
+          Tp->getElementType()->getPrimitiveSizeInBits() ||
+      LegalNumElts >= Tp->getElementCount().getFixedValue())
+    return InstructionCost::getInvalid();
+
+  unsigned VecTySize = TTI.getDataLayout().getTypeStoreSize(Tp);
+  unsigned LegalVTSize = LegalVT.getStoreSize();
+  // Number of source vectors after legalization:
+  unsigned NumOfSrcs = divideCeil(VecTySize, LegalVTSize);
+
+  auto *SingleOpTy = FixedVectorType::get(Tp->getElementType(), LegalNumElts);
+
+  unsigned NormalizedVF = LegalNumElts * std::max(NumOfSrcs, NumOfDests);
+  unsigned NumOfSrcRegs = NormalizedVF / LegalNumElts;
+  unsigned NumOfDestRegs = NormalizedVF / LegalNumElts;
+  SmallVector<int> NormalizedMask(NormalizedVF, PoisonMaskElem);
+  assert(NormalizedVF >= Mask.size() &&
+         "Normalized mask expected to be not shorter than original mask.");
+  copy(Mask, NormalizedMask.begin());
+  InstructionCost Cost = 0;
+  SmallDenseSet<std::pair<ArrayRef<int>, unsigned>> ReusedSingleSrcShuffles;
+  processShuffleMasks(
+      NormalizedMask, NumOfSrcRegs, NumOfDestRegs, NumOfDestRegs, []() {},
+      [&](ArrayRef<int> RegMask, unsigned SrcReg, unsigned DestReg) {
+        if (ShuffleVectorInst::isIdentityMask(RegMask, RegMask.size()))
+          return;
+        if (!ReusedSingleSrcShuffles.insert(std::make_pair(RegMask, SrcReg))
+                 .second)
+          return;
+        Cost += TTI.getShuffleCost(TTI::SK_PermuteSingleSrc, SingleOpTy,
+                                   RegMask, CostKind, 0, nullptr);
+      },
+      [&](ArrayRef<int> RegMask, unsigned Idx1, unsigned Idx2, bool NewReg) {
+        Cost += TTI.getShuffleCost(TTI::SK_PermuteTwoSrc, SingleOpTy, RegMask,
+                                   CostKind, 0, nullptr);
+      });
+  return Cost;
+}
+
 /// Try to perform better estimation of the permutation.
 /// 1. Split the source/destination vectors into real registers.
 /// 2. Do the mask analysis to identify which real registers are
@@ -647,48 +703,13 @@ InstructionCost RISCVTTIImpl::getShuffleCost(TTI::ShuffleKind Kind,
         return true;
       }
     };
-    // We are going to permute multiple sources and the result will be in
-    // multiple destinations. Providing an accurate cost only for splits where
-    // the element type remains the same.
+
     if (!Mask.empty() && LT.first.isValid() && LT.first != 1 &&
-        shouldSplit(Kind) &&
-        LT.second.getVectorElementType().getSizeInBits() ==
-        Tp->getElementType()->getPrimitiveSizeInBits() &&
-        LT.second.getVectorNumElements() <
-        cast<FixedVectorType>(Tp)->getNumElements() &&
-        divideCeil(Mask.size(),
-                   cast<FixedVectorType>(Tp)->getNumElements()) ==
-        static_cast<unsigned>(*LT.first.getValue())) {
-      unsigned NumRegs = *LT.first.getValue();
-      unsigned VF = cast<FixedVectorType>(Tp)->getNumElements();
-      unsigned SubVF = PowerOf2Ceil(VF / NumRegs);
-      auto *SubVecTy = FixedVectorType::get(Tp->getElementType(), SubVF);
-
-      InstructionCost Cost = 0;
-      for (unsigned I = 0, NumSrcRegs = divideCeil(Mask.size(), SubVF);
-           I < NumSrcRegs; ++I) {
-        bool IsSingleVector = true;
-        SmallVector<int> SubMask(SubVF, PoisonMaskElem);
-        transform(
-                  Mask.slice(I * SubVF,
-                             I == NumSrcRegs - 1 ? Mask.size() % SubVF : SubVF),
-                  SubMask.begin(), [&](int I) -> int {
-                    if (I == PoisonMaskElem)
-                      return PoisonMaskElem;
-                    bool SingleSubVector = I / VF == 0;
-                    IsSingleVector &= SingleSubVector;
-                    return (SingleSubVector ? 0 : 1) * SubVF + (I % VF) % SubVF;
-                  });
-        if (all_of(enumerate(SubMask), [](auto &&P) {
-          return P.value() == PoisonMaskElem ||
-            static_cast<unsigned>(P.value()) == P.index();
-        }))
-          continue;
-        Cost += getShuffleCost(IsSingleVector ? TTI::SK_PermuteSingleSrc
-                               : TTI::SK_PermuteTwoSrc,
-                               SubVecTy, SubMask, CostKind, 0, nullptr);
-      }
-      return Cost;
+        shouldSplit(Kind)) {
+      InstructionCost SplitCost =
+          costShuffleViaSplitting(*this, LT.second, FVTp, Mask, CostKind);
+      if (SplitCost.isValid())
+        return SplitCost;
     }
   }