[DAG] use isConstOrConstSplat in ComputeNumSignBits to optimize SRA

[spatel]

Summary

This is the minimum patch needed to solve the missed SRA fold that was noted in:
https://reviews.llvm.org/D25485

and partially fixed with:
https://reviews.llvm.org/rL284395

Of course, we could move more of the related analysis functions in DAGCombiner that are currently static/local to have a more unified approach.

Ie, move all of these (some were recently added in D25374)?
isConstantFPBuildVectorOrConstantFP
isConstOrConstSplatFP
isConstantOrConstantVector
isNullConstantOrNullSplatConstant
isOneConstantOrOneSplatConstant
isAllOnesConstantOrAllOnesSplatConstant

It's not clear to me exactly where the handling of undef elements is different and where that is intentional.

[spatel]

More consolidation is possible with the ISD namespace helpers, or should the previously listed functions go here too?

namespace ISD {
/// Node predicates

/// If N is a BUILD_VECTOR node whose elements are all the same constant or
/// undefined, return true and return the constant value in \p SplatValue.
bool isConstantSplatVector(const SDNode *N, APInt &SplatValue);

/// Return true if the specified node is a BUILD_VECTOR where all of the
/// elements are ~0 or undef.
bool isBuildVectorAllOnes(const SDNode *N);

/// Return true if the specified node is a BUILD_VECTOR where all of the
/// elements are 0 or undef.
bool isBuildVectorAllZeros(const SDNode *N);

/// Return true if the specified node is a BUILD_VECTOR node of all
/// ConstantSDNode or undef.
bool isBuildVectorOfConstantSDNodes(const SDNode *N);

/// Return true if the specified node is a BUILD_VECTOR node of all
/// ConstantFPSDNode or undef.
bool isBuildVectorOfConstantFPSDNodes(const SDNode *N);

/// Return true if the node has at least one operand and all operands of the
/// specified node are ISD::UNDEF.
bool allOperandsUndef(const SDNode *N);
}  // end llvm:ISD namespace

[RKSimon]

I agree that all those similar predicates is a bit of pain!

We have issues with handling whether we permit opaque constants and whether UNDEF elements are allowable.

include/llvm/CodeGen/SelectionDAGNodes.h

1419

Might this be a good time to add an 'AllowUndefs' argument so that we permit cases where there is at least one ConstantSDNode in the splat but other elements can be UNDEF.

[RKSimon]

LGTM - I think isConstOrConstSplatFP should probably moved over at the same time as isConstOrConstSplat (separate preliminary commit) followed by the SRA fix - make sense?

We can then investigate adding a "AllowUndefs" args to both of them in a followup commit and start the fun job of consolidating all those predicates.....

include/llvm/CodeGen/SelectionDAGNodes.h

public:
   static bool classof(const SDNode *N) {
     return N->getOpcode() == ISD::ConstantFP ||
            N->getOpcode() == ISD::TargetConstantFP;
   }
 };
 
 /// Returns true if \p V is a constant integer zero.
 bool isNullConstant(SDValue V);
+
 /// Returns true if \p V is an FP constant with a value of positive zero.
 bool isNullFPConstant(SDValue V);
+
 /// Returns true if \p V is an integer constant with all bits set.
 bool isAllOnesConstant(SDValue V);
+
 /// Returns true if \p V is a constant integer one.
 bool isOneConstant(SDValue V);
+
 /// Returns true if \p V is a bitwise not operation. Assumes that an all ones
 /// constant is canonicalized to be operand 1.
 bool isBitwiseNot(SDValue V);
 
