[ConstantFolding] Fold scalable shufflevector of poison/undef. (#143475)
diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index b9db402..d4ad21e 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -463,10 +463,9 @@
Constant *Elt =
ConstantExpr::getExtractElement(V1, ConstantInt::get(Ty, 0));
- if (Elt->isNullValue()) {
- auto *VTy = VectorType::get(EltTy, MaskEltCount);
- return ConstantAggregateZero::get(VTy);
- } else if (!MaskEltCount.isScalable())
+ // For scalable vectors, make sure this doesn't fold back into a
+ // shufflevector.
+ if (!MaskEltCount.isScalable() || Elt->isNullValue() || isa<UndefValue>(Elt))
return ConstantVector::getSplat(MaskEltCount, Elt);
}
diff --git a/llvm/test/Transforms/InstSimplify/shufflevector.ll b/llvm/test/Transforms/InstSimplify/shufflevector.ll
index 0442bd7..feab024 100644
--- a/llvm/test/Transforms/InstSimplify/shufflevector.ll
+++ b/llvm/test/Transforms/InstSimplify/shufflevector.ll
@@ -337,3 +337,19 @@
%revshuf = shufflevector <4 x i32> %shuf, <4 x i32> poison, <4 x i32> <i32 3, i32 2, i32 1, i32 0>
ret <4 x i32> %revshuf
}
+
+define <vscale x 2 x i32> @scalable_splat_undef() {
+; CHECK-LABEL: @scalable_splat_undef(
+; CHECK-NEXT: ret <vscale x 2 x i32> undef
+;
+ %shuf = shufflevector <vscale x 2 x i32> undef, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer
+ ret <vscale x 2 x i32> %shuf
+}
+
+define <vscale x 2 x i32> @scalable_splat_poison() {
+; CHECK-LABEL: @scalable_splat_poison(
+; CHECK-NEXT: ret <vscale x 2 x i32> poison
+;
+ %shuf = shufflevector <vscale x 2 x i32> poison, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
+ ret <vscale x 2 x i32> %shuf
+}