[llvm] [VectorCombine] Fold ZExt/SExt (Shuffle (ZExt/SExt %src)) to ZExt/SExt (Shuffle %src). (PR #141109)

[Philip Reames via llvm-commits]

================
@@ -2791,6 +2792,60 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
   return true;
 }
 
+bool VectorCombine::foldShuffleExt(Instruction &I) {
+  // Try to fold vector zero- and sign-extends split across multiple operations
+  // into a single extend.
+
+  // Check if we have ZEXT/SEXT (SHUFFLE (ZEXT/SEXT %src), _, identity-mask),
+  // with an identity mask extracting the first sub-vector.
+  Value *Src;
+  ArrayRef<int> Mask;
+  if (!match(&I, m_OneUse(m_Shuffle(m_OneUse(m_ZExtOrSExt(m_Value(Src))),
+                                    m_Value(), m_Mask(Mask)))) ||
+      !cast<ShuffleVectorInst>(&I)->isIdentityWithExtract())
+    return false;
+  auto *InnerExt = cast<Instruction>(I.getOperand(0));
+  auto *OuterExt = cast<Instruction>(*I.user_begin());
+  if (!isa<SExtInst, ZExtInst>(OuterExt))
+    return false;
+
+  // If the inner extend is a sign extend and the outer one isnt (i.e. a
+  // zero-extend), don't fold. If the first one is zero-extend, it doesn't
+  // matter if the second one is a sign- or zero-extend.
+  if (isa<SExtInst>(InnerExt) && !isa<SExtInst>(OuterExt))
+    return false;
+
+  auto *DstTy = cast<FixedVectorType>(OuterExt->getType());
+  auto *SrcTy =
+      FixedVectorType::get(InnerExt->getOperand(0)->getType()->getScalarType(),
+                           DstTy->getNumElements());
+
+  // Don't perform the fold if the cost of the new extend is worse than the cost
+  // of the 2 original extends.
+  InstructionCost OriginalCost =
+      TTI.getCastInstrCost(InnerExt->getOpcode(), SrcTy, InnerExt->getType(),
----------------
preames wrote:

It doesn't look like you have the right params for costing the original outer param - the inner params are repeated twice.  

https://github.com/llvm/llvm-project/pull/141109