From 5e9107066d9a283f2e9d040747a9b4e0537a6359 Mon Sep 17 00:00:00 2001 From: Saurabh Saxena Date: Wed, 21 Oct 2020 20:05:23 -0700 Subject: [PATCH] Add read variable op to function control outputs After inlining, we convert control dependencies to the function call to depend on all control outputs. Although read variable ops doesn't have side effect itself, it must be executed before variable updates after the function call. Note that it's not enough to convert control dependencies to the function call to depend on both control outputs and data outputs of the function call. The read variable ops can be inputs to ops that have side effects, e.g. assert and print, which are not the function data outputs. PiperOrigin-RevId: 338391537 Change-Id: Ibc843c1a584088e54010685ebce678b047f7d94d --- tensorflow/python/eager/function_test.py | 41 +++++++++++++++++++ .../python/framework/auto_control_deps.py | 11 +++-- .../framework/auto_control_deps_test.py | 34 ++++++++------- .../python/grappler/constant_folding_test.py | 8 +--- 4 files changed, 66 insertions(+), 28 deletions(-) diff --git a/tensorflow/python/eager/function_test.py b/tensorflow/python/eager/function_test.py index bd7a5060f39..7ebcf77ff61 100644 --- a/tensorflow/python/eager/function_test.py +++ b/tensorflow/python/eager/function_test.py @@ -4383,6 +4383,47 @@ class FunctionTest(test.TestCase, parameterized.TestCase): self.assertAllEqual(g(), array_ops.ones([1, 2, 3])) + @test_util.run_v2_only + def testControlDependencyAfterInline(self): + v = variables.Variable(0.) + + @def_function.function + def assign(): + return v.assign(1.) + + @def_function.function + def assign_add(): + return v.assign_add(1.) + + @def_function.function + def f(): + check_ops.assert_equal_v2(assign(), 1.) + check_ops.assert_equal_v2(assign_add(), 2.) + + # We don't have a way to inspect the inlined graph in Python, so we run it + # multiple times to have more confidence the dependency is correct. + for _ in range(30): + f() + + @test_util.run_v2_only + def testReadInFuncWriteOutside(self): + # Run many times since we are testing for a potential race condition. + for _ in range(30): + # pylint: disable=cell-var-from-loop + v = variables.Variable(1.) + + @def_function.function + def add_one(): + return v + 1. + + @def_function.function + def get_v_plus_one(): + v_plus_one = add_one() + v.assign_add(2.0) + return v_plus_one + + self.assertAllEqual(get_v_plus_one(), 2.0) + class MultiDeviceTest(test.TestCase, parameterized.TestCase): diff --git a/tensorflow/python/framework/auto_control_deps.py b/tensorflow/python/framework/auto_control_deps.py index 12f35a494b1..3f6ab98605b 100644 --- a/tensorflow/python/framework/auto_control_deps.py +++ b/tensorflow/python/framework/auto_control_deps.py @@ -373,11 +373,14 @@ class AutomaticControlDependencies(object): if control_flow_util.IsInWhileLoop(op): continue control_inputs = set() - # Ensure stateful ops run + # Ensure stateful ops run. + # Read-only ops are added to control outputs if the read value is + # consumed. This covers the case when the read value is returned from + # the function since that goes through a tf.identity in mark_as_return. if (op_def_registry.get(op.type) is None or - (op_is_stateful(op) and op.type not in utils.RESOURCE_READ_OPS)): - # TODO(srbs): Do not add functional ops to `ops_which_must_run` if - # they only have variable reads and are otherwise stateless. + (op_is_stateful(op) and + (op.type not in utils.RESOURCE_READ_OPS or + any(output.consumers() for output in op.outputs)))): ops_which_must_run.add(op) # Make a note of all opened manager_ids. if op.type == "NoOp": diff --git a/tensorflow/python/framework/auto_control_deps_test.py b/tensorflow/python/framework/auto_control_deps_test.py index dc5d8986958..a7b238c31b8 100644 --- a/tensorflow/python/framework/auto_control_deps_test.py +++ b/tensorflow/python/framework/auto_control_deps_test.py @@ -102,23 +102,16 @@ class AutomaticControlDependenciesTest(test.TestCase): self.assertNotIn(read_op1, read_op2.control_inputs) self.assertNotIn(read_op2, read_op1.control_inputs) - def testVariableReadsNotInOpsWithMustRun(self): + def testVariableReadsInOpsWithMustRun(self): with context.graph_mode(), self.cached_session(): v = resource_variable_ops.ResourceVariable(1.0) self.evaluate(variables.global_variables_initializer()) with acd.AutomaticControlDependencies() as c: - read_op1 = gen_resource_variable_ops.read_variable_op( - v.handle, v.dtype).op - read_op2 = gen_resource_variable_ops.read_variable_op( - v.handle, v.dtype).op - assign_op = gen_resource_variable_ops.assign_variable_op( - v.handle, v + 1) - # Reads must not be in `ops_which_must_run` since those get added to the - # `control_outputs`. - self.assertNotIn(read_op1, c.ops_which_must_run) - self.assertNotIn(read_op2, c.ops_which_must_run) - # Last write must be in `ops_which_must_run`. - self.assertIn(assign_op, c.ops_which_must_run) + read_op = gen_resource_variable_ops.read_variable_op(v.handle, + v.dtype).op + # Read ops get added to control outputs only if they have consumers. + c.mark_as_return(read_op.outputs[0]) + self.assertIn(read_op, c.ops_which_must_run) def testVariableMultipleReadsAndWrites(self): with context.graph_mode(), self.cached_session(): @@ -142,6 +135,11 @@ class AutomaticControlDependenciesTest(test.TestCase): v.handle, v + 1) assign_op4 = gen_resource_variable_ops.assign_variable_op( v.handle, v + 1) + # Read ops get added to control outputs only if they have consumers. + c.mark_as_return(read_op1.outputs[0]) + c.mark_as_return(read_op2.outputs[0]) + c.mark_as_return(read_op3.outputs[0]) + c.mark_as_return(read_op4.outputs[0]) # Verify the control edges. self.assertIn(read_op1, assign_op1.control_inputs) @@ -158,11 +156,11 @@ class AutomaticControlDependenciesTest(test.TestCase): for src_op, tgt_op in itertools.product(read_ops, read_ops): self.assertNotIn(src_op, tgt_op.control_inputs) - # Reads must not be in `ops_which_must_run`. - self.assertNotIn(read_op1, c.ops_which_must_run) - self.assertNotIn(read_op2, c.ops_which_must_run) - self.assertNotIn(read_op3, c.ops_which_must_run) - self.assertNotIn(read_op4, c.ops_which_must_run) + # Reads must be in `ops_which_must_run`. + self.assertIn(read_op1, c.ops_which_must_run) + self.assertIn(read_op2, c.ops_which_must_run) + self.assertIn(read_op3, c.ops_which_must_run) + self.assertIn(read_op4, c.ops_which_must_run) # Last write must be in `ops_which_must_run`. self.assertIn(assign_op4, c.ops_which_must_run) diff --git a/tensorflow/python/grappler/constant_folding_test.py b/tensorflow/python/grappler/constant_folding_test.py index 3336d3f7e8f..b5a9e7c8e48 100644 --- a/tensorflow/python/grappler/constant_folding_test.py +++ b/tensorflow/python/grappler/constant_folding_test.py @@ -96,19 +96,15 @@ class ConstantFoldingTest(test.TestCase): f(x, y).numpy() self.assertLen(graphs, 1) assign_count = 0 - read_count = 0 for node in graphs[0].node: if node.op == 'AssignAddVariableOp': self.assertEqual(node.input[0], 'y') assign_count += 1 - if node.op == 'ReadVariableOp': - read_count += 1 # Make sure that the only variable update that remains after - # grappler optimization is that of y, and that we prune all - # but the 2 necessary variable reads. + # grappler optimization is that of y. self.assertEqual(assign_count, 1) - self.assertEqual(read_count, 2) + self.assertLen(graphs[0].node, 11) if __name__ == '__main__':