| 528 | } |
| 529 | |
| 530 | Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 531 | GraphDef* optimized_graph) { |
| 532 | const auto preserve_set = item.NodesToPreserve(); |
| 533 | const absl::flat_hash_set<absl::string_view> nodes_to_preserve( |
| 534 | preserve_set.begin(), preserve_set.end()); |
| 535 | |
| 536 | // Prune all the nodes that won't be executed, ie all the nodes that aren't in |
| 537 | // the fanin of a fetch node. If fetch nodes aren't specified, we'll assume |
| 538 | // the whole graph might be executed. |
| 539 | GraphDef graph = item.graph; |
| 540 | Status status; |
| 541 | utils::MutableGraphView graph_view(&graph, &status); |
| 542 | TF_RETURN_IF_ERROR(status); |
| 543 | if (!nodes_to_preserve.empty()) { |
| 544 | TF_RETURN_IF_ERROR(PruneUnreachableNodes(&graph_view, nodes_to_preserve)); |
| 545 | bool did_split_identity_n = false; |
| 546 | TF_RETURN_IF_ERROR(SplitIdentityNInputs(&graph_view, nodes_to_preserve, |
| 547 | &did_split_identity_n)); |
| 548 | if (did_split_identity_n) { |
| 549 | TF_RETURN_IF_ERROR(PruneUnreachableNodes(&graph_view, nodes_to_preserve)); |
| 550 | } |
| 551 | // TODO(lyandy): Remove sorting once ArithmeticOptimizer |
| 552 | // (MinimizeBroadcasts) is migrated over to using utils::GraphView. |
| 553 | TF_RETURN_IF_ERROR( |
| 554 | graph_view.SortTopologically(/*ignore_cycles=*/true, {})); |
| 555 | } |
| 556 | |
| 557 | absl::flat_hash_set<absl::string_view> function_names; |
| 558 | function_names.reserve(graph.library().function_size()); |
| 559 | for (const auto& function : graph.library().function()) { |
| 560 | function_names.insert(function.signature().name()); |
| 561 | } |
| 562 | OpRegistryInterface* op_registry = OpRegistry::Global(); |
| 563 | |
| 564 | // Check if we can further prune the graph, by removing the trivial ops. |
| 565 | const int num_nodes = graph_view.NumNodes(); |
| 566 | std::vector<int> nodes_to_delete; |
| 567 | nodes_to_delete.reserve(num_nodes); |
| 568 | for (auto& node : graph_view.GetNodes()) { |
| 569 | if (!IsTrivialOp(node)) { |
| 570 | continue; |
| 571 | } |
| 572 | |
| 573 | // Don't remove nodes that must be preserved. |
| 574 | if (nodes_to_preserve.contains(node.GetName())) { |
| 575 | continue; |
| 576 | } |
| 577 | |
| 578 | // - Don't remove nodes that drive control dependencies. |
| 579 | // - Don't remove nodes that are driven by control dependencies either since |
| 580 | // we can't ensure (yet) that we won't increase the number of control |
| 581 | // dependency edges by deleting them (for example, removing a node driven |
| 582 | // by 10 control edges and driving 10 control edges would result in the |
| 583 | // creation of 100 edges). |
| 584 | // - Don't modify nodes that are connected to functions since that can |
| 585 | // result in inlining failures later on. |
| 586 | // - Don't prune nodes that are driven by another device since these could |
| 587 | // be used to reduce cross device communication. |
nothing calls this directly
no test coverage detected