experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use the zero-copy implementation of GraphConstructor in more places.

Browse Source 
Many uses of GraphConstructor take a `const GraphDef&` to a locally-defined GraphDef that is subsequently destroyed. We can move the GraphDef into GraphConstructor to avoid copying the graph nodes repeatedly. In some cases with large GraphDefs (e.g. with large embedded constant tensors) this optimization will reduce peak memory consumption.

PiperOrigin-RevId: 259809688
This commit is contained in:
Derek Murray 2019-07-24 13:50:36 -07:00 committed by TensorFlower Gardener
parent 3a72de3a1b
commit 5d37c2b785
11 changed files with 27 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/cc/framework/scope.cc
View File

@ -318,7 +318,7 @@ Status Scope::ToGraph(Graph* g, GraphConstructorOptions opts) const {
if (ok()) {
GraphDef graph_def;
graph()->ToGraphDef(&graph_def);
UpdateStatus(ConvertGraphDefToGraph(opts, graph_def, g));
UpdateStatus(ConvertGraphDefToGraph(opts, std::move(graph_def), g));
}
return *impl()->status_;
}

8
tensorflow/compiler/mlir/tensorflow/translate/import_graphdef.cc
View File

@ -300,8 +300,8 @@ Status Importer::RemoveBackedges(const Graph& graph) {
graph_ = absl::make_unique<Graph>(graph.flib_def());
GraphConstructorOptions opts;
opts.allow_internal_ops = true;
TF_RETURN_IF_ERROR(
::tensorflow::ConvertGraphDefToGraph(opts, graph_def, graph_.get()));
TF_RETURN_IF_ERROR(::tensorflow::ConvertGraphDefToGraph(
opts, std::move(graph_def), graph_.get()));
// Remove all the backedges. So the nodes can be added to the shape refiner.
TF_RETURN_IF_ERROR(back_edge_helper_.Remove(graph_.get()));
@ -1394,8 +1394,8 @@ StatusOr<mlir::OwningModuleRef> ConvertGraphdefToMlir(
if (add_default_attributes) {
TF_RETURN_IF_ERROR(AddDefaultsToNodeDef(&preprocessed_graphdef));
}
TF_RETURN_IF_ERROR(
ConvertGraphDefToGraph(options, preprocessed_graphdef, &graph));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
options, std::move(preprocessed_graphdef), &graph));
return ConvertGraphToMlir(graph, debug_info, graph.flib_def(), specs,
context);

4
tensorflow/compiler/tf2xla/tf2xla.cc
View File

@ -384,8 +384,8 @@ Status InitGraph(const GraphDef& graph_def, const tf2xla::Config& config,
TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(
&second_copy_def, *g->op_registry(), /*node_offset=*/0));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(GraphConstructorOptions(),
second_copy_def, g.get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
GraphConstructorOptions(), std::move(second_copy_def), g.get()));
TF_RETURN_IF_ERROR(RewriteAndPruneGraph(g.get(), config, feed_remapping));
// Functionalize control flow.

6
tensorflow/core/common_runtime/direct_session.cc
View File

@ -1614,15 +1614,15 @@ Status DirectSession::CreateGraphs(
}
}
for (const auto& partition : partitions) {
for (auto& partition : partitions) {
std::unique_ptr<Graph> device_graph(
new Graph(client_graph->flib_def.get()));
GraphConstructorOptions device_opts;
// There are internal operations (e.g., send/recv) that we now allow.
device_opts.allow_internal_ops = true;
device_opts.expect_device_spec = true;
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(device_opts, partition.second,
device_graph.get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
device_opts, std::move(partition.second), device_graph.get()));
outputs->emplace(partition.first, std::move(device_graph));
}

4
tensorflow/core/common_runtime/graph_execution_state.cc
View File

