[PATCH] D33775: [SelectionDAG] Get rid of recursion in findNonImmUse

[Justin Bogner via llvm-commits]

Max Kazantsev via Phabricator <reviews at reviews.llvm.org> writes:
> mkazantsev created this revision.
>
> The recursive implementation of findNonImmUse may overflow stack
> on extremely long use chains. This patch replaces it with an equivalent
> iterative implementation.

LGTM.

> https://reviews.llvm.org/D33775
>
> Files:
>   lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
>
> Index: lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> +++ lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> @@ -2022,7 +2022,7 @@
>  }
>  
>  /// findNonImmUse - Return true if "Use" is a non-immediate use of "Def".
> -/// This function recursively traverses up the operand chain, ignoring
> +/// This function iteratively traverses up the operand chain, ignoring
>  /// certain nodes.
>  static bool findNonImmUse(SDNode *Use, SDNode* Def, SDNode *ImmedUse,
>                            SDNode *Root, SmallPtrSetImpl<SDNode*> &Visited,
> @@ -2035,30 +2035,36 @@
>    // The Use may be -1 (unassigned) if it is a newly allocated node.  This can
>    // happen because we scan down to newly selected nodes in the case of glue
>    // uses.
> -  if ((Use->getNodeId() < Def->getNodeId() && Use->getNodeId() != -1))
> -    return false;
> +  std::vector<SDNode *> WorkList;
> +  WorkList.push_back(Use);
>  
> -  // Don't revisit nodes if we already scanned it and didn't fail, we know we
> -  // won't fail if we scan it again.
> -  if (!Visited.insert(Use).second)
> -    return false;
> +  while (!WorkList.empty()) {
> +    Use = WorkList.back();
> +    WorkList.pop_back();
> +    if (Use->getNodeId() < Def->getNodeId() && Use->getNodeId() != -1)
> +      continue;
>  
> -  for (const SDValue &Op : Use->op_values()) {
> -    // Ignore chain uses, they are validated by HandleMergeInputChains.
> -    if (Op.getValueType() == MVT::Other && IgnoreChains)
> +    // Don't revisit nodes if we already scanned it and didn't fail, we know we
> +    // won't fail if we scan it again.
> +    if (!Visited.insert(Use).second)
>        continue;
>  
> -    SDNode *N = Op.getNode();
> -    if (N == Def) {
> -      if (Use == ImmedUse || Use == Root)
> -        continue;  // We are not looking for immediate use.
> -      assert(N != Root);
> -      return true;
> -    }
> +    for (const SDValue &Op : Use->op_values()) {
> +      // Ignore chain uses, they are validated by HandleMergeInputChains.
> +      if (Op.getValueType() == MVT::Other && IgnoreChains)
> +        continue;
>  
> -    // Traverse up the operand chain.
> -    if (findNonImmUse(N, Def, ImmedUse, Root, Visited, IgnoreChains))
> -      return true;
> +      SDNode *N = Op.getNode();
> +      if (N == Def) {
> +        if (Use == ImmedUse || Use == Root)
> +          continue;  // We are not looking for immediate use.
> +        assert(N != Root);
> +        return true;
> +      }
> +
> +      // Traverse up the operand chain.
> +      WorkList.push_back(N);
> +    }
>    }
>    return false;
>  }
>
>
> Index: lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> +++ lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
> @@ -2022,7 +2022,7 @@
>  }
>  
>  /// findNonImmUse - Return true if "Use" is a non-immediate use of "Def".
> -/// This function recursively traverses up the operand chain, ignoring
> +/// This function iteratively traverses up the operand chain, ignoring
>  /// certain nodes.
>  static bool findNonImmUse(SDNode *Use, SDNode* Def, SDNode *ImmedUse,
>                            SDNode *Root, SmallPtrSetImpl<SDNode*> &Visited,
> @@ -2035,30 +2035,36 @@
>    // The Use may be -1 (unassigned) if it is a newly allocated node.  This can
>    // happen because we scan down to newly selected nodes in the case of glue
>    // uses.
> -  if ((Use->getNodeId() < Def->getNodeId() && Use->getNodeId() != -1))
> -    return false;
> +  std::vector<SDNode *> WorkList;
> +  WorkList.push_back(Use);
>  
> -  // Don't revisit nodes if we already scanned it and didn't fail, we know we
> -  // won't fail if we scan it again.
> -  if (!Visited.insert(Use).second)
> -    return false;
> +  while (!WorkList.empty()) {
> +    Use = WorkList.back();
> +    WorkList.pop_back();
> +    if (Use->getNodeId() < Def->getNodeId() && Use->getNodeId() != -1)
> +      continue;
>  
> -  for (const SDValue &Op : Use->op_values()) {
> -    // Ignore chain uses, they are validated by HandleMergeInputChains.
> -    if (Op.getValueType() == MVT::Other && IgnoreChains)
> +    // Don't revisit nodes if we already scanned it and didn't fail, we know we
> +    // won't fail if we scan it again.
> +    if (!Visited.insert(Use).second)
>        continue;
>  
> -    SDNode *N = Op.getNode();
> -    if (N == Def) {
> -      if (Use == ImmedUse || Use == Root)
> -        continue;  // We are not looking for immediate use.
> -      assert(N != Root);
> -      return true;
> -    }
> +    for (const SDValue &Op : Use->op_values()) {
> +      // Ignore chain uses, they are validated by HandleMergeInputChains.
> +      if (Op.getValueType() == MVT::Other && IgnoreChains)
> +        continue;
>  
> -    // Traverse up the operand chain.
> -    if (findNonImmUse(N, Def, ImmedUse, Root, Visited, IgnoreChains))
> -      return true;
> +      SDNode *N = Op.getNode();
> +      if (N == Def) {
> +        if (Use == ImmedUse || Use == Root)
> +          continue;  // We are not looking for immediate use.
> +        assert(N != Root);
> +        return true;
> +      }
> +
> +      // Traverse up the operand chain.
> +      WorkList.push_back(N);
> +    }
>    }
>    return false;
>  }