TensorFlow: implement placement heuristic that takes generator nodes

(nodes with one non-ref output and one consumer), and places it
preferentially with its consumer.

For example:
    assign
   /      \
var       input

In the above graph, assign is bound to the device of 'var' due to the
reference edge.  This heuristic binds 'input' to the same device as
the assign, because it has only one consumer.

This addresses the general problem of colocating initializers with
their variables, and similar other cases.  There are very few reasons
to want to place the 'input' on a node other than its consumer (there
are some contrived cases, but that's why this is a heuristic).

This CL adds a test case for this small example above, illustrative
of the general problem.

An extension of this CL would be to do the same thing not just for
single output / single consumer nodes, but whenever all out edges of
a node connect to the same 'colcoation group'.
Change: 123896863

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/core/common_runtime/simple_placer.h b/tensorflow/core/common_runtime/simple_placer.h
index 7bfd5a1311..1dfb373307 100644
--- a/tensorflow/core/common_runtime/simple_placer.h
+++ b/tensorflow/core/common_runtime/simple_placer.h
@@ -79,6 +79,15 @@ class SimplePlacer {
   Status Run();
 
  private:
+  // Returns true if the device type of 'candidate_device_name' is
+  // found in 'devices'.
+  bool CanAssignToDevice(const string& candidate_device_name,
+                         const std::vector<Device*> devices) const;
+
+  // Assigns 'node's devices to 'assigned_device', and logs the
+  // placement if the SessionOptions entry in 'options_' requests it.
+  void AssignAndLog(const string& assigned_device, Node* node) const;
+
   Graph* const graph_;                           // Not owned.
   const DeviceSet* const devices_;               // Not owned.
   const SessionOptions* options_;                // Not owned.