[InstCombine] Avoid moving ops that do restrict undef across shuffles.

[fhahn]

Summary

I think we have to be a bit more careful when it comes to moving
ops across shuffles, if the op does restrict undef. For example, without
this patch, we would move 'and %v, <0, 0, -1, -1>' over a
'shufflevector %a, undef, <undef, undef, 1, 2>'. As a result, the first
2 lanes of the result are undef after the combine, but they really
should be 0, unless I am missing something.

For ops that do fold to undef on undef operands, the current behavior
should be fine. I've add conservative check OpDoesRestrictUndef, maybe
there's a better existing utility?

[fhahn]

clang-format

[lebedev.ri]

llvm/test/Transforms/InstCombine/vec_shuffle.ll

1010–1016

Hmm yes, i'd agree this is a miscompile..

[RKSimon]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1513

@spatel Do we have an existing helper for something like this? I know you added getKnownUndefForVectorBinop to the DAG

[spatel]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1513

I don't think there's a twin of that for IR. Unfortunately, we have bits and pieces of the puzzle in both places...

1552

What if we extend this check instead:

if (I >= SrcVecNumElts || ShMask[I] < 0) {

Currently, we're transferring an undef element from the original shuffle to the newly created shuffle, but that's not safe in general. It is safe, however, if the binop would constant fold to undef given an undef operand.

(Sorry this code got so complicated...)

llvm/test/Transforms/InstCombine/vec_shuffle.ll

1010–1016

This test is not committed to trunk? Add that along with similar patterns for other binop opcodes like the add/sub/shl noted in this patch? That way we'll know if this or my alternate suggestion is behaving as intended.

[fhahn]

llvm/test/Transforms/InstCombine/vec_shuffle.ll

1010–1016

I did not commit the test yet, but the diff here just shows the difference in output. I'll check with more positive and negative scenarios

[fhahn]

Update to use code that uses ConstantExpr to check if things fold to undef again.

[fhahn]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

1552

That's indeed a better place. I've updated the code to use the existing checks here.

[fhahn]

The FIXMEs are addressed by D70169

[spatel]

LGTM - thanks!

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombiner::foldVectorBinop(BinaryOperator &Inst) {
   Constant *C;
   if (match(&Inst, m_c_BinOp(
           m_OneUse(m_ShuffleVector(m_Value(V1), m_Undef(), m_Constant(Mask))),
           m_Constant(C))) &&
       V1->getType()->getVectorNumElements() <= NumElts) {
     assert(Inst.getType()->getScalarType() == V1->getType()->getScalarType() &&
            "Shuffle should not change scalar type");
 
     // Find constant NewC that has property:
     //   shuffle(NewC, ShMask) = C
     // If such constant does not exist (example: ShMask=<0,0> and C=<1,2>)
     // reorder is not possible. A 1-to-1 mapping is not required. Example:
     // ShMask = <1,1,2,2> and C = <5,5,6,6> --> NewC = <undef,5,6,undef>
     bool ConstOp1 = isa<Constant>(RHS);
     SmallVector<int, 16> ShMask;
     ShuffleVectorInst::getShuffleMask(Mask, ShMask);
     unsigned SrcVecNumElts = V1->getType()->getVectorNumElements();
for (unsigned I = 0; I < NumElts; ++I) {
           MayChange = false;
           break;
         }
         NewVecC[ShMask[I]] = CElt;
       }
       // If this is a widening shuffle, we must be able to extend with undef
       // elements. If the original binop does not produce an undef in the high
       // lanes, then this transform is not safe.
+      // Similarly for undef lanes due to the shuffle mask, we can only
+      // transform binops that preserve undef.
       // TODO: We could shuffle those non-undef constant values into the
       //       result by using a constant vector (rather than an undef vector)
       //       as operand 1 of the new binop, but that might be too aggressive
       //       for target-independent shuffle creation.
-      if (I >= SrcVecNumElts) {
+      if (I >= SrcVecNumElts || ShMask[I] < 0) {
         Constant *MaybeUndef =
             ConstOp1 ? ConstantExpr::get(Opcode, UndefScalar, CElt)
                      : ConstantExpr::get(Opcode, CElt, UndefScalar);
         if (!isa<UndefValue>(MaybeUndef)) {
           MayChange = false;
           break;
         }
       }

llvm/test/Transforms/InstCombine/vec_shuffle.ll

 ; CHECK-NEXT:    [[R:%.*]] = shufflevector <2 x i32> [[TMP1]], <2 x i32> undef, <2 x i32> zeroinitializer
 ; CHECK-NEXT:    ret <2 x i32> [[R]]
 ;
   %splat = shufflevector <2 x i32> %x, <2 x i32> undef, <2 x i32> zeroinitializer
   %r = and <2 x i32> %splat, <i32 42, i32 42>
   ret <2 x i32> %r
 }
 
+; AND does not fold to undef for undef operands, we cannot move it
+; across a shuffle with undef masks.
 define <4 x i16> @and_constant_mask_undef(<4 x i16> %add) {
 ; CHECK-LABEL: @and_constant_mask_undef(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[AND:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[AND:%.*]] = and <4 x i16> [[SHUFFLE]], <i16 0, i16 0, i16 -1, i16 -1>
 ; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %add, <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
   %and = and <4 x i16> %shuffle, <i16 0, i16 0, i16 -1, i16 -1>
   ret <4 x i16> %and
 }
 
+; AND does not fold to undef for undef operands, we cannot move it
+; across a shuffle with undef masks.
 define <4 x i16> @and_constant_mask_undef_2(<4 x i16> %add) {
 ; CHECK-LABEL: @and_constant_mask_undef_2(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[AND:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 1, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 1, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[AND:%.*]] = and <4 x i16> [[SHUFFLE]], <i16 -1, i16 -1, i16 -1, i16 0>
 ; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %add, <4 x i16> undef, <4 x i32> <i32 1, i32 1, i32 1, i32 undef>
   %and = and <4 x i16> %shuffle, <i16 -1, i16 -1, i16 -1, i16 -0>
   ret <4 x i16> %and
 }
 
+; FIXME: We should be able to move the AND across the shuffle, as -1 (AND identity value) is used for undef lanes.
 define <4 x i16> @and_constant_mask_undef_3(<4 x i16> %add) {
 ; CHECK-LABEL: @and_constant_mask_undef_3(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    ret <4 x i16> <i16 0, i16 0, i16 0, i16 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[AND:%.*]] = and <4 x i16> [[SHUFFLE]], <i16 0, i16 0, i16 0, i16 -1>
+; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %add, <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 1, i32 undef>
   %and = and <4 x i16> %shuffle, <i16 0, i16 0, i16 0, i16 -1>
   ret <4 x i16> %and
 }
 
+; FIXME: We should be able to move the AND across the shuffle, as -1 (AND identity value) is used for undef lanes.
 define <4 x i16> @and_constant_mask_undef_4(<4 x i16> %add) {
 ; CHECK-LABEL: @and_constant_mask_undef_4(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = and <4 x i16> [[ADD:%.*]], <i16 9, i16 20, i16 undef, i16 undef>
-; CHECK-NEXT:    [[AND:%.*]] = shufflevector <4 x i16> [[TMP0]], <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[ADD:%.*]], <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[AND:%.*]] = and <4 x i16> [[SHUFFLE]], <i16 9, i16 20, i16 20, i16 -1>
 ; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %add, <4 x i16> undef, <4 x i32> <i32 0, i32 1, i32 1, i32 undef>
   %and = and <4 x i16> %shuffle, <i16 9, i16 20, i16 20, i16 -1>
   ret <4 x i16> %and
 }
 
-
 define <4 x i16> @and_constant_mask_not_undef(<4 x i16> %add) {
 ; CHECK-LABEL: @and_constant_mask_not_undef(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = and <4 x i16> [[ADD:%.*]], <i16 undef, i16 -1, i16 0, i16 0>
 ; CHECK-NEXT:    [[AND:%.*]] = shufflevector <4 x i16> [[TMP0]], <4 x i16> undef, <4 x i32> <i32 2, i32 3, i32 1, i32 1>
 ; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %add, <4 x i16> undef, <4 x i32> <i32 2, i32 3, i32 1, i32 1>
   %and = and <4 x i16> %shuffle, <i16 0, i16 0, i16 -1, i16 -1>
   ret <4 x i16> %and
 }
 
+; OR does not fold to undef for undef operands, we cannot move it
+; across a shuffle with undef masks.
 define <4 x i16> @or_constant_mask_undef(<4 x i16> %in) {
 ; CHECK-LABEL: @or_constant_mask_undef(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[OR:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[OR:%.*]] = or <4 x i16> [[SHUFFLE]], <i16 -1, i16 -1, i16 0, i16 0>
 ; CHECK-NEXT:    ret <4 x i16> [[OR]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 undef, i32 undef, i32 1, i32 1>
   %or = or <4 x i16> %shuffle, <i16 -1, i16 -1, i16 0, i16 0>
   ret <4 x i16> %or
 }
 
+; OR does not fold to undef for undef operands, we cannot move it
+; across a shuffle with undef masks.
 define <4 x i16> @or_constant_mask_undef_2(<4 x i16> %in) {
 ; CHECK-LABEL: @or_constant_mask_undef_2(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[OR:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[OR:%.*]] = or <4 x i16> [[SHUFFLE]], <i16 -1, i16 0, i16 0, i16 -1>
 ; CHECK-NEXT:    ret <4 x i16> [[OR]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
   %or = or <4 x i16> %shuffle, <i16 -1, i16 0, i16 0, i16 -1>
   ret <4 x i16> %or
 }
 
+; FIXME: We should be able to move the OR across the shuffle, as 0 (OR identity value) is used for undef lanes.
 define <4 x i16> @or_constant_mask_undef_3(<4 x i16> %in) {
 ; CHECK-LABEL: @or_constant_mask_undef_3(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    ret <4 x i16> <i16 undef, i16 -1, i16 -1, i16 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[OR:%.*]] = or <4 x i16> [[SHUFFLE]], <i16 0, i16 -1, i16 -1, i16 0>
+; CHECK-NEXT:    ret <4 x i16> [[OR]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
   %or = or <4 x i16> %shuffle, <i16 0, i16 -1, i16 -1, i16 0>
   ret <4 x i16> %or
 }
 
+; FIXME: We should be able to move the OR across the shuffle, as 0 (OR identity value) is used for undef lanes.
 define <4 x i16> @or_constant_mask_undef_4(<4 x i16> %in) {
 ; CHECK-LABEL: @or_constant_mask_undef_4(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = or <4 x i16> [[IN:%.*]], <i16 undef, i16 99, i16 undef, i16 undef>
-; CHECK-NEXT:    [[OR:%.*]] = shufflevector <4 x i16> [[TMP0]], <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
+; CHECK-NEXT:    [[OR:%.*]] = or <4 x i16> [[SHUFFLE]], <i16 0, i16 99, i16 99, i16 0>
 ; CHECK-NEXT:    ret <4 x i16> [[OR]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 undef, i32 1, i32 1, i32 undef>
   %or = or <4 x i16> %shuffle, <i16 0, i16 99, i16 99, i16 0>
   ret <4 x i16> %or
 }
 
 define <4 x i16> @or_constant_mask_not_undef(<4 x i16> %in) {
 ; CHECK-LABEL: @or_constant_mask_not_undef(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = or <4 x i16> [[IN:%.*]], <i16 undef, i16 -1, i16 0, i16 0>
 ; CHECK-NEXT:    [[AND:%.*]] = shufflevector <4 x i16> [[TMP0]], <4 x i16> undef, <4 x i32> <i32 2, i32 3, i32 1, i32 1>
 ; CHECK-NEXT:    ret <4 x i16> [[AND]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 2, i32 3, i32 1, i32 1>
   %and = or <4 x i16> %shuffle, <i16 0, i16 0, i16 -1, i16 -1>
   ret <4 x i16> %and
 }
 
 define <4 x i16> @shl_constant_mask_undef(<4 x i16> %in) {
 ; CHECK-LABEL: @shl_constant_mask_undef(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = shl <4 x i16> [[IN:%.*]], <i16 10, i16 0, i16 0, i16 0>
-; CHECK-NEXT:    [[SHL:%.*]] = shufflevector <4 x i16> [[TMP0]], <4 x i16> undef, <4 x i32> <i32 0, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[SHUFFLE:%.*]] = shufflevector <4 x i16> [[IN:%.*]], <4 x i16> undef, <4 x i32> <i32 0, i32 undef, i32 1, i32 1>
+; CHECK-NEXT:    [[SHL:%.*]] = shl <4 x i16> [[SHUFFLE]], <i16 10, i16 3, i16 0, i16 0>
 ; CHECK-NEXT:    ret <4 x i16> [[SHL]]
 ;
 entry:
   %shuffle = shufflevector <4 x i16> %in, <4 x i16> undef, <4 x i32> <i32 0, i32 undef, i32 1, i32 1>
   %shl = shl <4 x i16> %shuffle, <i16 10, i16 3, i16 0, i16 0>
   ret <4 x i16> %shl
 }