[InstCombine] Fix poison safety of folding shufflevector into select

[dtcxzyw]

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 #115465.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

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 #115465.

---

Full diff: https://github.com/llvm/llvm-project/pull/115483.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+3-1)
• (modified) llvm/test/Transforms/InstCombine/vec_shuffle.ll (+12)

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 163d9c9557b239..9fb68b5399c845 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:    [[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>
+  ret <4 x i8> %shuf
+}
+
 declare i1 @cond()
 declare <4 x i32> @value()
 

[nikic]

LGTM