+/// Returns the SDNode if it is a constant splat BuildVector or constant int.
+ConstantSDNode *isConstOrConstSplat(SDValue N);
+
+
 class GlobalAddressSDNode : public SDNode {
   const GlobalValue *TheGlobal;
   int64_t Offset;
   unsigned char TargetFlags;
   friend class SelectionDAG;
   GlobalAddressSDNode(unsigned Opc, unsigned Order, const DebugLoc &DL,
                       const GlobalValue *GA, EVT VT, int64_t o,
                       unsigned char TargetFlags);

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDNode *isConstantFPBuildVectorOrConstantFP(SDValue N) {
   if (isa<ConstantFPSDNode>(N))
     return N.getNode();
   if (ISD::isBuildVectorOfConstantFPSDNodes(N.getNode()))
     return N.getNode();
   return nullptr;
 }
 
 // \brief Returns the SDNode if it is a constant splat BuildVector or constant
-// int.
-static ConstantSDNode *isConstOrConstSplat(SDValue N) {
-  if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
-    return CN;
-
-  if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
-    BitVector UndefElements;
-    ConstantSDNode *CN = BV->getConstantSplatNode(&UndefElements);
-
-    // BuildVectors can truncate their operands. Ignore that case here.
-    // FIXME: We blindly ignore splats which include undef which is overly
-    // pessimistic.
-    if (CN && UndefElements.none() &&
-        CN->getValueType(0) == N.getValueType().getScalarType())
-      return CN;
-  }
-
-  return nullptr;
-}
-
-// \brief Returns the SDNode if it is a constant splat BuildVector or constant
 // float.
 static ConstantFPSDNode *isConstOrConstSplatFP(SDValue N) {
   if (ConstantFPSDNode *CN = dyn_cast<ConstantFPSDNode>(N))
     return CN;
 
   if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
     BitVector UndefElements;
     ConstantFPSDNode *CN = BV->getConstantFPSplatNode(&UndefElements);

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::ADD:
     if (Tmp2 == 1) return 1;
     return std::min(Tmp, Tmp2)-1;
 
   case ISD::SUB:
     Tmp2 = ComputeNumSignBits(Op.getOperand(1), Depth+1);
     if (Tmp2 == 1) return 1;
 
     // Handle NEG.
-    if (ConstantSDNode *CLHS = dyn_cast<ConstantSDNode>(Op.getOperand(0)))
+    if (ConstantSDNode *CLHS = isConstOrConstSplat(Op.getOperand(0)))
       if (CLHS->isNullValue()) {
         APInt KnownZero, KnownOne;
         computeKnownBits(Op.getOperand(1), KnownZero, KnownOne, Depth+1);
         // If the input is known to be 0 or 1, the output is 0/-1, which is all
         // sign bits set.
         if ((KnownZero | APInt(VTBits, 1)).isAllOnesValue())
           return VTBits;
 
bool llvm::isOneConstant(SDValue V) {
   ConstantSDNode *Const = dyn_cast<ConstantSDNode>(V);
   return Const != nullptr && Const->isOne();
 }
 
 bool llvm::isBitwiseNot(SDValue V) {
   return V.getOpcode() == ISD::XOR && isAllOnesConstant(V.getOperand(1));
 }
 
+ConstantSDNode *llvm::isConstOrConstSplat(SDValue N) {
+  if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
+    return CN;
+
+  if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
+    BitVector UndefElements;
+    ConstantSDNode *CN = BV->getConstantSplatNode(&UndefElements);
+
+    // BuildVectors can truncate their operands. Ignore that case here.
+    // FIXME: We blindly ignore splats which include undef which is overly
+    // pessimistic.
+    if (CN && UndefElements.none() &&
+        CN->getValueType(0) == N.getValueType().getScalarType())
+      return CN;
+  }
+
+  return nullptr;
+}
+
 HandleSDNode::~HandleSDNode() {
   DropOperands();
 }
 
 GlobalAddressSDNode::GlobalAddressSDNode(unsigned Opc, unsigned Order,
                                          const DebugLoc &DL,
                                          const GlobalValue *GA, EVT VT,
                                          int64_t o, unsigned char TF)

test/CodeGen/X86/sar_fold64.ll

 
 define <4 x i32> @all_sign_bit_ashr_vec(<4 x i32> %x) {
 ; CHECK-LABEL: all_sign_bit_ashr_vec:
 ; CHECK:       # BB#0:
 ; CHECK-NEXT:    pand {{.*}}(%rip), %xmm0
 ; CHECK-NEXT:    pxor %xmm1, %xmm1
 ; CHECK-NEXT:    psubd %xmm0, %xmm1
 ; CHECK-NEXT:    movdqa %xmm1, %xmm0
-; CHECK-NEXT:    psrad $31, %xmm0
-; CHECK-NEXT:    movdqa %xmm1, %xmm2
-; CHECK-NEXT:    movsd {{.*#+}} xmm2 = xmm0[0],xmm2[1]
-; CHECK-NEXT:    pshufd {{.*#+}} xmm2 = xmm2[1,3,2,3]
-; CHECK-NEXT:    movdqa %xmm1, %xmm0
-; CHECK-NEXT:    psrad $5, %xmm0
-; CHECK-NEXT:    psrad $1, %xmm1
-; CHECK-NEXT:    movsd {{.*#+}} xmm0 = xmm1[0],xmm0[1]
-; CHECK-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[0,2,2,3]
-; CHECK-NEXT:    punpckldq {{.*#+}} xmm0 = xmm0[0],xmm2[0],xmm0[1],xmm2[1]
 ; CHECK-NEXT:    retq
 ;
   %and = and <4 x i32> %x, <i32 1, i32 1, i32 1 , i32 1>
   %neg = sub <4 x i32> zeroinitializer, %and
   %sar = ashr <4 x i32> %neg, <i32 1, i32 31, i32 5, i32 0>
   ret <4 x i32> %sar
 }
 
 attributes #0 = { nounwind }