[InstCombine] Do not perform binop-of-shuffle when mask is poison

[antoniofrighetto]

Shuffle masks operand may be turn into poison if this does not lead to observable changes. This however may not guarantee binop to binop-of-shuffle replacement to be sound anymore. IIUC the transformation, we shouldn't need to check for all the mask operands being poisoned.

Fixes: #82052.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Antonio Frighetto (antoniofrighetto)

Changes

Shuffle masks operand may be turn into poison if this does not lead to observable changes. This however may not guarantee binop to binop-of-shuffle replacement to be sound anymore. IIUC the transformation, we shouldn't need to check for all the mask operands being poisoned.

Fixes: #82052.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp (+5-3)
• (modified) llvm/test/Transforms/InstCombine/vec_demanded_elts.ll (+50)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
index 5f13454089e515..97ae980a7cba70 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp
@@ -1870,14 +1870,16 @@ Value *InstCombinerImpl::SimplifyDemandedVectorElts(Value *V,
         Value *ShufOp = MatchShufAsOp0 ? X : Y;
         Value *OtherOp = MatchShufAsOp0 ? Y : X;
         for (User *U : OtherOp->users()) {
-          auto Shuf = m_Shuffle(m_Specific(ShufOp), m_Value(), m_ZeroMask());
+          ArrayRef<int> Mask;
+          auto Shuf = m_Shuffle(m_Specific(ShufOp), m_Value(), m_Mask(Mask));
           if (BO->isCommutative()
                   ? match(U, m_c_BinOp(Opcode, Shuf, m_Specific(OtherOp)))
                   : MatchShufAsOp0
                         ? match(U, m_BinOp(Opcode, Shuf, m_Specific(OtherOp)))
                         : match(U, m_BinOp(Opcode, m_Specific(OtherOp), Shuf)))
-            if (DT.dominates(U, I))
-              return U;
+            if (match(Mask, m_ZeroMask()) && Mask[0] != PoisonMaskElem)
+              if (DT.dominates(U, I))
+                return U;
         }
         return nullptr;
       };
diff --git a/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll b/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
index 4e5fdbb1eab991..c2c7d264c66491 100644
--- a/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
+++ b/llvm/test/Transforms/InstCombine/vec_demanded_elts.ll
@@ -1159,3 +1159,53 @@ define i4 @common_binop_demand_via_extelt_op0_mismatch_elt1(<2 x i4> %x, <2 x i4
   call void @use(<2 x i4> %b_xshuf_y)
   ret i4 %b_xy0
 }
+
+define <2 x i8> @common_binop_demand_via_splat_mask_poison(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @common_binop_demand_via_splat_mask_poison(
+; CHECK-NEXT:    [[YSPLAT:%.*]] = shufflevector <2 x i8> [[Y:%.*]], <2 x i8> poison, <2 x i32> <i32 0, i32 poison>
+; CHECK-NEXT:    [[VV:%.*]] = add <2 x i8> [[YSPLAT]], [[X:%.*]]
+; CHECK-NEXT:    [[MSPLAT:%.*]] = shufflevector <2 x i8> [[VV]], <2 x i8> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[RES:%.*]] = add <2 x i8> [[VV]], [[MSPLAT]]
+; CHECK-NEXT:    ret <2 x i8> [[RES]]
+;
+  %ysplat = shufflevector <2 x i8> %y, <2 x i8> poison, <2 x i32> <i32 0, i32 poison>                       ; <y0, poison>
+  %vv = add <2 x i8> %x, %ysplat                                                                            ; <x0+y0, poison>
+  %m = add <2 x i8> %x, %y                                                                                  ; <x0+y0, x1+y1>
+  %msplat = shufflevector <2 x i8> %m, <2 x i8> poison, <2 x i32> <i32 0, i32 0>      ; LeftDemanded = 1    ; <x0+y0, x0+y0>
+  %res = add <2 x i8> %vv, %msplat                                                                          ; <x0+y0+x0+y0, poison>
+  ret <2 x i8> %res
+}
+
+define <2 x i8> @common_binop_demand_via_splat_mask_poison_2(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @common_binop_demand_via_splat_mask_poison_2(
+; CHECK-NEXT:    [[YSPLAT:%.*]] = shufflevector <2 x i8> [[Y:%.*]], <2 x i8> poison, <2 x i32> <i32 poison, i32 0>
+; CHECK-NEXT:    [[VV:%.*]] = add <2 x i8> [[YSPLAT]], [[X:%.*]]
+; CHECK-NEXT:    [[M:%.*]] = add <2 x i8> [[X]], [[Y]]
+; CHECK-NEXT:    [[MSPLAT:%.*]] = shufflevector <2 x i8> [[M]], <2 x i8> [[Y]], <2 x i32> <i32 0, i32 2>
+; CHECK-NEXT:    [[RES:%.*]] = add <2 x i8> [[VV]], [[MSPLAT]]
+; CHECK-NEXT:    ret <2 x i8> [[RES]]
+;
+  %ysplat = shufflevector <2 x i8> %y, <2 x i8> poison, <2 x i32> <i32 poison, i32 0>
+  %vv = add <2 x i8> %x, %ysplat
+  %m = add <2 x i8> %x, %y
+  %msplat = shufflevector <2 x i8> %m, <2 x i8> %y, <2 x i32> <i32 0, i32 2>           ; LeftDemanded = 1, RightDemanded = 1
+  %res = add <2 x i8> %vv, %msplat
+  ret <2 x i8> %res
+}
+
+define <2 x i8> @common_binop_demand_via_splat_mask_poison_3(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
+; CHECK-LABEL: @common_binop_demand_via_splat_mask_poison_3(
+; CHECK-NEXT:    [[YSPLAT:%.*]] = shufflevector <2 x i8> [[Y:%.*]], <2 x i8> poison, <2 x i32> <i32 poison, i32 0>
+; CHECK-NEXT:    [[VV:%.*]] = add <2 x i8> [[YSPLAT]], [[X:%.*]]
+; CHECK-NEXT:    [[M:%.*]] = add <2 x i8> [[X]], [[Y]]
+; CHECK-NEXT:    [[MSPLAT:%.*]] = shufflevector <2 x i8> [[M]], <2 x i8> poison, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[RES:%.*]] = add <2 x i8> [[VV]], [[MSPLAT]]
+; CHECK-NEXT:    ret <2 x i8> [[RES]]
+;
+  %ysplat = shufflevector <2 x i8> %y, <2 x i8> poison, <2 x i32> <i32 poison, i32 0>
+  %vv = add <2 x i8> %x, %ysplat
+  %m = add <2 x i8> %x, %y
+  %msplat = shufflevector <2 x i8> %m, <2 x i8> poison, <2 x i32> <i32 0, i32 0>       ; LeftDemanded = 1
+  %res = add <2 x i8> %vv, %msplat
+  ret <2 x i8> %res
+}

[goldsteinn]

It looks like only your third test is problematic: https://alive2.llvm.org/ce/z/WrH49f
can the fix be made more precise?

[antoniofrighetto]

It looks like only your third test is problematic: https://alive2.llvm.org/ce/z/WrH49f
can the fix be made more precise?

Wouldn't that require to know right / left demanded of the mask of the use of the operand passed to SimplifyDemandedVectorElts (%msplat)? Not sure if there's a better / worth way to do this.

[antoniofrighetto]

@goldsteinn, don't have a strong opinion, but considering this looks quite of a corner case, it is fuzzer-generated, and in the general case it may not be safe, I'd rather favour readability / not make the code uglier, and rather being more conservative (the input space looks already quite small anyways).

[nikic]

LGTM