| 517 | } |
| 518 | |
| 519 | Status RemoveStackOps(const std::unordered_set<string>& nodes_to_preserve, |
| 520 | GraphDef* optimized_graph) { |
| 521 | NodeMap node_map(optimized_graph); |
| 522 | GraphTopologyView graph_view; |
| 523 | TF_RETURN_IF_ERROR(graph_view.InitializeFromGraph(*optimized_graph)); |
| 524 | |
| 525 | for (int node_idx = 0; node_idx < optimized_graph->node_size(); ++node_idx) { |
| 526 | if (IsStackOp(optimized_graph->node(node_idx))) { |
| 527 | for (int push_node_idx : GetStackPushNodesToConvert( |
| 528 | graph_view, nodes_to_preserve, node_idx)) { |
| 529 | // We found push nodes without corresponding pops. Convert them to |
| 530 | // Identity passing the data through and add a control dependency from |
| 531 | // the op supplying the stack handle. |
| 532 | NodeDef* push_node = optimized_graph->mutable_node(push_node_idx); |
| 533 | VLOG(1) << "Converting " << push_node_idx << " : " |
| 534 | << push_node->DebugString(); |
| 535 | if (push_node->attr().count("swap_memory") != 0) { |
| 536 | push_node->mutable_attr()->erase("swap_memory"); |
| 537 | } |
| 538 | push_node->set_op("Identity"); |
| 539 | push_node->mutable_input()->SwapElements(0, 1); |
| 540 | const string ctrl_dep = ConstantFolding::AddControlDependency( |
| 541 | push_node->input(1), optimized_graph, &node_map); |
| 542 | push_node->set_input(1, ctrl_dep); |
| 543 | VLOG(1) << "After converting: " << push_node->DebugString(); |
| 544 | } |
| 545 | } |
| 546 | } |
| 547 | return Status::OK(); |
| 548 | } |
| 549 | |
| 550 | bool IsSimpleBinaryOperator(const NodeDef& node) { |
| 551 | return (IsLess(node) || IsLessEqual(node) || IsGreater(node) || |
no test coverage detected