[llvm] 192cb71 - [InstCombine] avoid crashing on select-shuffle detection

[Sanjay Patel via llvm-commits]

Author: Sanjay Patel
Date: 2020-06-04T17:27:14-04:00
New Revision: 192cb718361dbd7be082bc0893f43bbc9782288f

URL: https://github.com/llvm/llvm-project/commit/192cb718361dbd7be082bc0893f43bbc9782288f
DIFF: https://github.com/llvm/llvm-project/commit/192cb718361dbd7be082bc0893f43bbc9782288f.diff

LOG: [InstCombine] avoid crashing on select-shuffle detection

As mentioned in the post-commit comments of D81013 -
the mask check API has to assume the shuffle is
not length-changing, but we have not ruled that out
in this code. Use the ShuffleVectorInst call instead.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
    llvm/test/Transforms/InstCombine/select-select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 4b3b3c109543..f406eb40e16b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -2413,7 +2413,7 @@ Instruction *InstCombiner::foldVectorSelect(SelectInst &Sel) {
   ArrayRef<int> Mask;
   if (match(TVal, m_OneUse(m_Shuffle(m_Value(X), m_Value(Y), m_Mask(Mask)))) &&
       !is_contained(Mask, UndefMaskElem) &&
-      ShuffleVectorInst::isSelectMask(Mask)) {
+      cast<ShuffleVectorInst>(TVal)->isSelect()) {
     if (X == FVal) {
       // select Cond, (shuf_sel X, Y), X --> shuf_sel X, (select Cond, Y, X)
       Value *NewSel = Builder.CreateSelect(Cond, Y, X, "sel", &Sel);
@@ -2427,7 +2427,7 @@ Instruction *InstCombiner::foldVectorSelect(SelectInst &Sel) {
   }
   if (match(FVal, m_OneUse(m_Shuffle(m_Value(X), m_Value(Y), m_Mask(Mask)))) &&
       !is_contained(Mask, UndefMaskElem) &&
-      ShuffleVectorInst::isSelectMask(Mask)) {
+      cast<ShuffleVectorInst>(FVal)->isSelect()) {
     if (X == TVal) {
       // select Cond, X, (shuf_sel X, Y) --> shuf_sel X, (select Cond, X, Y)
       Value *NewSel = Builder.CreateSelect(Cond, X, Y, "sel", &Sel);

diff  --git a/llvm/test/Transforms/InstCombine/select-select.ll b/llvm/test/Transforms/InstCombine/select-select.ll
index c78326e7f7b6..f6662763012f 100644
--- a/llvm/test/Transforms/InstCombine/select-select.ll
+++ b/llvm/test/Transforms/InstCombine/select-select.ll
@@ -149,3 +149,29 @@ define <4 x i8> @sel_shuf_no_common_operand(<4 x i8> %x, <4 x i8> %y, <4 x i1> %
   %r = select <4 x i1> %cmp, <4 x i8> %z, <4 x i8> %blend
   ret <4 x i8> %r
 }
+
+; Negative test - don't crash (this is not a select shuffle because it changes vector length)
+
+define <2 x i8> @sel_shuf_narrowing_commute1(<4 x i8> %x, <4 x i8> %y, <2 x i8> %x2, <2 x i1> %cmp) {
+; CHECK-LABEL: @sel_shuf_narrowing_commute1(
+; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <2 x i32> <i32 0, i32 5>
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP:%.*]], <2 x i8> [[BLEND]], <2 x i8> [[X2:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[R]]
+;
+  %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <2 x i32> <i32 0, i32 5>
+  %r = select <2 x i1> %cmp, <2 x i8> %blend, <2 x i8> %x2
+  ret <2 x i8> %r
+}
+
+; Negative test - don't crash (this is not a select shuffle because it changes vector length)
+
+define <2 x i8> @sel_shuf_narrowing_commute2(<4 x i8> %x, <4 x i8> %y, <2 x i8> %x2, <2 x i1> %cmp) {
+; CHECK-LABEL: @sel_shuf_narrowing_commute2(
+; CHECK-NEXT:    [[BLEND:%.*]] = shufflevector <4 x i8> [[X:%.*]], <4 x i8> [[Y:%.*]], <2 x i32> <i32 0, i32 5>
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[CMP:%.*]], <2 x i8> [[X2:%.*]], <2 x i8> [[BLEND]]
+; CHECK-NEXT:    ret <2 x i8> [[R]]
+;
+  %blend = shufflevector <4 x i8> %x, <4 x i8> %y, <2 x i32> <i32 0, i32 5>
+  %r = select <2 x i1> %cmp, <2 x i8> %x2, <2 x i8> %blend
+  ret <2 x i8> %r
+}