| 1564 | } |
| 1565 | |
| 1566 | Status ConstantFolding::FoldGraph( |
| 1567 | const GraphProperties& properties, GraphDef* output, |
| 1568 | absl::flat_hash_set<string>* nodes_to_not_simplify) { |
| 1569 | std::unordered_set<string> processed_nodes; |
| 1570 | std::deque<NodeDef*> queue; |
| 1571 | for (int i = 0; i < graph_->node_size(); i++) { |
| 1572 | if (IsFoldable(graph_->node(i), &properties)) { |
| 1573 | queue.push_back(graph_->mutable_node(i)); |
| 1574 | } |
| 1575 | } |
| 1576 | while (!queue.empty()) { |
| 1577 | NodeDef* node = queue.front(); |
| 1578 | queue.pop_front(); |
| 1579 | if (processed_nodes.count(node->name())) { |
| 1580 | continue; |
| 1581 | } |
| 1582 | // We need to record a copy of output nodes before FoldNode() modifies it. |
| 1583 | // We also need to ensure that the fanout is sorted deterministically. |
| 1584 | const std::set<NodeDef*>& outputs = node_map_->GetOutputs(node->name()); |
| 1585 | std::vector<NodeDef*> fanout(outputs.begin(), outputs.end()); |
| 1586 | std::sort(fanout.begin(), fanout.end(), |
| 1587 | [](const NodeDef* n1, const NodeDef* n2) { |
| 1588 | return n1->name() < n2->name(); |
| 1589 | }); |
| 1590 | |
| 1591 | bool result_too_large = false; |
| 1592 | Status s = FoldNode(node, output, &result_too_large); |
| 1593 | processed_nodes.insert(node->name()); |
| 1594 | if (!s.ok()) { |
| 1595 | VLOG(1) << "Failed to fold node " << node->DebugString() |
| 1596 | << "\nError message: " << s; |
| 1597 | if (result_too_large) { |
| 1598 | nodes_to_not_simplify->emplace(node->name()); |
| 1599 | } |
| 1600 | } else { |
| 1601 | for (auto& output : fanout) { |
| 1602 | if (IsFoldable(*output, &properties)) { |
| 1603 | queue.push_back(output); |
| 1604 | } |
| 1605 | } |
| 1606 | } |
| 1607 | } |
| 1608 | |
| 1609 | // Delete the newly created nodes that don't feed anything. |
| 1610 | std::vector<int> nodes_to_delete; |
| 1611 | for (int i = 0; i < output->node_size(); i++) { |
| 1612 | auto fanout = node_map_->GetOutputs(output->node(i).name()); |
| 1613 | if (fanout.empty()) nodes_to_delete.push_back(i); |
| 1614 | } |
| 1615 | EraseNodesFromGraph(std::move(nodes_to_delete), output); |
| 1616 | |
| 1617 | for (const auto& node : graph_->node()) { |
| 1618 | // If no fetch nodes is provided, we conservatively |
| 1619 | // keep all nodes in the original graph in case users need to fetch |
| 1620 | // their values. |
| 1621 | auto fanout = node_map_->GetOutputs(node.name()); |
| 1622 | if (!fanout.empty() || !has_fetch_ || |
| 1623 | nodes_to_preserve_.find(node.name()) != nodes_to_preserve_.end()) { |
nothing calls this directly
no test coverage detected