[llvm] [InstCombine] Fix poison safety of folding shufflevector into select (PR #115483)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/115483

We are allowed to fold shufflevector into select iff the condition is guaranteed not to be poison or the RHS is a poison.
Alive2: https://alive2.llvm.org/ce/z/28zEWR

Closes https://github.com/llvm/llvm-project/issues/115465.


>From 3f2be7a06aa58df551bf71e16d98490951a6be57 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 8 Nov 2024 20:45:35 +0800
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests. NFC.

---
 llvm/test/Transforms/InstCombine/vec_shuffle.ll | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index 163d9c9557b239..49514681ef8ac3 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -2411,6 +2411,18 @@ define <4 x i32> @shuf_same_length_vec_select(<4 x i1> %cond) {
   ret <4 x i32> %shuf
 }
 
+; Make sure we do not fold in this case.
+define <4 x i8> @shuf_cmp_may_be_poison(<4 x i8> %x, <4 x i8> %y, i1 %cmp) {
+; CHECK-LABEL: @shuf_cmp_may_be_poison(
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i8> [[Y:%.*]], <4 x i8> <i8 poison, i8 1, i8 poison, i8 3>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    [[SHUF:%.*]] = select i1 [[CMP:%.*]], <4 x i8> [[TMP1]], <4 x i8> <i8 0, i8 1, i8 0, i8 3>
+; CHECK-NEXT:    ret <4 x i8> [[SHUF]]
+;
+  %sel = select i1 %cmp, <4 x i8> %y, <4 x i8> <i8 0, i8 poison, i8 0, i8 poison>
+  %shuf = shufflevector <4 x i8> %sel, <4 x i8> <i8 poison, i8 1, i8 poison, i8 3>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+  ret <4 x i8> %shuf
+}
+
 declare i1 @cond()
 declare <4 x i32> @value()
 

>From 24ee5fcdbca3e896ec7e739f82f8d0f28df125de Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Fri, 8 Nov 2024 21:38:31 +0800
Subject: [PATCH 2/2] [InstCombine] Fix poison safety of folding shufflevector
 into select

---
 llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp | 4 +++-
 llvm/test/Transforms/InstCombine/vec_shuffle.ll          | 6 +++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 454fe5a91d375a..ede89b099e8deb 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2904,7 +2904,9 @@ Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
     if (auto *SI = dyn_cast<SelectInst>(LHS)) {
       // We cannot do this fold for elementwise select since ShuffleVector is
       // not elementwise.
-      if (SI->getCondition()->getType()->isIntegerTy()) {
+      if (SI->getCondition()->getType()->isIntegerTy() &&
+          (isa<PoisonValue>(RHS) ||
+           isGuaranteedNotToBePoison(SI->getCondition()))) {
         if (Instruction *I = FoldOpIntoSelect(SVI, SI))
           return I;
       }
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index 49514681ef8ac3..9fb68b5399c845 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -2414,9 +2414,9 @@ define <4 x i32> @shuf_same_length_vec_select(<4 x i1> %cond) {
 ; Make sure we do not fold in this case.
 define <4 x i8> @shuf_cmp_may_be_poison(<4 x i8> %x, <4 x i8> %y, i1 %cmp) {
 ; CHECK-LABEL: @shuf_cmp_may_be_poison(
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i8> [[Y:%.*]], <4 x i8> <i8 poison, i8 1, i8 poison, i8 3>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
-; CHECK-NEXT:    [[SHUF:%.*]] = select i1 [[CMP:%.*]], <4 x i8> [[TMP1]], <4 x i8> <i8 0, i8 1, i8 0, i8 3>
-; CHECK-NEXT:    ret <4 x i8> [[SHUF]]
+; CHECK-NEXT:    [[Y:%.*]] = select i1 [[CMP:%.*]], <4 x i8> [[Y1:%.*]], <4 x i8> <i8 0, i8 poison, i8 0, i8 poison>
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i8> [[Y]], <4 x i8> <i8 poison, i8 1, i8 poison, i8 3>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>
+; CHECK-NEXT:    ret <4 x i8> [[TMP1]]
 ;
   %sel = select i1 %cmp, <4 x i8> %y, <4 x i8> <i8 0, i8 poison, i8 0, i8 poison>
   %shuf = shufflevector <4 x i8> %sel, <4 x i8> <i8 poison, i8 1, i8 poison, i8 3>, <4 x i32> <i32 0, i32 5, i32 2, i32 7>