| 649 | } |
| 650 | |
| 651 | Status MetaOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 652 | GraphDef* optimized_graph) { |
| 653 | VLOG(1) << "Starting optimization for grappler item: " << item.id; |
| 654 | optimization_results_.clear(); |
| 655 | |
| 656 | // Constructs a FunctionLibraryDefinition with functions that are reachable |
| 657 | // from the nodes of the graph. |
| 658 | const auto minimized_flib = |
| 659 | [](const GraphDef& graph) -> FunctionLibraryDefinition { |
| 660 | return FunctionLibraryDefinition(OpRegistry::Global(), graph.library()) |
| 661 | .ReachableDefinitions(graph); |
| 662 | }; |
| 663 | |
| 664 | // 0. Original graph might contain a huge function library, that is mostly |
| 665 | // unused. This library copied over by each individual Grappler optimizer, |
| 666 | // which adds a huge overhead. Before starting optimization passes we just |
| 667 | // remove all the unreachable functions. |
| 668 | // TODO(ezhulenev): Construct reachable function library definition directly |
| 669 | // from the proto without constructing temporary FunctionLibraryDefinition. |
| 670 | GraphDef trimmed_graph; // do not copy graph with a potentially huge library |
| 671 | *trimmed_graph.mutable_node() = item.graph.node(); |
| 672 | *trimmed_graph.mutable_versions() = item.graph.versions(); |
| 673 | *trimmed_graph.mutable_library() = minimized_flib(item.graph).ToProto(); |
| 674 | |
| 675 | GrapplerItem trimmed_item = item.WithGraph(std::move(trimmed_graph)); |
| 676 | |
| 677 | VLOG(1) << absl::Substitute( |
| 678 | "Deleted $0 unreachable functions from the graph (library size = $1)", |
| 679 | item.graph.library().function_size() - |
| 680 | trimmed_item.graph.library().function_size(), |
| 681 | trimmed_item.graph.library().function_size()); |
| 682 | |
| 683 | // 1. Optimize main graph |
| 684 | TF_RETURN_IF_ERROR(OptimizeGraph(cluster, trimmed_item, optimized_graph)); |
| 685 | VLOG(1) << "Optimized main graph."; |
| 686 | GRAPPLER_RETURN_IF_DEADLINE_EXCEEDED(); |
| 687 | |
| 688 | // 2. Optimize functions reachable from the optimized graph. |
| 689 | FunctionLibraryDefinition flib = minimized_flib(*optimized_graph); |
| 690 | using NodeDefs = protobuf::RepeatedPtrField<NodeDef>; |
| 691 | |
| 692 | // Find functions for which we might need to compute a gradient at runtime. |
| 693 | absl::flat_hash_set<string> differentiable_functions; |
| 694 | |
| 695 | const auto find_differentiable_functions = |
| 696 | [&](const NodeDefs& nodes) -> void { |
| 697 | for (const NodeDef& node : nodes) { |
| 698 | if (IsSymbolicGradient(node)) { |
| 699 | const auto* f_attr = gtl::FindOrNull(node.attr(), "f"); |
| 700 | if (f_attr) differentiable_functions.insert(f_attr->func().name()); |
| 701 | } |
| 702 | } |
| 703 | }; |
| 704 | |
| 705 | // SymbolicGradient nodes inside the main graph. |
| 706 | find_differentiable_functions(optimized_graph->node()); |
| 707 | // SymbolicGradient nodes inside the function library. |
| 708 | for (const FunctionDef& function : optimized_graph->library().function()) { |
no test coverage detected