| 566 | } // namespace |
| 567 | |
| 568 | Status ConstantFold(const ConstantFoldingOptions& opts, |
| 569 | FunctionLibraryRuntime* function_library, Env* env, |
| 570 | const Device* partition_device, Graph* graph, |
| 571 | bool* was_mutated) { |
| 572 | // TensorFlow flushes denormals to zero and rounds to nearest, so we do |
| 573 | // the same here. |
| 574 | port::ScopedFlushDenormal flush; |
| 575 | port::ScopedSetRound round(FE_TONEAREST); |
| 576 | |
| 577 | DumpGraph("Before", graph); |
| 578 | ConstantFoldNameGenerator generate_new_name = opts.generate_new_name; |
| 579 | if (generate_new_name == nullptr) { |
| 580 | generate_new_name = [](Graph* graph, string old_name) { |
| 581 | return strings::StrCat(graph->NewName(old_name), "__cf__", |
| 582 | UniqueConstantId()); |
| 583 | }; |
| 584 | } |
| 585 | |
| 586 | std::vector<Node*> constant_foldable_nodes; |
| 587 | std::unordered_map<const Node*, gtl::FlatSet<Node*>> constant_control_deps; |
| 588 | std::unordered_map<const Node*, std::vector<Tensor>> shape_replacement_map; |
| 589 | FindConstantFoldableNodes(graph, opts, &constant_foldable_nodes, |
| 590 | &constant_control_deps, &shape_replacement_map); |
| 591 | if (constant_foldable_nodes.empty()) { |
| 592 | VLOG(1) << "No constant foldable nodes found"; |
| 593 | *was_mutated = false; |
| 594 | // This is not an error, so return the status as OK. |
| 595 | return Status::OK(); |
| 596 | } |
| 597 | |
| 598 | std::map<NodeAndOutput, NodeAndOutput> tensors_to_fetch; |
| 599 | std::unique_ptr<Graph> constant_graph( |
| 600 | GetConstantGraph(graph, constant_foldable_nodes, shape_replacement_map, |
| 601 | &tensors_to_fetch, generate_new_name)); |
| 602 | DumpGraph("Constant graph", constant_graph.get()); |
| 603 | |
| 604 | if (tensors_to_fetch.empty()) { |
| 605 | VLOG(1) << "No constant nodes found that feed into the original graph."; |
| 606 | *was_mutated = false; |
| 607 | // This is not an error, so return the status as OK. |
| 608 | return Status::OK(); |
| 609 | } |
| 610 | VLOG(1) << "Constant foldable " << constant_graph->num_node_ids() << " : " |
| 611 | << graph->num_node_ids(); |
| 612 | |
| 613 | std::vector<string> tensors_to_fetch_names; |
| 614 | std::vector<NodeAndOutput> tensors_to_replace; |
| 615 | // Sorting the nodes based on the name gives us a stable ordering between runs |
| 616 | // for the same graph. |
| 617 | std::vector<std::pair<NodeAndOutput, NodeAndOutput>> tensors_to_fetch_sorted( |
| 618 | tensors_to_fetch.begin(), tensors_to_fetch.end()); |
| 619 | std::sort(tensors_to_fetch_sorted.begin(), tensors_to_fetch_sorted.end(), |
| 620 | [](const std::pair<NodeAndOutput, NodeAndOutput>& n1, |
| 621 | const std::pair<NodeAndOutput, NodeAndOutput>& n2) { |
| 622 | return std::tie(n1.first.first->name(), n1.first.second) < |
| 623 | std::tie(n2.first.first->name(), n2.first.second); |
| 624 | }); |
| 625 | for (auto n : tensors_to_fetch_sorted) { |