[SelectionDAG] Optional handling of UNDEF elements in matchUnaryPredicate

[RKSimon]

Summary

Now that SimplifyDemandedBits/SimplifyDemandedVectorElts is simplifying vector elements, we're seeing more constant BUILD_VECTOR containing UNDEFs.

This patch provides opt-in handling of UNDEF elements in matchUnaryPredicate, passing NULL instead of the result ConstantSDNode* argument.

I've updated SelectionDAG::simplifyShift to demonstrate its use.

An equivalent matchBinaryPredicate patch is in progress.

[andreadb]

I think this is the right way to go.

SimplifyDemandedVectorElts is getting better and better at identifying unused lanes of a vector. The more we improve that logic, the higher the risk of propagating undefs to unused lanes of a constant build_vector. Teaching matchUnaryPredicate (and, in future, matchBinaryPredicate too) how/when to ignore undef elements makes a lot of sense to me. It makes the logic resilient to the presence of undefs, and it avoids unexpected regressions (indirectly caused by the aggressive propagation of undefs).

[craig.topper]

LGTM

llvm/trunk/include/llvm/CodeGen/SelectionDAGNodes.h

return (Op == ISD::ADD || Op == ISD::SUB || Op == ISD::MUL ||
             Op == ISD::SHL || Op == ISD::SRL || Op == ISD::SRA ||
             Op == ISD::SDIV || Op == ISD::UDIV || Op == ISD::SREM ||
             Op == ISD::UREM || Op == ISD::FADD || Op == ISD::FSUB ||
             Op == ISD::FMUL || Op == ISD::FDIV || Op == ISD::FREM);
   }
 
   /// Attempt to match a unary predicate against a scalar/splat constant or
   /// every element of a constant BUILD_VECTOR.
+  /// If AllowUndef is true, then UNDEF element will pass nullptr to Match.
   bool matchUnaryPredicate(SDValue Op,
-                           std::function<bool(ConstantSDNode *)> Match);
+                           std::function<bool(ConstantSDNode *)> Match,
+                           bool AllowUndefs = false);
 
   /// Attempt to match a binary predicate against a pair of scalar/splat
   /// constants or every element of a pair of constant BUILD_VECTORs.
   bool matchBinaryPredicate(
       SDValue LHS, SDValue RHS,
       std::function<bool(ConstantSDNode *, ConstantSDNode *)> Match);
 
 } // end namespace ISD
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_SELECTIONDAGNODES_H