@ -757,8 +757,8 @@ Status GraphExecutionState::OptimizeGraph(
GraphConstructorOptions opts;
opts.allow_internal_ops = true;
TF_RETURN_IF_ERROR(
ConvertGraphDefToGraph(opts, new_graph, optimized_graph->get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(opts, std::move(new_graph),
optimized_graph->get()));
// The graph conversion sets the requested device names but not the
// assigned device names. However, since at this point the graph is placed
// TF expects an assigned device name for every node. Therefore we copy

6
tensorflow/core/distributed_runtime/graph_mgr.cc
View File

@ -179,14 +179,14 @@ Status GraphMgr::InitItem(const string& handle, const GraphDef& gdef,
}
std::unordered_map<string, std::unique_ptr<Graph>> partition_graphs;
for (const auto& partition : partitions) {
for (auto& partition : partitions) {
std::unique_ptr<Graph> device_graph(new Graph(OpRegistry::Global()));
GraphConstructorOptions device_opts;
// There are internal operations (e.g., send/recv) that we now allow.
device_opts.allow_internal_ops = true;
device_opts.expect_device_spec = true;
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(device_opts, partition.second,
device_graph.get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
device_opts, std::move(partition.second), device_graph.get()));
partition_graphs.emplace(partition.first, std::move(device_graph));
}

2
tensorflow/core/graph/graph_def_builder_util.cc
View File

@ -22,7 +22,7 @@ Status GraphDefBuilderToGraph(const GraphDefBuilder& builder, Graph* graph) {
GraphDef graph_def;
TF_RETURN_IF_ERROR(builder.ToGraphDef(&graph_def));
GraphConstructorOptions opts;
return ConvertGraphDefToGraph(opts, graph_def, graph);
return ConvertGraphDefToGraph(opts, std::move(graph_def), graph);
}
} // namespace tensorflow

4
tensorflow/core/grappler/grappler_item_builder.cc
View File

@ -267,8 +267,8 @@ Status RuntimeGraphOptimizer(const GraphDef& graph_def_arg,
graph_ctor_opts.expect_device_spec = false;
std::unique_ptr<Graph> graphptr(new Graph(function_library));
TF_RETURN_IF_ERROR(
ConvertGraphDefToGraph(graph_ctor_opts, graph_def, graphptr.get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
graph_ctor_opts, std::move(graph_def), graphptr.get()));
// Optimize the graph.
::tensorflow::GraphOptimizer optimizer(*optimizer_opts);

2
tensorflow/core/grappler/optimizers/function_optimizer.cc
View File

@ -784,7 +784,7 @@ constexpr const char* const kLowerAsMultiDeviceFunctionAttr =
using KeepCallerNode = InlineFunctionBodyOptions::KeepCallerNode;
using OutputControlSource = InlineFunctionBodyOptions::OutputControlSource;
// Checks if boolean attribute is defined and it's value is 'true'.
// Checks if boolean attribute is defined and its value is 'true'.
bool CheckBoolAttr(const Node* n, absl::string_view attr_name) {
bool match;
Status s = GetNodeAttr(n->attrs(), attr_name, &match);

11
tensorflow/core/grappler/optimizers/meta_optimizer.cc
View File

@ -802,8 +802,6 @@ Status OptimizeGraph(
std::unique_ptr<tensorflow::Graph> optimized_graph(
new tensorflow::Graph(OpRegistry::Global()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(GraphConstructorOptions(),
out_graph, optimized_graph.get()));
// Copy optimized functions back to the overlay lib.
if (flib) {
@ -817,25 +815,28 @@ Status OptimizeGraph(
}
}
*g = std::move(optimized_graph);
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(
GraphConstructorOptions(), std::move(out_graph), optimized_graph.get()));
// The graph conversion sets the requested device names but not the
// assigned device names. However, since at this point the graph is
// placed TF expects an assigned device name for every node. Therefore
// we copy the requested device into the assigned device field.
for (Node* node : (*g)->nodes()) {
for (Node* node : optimized_graph->nodes()) {
if (node->IsOp() && node->assigned_device_name().empty()) {
if (node->requested_device().empty()) {
return errors::Internal(
"Either placer did not place the node or Grappler did not "
"copy the assigned device. Contact Grappler team since latter "
"is more likely. Node=",
node->name(), " Graph: ", (*g)->ToGraphDefDebug().DebugString());
node->name(),
" Graph: ", optimized_graph->ToGraphDefDebug().DebugString());
}
node->set_assigned_device_name(node->requested_device());
}
}
*g = std::move(optimized_graph);
return Status::OK();
}

4
tensorflow/tools/optimization/optimization_pass_runner.cc
View File

@ -111,8 +111,8 @@ Status OptimizationPassRunner::Run(absl::string_view pass_to_run,
GraphConstructorOptions graph_opts;
graph_opts.expect_device_spec = true;
graph_opts.allow_internal_ops = true;
TF_RETURN_IF_ERROR(
ConvertGraphDefToGraph(graph_opts, input, options.graph->get()));
TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(graph_opts, std::move(input),
options.graph->get()));
// Add all devices that were previously configured with AddDevice.
DeviceSet device_set;