From 05e840827084935c75238c234643ce59efba12f4 Mon Sep 17 00:00:00 2001 From: Kuangyuan Chen Date: Mon, 1 Feb 2021 10:21:15 -0800 Subject: [PATCH] Make GraphExecutionState::OptimizeGraph() a public method so that it can reused by other components. This is fine as OptimizeGraph() does not change the internal states of GraphExecutionState. PiperOrigin-RevId: 354969297 Change-Id: Ia37c8a0366c9f6263f7c4a4d53161a622c2fc6ce --- .../core/common_runtime/graph_execution_state.h | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tensorflow/core/common_runtime/graph_execution_state.h b/tensorflow/core/common_runtime/graph_execution_state.h index f9a1e35b476..ba16fb030b5 100644 --- a/tensorflow/core/common_runtime/graph_execution_state.h +++ b/tensorflow/core/common_runtime/graph_execution_state.h @@ -137,10 +137,18 @@ class GraphExecutionState { Status BuildGraph(const BuildGraphOptions& options, std::unique_ptr* out); + // Optimize the graph with the node set specified in `options`. + Status OptimizeGraph( + const BuildGraphOptions& options, std::unique_ptr* optimized_graph, + std::unique_ptr* optimized_flib); + // The graph returned by BuildGraph may contain only the pruned // graph, whereas some clients may want access to the full graph. const Graph* full_graph() { return graph_; } + // The original function library of this graph. + const FunctionLibraryDefinition& flib_def() const { return *flib_def_; } + // Returns the node with the given name, or null if it does not exist. const Node* get_node_by_name(const string& name) const { NodeNameToCostIdMap::const_iterator iter = @@ -179,10 +187,6 @@ class GraphExecutionState { Status PruneGraph(const BuildGraphOptions& options, Graph* graph, subgraph::RewriteGraphMetadata* out_rewrite_metadata); - Status OptimizeGraph( - const BuildGraphOptions& options, std::unique_ptr* optimized_graph, - std::unique_ptr* optimized_flib); - // The GraphExecutionState must store a copy of the original GraphDef if // either of the following conditions holds: //