| 3947 | } |
| 3948 | |
| 3949 | std::tuple<bool, std::vector<Node*>, const MklLayoutRewritePass::FusionInfo> |
| 3950 | MklLayoutRewritePass::CheckForNodeFusion(Node* a) const { |
| 3951 | // Stores matched nodes, in the same order as node_checkers. |
| 3952 | std::vector<Node*> nodes; |
| 3953 | |
| 3954 | for (auto fi = finfo_.begin(); fi != finfo_.end(); ++fi) { |
| 3955 | // |
| 3956 | // Make sure node "a" and its succeding nodes (b, c ...), match the pattern |
| 3957 | // defined in fusion info (ops[0], ops[1], ...), |
| 3958 | // a.k.a. "a->b->c" matches "op1->op2->op3" |
| 3959 | // |
| 3960 | |
| 3961 | // Stores the first unvisited outgoing edge of each matched node in "nodes". |
| 3962 | std::stack<EdgeSet::const_iterator> current_neighbor_stack; |
| 3963 | nodes.clear(); |
| 3964 | |
| 3965 | auto node_checker = fi->node_checkers.begin(); |
| 3966 | if (a != nullptr && (*node_checker)(a)) { |
| 3967 | nodes.push_back(a); |
| 3968 | current_neighbor_stack.push(a->out_edges().begin()); |
| 3969 | ++node_checker; |
| 3970 | } |
| 3971 | |
| 3972 | while (!nodes.empty()) { |
| 3973 | auto& current_neighbor_iter = current_neighbor_stack.top(); |
| 3974 | |
| 3975 | if (current_neighbor_iter != nodes.back()->out_edges().end()) { |
| 3976 | // Found an unvisited edge. Goes through the edge to get the neighbor. |
| 3977 | Node* neighbor_node = (*current_neighbor_iter)->dst(); |
| 3978 | ++current_neighbor_stack.top(); // Retrieves the next unvisited edge. |
| 3979 | |
| 3980 | if ((*node_checker)(neighbor_node)) { |
| 3981 | // Found a match. Stores the node and moves to the next checker. |
| 3982 | nodes.push_back(neighbor_node); |
| 3983 | current_neighbor_stack.push(neighbor_node->out_edges().begin()); |
| 3984 | if (++node_checker == fi->node_checkers.end()) { |
| 3985 | return make_tuple(true, nodes, *fi); |
| 3986 | } |
| 3987 | } |
| 3988 | } else { |
| 3989 | // Removes the current node since none of its neighbor leads to a |
| 3990 | // further match. |
| 3991 | nodes.pop_back(); |
| 3992 | current_neighbor_stack.pop(); |
| 3993 | --node_checker; |
| 3994 | } |
| 3995 | } |
| 3996 | } |
| 3997 | |
| 3998 | return make_tuple(false, std::vector<Node*>(), FusionInfo()); |
| 3999 | } |
| 4000 | |
| 4001 | /////////////////////////////////////////////////////////////////////////////// |
| 4002 | // Post-rewrite OneDNN metadata fixup pass |