llvm/trunk/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool ISD::allOperandsUndef(const SDNode *N) {
   for (const SDValue &Op : N->op_values())
     if (!Op.isUndef())
       return false;
 
   return true;
 }
 
 bool ISD::matchUnaryPredicate(SDValue Op,
-                              std::function<bool(ConstantSDNode *)> Match) {
+                              std::function<bool(ConstantSDNode *)> Match,
+                              bool AllowUndefs) {
+  // FIXME: Add support for scalar UNDEF cases?
   if (auto *Cst = dyn_cast<ConstantSDNode>(Op))
     return Match(Cst);
 
+  // FIXME: Add support for vector UNDEF cases?
   if (ISD::BUILD_VECTOR != Op.getOpcode())
     return false;
 
   EVT SVT = Op.getValueType().getScalarType();
   for (unsigned i = 0, e = Op.getNumOperands(); i != e; ++i) {
+    if (AllowUndefs && Op.getOperand(i).isUndef()) {
+      if (!Match(nullptr))
+        return false;
+      continue;
+    }
+
     auto *Cst = dyn_cast<ConstantSDNode>(Op.getOperand(i));
     if (!Cst || Cst->getValueType(0) != SVT || !Match(Cst))
       return false;
   }
   return true;
 }
 
 bool ISD::matchBinaryPredicate(
if (Y.isUndef())
     return getUNDEF(X.getValueType());
 
   // shift 0, Y --> 0
   // shift X, 0 --> X
   if (isNullOrNullSplat(X) || isNullOrNullSplat(Y))
     return X;
 
   // shift X, C >= bitwidth(X) --> undef
-  // All vector elements must be too big to avoid partial undefs.
+  // All vector elements must be too big (or undef) to avoid partial undefs.
   auto isShiftTooBig = [X](ConstantSDNode *Val) {
-    return Val->getAPIntValue().uge(X.getScalarValueSizeInBits());
+    return !Val || Val->getAPIntValue().uge(X.getScalarValueSizeInBits());
   };
-  if (ISD::matchUnaryPredicate(Y, isShiftTooBig))
+  if (ISD::matchUnaryPredicate(Y, isShiftTooBig, true))
     return getUNDEF(X.getValueType());
 
   return SDValue();
 }
 
 SDValue SelectionDAG::getVAArg(EVT VT, const SDLoc &dl, SDValue Chain,
                                SDValue Ptr, SDValue SV, unsigned Align) {
   SDValue Ops[] = { Chain, Ptr, SV, getTargetConstant(Align, dl, MVT::i32) };

llvm/trunk/test/CodeGen/X86/combine-shl.ll

 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    retq
   %1 = and <4 x i32> %a0, <i32 2147483647, i32 2147483647, i32 2147483647, i32 2147483647>
   %2 = shl <4 x i32> %1, <i32 33, i32 33, i32 33, i32 33>
   ret <4 x i32> %2
 }
 
 define <4 x i32> @combine_vec_shl_outofrange3(<4 x i32> %a0) {
-; SSE-LABEL: combine_vec_shl_outofrange3:
-; SSE:       # %bb.0:
-; SSE-NEXT:    xorps %xmm0, %xmm0
-; SSE-NEXT:    retq
-;
-; AVX-LABEL: combine_vec_shl_outofrange3:
-; AVX:       # %bb.0:
-; AVX-NEXT:    vpsllvd {{.*}}(%rip), %xmm0, %xmm0
-; AVX-NEXT:    retq
+; CHECK-LABEL: combine_vec_shl_outofrange3:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    retq
   %1 = shl <4 x i32> %a0, <i32 33, i32 34, i32 35, i32 undef>
   ret <4 x i32> %1
 }
 
 ; fold (shl x, 0) -> x
 define <4 x i32> @combine_vec_shl_by_zero(<4 x i32> %x) {
 ; CHECK-LABEL: combine_vec_shl_by_zero:
 ; CHECK:       # %bb.0:

llvm/trunk/test/CodeGen/X86/combine-sra.ll

 ; CHECK-LABEL: combine_vec_ashr_outofrange1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    retq
   %1 = ashr <4 x i32> %x, <i32 33, i32 34, i32 35, i32 36>
   ret <4 x i32> %1
 }
 
 define <4 x i32> @combine_vec_ashr_outofrange2(<4 x i32> %x) {
-; SSE-LABEL: combine_vec_ashr_outofrange2:
-; SSE:       # %bb.0:
-; SSE-NEXT:    retq
-;
-; AVX-LABEL: combine_vec_ashr_outofrange2:
-; AVX:       # %bb.0:
-; AVX-NEXT:    vpsravd {{.*}}(%rip), %xmm0, %xmm0
-; AVX-NEXT:    retq
+; CHECK-LABEL: combine_vec_ashr_outofrange2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    retq
   %1 = ashr <4 x i32> %x, <i32 33, i32 34, i32 35, i32 undef>
   ret <4 x i32> %1
 }
 
 ; fold (sra x, 0) -> x
 define <4 x i32> @combine_vec_ashr_by_zero(<4 x i32> %x) {
 ; CHECK-LABEL: combine_vec_ashr_by_zero:
 ; CHECK:       # %bb.0:

llvm/trunk/test/CodeGen/X86/combine-srl.ll

 ; CHECK-LABEL: combine_vec_lshr_outofrange1:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    retq
   %1 = lshr <4 x i32> %x, <i32 33, i32 34, i32 35, i32 36>
   ret <4 x i32> %1
 }
 
 define <4 x i32> @combine_vec_lshr_outofrange2(<4 x i32> %x) {
-; SSE-LABEL: combine_vec_lshr_outofrange2:
-; SSE:       # %bb.0:
-; SSE-NEXT:    retq
-;
-; AVX-LABEL: combine_vec_lshr_outofrange2:
-; AVX:       # %bb.0:
-; AVX-NEXT:    vpsrlvd {{.*}}(%rip), %xmm0, %xmm0
-; AVX-NEXT:    retq
+; CHECK-LABEL: combine_vec_lshr_outofrange2:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    retq
   %1 = lshr <4 x i32> %x, <i32 33, i32 34, i32 35, i32 undef>
   ret <4 x i32> %1
 }
 
 ; fold (srl x, 0) -> x
 define <4 x i32> @combine_vec_lshr_by_zero(<4 x i32> %x) {
 ; CHECK-LABEL: combine_vec_lshr_by_zero:
 ; CHECK:       # %bb.0: