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
This commit is contained in:
Kuangyuan Chen 2021-02-01 10:21:15 -08:00 committed by TensorFlower Gardener
parent 30f2b5fa02
commit 05e8408270

View File

@ -137,10 +137,18 @@ class GraphExecutionState {
Status BuildGraph(const BuildGraphOptions& options,
std::unique_ptr<ClientGraph>* out);
// Optimize the graph with the node set specified in `options`.
Status OptimizeGraph(
const BuildGraphOptions& options, std::unique_ptr<Graph>* optimized_graph,
std::unique_ptr<FunctionLibraryDefinition>* 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<Graph>* optimized_graph,
std::unique_ptr<FunctionLibraryDefinition>* optimized_flib);
// The GraphExecutionState must store a copy of the original GraphDef if
// either of the following conditions holds:
//