[InstCombine] Do not fold shufflevector(select) if the select condition is a vector

[dtcxzyw]

Since shufflevector is not element-wise, we cannot do fold it into select when the select condition is a vector.
For shufflevector that doesn't change the length, it doesn't crash, but it is still a miscompilation: https://alive2.llvm.org/ce/z/s8saCx

Fixes #113986.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Since shufflevector is not element-wise, we cannot do fold it into select when the select condition is a vector.
For shufflevector that doesn't change the length, it doesn't crash, but it is still a miscompilation: https://alive2.llvm.org/ce/z/s8saCx

Fixes #113986.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+6-2)
• (modified) llvm/test/Transforms/InstCombine/vec_shuffle.ll (+24)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 75e7c1c97018cb..454fe5a91d375a 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -2902,8 +2902,12 @@ Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
 
   if (match(RHS, m_Constant())) {
     if (auto *SI = dyn_cast<SelectInst>(LHS)) {
-      if (Instruction *I = FoldOpIntoSelect(SVI, SI))
-        return I;
+      // We cannot do this fold for elementwise select since ShuffleVector is
+      // not elementwise.
+      if (SI->getCondition()->getType()->isIntegerTy()) {
+        if (Instruction *I = FoldOpIntoSelect(SVI, SI))
+          return I;
+      }
     }
     if (auto *PN = dyn_cast<PHINode>(LHS)) {
       if (Instruction *I = foldOpIntoPhi(SVI, PN))
diff --git a/llvm/test/Transforms/InstCombine/vec_shuffle.ll b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
index d050cf10849e3c..39a9db02eef293 100644
--- a/llvm/test/Transforms/InstCombine/vec_shuffle.ll
+++ b/llvm/test/Transforms/InstCombine/vec_shuffle.ll
@@ -2387,6 +2387,30 @@ define <2 x i32> @foldselect0(i1 %c) {
   ret <2 x i32> %shuf
 }
 
+; Make sure we do not crash in this case.
+define <4 x float> @shuf_larger_length_vec_select(<2 x i1> %cond) {
+; CHECK-LABEL: @shuf_larger_length_vec_select(
+; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND:%.*]], <2 x float> zeroinitializer, <2 x float> <float 1.000000e+00, float 1.000000e+00>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <2 x float> [[SEL]], <2 x float> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x float> [[SHUF]]
+;
+  %sel = select <2 x i1> %cond, <2 x float> zeroinitializer, <2 x float> splat(float 1.000000e+00)
+  %shuf = shufflevector <2 x float> %sel, <2 x float> zeroinitializer, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  ret <4 x float> %shuf
+}
+
+; Make sure we do not fold in this case.
+define <4 x i32> @shuf_same_length_vec_select(<4 x i1> %cond) {
+; CHECK-LABEL: @shuf_same_length_vec_select(
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[COND:%.*]], <4 x i32> <i32 poison, i32 1, i32 2, i32 3>, <4 x i32> <i32 poison, i32 5, i32 6, i32 7>
+; CHECK-NEXT:    [[SHUF:%.*]] = shufflevector <4 x i32> [[SEL]], <4 x i32> <i32 poison, i32 9, i32 poison, i32 poison>, <4 x i32> <i32 2, i32 1, i32 3, i32 5>
+; CHECK-NEXT:    ret <4 x i32> [[SHUF]]
+;
+  %sel = select <4 x i1> %cond, <4 x i32> <i32 0, i32 1, i32 2, i32 3>, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
+  %shuf = shufflevector <4 x i32> %sel, <4 x i32> <i32 8, i32 9, i32 10, i32 11>, <4 x i32> <i32 2, i32 1, i32 3, i32 5>
+  ret <4 x i32> %shuf
+}
+
 declare i1 @cond()
 declare <4 x i32> @value()
 

[MatzeB]

LGTM, thanks for fixing!

[MatzeB]

I think we could still perform this transformation when the shuffle operand types are the same as the shuffle output type? Regardless it is probably best to just land this change as quick as possible to unbreak things. I can look into re-enabling it for same-type myself then.

[dtcxzyw]

There is a counterexample: https://alive2.llvm.org/ce/z/s8saCx.