From 3f43965a44f1bbf897c51f254ababd083087979d Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Wed, 5 Dec 2018 14:55:02 -0800 Subject: [PATCH] [Grappler] Cleanup GrapplerItem constructors PiperOrigin-RevId: 224226461 --- tensorflow/core/grappler/grappler_item.cc | 30 +++++++++-------- tensorflow/core/grappler/grappler_item.h | 11 ++++--- .../optimizers/arithmetic_optimizer.cc | 3 +- ...perimental_implementation_selector_test.cc | 4 +-- .../optimizers/function_optimizer_test.cc | 32 +++++++++---------- .../grappler/optimizers/memory_optimizer.cc | 7 ++-- .../optimizers/memory_optimizer_test.cc | 14 ++++---- .../grappler/optimizers/meta_optimizer.cc | 2 +- .../optimizers/meta_optimizer_test.cc | 4 +-- .../core/grappler/optimizers/remapper.cc | 2 +- 10 files changed, 56 insertions(+), 53 deletions(-) diff --git a/tensorflow/core/grappler/grappler_item.cc b/tensorflow/core/grappler/grappler_item.cc index f7cda35368e..74bde67f198 100644 --- a/tensorflow/core/grappler/grappler_item.cc +++ b/tensorflow/core/grappler/grappler_item.cc @@ -30,20 +30,22 @@ limitations under the License. namespace tensorflow { namespace grappler { -GrapplerItem::GrapplerItem(const GrapplerItem& other, GraphDef* graph_def) { - id = other.id; - feed = other.feed; - fetch = other.fetch; - init_ops = other.init_ops; - keep_ops = other.keep_ops; - expected_init_time = other.expected_init_time; - save_op = other.save_op; - restore_op = other.restore_op; - save_restore_loc_tensor = other.save_restore_loc_tensor; - queue_runners = other.queue_runners; - devices_ = other.devices_; - allowed_optimizations_ = other.allowed_optimizations_; - graph.Swap(graph_def); +GrapplerItem GrapplerItem::WithGraph(GraphDef&& graph_def) const { + GrapplerItem item; + item.id = id; + item.feed = feed; + item.fetch = fetch; + item.init_ops = init_ops; + item.keep_ops = keep_ops; + item.expected_init_time = expected_init_time; + item.save_op = save_op; + item.restore_op = restore_op; + item.save_restore_loc_tensor = save_restore_loc_tensor; + item.queue_runners = queue_runners; + item.devices_ = devices_; + item.allowed_optimizations_ = allowed_optimizations_; + item.graph.Swap(&graph_def); + return item; } std::vector GrapplerItem::MainOpsFanin() const { diff --git a/tensorflow/core/grappler/grappler_item.h b/tensorflow/core/grappler/grappler_item.h index 6ef4f14247b..9051542988c 100644 --- a/tensorflow/core/grappler/grappler_item.h +++ b/tensorflow/core/grappler/grappler_item.h @@ -35,12 +35,15 @@ namespace grappler { // nodes, and potentially a set of nodes to feed. struct GrapplerItem { GrapplerItem() = default; - GrapplerItem(const GrapplerItem& other, GraphDef&& graph_def) - : GrapplerItem(other, &graph_def) {} - // Swaps *graph_def with an empty GraphDef. - GrapplerItem(const GrapplerItem& other, GraphDef* graph_def); + GrapplerItem(const GrapplerItem& other) = default; + GrapplerItem(GrapplerItem&& other) = default; + GrapplerItem& operator=(const GrapplerItem& other) = default; + GrapplerItem& operator=(GrapplerItem&& other) = default; virtual ~GrapplerItem() = default; + // Create a copy of this GrapplerItem with graph swapped with the argument. + GrapplerItem WithGraph(GraphDef&& graph) const; + string id; // A unique id for this item // Inputs diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc index e41b1cf6840..d35c00f29ec 100644 --- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.cc @@ -3561,8 +3561,7 @@ Status ArithmeticOptimizer::Optimize(Cluster* /*cluster*/, // Set up helper data structures. nodes_to_preserve_ = item.NodesToPreserve(); fetch_nodes_known_ = !item.fetch.empty(); - *optimized_graph = item.graph; - GrapplerItem optimized_item(item, optimized_graph); + GrapplerItem optimized_item(item); optimized_graph_ = &optimized_item.graph; node_map_.reset(new NodeMap(optimized_graph_)); diff --git a/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc b/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc index e1ac7766d34..e330835e9bc 100644 --- a/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc +++ b/tensorflow/core/grappler/optimizers/experimental_implementation_selector_test.cc @@ -127,7 +127,7 @@ TEST_F(ExperimentalImplementationSelectorTest, SwapImplementationEval) { test::AsScalar(4.0f)); TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output)); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); const auto twice_boosted_tensor = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(twice_boosted_tensor[0], test::AsScalar(2.0f)); @@ -223,7 +223,7 @@ TEST_F(ExperimentalImplementationSelectorTest, SwapImplementationWithGradient) { test::AsScalar(4.0f)); TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output)); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); const auto twice_boosted_tensor = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(twice_boosted_tensor[0], test::AsScalar(2.0f)); diff --git a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc index 935fdc64095..c971eec3f4d 100644 --- a/tensorflow/core/grappler/optimizers/function_optimizer_test.cc +++ b/tensorflow/core/grappler/optimizers/function_optimizer_test.cc @@ -108,7 +108,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_SimpleFunction) { item.fetch = {"z"}; item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -184,7 +184,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_SkipErrorsIfGraphNotModified) { item.fetch = {"z1"}; item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -284,7 +284,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_FixedTypeFunction) { item.fetch = {"z"}; item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -368,7 +368,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_FunctionWithOutputMapping) { item.fetch = {"z"}; item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -418,7 +418,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_FunctionWithInputForwarding) { item.feed.emplace_back("x4", test::AsScalar(-1.0f)); item.feed.emplace_back("x3", test::AsScalar(1234)); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); test::ExpectTensorEqual(tensors_expected[1], tensors[1]); @@ -549,7 +549,7 @@ TEST_F(FunctionOptimizerTest, InlineFunction_FunctionWithNestedFunctionCall) { item.feed.emplace_back("a", test::AsScalar(2.0f)); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); @@ -748,7 +748,7 @@ TEST_F(FunctionOptimizerTest, InlineIndirectFunctionSimpleFunction) { item.feed.emplace_back("a", pi); item.feed.emplace_back("b", pi); - GrapplerItem optimized(item, std::move(optimized_graph)); + GrapplerItem optimized = item.WithGraph(std::move(optimized_graph)); auto tensors_expected = EvaluateFetchNodes(item); auto tensors = EvaluateFetchNodes(optimized); ASSERT_EQ(tensors_expected.size(), 1); @@ -876,7 +876,7 @@ TEST_F(FunctionOptimizerTest, InlineIndirectFunctionWithControlDependencies) { EXPECT_EQ(tensors_expected[0].flat()(0), 4.0); // mul EXPECT_EQ(tensors_expected[1].flat()(0), 3.0); // read variable - GrapplerItem optimized(item, std::move(optimized_graph)); + GrapplerItem optimized = item.WithGraph(std::move(optimized_graph)); auto tensors = EvaluateFetchNodes(optimized); ASSERT_EQ(tensors.size(), 2); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); @@ -1019,7 +1019,7 @@ TEST_F(FunctionOptimizerTest, InlineIndirectFunctionWithoutSideEffects) { auto tensors_expected = EvaluateFetchNodes(item); ASSERT_EQ(tensors_expected.size(), 1); - GrapplerItem optimized(item, std::move(optimized_graph)); + GrapplerItem optimized = item.WithGraph(std::move(optimized_graph)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -1067,7 +1067,7 @@ TEST_F(FunctionOptimizerTest, SpecializeFunctionXTimesTwo) { item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -1131,7 +1131,7 @@ TEST_F(FunctionOptimizerTest, SpecializeIndirectFunctionXTimesTwo) { item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -1194,7 +1194,7 @@ TEST_F(FunctionOptimizerTest, SpecializeFunctionPushDownConstInput) { item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -1274,7 +1274,7 @@ TEST_F(FunctionOptimizerTest, SpecializeIndirectFunctionPushDownConstInput) { item.feed.emplace_back("x", pi); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -1390,7 +1390,7 @@ TEST_F(FunctionOptimizerTest, SpecializeFunction_OncePerUniqueContext) { item.feed = {{"xf", pi}, {"yf", pi}, {"xi", four}, {"yi", four}}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); @@ -1499,7 +1499,7 @@ TEST_F(FunctionOptimizerTest, SpecializeFunctionForUsedOutputTensors) { item.feed = {{"xf", pi}, {"yf", pi}}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); ASSERT_EQ(tensors_expected.size(), tensors.size()); @@ -1660,7 +1660,7 @@ TEST_F(FunctionOptimizerTest, SpecializeIndirectFunctionForUsedOutputTensors) { item.feed = {{"xf", pi}, {"yf", pi}}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); ASSERT_EQ(tensors_expected.size(), tensors.size()); diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc index 453db5d91e7..227c2bb8b0f 100644 --- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc @@ -1306,13 +1306,12 @@ Status RelaxAllocatorConstraints(GraphDef* optimized_graph) { Status MemoryOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, GraphDef* optimized_graph) { - *optimized_graph = item.graph; + GrapplerItem optimized_item(item); RecomputationRewritingPass(optimization_level_, - recomputation_targets_name_scope_, optimized_graph, - item); + recomputation_targets_name_scope_, + &optimized_item.graph, item); - GrapplerItem optimized_item(item, optimized_graph); std::unordered_set skip_list; // Bound the number of rewrite passes to avoid long processing times on graphs // that simply won't fit in memory. diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc index 75285b07bb9..356b23dec0d 100644 --- a/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc +++ b/tensorflow/core/grappler/optimizers/memory_optimizer_test.cc @@ -279,7 +279,7 @@ TEST_F(MemoryOptimizerTest, SimpleSwapping) { EXPECT_EQ("^swap_out_e_0", new_c.input(1)); // Run the optimizer a second time to ensure it's idempotent. - GrapplerItem item_copy(item, std::move(output)); + GrapplerItem item_copy = item.WithGraph(std::move(output)); status = optimizer.Optimize(cluster.get(), item_copy, &output); TF_EXPECT_OK(status); @@ -287,7 +287,7 @@ TEST_F(MemoryOptimizerTest, SimpleSwapping) { item.fetch = {"e"}; item.init_ops = {init.name()}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); #endif @@ -337,7 +337,7 @@ TEST_F(MemoryOptimizerTest, SwappingHeuristics) { #if GOOGLE_CUDA auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); for (int i = 0; i < item.fetch.size(); ++i) { test::ExpectTensorEqual(tensors_expected[i], tensors[i]); @@ -386,7 +386,7 @@ TEST_F(MemoryOptimizerTest, UnswappableInputs) { #if GOOGLE_CUDA auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); #endif @@ -474,7 +474,7 @@ TEST_F(RelaxAllocatorConstraintsTest, SameDevice) { item.fetch = {"exp"}; item.init_ops = {"variable"}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); } @@ -505,7 +505,7 @@ TEST_F(RelaxAllocatorConstraintsTest, DifferentDevice) { item.fetch = {"exp"}; item.init_ops = {"variable"}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); #endif @@ -598,7 +598,7 @@ TEST_F(RelaxAllocatorConstraintsTest, AssignNodeInFanout) { item.fetch = {"assign0", "assign1"}; item.init_ops = {"exp1", "variable1"}; auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); for (int i = 0; i < tensors_expected.size(); ++i) { test::ExpectTensorEqual(tensors_expected[i], tensors[i]); diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc index 7c83036341c..572cc41d765 100644 --- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc @@ -440,7 +440,7 @@ Status MetaOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, item.graph) .ToProto(); - GrapplerItem trimmed_item(item, std::move(trimmed_graph)); + GrapplerItem trimmed_item = item.WithGraph(std::move(trimmed_graph)); VLOG(1) << absl::Substitute( "Deleted $0 unreachable functions from the graph (library size = $1)", diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc index 42b867b6ac1..12db5d6ca9b 100644 --- a/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc +++ b/tensorflow/core/grappler/optimizers/meta_optimizer_test.cc @@ -396,7 +396,7 @@ TEST_F(MetaOptimizerTest, OptimizeFunctionLibrary) { item.feed.emplace_back("b", test::AsScalar(4)); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); @@ -502,7 +502,7 @@ TEST_F(MetaOptimizerTest, OptimizeFunctionLibraryPruneFunctionBody) { item.feed.emplace_back("b", test::AsScalar(3.123f)); auto tensors_expected = EvaluateFetchNodes(item); - GrapplerItem optimized(item, std::move(output)); + GrapplerItem optimized = item.WithGraph(std::move(output)); auto tensors = EvaluateFetchNodes(optimized); test::ExpectTensorEqual(tensors_expected[0], tensors[0]); diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc index d8e62e0b24e..3fb3f2b0ec7 100644 --- a/tensorflow/core/grappler/optimizers/remapper.cc +++ b/tensorflow/core/grappler/optimizers/remapper.cc @@ -665,7 +665,7 @@ Status Remapper::Optimize(Cluster* /*cluster*/, const GrapplerItem& item, std::reverse(topo_sorted_graph.mutable_node()->begin(), topo_sorted_graph.mutable_node()->end()); - GrapplerItem topo_sorted_item(item, std::move(topo_sorted_graph)); + GrapplerItem topo_sorted_item = item.WithGraph(std::move(topo_sorted_graph)); RemapperContext ctx(topo_sorted_item); // Skip nodes that were invalidated by a remapper, e.g. do not process BiasAdd