[SelectionDAG] Eliminate exponential behavior in WalkChainUsers

[timshen]

Summary

This patch fixes a potential exponential behavior due to no memoization in a DFS.

For example, consider such a graph:
ret -> a0, ret -> b0,
a0 -> a1, a0-> b1, b0 -> a1, b0 -> b1,
a1 -> a2, a1-> b2, b1 -> a2, b1 -> b2,
a2 -> a3, a2-> b3, b2 -> a3, b2 -> b3,
...

If this DAG has n layers, the all paths WalkChainUsers will walk through is O(2^n); after memoization it will be O(n).

There is no functional change.

[resistor]

Functionality LGTM. I would add a comment at the use site of TokenFactorResult to point out that the walk is being memoized.

[timshen]

Updated comments.

Also use the DenseMap iterator correctly.

[timshen]

Committed as r259315: http://reviews.llvm.org/rL259315

lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

 /// non-pattern node in between two pattern nodes (which prevent folding because
 /// it would induce a cycle) and whether we have a TokenFactor node sandwiched
 /// between pattern nodes (in which case the TF becomes part of the pattern).
 ///
 /// The walk we do here is guaranteed to be small because we quickly get down to
 /// already selected nodes "below" us.
 static ChainResult
 WalkChainUsers(const SDNode *ChainedNode,
-               SmallVectorImpl<SDNode*> &ChainedNodesInPattern,
-               SmallVectorImpl<SDNode*> &InteriorChainedNodes) {
+               SmallVectorImpl<SDNode *> &ChainedNodesInPattern,
+               DenseMap<const SDNode *, ChainResult> &TokenFactorResult,
+               SmallVectorImpl<SDNode *> &InteriorChainedNodes) {
   ChainResult Result = CR_Simple;
 
   for (SDNode::use_iterator UI = ChainedNode->use_begin(),
          E = ChainedNode->use_end(); UI != E; ++UI) {
     // Make sure the use is of the chain, not some other value we produce.
     if (UI.getUse().getValueType() != MVT::Other) continue;
 
     SDNode *User = *UI;
for (SDNode::use_iterator UI = ChainedNode->use_begin(),
     //      \        /
     //       \      /
     //       [Store]
     //
     // In this case, the TokenFactor becomes part of our match and we rewrite it
     // as a new TokenFactor.
     //
     // To distinguish these two cases, do a recursive walk down the uses.
-    switch (WalkChainUsers(User, ChainedNodesInPattern, InteriorChainedNodes)) {
+    auto MemoizeResult = TokenFactorResult.find(User);
+    bool Visited = MemoizeResult != TokenFactorResult.end();
+    // Recursively walk chain users only if the result is not memoized.
+    if (!Visited) {
+      auto Res = WalkChainUsers(User, ChainedNodesInPattern, TokenFactorResult,
+                                InteriorChainedNodes);
+      MemoizeResult = TokenFactorResult.insert(std::make_pair(User, Res)).first;
+    }
+    switch (MemoizeResult->second) {
     case CR_Simple:
       // If the uses of the TokenFactor are just already-selected nodes, ignore
       // it, it is "below" our pattern.
       continue;
     case CR_InducesCycle:
       // If the uses of the TokenFactor lead to nodes that are not part of our
       // pattern that are not selected, folding would turn this into a cycle,
       // bail out now.
       return CR_InducesCycle;
     case CR_LeadsToInteriorNode:
       break;  // Otherwise, keep processing.
     }
 
     // Okay, we know we're in the interesting interior case.  The TokenFactor
     // is now going to be considered part of the pattern so that we rewrite its
     // uses (it may have uses that are not part of the pattern) with the
     // ultimate chain result of the generated code.  We will also add its chain
     // inputs as inputs to the ultimate TokenFactor we create.
     Result = CR_LeadsToInteriorNode;
+    if (!Visited) {
       ChainedNodesInPattern.push_back(User);
       InteriorChainedNodes.push_back(User);
+    }
     continue;
   }
 
   return Result;
 }
 
 /// HandleMergeInputChains - This implements the OPC_EmitMergeInputChains
 /// operation for when the pattern matched at least one node with a chains.  The
 /// input vector contains a list of all of the chained nodes that we match.  We
 /// must determine if this is a valid thing to cover (i.e. matching it won't
 /// induce cycles in the DAG) and if so, creating a TokenFactor node. that will
 /// be used as the input node chain for the generated nodes.
 static SDValue
 HandleMergeInputChains(SmallVectorImpl<SDNode*> &ChainNodesMatched,
                        SelectionDAG *CurDAG) {
+  // Used for memoization. Without it WalkChainUsers could take exponential
+  // time to run.
+  DenseMap<const SDNode *, ChainResult> TokenFactorResult;
   // Walk all of the chained nodes we've matched, recursively scanning down the
   // users of the chain result. This adds any TokenFactor nodes that are caught
   // in between chained nodes to the chained and interior nodes list.
   SmallVector<SDNode*, 3> InteriorChainedNodes;
   for (unsigned i = 0, e = ChainNodesMatched.size(); i != e; ++i) {
     if (WalkChainUsers(ChainNodesMatched[i], ChainNodesMatched,
+                       TokenFactorResult,
                        InteriorChainedNodes) == CR_InducesCycle)
       return SDValue(); // Would induce a cycle.
   }
 
   // Okay, we have walked all the matched nodes and collected TokenFactor nodes
   // that we are interested in.  Form our input TokenFactor node.
   SmallVector<SDValue, 3> InputChains;
   for (unsigned i = 0, e = ChainNodesMatched.size(); i != e; ++i) {