[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...

1568

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

1568

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

if (LShuf->isSelect() && !LShuf->getMask()->containsUndefElement() &&
       // RHS = shuffle V2, V1, <0, 5, 6, 3>
       // LHS + RHS --> (V10+V20, V21+V11, V22+V12, V13+V23) --> V1 + V2
       Instruction *NewBO = BinaryOperator::Create(Opcode, V1, V2);
       NewBO->copyIRFlags(&Inst);
       return NewBO;
     }
   }
 
+  // Returns true, if the binary operator restricts undef operands, i.e. may no fold into
+  // undef for undef operands, like 'and undef, %x' != undef.
+  auto OpDoesRestrictUndef = [](BinaryOperator::BinaryOps Opcode) {
+    // Conservatively check opcodes that are guaranteed to fold to undef,
+    // for an undef operand.
+    switch (Opcode) {
+      case Instruction::Add:
+      case Instruction::Sub:
+      case Instruction::Shl:
+        return false;
+    default:
+      return true;
+    }
+  };
+
   // If one argument is a shuffle within one vector and the other is a constant,
   // try moving the shuffle after the binary operation. This canonicalization
   // intends to move shuffles closer to other shuffles and binops closer to
   // other binops, so they can be folded. It may also enable demanded elements
   // transforms.
   Constant *C;
   if (match(&Inst, m_c_BinOp(
           m_OneUse(m_ShuffleVector(m_Value(V1), m_Undef(), m_Constant(Mask))),
for (unsigned I = 0; I < NumElts; ++I) {
         // 3. This is a widening shuffle that copies elements of V1 into the
         //    extended elements (extending with undef is allowed).
         if (!CElt || (!isa<UndefValue>(NewCElt) && NewCElt != CElt) ||
             I >= SrcVecNumElts) {
           MayChange = false;
           break;
         }
         NewVecC[ShMask[I]] = CElt;
+      } else if (OpDoesRestrictUndef(Opcode)) {
+        MayChange = false;
+        break;
       }
       // 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.
       // 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) {
         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
 }
 
 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]]