| 927 | } |
| 928 | |
| 929 | Status FindGatherNode(Node* n, const std::unordered_set<std::string>& stop_nodes, |
| 930 | Node** gather_node) { |
| 931 | *gather_node = nullptr; |
| 932 | std::unordered_set<Node*> pushed; |
| 933 | std::queue<Node*> q; |
| 934 | q.push(n); |
| 935 | pushed.insert(n); |
| 936 | while (!q.empty()) { |
| 937 | Node* curr = q.front(); |
| 938 | q.pop(); |
| 939 | if (curr->op_def().name() == "KvResourceGather") { |
| 940 | if (*gather_node) { |
| 941 | LOG(FATAL) << "Find many KvResourceGather op from variable: " << n->DebugString() |
| 942 | << ", gather1: " << (*gather_node)->DebugString() |
| 943 | << ", gather2: " << curr->DebugString(); |
| 944 | } |
| 945 | *gather_node = curr; |
| 946 | continue; |
| 947 | } |
| 948 | |
| 949 | if (stop_nodes.find(curr->op_def().name()) == stop_nodes.end()) { |
| 950 | for (const Edge* edge : curr->out_edges()) { |
| 951 | if (pushed.find(edge->dst()) == pushed.end()) { |
| 952 | pushed.insert(edge->dst()); |
| 953 | q.push(edge->dst()); |
| 954 | } |
| 955 | } |
| 956 | } |
| 957 | } |
| 958 | |
| 959 | return Status::OK(); |
| 960 | } |
| 961 | |
| 962 | bool IsDynamicStitchOp(Node* n) { |
| 963 | return n->op_def().name() == "DynamicStitch" || |