[llvm] [VectorCombine] Fold permute of intrinsics into intrinsic of permutes: shuffle(intrinsic, poison/undef) -> intrinsic(shuffle) (PR #170052)

[Jerry Dang via llvm-commits]

================
@@ -2961,6 +2962,94 @@ bool VectorCombine::foldShuffleOfIntrinsics(Instruction &I) {
   return true;
 }
 
+/// Try to convert
+/// "shuffle (intrinsic), (poison/undef)" into "intrinsic (shuffle)".
+bool VectorCombine::foldPermuteOfIntrinsic(Instruction &I) {
+  Value *V0, *V1;
+  ArrayRef<int> Mask;
+  if (!match(&I, m_Shuffle(m_OneUse(m_Value(V0)), m_Value(V1), m_Mask(Mask))))
+    return false;
+
+  // Check for permute
+  if (!match(V1, m_Poison()) && !match(V1, m_Undef())) {
+    LLVM_DEBUG(dbgs() << "not a permute\n");
+    return false;
+  }
+
+  auto *II0 = dyn_cast<IntrinsicInst>(V0);
+  if (!II0)
+    return false;
+
+  auto *ShuffleDstTy = dyn_cast<FixedVectorType>(I.getType());
+  auto *IntrinsicSrcTy = dyn_cast<FixedVectorType>(II0->getType());
+  if (!ShuffleDstTy || !IntrinsicSrcTy)
+    return false;
+
+  // Validate it's a pure permute, mask should only reference the first vector
+  unsigned NumSrcElts = IntrinsicSrcTy->getNumElements();
+  for (int Idx : Mask) {
+    if (Idx > 0 && Idx >= (int)NumSrcElts)
+      return false;
+  }
----------------
kuroyukiasuna wrote:

The `M < 0` condition would cause the optimization failing to trigger because poison values in the mask produces -1.

E.g.:
```
%1 = shufflevector <2 x i16> %l266, <2 x i16> poison, <4 x i32> <i32 0, i32 1, i32 poison, i32 poison>
```
that is `<0, 1, -1, -1>`, and we want the optimization to run in this case.

I ended up using:
```
if (any_of(Mask, [NumSrcElts](int M){ return M >= (int)NumSrcElts; }))
  return false;
```
Let me know if you have other suggestions. Thanks!

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