[llvm] 1831109 - [InstCombine] Do not fold `shufflevector(select)` if the select condition is a vector (#113993)

[via llvm-commits]

Author: Yingwei Zheng
Date: 2024-10-29T10:39:07+08:00
New Revision: 18311093abe6481388a0d963d58438d743b47569

URL: https://github.com/llvm/llvm-project/commit/18311093abe6481388a0d963d58438d743b47569
DIFF: https://github.com/llvm/llvm-project/commit/18311093abe6481388a0d963d58438d743b47569.diff

LOG: [InstCombine] Do not fold `shufflevector(select)` if the select condition is a vector (#113993)

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 https://github.com/llvm/llvm-project/issues/113986.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
    llvm/test/Transforms/InstCombine/vec_shuffle.ll

Removed: 
    


################################################################################
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()