[llvm] [ConstantFolding] Fold scalable shufflevector of poison/undef. (PR #143475)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/143475

None

>From 8df9851e9935eca431b6b10a41a1beb2b5f37ece Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Mon, 9 Jun 2025 22:04:29 -0700
Subject: [PATCH 1/2] Pre-commit test

---
 .../Transforms/InstSimplify/shufflevector.ll     | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/llvm/test/Transforms/InstSimplify/shufflevector.ll b/llvm/test/Transforms/InstSimplify/shufflevector.ll
index 0442bd721974b..f588bf950b9ee 100644
--- a/llvm/test/Transforms/InstSimplify/shufflevector.ll
+++ b/llvm/test/Transforms/InstSimplify/shufflevector.ll
@@ -337,3 +337,19 @@ define <4 x i32> @not_fold_identity2(<4 x i32> %x) {
   %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> shufflevector (<vscale x 2 x i32> undef, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer)
+;
+  %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> shufflevector (<vscale x 2 x i32> poison, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
+;
+  %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
+}

>From 8b85e7463084ef26445197a944f195c58d291c4e Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Mon, 9 Jun 2025 22:06:46 -0700
Subject: [PATCH 2/2] [ConstantFolding] Fold scalable shufflevector of
 poison/undef.

---
 llvm/lib/IR/ConstantFold.cpp                       | 5 +----
 llvm/test/Transforms/InstSimplify/shufflevector.ll | 4 ++--
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/IR/ConstantFold.cpp b/llvm/lib/IR/ConstantFold.cpp
index b9db402fe9562..6c8850aae0476 100644
--- a/llvm/lib/IR/ConstantFold.cpp
+++ b/llvm/lib/IR/ConstantFold.cpp
@@ -463,10 +463,7 @@ Constant *llvm::ConstantFoldShuffleVectorInstruction(Constant *V1, Constant *V2,
     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())
+    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 f588bf950b9ee..feab024a29213 100644
--- a/llvm/test/Transforms/InstSimplify/shufflevector.ll
+++ b/llvm/test/Transforms/InstSimplify/shufflevector.ll
@@ -340,7 +340,7 @@ define <4 x i32> @not_fold_identity2(<4 x i32> %x) {
 
 define <vscale x 2 x i32> @scalable_splat_undef() {
 ; CHECK-LABEL: @scalable_splat_undef(
-; CHECK-NEXT:    ret <vscale x 2 x i32> shufflevector (<vscale x 2 x i32> undef, <vscale x 2 x i32> undef, <vscale x 2 x i32> zeroinitializer)
+; 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
@@ -348,7 +348,7 @@ define <vscale x 2 x i32> @scalable_splat_undef() {
 
 define <vscale x 2 x i32> @scalable_splat_poison() {
 ; CHECK-LABEL: @scalable_splat_poison(
-; CHECK-NEXT:    ret <vscale x 2 x i32> shufflevector (<vscale x 2 x i32> poison, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer)
+; 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