Limit the recursion depth of SelectionDAG::isSplatValue()

[capn]

Summary

This method previously always recursively checked both the left-hand
side and right-hand side of binary operations for splatted (broadcast)
vector values to determine if the parent DAG node is a splat.

Like several other SelectionDAG methods, limit the recursion depth to
MaxRecursionDepth (6). This prevents stack overflow.
See also https://issuetracker.google.com/173785481

[capn]

Fix default argument

[bogner]

LGTM once it actually compiles (we'll certainly need to call isSplatValue with a depth of zero somewhere)

[RKSimon]

makes sense - thanks for dealing with this

[capn]

Thanks for the quick review! I do not have commit access, please commit it for me.

[capn]

Friendly ping. Can someone with commit privileges land this please?

[bogner]

In D92421#2438963, @capn wrote:

Friendly ping. Can someone with commit privileges land this please?

Oh, sorry about that. I'll commit this for you today.

llvm/include/llvm/CodeGen/SelectionDAG.h

public:
 
   /// Test whether \p V has a splatted value for all the demanded elements.
   ///
   /// On success \p UndefElts will indicate the elements that have UNDEF
   /// values instead of the splat value, this is only guaranteed to be correct
   /// for \p DemandedElts.
   ///
   /// NOTE: The function will return true for a demanded splat of UNDEF values.
-  bool isSplatValue(SDValue V, const APInt &DemandedElts, APInt &UndefElts);
+  bool isSplatValue(SDValue V, const APInt &DemandedElts, APInt &UndefElts,
+                    unsigned Depth);
 
   /// Test whether \p V has a splatted value.
   bool isSplatValue(SDValue V, bool AllowUndefs = false);
 
   /// If V is a splatted value, return the source vector and its splat index.
   SDValue getSplatSourceVector(SDValue V, int &SplatIndex);
 
   /// If V is a splat vector, return its scalar source operand by extracting

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool SelectionDAG::MaskedValueIsAllOnes(SDValue V, const APInt &Mask,
   return Mask.isSubsetOf(computeKnownBits(V, Depth).One);
 }
 
 /// isSplatValue - Return true if the vector V has the same value
 /// across all DemandedElts. For scalable vectors it does not make
 /// sense to specify which elements are demanded or undefined, therefore
 /// they are simply ignored.
 bool SelectionDAG::isSplatValue(SDValue V, const APInt &DemandedElts,
-                                APInt &UndefElts) {
+                                APInt &UndefElts, unsigned Depth) {
   EVT VT = V.getValueType();
   assert(VT.isVector() && "Vector type expected");
 
   if (!VT.isScalableVector() && !DemandedElts)
     return false; // No demanded elts, better to assume we don't know anything.
 
+  if (Depth >= MaxRecursionDepth)
+    return false; // Limit search depth.
+
   // Deal with some common cases here that work for both fixed and scalable
   // vector types.
   switch (V.getOpcode()) {
   case ISD::SPLAT_VECTOR:
     UndefElts = V.getOperand(0).isUndef()
                     ? APInt::getAllOnesValue(DemandedElts.getBitWidth())
                     : APInt(DemandedElts.getBitWidth(), 0);
     return true;
   case ISD::ADD:
   case ISD::SUB:
   case ISD::AND: {
     APInt UndefLHS, UndefRHS;
     SDValue LHS = V.getOperand(0);
     SDValue RHS = V.getOperand(1);
-    if (isSplatValue(LHS, DemandedElts, UndefLHS) &&
-        isSplatValue(RHS, DemandedElts, UndefRHS)) {
+    if (isSplatValue(LHS, DemandedElts, UndefLHS, Depth + 1) &&
+        isSplatValue(RHS, DemandedElts, UndefRHS, Depth + 1)) {
       UndefElts = UndefLHS | UndefRHS;
       return true;
     }
     break;
   }
   case ISD::TRUNCATE:
   case ISD::SIGN_EXTEND:
   case ISD::ZERO_EXTEND:
-    return isSplatValue(V.getOperand(0), DemandedElts, UndefElts);
+    return isSplatValue(V.getOperand(0), DemandedElts, UndefElts, Depth + 1);
   }
 
   // We don't support other cases than those above for scalable vectors at
   // the moment.
   if (VT.isScalableVector())
     return false;
 
   unsigned NumElts = VT.getVectorNumElements();
bool SelectionDAG::isSplatValue(SDValue V, const APInt &DemandedElts,
   }
   case ISD::EXTRACT_SUBVECTOR: {
     // Offset the demanded elts by the subvector index.
     SDValue Src = V.getOperand(0);
     uint64_t Idx = V.getConstantOperandVal(1);
     unsigned NumSrcElts = Src.getValueType().getVectorNumElements();
     APInt UndefSrcElts;
     APInt DemandedSrcElts = DemandedElts.zextOrSelf(NumSrcElts).shl(Idx);
-    if (isSplatValue(Src, DemandedSrcElts, UndefSrcElts)) {
+    if (isSplatValue(Src, DemandedSrcElts, UndefSrcElts, Depth + 1)) {
       UndefElts = UndefSrcElts.extractBits(NumElts, Idx);
       return true;
     }
     break;
   }
   }
 
   return false;
 }
 
 /// Helper wrapper to main isSplatValue function.
 bool SelectionDAG::isSplatValue(SDValue V, bool AllowUndefs) {
   EVT VT = V.getValueType();
   assert(VT.isVector() && "Vector type expected");
 
   APInt UndefElts;
   APInt DemandedElts;
 
   // For now we don't support this with scalable vectors.
   if (!VT.isScalableVector())
     DemandedElts = APInt::getAllOnesValue(VT.getVectorNumElements());
   return isSplatValue(V, DemandedElts, UndefElts) &&
          (AllowUndefs || !UndefElts);
 }
 
 SDValue SelectionDAG::getSplatSourceVector(SDValue V, int &SplatIdx) {
   V = peekThroughExtractSubvectors(V);
 
   EVT VT = V.getValueType();
   unsigned Opcode = V.getOpcode();
   switch (Opcode) {
   default: {
     APInt UndefElts;
     APInt DemandedElts;
 
     if (!VT.isScalableVector())
       DemandedElts = APInt::getAllOnesValue(VT.getVectorNumElements());
 
     if (isSplatValue(V, DemandedElts, UndefElts)) {
       if (VT.isScalableVector()) {
         // DemandedElts and UndefElts are ignored for scalable vectors, since
         // the only supported cases are SPLAT_VECTOR nodes.
         SplatIdx = 0;
       } else {
         // Handle case where all demanded elements are UNDEF.
         if (DemandedElts.isSubsetOf(UndefElts)) {
           SplatIdx = 0;