[ConstantFolding] Fold intrinsics of scalable vectors with splatted operands (#141845)

As noted in
https://github.com/llvm/llvm-project/pull/141821#issuecomment-2917328924,
whilst we currently constant fold intrinsics of fixed-length vectors via
their scalar counterpart, we don't do the same for scalable vectors.

This handles the scalable vector case when the operands are splats.

One weird snag in ConstantVector::getSplat was that it produced a undef
if passed in poison, so this also contains a fix by checking for
PoisonValue before UndefValue.

1 files changed, 29 insertions, 1 deletions

diff --git a/llvm/lib/Analysis/ConstantFolding.cpp b/llvm/lib/Analysis/ConstantFolding.cpp
index 412a0e8..40302fb 100644
--- a/llvm/lib/Analysis/ConstantFolding.cpp
+++ b/llvm/lib/Analysis/ConstantFolding.cpp
@@ -3780,7 +3780,35 @@ static Constant *ConstantFoldScalableVectorCall(
   default:
     break;
   }
-  return nullptr;
+
+  // If trivially vectorizable, try folding it via the scalar call if all
+  // operands are splats.
+
+  // TODO: ConstantFoldFixedVectorCall should probably check this too?
+  if (!isTriviallyVectorizable(IntrinsicID))
+    return nullptr;
+
+  SmallVector<Constant *, 4> SplatOps;
+  for (auto [I, Op] : enumerate(Operands)) {
+    if (isVectorIntrinsicWithScalarOpAtArg(IntrinsicID, I, /*TTI=*/nullptr)) {
+      SplatOps.push_back(Op);
+      continue;
+    }
+    // TODO: Should getSplatValue return a poison scalar for a poison vector?
+    if (isa<PoisonValue>(Op)) {
+      SplatOps.push_back(PoisonValue::get(Op->getType()->getScalarType()));
+      continue;
+    }
+    Constant *Splat = Op->getSplatValue();
+    if (!Splat)
+      return nullptr;
+    SplatOps.push_back(Splat);
+  }
+  Constant *Folded = ConstantFoldScalarCall(
+      Name, IntrinsicID, SVTy->getElementType(), SplatOps, TLI, Call);
+  if (!Folded)
+    return nullptr;
+  return ConstantVector::getSplat(SVTy->getElementCount(), Folded);
 }
 
 static std::pair<Constant *, Constant *>