1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/common_runtime/placer.h"
17
18#include <memory>
19#include <string>
20#include <utility>
21#include <vector>
22
23#include "tensorflow/core/common_runtime/device.h"
24#include "tensorflow/core/common_runtime/device_factory.h"
25#include "tensorflow/core/common_runtime/device_set.h"
26#include "tensorflow/core/framework/device_attributes.pb.h"
27#include "tensorflow/core/framework/kernel_def_builder.h"
28#include "tensorflow/core/framework/op.h"
29#include "tensorflow/core/framework/op_def_builder.h"
30#include "tensorflow/core/framework/op_kernel.h"
31#include "tensorflow/core/graph/graph.h"
32#include "tensorflow/core/graph/graph_def_builder.h"
33#include "tensorflow/core/graph/graph_def_builder_util.h"
34#include "tensorflow/core/lib/core/error_codes.pb.h"
35#include "tensorflow/core/lib/core/errors.h"
36#include "tensorflow/core/lib/core/status_test_util.h"
37#include "tensorflow/core/lib/strings/strcat.h"
38#include "tensorflow/core/platform/test.h"
39
40namespace tensorflow {
41
42namespace {
43
44////////////////////////////////////////////////////////////////////////////////
45//
46// Op, kernel, and device registrations to set up the environment.
47//
48// The Placer uses information about the op (input types),
49// kernel (device constraints), and available devices to make
50// placement decisions. To avoid depending on the full runtime, we
51// define dummy implementations of these, and register them with the
52// runtime.
53//
54////////////////////////////////////////////////////////////////////////////////
55
56// A dummy OpKernel that is used to register ops on different devices.
57class DummyOp : public OpKernel {
58 public:
59  explicit DummyOp(OpKernelConstruction* context) : OpKernel(context) {}
60  void Compute(OpKernelContext* context) override {}
61};
62
63// A fake device that has specific device attributes, used to simulate
64// the presence of a CPU or a GPU (without depending on that part of
65// the runtime.
66class FakeDevice : public Device {
67 private:
68  explicit FakeDevice(const DeviceAttributes& device_attributes)
69      : Device(nullptr, device_attributes) {}
70
71 public:
72  Status Sync() override { return errors::Unimplemented("FakeDevice::Sync()"); }
73
74  Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; }
75
76  static std::unique_ptr<Device> MakeCPU(const string& name) {
77    DeviceAttributes device_attributes;
78    device_attributes.set_name(name);
79    device_attributes.set_device_type(DeviceType("FakeCPU").type());
80    return std::unique_ptr<Device>(new FakeDevice(device_attributes));
81  }
82
83  static std::unique_ptr<Device> MakeGPU(const string& name) {
84    DeviceAttributes device_attributes;
85    device_attributes.set_name(name);
86    device_attributes.set_device_type(DeviceType("FakeGPU").type());
87    return std::unique_ptr<Device>(new FakeDevice(device_attributes));
88  }
89};
90
91class DummyFactory : public DeviceFactory {
92 public:
93  Status CreateDevices(const SessionOptions& options, const string& name_prefix,
94                       std::vector<Device*>* devices) override {
95    return Status::OK();
96  }
97};
98
99// Device order now depends on the registration of devices, not a fixed
100// value in device_set.cc.  To avoid the need to link in the real CPU and GPU
101// devices into this test, we create fake devices and registrations that
102// can stand-in for the real devices for the purposes of testing placement
103// and ordering.
104REGISTER_LOCAL_DEVICE_FACTORY("FakeCPU", DummyFactory);
105REGISTER_LOCAL_DEVICE_FACTORY("FakeGPU", DummyFactory, 51);
106
107// Register the following ops so they can be added to a Graph, and
108// kernels so that they can be placed on particular device types.
109REGISTER_OP("TestVariable").Output("o: Ref(float)");
110REGISTER_KERNEL_BUILDER(Name("TestVariable").Device("FakeCPU"), DummyOp);
111REGISTER_KERNEL_BUILDER(Name("TestVariable").Device("FakeGPU"), DummyOp);
112
113REGISTER_OP("VariableCPU").Output("o: Ref(float)");
114REGISTER_KERNEL_BUILDER(Name("VariableCPU").Device("FakeCPU"), DummyOp);
115
116REGISTER_OP("VariableGPU").Output("o: Ref(float)");
117REGISTER_KERNEL_BUILDER(Name("VariableGPU").Device("FakeGPU"), DummyOp);
118
119REGISTER_OP("VariableNoKernels").Output("o: Ref(float)");
120
121REGISTER_OP("TestAdd").Input("a: float").Input("b: float").Output("o: float");
122REGISTER_KERNEL_BUILDER(Name("TestAdd").Device("FakeCPU"), DummyOp);
123REGISTER_KERNEL_BUILDER(Name("TestAdd").Device("FakeGPU"), DummyOp);
124
125REGISTER_OP("TestRelu").Input("i: float").Output("o: float");
126REGISTER_KERNEL_BUILDER(Name("TestRelu").Device("FakeCPU"), DummyOp);
127REGISTER_KERNEL_BUILDER(Name("TestRelu").Device("FakeGPU"), DummyOp);
128
129REGISTER_OP("ReluCPU").Input("i: float").Output("o: float");
130REGISTER_KERNEL_BUILDER(Name("ReluCPU").Device("FakeCPU"), DummyOp);
131
132REGISTER_OP("ReluGPU").Input("i: float").Output("o: float");
133REGISTER_KERNEL_BUILDER(Name("ReluGPU").Device("FakeGPU"), DummyOp);
134
135REGISTER_OP("TestAssign").Input("i: Ref(float)").Input("v: float");
136REGISTER_KERNEL_BUILDER(Name("TestAssign").Device("FakeCPU"), DummyOp);
137REGISTER_KERNEL_BUILDER(Name("TestAssign").Device("FakeGPU"), DummyOp);
138
139REGISTER_OP("AssignCPU").Input("i: Ref(float)").Input("v: float");
140REGISTER_KERNEL_BUILDER(Name("AssignCPU").Device("FakeCPU"), DummyOp);
141
142REGISTER_OP("AssignGPU").Input("i: Ref(float)").Input("v: float");
143REGISTER_KERNEL_BUILDER(Name("AssignGPU").Device("FakeGPU"), DummyOp);
144
145REGISTER_OP("TestInput").Output("a: float").Output("b: float");
146REGISTER_KERNEL_BUILDER(Name("TestInput").Device("FakeCPU"), DummyOp);
147
148// Op producing an output that can be placed on CPU or GPU.
149REGISTER_OP("TestCPUGPUOutput").Output("a: float");
150REGISTER_KERNEL_BUILDER(Name("TestCPUGPUOutput").Device("FakeCPU"), DummyOp);
151REGISTER_KERNEL_BUILDER(Name("TestCPUGPUOutput").Device("FakeGPU"), DummyOp);
152
153REGISTER_OP("TestGPUOutput").Output("a: float");
154REGISTER_KERNEL_BUILDER(Name("TestGPUOutput").Device("FakeGPU"), DummyOp);
155
156REGISTER_OP("TestDevice").Output("a: float").Output("b: float");
157REGISTER_KERNEL_BUILDER(Name("TestDevice").Device("FakeGPU"), DummyOp);
158
159REGISTER_OP("TestDeviceEnforce").Input("a: Ref(float)").Output("b: float");
160REGISTER_KERNEL_BUILDER(Name("TestDeviceEnforce").Device("FakeCPU"), DummyOp);
161REGISTER_KERNEL_BUILDER(Name("TestDeviceEnforce").Device("FakeGPU"), DummyOp);
162
163REGISTER_KERNEL_BUILDER(Name("Shape").Device("FakeCPU"), DummyOp);
164REGISTER_KERNEL_BUILDER(Name("Shape").Device("FakeGPU"), DummyOp);
165
166////////////////////////////////////////////////////////////////////////////////
167//
168// A PlacerTest method has three phases:
169//
170// 1. Build a TensorFlow graph, with no (or partial) device assignments.
171// 2. Attempt to compute a placement using the Placer.
172// 3. EITHER: test that the constraints implied by the graph are respected;
173//    or that an appropriate error was reported.
174//
175////////////////////////////////////////////////////////////////////////////////
176class PlacerTest : public ::testing::Test {
177 protected:
178  PlacerTest() {
179    // Build a set of 10 GPU and 10 CPU devices.
180    // NOTE: this->local_devices_ owns the device objects;
181    // this->devices_ contains borrowed pointers to the device
182    // objects.
183    for (int i = 0; i < 10; ++i) {
184      local_devices_.emplace_back(FakeDevice::MakeCPU(
185          strings::StrCat("/job:a/replica:0/task:0/device:fakecpu:", i)));
186      devices_.AddDevice(local_devices_.back().get());
187      // Insert the GPUs in reverse order.
188      local_devices_.emplace_back(FakeDevice::MakeGPU(
189          strings::StrCat("/job:a/replica:0/task:0/device:fakegpu:", 9 - i)));
190      devices_.AddDevice(local_devices_.back().get());
191    }
192  }
193
194  // Builds the given graph, and (if successful) indexes the node
195  // names for use in placement, and later lookup.
196  Status BuildGraph(const GraphDefBuilder& builder, Graph* out_graph) {
197    TF_RETURN_IF_ERROR(GraphDefBuilderToGraph(builder, out_graph));
198    nodes_by_name_.clear();
199    for (Node* node : out_graph->nodes()) {
200      nodes_by_name_[node->name()] = node->id();
201    }
202    return Status::OK();
203  }
204
205  // Invokes the Placer on "graph". If no DeviceSet is specified, the
206  // placement will use the default DeviceSet (of 10 CPU and 10 GPU devices).
207  //
208  // REQUIRES: "*graph" was produced by the most recent call to BuildGraph.
209  Status Place(Graph* graph, DeviceSet* devices, SessionOptions* options) {
210    Placer placer(graph, devices, options);
211    return placer.Run();
212  }
213
214  Status Place(Graph* graph, DeviceSet* devices) {
215    return Place(graph, devices, nullptr);
216  }
217
218  Status Place(Graph* graph, SessionOptions* options) {
219    return Place(graph, &devices_, options);
220  }
221
222  Status Place(Graph* graph) { return Place(graph, &devices_, nullptr); }
223
224  // Returns the node in "graph" with the given name.
225  //
226  // REQUIRES: "graph" was produced by the most recent call to BuildGraph.
227  Node* GetNodeByName(const Graph& graph, const string& name) {
228    const auto search = nodes_by_name_.find(name);
229    CHECK(search != nodes_by_name_.end()) << "Unknown node name: " << name;
230    return graph.FindNodeId(search->second);
231  }
232
233 protected:
234  std::vector<std::unique_ptr<Device>> local_devices_;
235  DeviceSet devices_;
236  Placer::NodeNameToIdMap nodes_by_name_;
237
238  Status ReferenceTestHelper(const string& variable_op_type,
239                             const string& assign_op_type,
240                             const DeviceType& expected_device_type);
241};
242
243#define EXPECT_COLOCATED(g, name_a, name_b)                         \
244  do {                                                              \
245    Graph& g_ = (g);                                                \
246    EXPECT_EQ(GetNodeByName(g_, (name_a))->assigned_device_name(),  \
247              GetNodeByName(g_, (name_b))->assigned_device_name()); \
248  } while (0)
249
250#define EXPECT_NOT_COLOCATED(g, name_a, name_b)                     \
251  do {                                                              \
252    Graph& g_ = (g);                                                \
253    EXPECT_NE(GetNodeByName(g_, (name_a))->assigned_device_name(),  \
254              GetNodeByName(g_, (name_b))->assigned_device_name()); \
255  } while (0)
256
257#define EXPECT_DEVICE_TYPE(g, name, expected_device_type)               \
258  EXPECT_EQ(DeviceType(expected_device_type).type(),                    \
259            devices_                                                    \
260                .FindDeviceByName(                                      \
261                    GetNodeByName((g), (name))->assigned_device_name()) \
262                ->attributes()                                          \
263                .device_type())
264
265#define EXPECT_DEVICE_CONTAINS(g, name, device_substr)                        \
266  EXPECT_TRUE(StringPiece(GetNodeByName((g), (name))->assigned_device_name()) \
267                  .contains(device_substr))
268
269// Test that a graph with no constraints will successfully assign nodes to the
270// "best available" device (i.e. prefer GPU over CPU).
271TEST_F(PlacerTest, TestNoConstraints) {
272  Graph g(OpRegistry::Global());
273  {  // Scope for temporary variables used to construct g.
274    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
275    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
276    ops::UnaryOp("TestRelu", ops::NodeOut(input, 0), b.opts().WithName("n1"));
277    ops::UnaryOp("TestRelu", ops::NodeOut(input, 1), b.opts().WithName("n2"));
278    TF_EXPECT_OK(BuildGraph(b, &g));
279  }
280
281  TF_EXPECT_OK(Place(&g));
282  EXPECT_DEVICE_TYPE(g, "in", "FakeCPU");
283  EXPECT_DEVICE_TYPE(g, "n1", "FakeGPU");
284  EXPECT_DEVICE_TYPE(g, "n2", "FakeGPU");
285}
286
287// Test that a graph with device type and reference constraints on
288// some of the ops will successfully assign nodes to the constrained
289// device, and colocate nodes with reference connections.
290TEST_F(PlacerTest, TestDeviceTypeConstraints) {
291  Graph g(OpRegistry::Global());
292  {  // Scope for temporary variables used to construct g.
293    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
294    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
295    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
296    ops::BinaryOp("AssignCPU", var_cpu, input, b.opts().WithName("assign_cpu"));
297    Node* var_gpu = ops::SourceOp("VariableGPU", b.opts().WithName("var_gpu"));
298    ops::BinaryOp("AssignGPU", var_gpu, input, b.opts().WithName("assign_gpu"));
299    TF_EXPECT_OK(BuildGraph(b, &g));
300  }
301
302  TF_EXPECT_OK(Place(&g));
303  EXPECT_DEVICE_TYPE(g, "in", "FakeCPU");
304  EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU");
305  EXPECT_DEVICE_TYPE(g, "assign_cpu", "FakeCPU");
306  EXPECT_COLOCATED(g, "var_cpu", "assign_cpu");
307  EXPECT_DEVICE_TYPE(g, "var_gpu", "FakeGPU");
308  EXPECT_DEVICE_TYPE(g, "assign_gpu", "FakeGPU");
309  EXPECT_COLOCATED(g, "var_gpu", "assign_gpu");
310}
311
312TEST_F(PlacerTest, TestMetadataColocatedWithInput) {
313  Graph g(OpRegistry::Global());
314  {  // Scope for temporary variables used to construct g.
315    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
316    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
317
318    // Normally, shape has a GPU implementation and would be placed
319    // on GPU.  However, because it is a metadata operation, it is
320    // placed on CPU to avoid transferring the data from CPU to GPU.
321    ops::UnaryOp("Shape", var_cpu, b.opts().WithName("shape_op"));
322    TF_EXPECT_OK(BuildGraph(b, &g));
323  }
324
325  TF_EXPECT_OK(Place(&g));
326  EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU");
327  EXPECT_DEVICE_TYPE(g, "shape_op", "FakeCPU");
328  EXPECT_COLOCATED(g, "var_cpu", "shape_op");
329}
330
331// Heuristic A implements "Island fusing": if a node only generates
332// an output and it has only one consumer, we place the node
333// with its consumer.
334TEST_F(PlacerTest, TestHeuristicGeneratorFollowsSingleConsumer) {
335  Graph g(OpRegistry::Global());
336  {  // Scope for temporary variables used to construct g.
337    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
338
339    // A variable is only on CPU
340    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
341
342    // The constant to be assigned can be on both GPU or CPU.
343    //
344    // Because of the heuristic, it gets placed on CPU to avoid a
345    // copy.
346    Node* input = ops::SourceOp("TestCPUGPUOutput", b.opts().WithName("in"));
347
348    // The assign is bound to CPU by the reference edge.
349    ops::BinaryOp("TestAssign", var_cpu, input, b.opts().WithName("assign"));
350
351    TF_EXPECT_OK(BuildGraph(b, &g));
352  }
353
354  TF_EXPECT_OK(Place(&g));
355  EXPECT_COLOCATED(g, "var_cpu", "in");
356  EXPECT_COLOCATED(g, "assign", "in");
357}
358
359TEST_F(PlacerTest, TestIgnoreGeneratorHeuristicIfWrongDevice) {
360  Graph g(OpRegistry::Global());
361  {  // Scope for temporary variables used to construct g.
362    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
363
364    // A variable is only on CPU
365    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
366
367    // The constant to be assigned can only be on GPU.
368    //
369    // The heuristic to place the generator with its consumer does
370    // not apply since the consumer's device is not in the list
371    // of valid devices for the generator.
372    Node* input = ops::SourceOp("TestGPUOutput", b.opts().WithName("in"));
373
374    // The assign is bound to CPU by the reference edge.
375    ops::BinaryOp("TestAssign", var_cpu, input, b.opts().WithName("assign"));
376
377    TF_EXPECT_OK(BuildGraph(b, &g));
378  }
379
380  TF_EXPECT_OK(Place(&g));
381  EXPECT_DEVICE_TYPE(g, "in", "FakeGPU");
382  EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU");
383  EXPECT_COLOCATED(g, "var_cpu", "assign");
384}
385
386TEST_F(PlacerTest, TestIgnoreGeneratorHeuristicIfWrongPartialDevice) {
387  Graph g(OpRegistry::Global());
388  {  // Scope for temporary variables used to construct g.
389    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
390
391    // A variable is only on CPU
392    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
393
394    // The constant to be assigned can be on CPU or GPU, but is explicitly
395    // placed on CPU:1.
396    //
397    // The heuristic to place the generator with its consumer does
398    // not apply since the consumer's device is not in the list
399    // of valid devices for the generator.
400    Node* input =
401        ops::SourceOp("TestCPUGPUOutput",
402                      b.opts().WithName("in").WithDevice("/device:fakecpu:1"));
403
404    // The assign is bound to CPU by the reference edge.
405    ops::BinaryOp("TestAssign", var_cpu, input, b.opts().WithName("assign"));
406
407    TF_EXPECT_OK(BuildGraph(b, &g));
408  }
409
410  TF_EXPECT_OK(Place(&g));
411  EXPECT_DEVICE_TYPE(g, "in", "FakeCPU");
412  EXPECT_DEVICE_CONTAINS(g, "in", "/device:fakecpu:1");
413  EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU");
414  EXPECT_COLOCATED(g, "var_cpu", "assign");
415  EXPECT_DEVICE_CONTAINS(g, "var_cpu", "/device:fakecpu:0");
416}
417
418// Test that a graph with partial device specifications on the ops
419// will successfully
420TEST_F(PlacerTest, TestPartialSpec) {
421  Graph g(OpRegistry::Global());
422  {  // Scope for temporary variables used to construct g.
423    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
424    ops::SourceOp("TestInput", b.opts().WithName("in").WithDevice("/job:a"));
425    ops::SourceOp("TestVariable",
426                  b.opts().WithName("var").WithDevice("/job:a"));
427    TF_EXPECT_OK(BuildGraph(b, &g));
428  }
429
430  TF_EXPECT_OK(Place(&g));
431  EXPECT_DEVICE_TYPE(g, "in", "FakeCPU");
432  EXPECT_DEVICE_CONTAINS(g, "in", "/job:a");
433  EXPECT_DEVICE_TYPE(g, "var", "FakeGPU");
434  EXPECT_DEVICE_CONTAINS(g, "var", "/job:a");
435}
436
437// Test that a node with a pre-assigned device is not relocated.
438TEST_F(PlacerTest, TestAssignedDevicePreserved) {
439  Graph g(OpRegistry::Global());
440  {  // Scope for temporary variables used to construct g.
441    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
442    ops::SourceOp("TestInput", b.opts().WithName("in"));
443    TF_EXPECT_OK(BuildGraph(b, &g));
444  }
445
446  GetNodeByName(g, "in")->set_assigned_device_name(
447      "/job:a/replica:0/task:0/device:fakecpu:7");
448
449  TF_EXPECT_OK(Place(&g));
450  EXPECT_EQ("/job:a/replica:0/task:0/device:fakecpu:7",
451            GetNodeByName(g, "in")->assigned_device_name());
452}
453
454// Test that a graph with partial device specifications for CPU-only ops
455// will be relocated to CPU.
456TEST_F(PlacerTest, TestPartialSpecGpuToCpu) {
457  Graph g(OpRegistry::Global());
458  {  // Scope for temporary variables used to construct g.
459    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
460    ops::SourceOp("TestInput",
461                  b.opts().WithName("in").WithDevice("/device:fakegpu:0"));
462    ops::SourceOp("TestVariable",
463                  b.opts().WithName("var").WithDevice("/device:fakegpu:0"));
464    TF_EXPECT_OK(BuildGraph(b, &g));
465  }
466
467  SessionOptions options;
468  options.config.set_allow_soft_placement(true);
469  TF_EXPECT_OK(Place(&g, &options));
470  EXPECT_DEVICE_TYPE(g, "in", "FakeCPU");
471  EXPECT_DEVICE_CONTAINS(g, "in", "/device:fakecpu");
472  EXPECT_DEVICE_TYPE(g, "var", "FakeGPU");
473  EXPECT_DEVICE_CONTAINS(g, "var", "/device:fakegpu:0");
474}
475
476// Test that a node with an assigned GPU device but has not registered
477// OpKernel will fail.
478TEST_F(PlacerTest, TestAssignedGpuDeviceToCpuDevice) {
479  Graph g(OpRegistry::Global());
480  {  // Scope for temporary variables used to construct g.
481    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
482    ops::SourceOp("TestInput", b.opts().WithName("in"));
483    TF_EXPECT_OK(BuildGraph(b, &g));
484  }
485
486  GetNodeByName(g, "in")->set_assigned_device_name(
487      "/job:a/replica:0/task:0/device:fakegpu:0");
488
489  Status s = Place(&g);
490  EXPECT_EQ(error::INTERNAL, s.code());
491  EXPECT_TRUE(
492      StringPiece(s.error_message())
493          .contains(
494              "Assigned device '/job:a/replica:0/task:0/device:fakegpu:0' "
495              "does not have registered OpKernel support for TestInput"));
496}
497
498// Test that graphs with reference connections are correctly placed.
499
500// Build a graph containing a Variable op of "variable_op_type" and an
501// Assign op of "assign_op_type", and expect all of the ops to be
502// placed on a device of type "expected_device_type".
503Status PlacerTest::ReferenceTestHelper(const string& variable_op_type,
504                                       const string& assign_op_type,
505                                       const DeviceType& expected_device_type) {
506  Graph g(OpRegistry::Global());
507  {  // Scope for temporary variables used to construct g.
508    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
509    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
510    // Build ten variable-and-assignment pairs.
511    for (int i = 0; i < 10; ++i) {
512      Node* var = ops::SourceOp(variable_op_type,
513                                b.opts().WithName(strings::StrCat("var_", i)));
514      ops::BinaryOp(assign_op_type, var, input,
515                    b.opts().WithName(strings::StrCat("assign_", i)));
516    }
517    TF_EXPECT_OK(BuildGraph(b, &g));
518  }
519
520  TF_RETURN_IF_ERROR(Place(&g));
521
522  for (int i = 0; i < 10; ++i) {
523    EXPECT_COLOCATED(g, strings::StrCat("var_", i),
524                     strings::StrCat("assign_", i));
525    EXPECT_DEVICE_TYPE(g, strings::StrCat("var_", i), expected_device_type);
526    EXPECT_DEVICE_TYPE(g, strings::StrCat("assign_", i), expected_device_type);
527  }
528
529  return Status::OK();
530}
531
532// Test all 2^3 combinations of Variable and Assignment op types
533// (unconstrained, CPU-only, and GPU-only).
534TEST_F(PlacerTest, TestReferenceConnection) {
535  Status s;
536  TF_EXPECT_OK(ReferenceTestHelper("TestVariable", "TestAssign", "FakeGPU"));
537  TF_EXPECT_OK(ReferenceTestHelper("TestVariable", "AssignCPU", "FakeCPU"));
538  TF_EXPECT_OK(ReferenceTestHelper("TestVariable", "AssignGPU", "FakeGPU"));
539  TF_EXPECT_OK(ReferenceTestHelper("VariableCPU", "TestAssign", "FakeCPU"));
540  TF_EXPECT_OK(ReferenceTestHelper("VariableCPU", "AssignCPU", "FakeCPU"));
541  {
542    Status s = ReferenceTestHelper("VariableCPU", "AssignGPU", "FakeCPU");
543    EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
544    EXPECT_TRUE(StringPiece(s.error_message())
545                    .contains("no device type supports both of those nodes"));
546  }
547  TF_EXPECT_OK(ReferenceTestHelper("VariableGPU", "TestAssign", "FakeGPU"));
548  {
549    Status s = ReferenceTestHelper("VariableGPU", "AssignCPU", "FakeCPU");
550    EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
551    EXPECT_TRUE(StringPiece(s.error_message())
552                    .contains("no device type supports both of those nodes"));
553  }
554  TF_EXPECT_OK(ReferenceTestHelper("VariableGPU", "AssignGPU", "FakeGPU"));
555}
556
557// Handle-using dummy variable ops.
558REGISTER_OP("TestHandleVariable").Output("o: resource");
559REGISTER_KERNEL_BUILDER(Name("TestHandleVariable").Device("FakeCPU"), DummyOp);
560REGISTER_KERNEL_BUILDER(Name("TestHandleVariable").Device("FakeGPU"), DummyOp);
561
562REGISTER_OP("HandleVariableCPU").Output("o: resource");
563REGISTER_KERNEL_BUILDER(Name("HandleVariableCPU").Device("FakeCPU"), DummyOp);
564
565REGISTER_OP("HandleVariableGPU").Output("o: resource");
566REGISTER_KERNEL_BUILDER(Name("HandleVariableGPU").Device("FakeGPU"), DummyOp);
567
568REGISTER_OP("TestHandleAssign").Input("i: resource").Input("v: float");
569REGISTER_KERNEL_BUILDER(Name("TestHandleAssign").Device("FakeCPU"), DummyOp);
570REGISTER_KERNEL_BUILDER(Name("TestHandleAssign").Device("FakeGPU"), DummyOp);
571
572REGISTER_OP("HandleAssignCPU").Input("i: resource").Input("v: float");
573REGISTER_KERNEL_BUILDER(Name("HandleAssignCPU").Device("FakeCPU"), DummyOp);
574
575REGISTER_OP("HandleAssignGPU").Input("i: resource").Input("v: float");
576REGISTER_KERNEL_BUILDER(Name("HandleAssignGPU").Device("FakeGPU"), DummyOp);
577
578// Tests all combinations of resource handles and ops using them.
579TEST_F(PlacerTest, TestResourceHandle) {
580  auto handle_test = [this](const string& var_op_name,
581                            const string& use_op_name, DeviceType device) {
582    Graph g(OpRegistry::Global());
583    {  // Scope for temporary variables used to construct g.
584      GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
585      Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
586      Node* var = ops::SourceOp(var_op_name, b.opts().WithName("var"));
587      ops::BinaryOp(use_op_name, var, input, b.opts().WithName("assign"));
588      TF_EXPECT_OK(BuildGraph(b, &g));
589    }
590
591    TF_RETURN_IF_ERROR(Place(&g));
592
593    EXPECT_COLOCATED(g, "var", "assign");
594    EXPECT_DEVICE_TYPE(g, "var", device);
595    EXPECT_DEVICE_TYPE(g, "assign", device);
596    return Status::OK();
597  };
598  TF_EXPECT_OK(
599      handle_test("TestHandleVariable", "TestHandleAssign", "FakeGPU"));
600  TF_EXPECT_OK(handle_test("TestHandleVariable", "HandleAssignCPU", "FakeCPU"));
601  TF_EXPECT_OK(handle_test("TestHandleVariable", "HandleAssignGPU", "FakeGPU"));
602  TF_EXPECT_OK(handle_test("HandleVariableCPU", "TestHandleAssign", "FakeCPU"));
603  TF_EXPECT_OK(handle_test("HandleVariableCPU", "HandleAssignCPU", "FakeCPU"));
604  TF_EXPECT_OK(handle_test("HandleVariableGPU", "HandleAssignGPU", "FakeGPU"));
605  TF_EXPECT_OK(handle_test("HandleVariableGPU", "TestHandleAssign", "FakeGPU"));
606  EXPECT_FALSE(
607      handle_test("HandleVariableGPU", "HandleAssignCPU", "FakeCPU").ok());
608  EXPECT_FALSE(
609      handle_test("HandleVariableCPU", "HandleAssignGPU", "FakeCPU").ok());
610}
611
612// Test that an assignment of an operator to the wrong device
613// is ignored when it could never be satisfied (due to reference
614// edges, for example).
615TEST_F(PlacerTest, TestReferenceConnectionIgnoreInfeasible) {
616  Status s;
617  Graph g(OpRegistry::Global());
618  {
619    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
620    Node* input = ops::SourceOp(
621        "TestDevice",
622        b.opts().WithName("in").WithDevice("/job:a/task:0/device:fakegpu:0"));
623    Node* var =
624        ops::SourceOp("TestVariable", b.opts().WithName("var_0").WithDevice(
625                                          "/job:a/task:0/device:fakegpu:0"));
626
627    // This op is specified on CPU, but in practice will be ignored,
628    // because the reference edges forces it on GPU.
629    ops::BinaryOp("TestAssign", var, input,
630                  b.opts().WithName("assign").WithDevice(
631                      "/job:a/task:0/device:fakecpu:0"));
632    TF_EXPECT_OK(BuildGraph(b, &g));
633  }
634
635  SessionOptions options;
636  s = Place(&g, &options);
637  TF_EXPECT_OK(s);
638  EXPECT_DEVICE_TYPE(g, "var_0", "FakeGPU");
639  EXPECT_DEVICE_TYPE(g, "assign", "FakeGPU");
640}
641
642// Test that an assignment of an operator to the a more specified device
643// causes the device to maintain its more specific placement.
644TEST_F(PlacerTest, TestReferenceConnectionMoreSpecificDestinationSourceWins) {
645  Status s;
646  Graph g(OpRegistry::Global());
647  {
648    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
649    // Input can be on either device
650    Node* input =
651        ops::SourceOp("TestCPUGPUOutput",
652                      b.opts().WithName("in").WithDevice("/job:a/task:0"));
653
654    // Variable can be on either device
655    Node* var = ops::SourceOp(
656        "TestVariable", b.opts().WithName("var_0").WithDevice("/job:a/task:0"));
657
658    // This op is specified on CPU and is more specific than the variable.
659    // Because the variable is less specified, the variable will be
660    // assigned to CPU.
661    ops::BinaryOp("TestAssign", var, input,
662                  b.opts().WithName("assign").WithDevice(
663                      "/job:a/task:0/device:fakecpu:0"));
664    TF_EXPECT_OK(BuildGraph(b, &g));
665  }
666
667  SessionOptions options;
668  s = Place(&g, &options);
669  TF_EXPECT_OK(s);
670  EXPECT_DEVICE_TYPE(g, "var_0", "FakeCPU");
671  EXPECT_DEVICE_TYPE(g, "assign", "FakeCPU");
672}
673
674// A reference connection exists between a variable and an assign,
675// where the assign has a device but the variable does not.  In this
676// case, the variable gets placed on the location of the assign
677// operation.
678TEST_F(PlacerTest, TestReferenceConnectionNoSourceDevice) {
679  Status s;
680  Graph g(OpRegistry::Global());
681  {
682    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
683    Node* input = ops::SourceOp(
684        "TestDevice",
685        b.opts().WithName("in").WithDevice("/job:a/task:0/device:fakegpu:0"));
686    Node* var = ops::SourceOp("TestVariable", b.opts().WithName("var_0"));
687    ops::BinaryOp("TestAssign", var, input,
688                  b.opts().WithName("assign").WithDevice(
689                      "/job:a/task:0/device:fakecpu:0"));
690    TF_EXPECT_OK(BuildGraph(b, &g));
691  }
692
693  SessionOptions options;
694  s = Place(&g, &options);
695  TF_EXPECT_OK(s);
696  EXPECT_DEVICE_TYPE(g, "var_0", "FakeCPU");
697  EXPECT_DEVICE_TYPE(g, "assign", "FakeCPU");
698}
699
700TEST_F(PlacerTest, TestColocationGroup) {
701  Graph g(OpRegistry::Global());
702  {  // Scope for temporary variables used to construct g.
703    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
704    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
705    Node* colocated_with_input = ops::UnaryOp(
706        "TestRelu", input,
707        b.opts().WithName("colocated_1").WithAttr("_class", {"loc:@in"}));
708
709    // This will not be colocated with the input because TestInput is
710    // only availbale on CPU and TestRelu will default to GPU.
711    Node* not_colocated_with_input =
712        ops::UnaryOp("TestRelu", input, b.opts().WithName("foo"));
713    CHECK(colocated_with_input);
714    CHECK(not_colocated_with_input);
715    TF_EXPECT_OK(BuildGraph(b, &g));
716  }
717
718  TF_EXPECT_OK(Place(&g));
719  EXPECT_COLOCATED(g, "in", "colocated_1");
720  EXPECT_NOT_COLOCATED(g, "in", "foo");
721}
722
723TEST_F(PlacerTest, TestMultipleColocationGroups) {
724  Graph g(OpRegistry::Global());
725  {  // Scope for temporary variables used to construct g.
726    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
727    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
728    Node* colocated_with_input = ops::UnaryOp(
729        "TestRelu", input,
730        b.opts().WithName("colocated_1").WithAttr("_class", {"loc:@in"}));
731    Node* colocated_with_input_and_other =
732        ops::UnaryOp("TestRelu", input,
733                     b.opts().WithName("foo").WithAttr(
734                         "_class", {"loc:@in", "loc:@colocated_1"}));
735    CHECK(colocated_with_input);
736    CHECK(colocated_with_input_and_other);
737    TF_EXPECT_OK(BuildGraph(b, &g));
738  }
739
740  TF_EXPECT_OK(Place(&g));
741  EXPECT_COLOCATED(g, "in", "colocated_1");
742  EXPECT_COLOCATED(g, "in", "foo");
743}
744
745TEST_F(PlacerTest, TestInvalidMultipleColocationGroups) {
746  Graph g(OpRegistry::Global());
747  {  // Scope for temporary variables used to construct g.
748    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
749    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
750    Node* colocated_with_input = ops::UnaryOp(
751        "ReluCPU", input,
752        b.opts().WithName("colocated_1").WithAttr("_class", {"loc:@in"}));
753    Node* colocated_with_input_and_other =
754        ops::UnaryOp("ReluGPU", input,
755                     b.opts().WithName("foo").WithAttr(
756                         "_class", {"loc:@in", "loc:@colocated_1"}));
757    CHECK(colocated_with_input);
758    CHECK(colocated_with_input_and_other);
759    TF_EXPECT_OK(BuildGraph(b, &g));
760  }
761
762  Status s = Place(&g);
763  EXPECT_TRUE(StringPiece(s.error_message())
764                  .contains("Cannot colocate nodes 'foo' and 'in' because no "
765                            "device type supports both of those nodes and the "
766                            "other nodes colocated with them"));
767}
768
769TEST_F(PlacerTest, TestColocationGroupWithReferenceConnections) {
770  Graph g(OpRegistry::Global());
771  {  // Scope for temporary variables used to construct g.
772    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
773    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
774    Node* var1 = ops::SourceOp("VariableCPU", b.opts().WithName("var1"));
775    Node* var2 = ops::SourceOp("VariableCPU", b.opts().WithName("var2"));
776
777    // Two assigns (reference connections) with two different
778    // colocation groups. Because their colocation groups all map to the
779    // same device, this is a valid assignment.
780    ops::BinaryOp(
781        "TestAssign", var1, input,
782        b.opts().WithName("assign1").WithAttr("_class", {"loc:@var1"}));
783    ops::BinaryOp(
784        "TestAssign", var2, input,
785        b.opts().WithName("assign2").WithAttr("_class", {"loc:@var2"}));
786    TF_EXPECT_OK(BuildGraph(b, &g));
787  }
788
789  TF_EXPECT_OK(Place(&g));
790  EXPECT_COLOCATED(g, "in", "var1");
791  EXPECT_COLOCATED(g, "in", "var2");
792  EXPECT_COLOCATED(g, "var1", "assign2");
793  EXPECT_COLOCATED(g, "var2", "assign1");
794}
795
796TEST_F(PlacerTest, TestColocationGroupWithUnsatisfiableReferenceConnections) {
797  Graph g(OpRegistry::Global());
798  {  // Scope for temporary variables used to construct g.
799    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
800    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
801
802    Node* var1 = ops::SourceOp("VariableCPU", b.opts().WithName("var1"));
803    Node* var2 = ops::SourceOp("VariableCPU", b.opts().WithName("var2"));
804    // Var 3 is on GPU
805    Node* var3 = ops::SourceOp("VariableGPU", b.opts().WithName("var3"));
806
807    // Two assigns (reference connections) with two different
808    // colocation groups. Because their colocation groups all map to the
809    // same device, this is a valid assignment.
810    ops::BinaryOp(
811        "TestAssign", var1, input,
812        b.opts().WithName("assign1").WithAttr("_class", {"loc:@var1"}));
813    ops::BinaryOp(
814        "TestAssign", var2, input,
815        b.opts().WithName("assign2").WithAttr("_class", {"loc:@var2"}));
816    // Assign to var3, but try to use a colocation group that matches
817    // the assign of var2.  This should fail because assign2 must be on CPU
818    // (it has a reference edge on var2), and assign3 must be on GPU,
819    // hence the conflict.
820    ops::BinaryOp(
821        "TestAssign", var3, input,
822        b.opts().WithName("assign3").WithAttr("_class", {"loc:@var2"}));
823    TF_EXPECT_OK(BuildGraph(b, &g));
824  }
825
826  Status s = Place(&g);
827  EXPECT_TRUE(
828      StringPiece(s.error_message())
829          .contains("Cannot colocate nodes 'var3' and 'assign3' because no "
830                    "device type supports both of those nodes and the other "
831                    "nodes colocated with them."));
832}
833
834TEST_F(PlacerTest, TestColocationAndReferenceConnections) {
835  Graph g(OpRegistry::Global());
836  {  // Scope for temporary variables used to construct g.
837    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
838    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
839    for (int i = 0; i < 10; ++i) {
840      // Declare ten variable and assignment pairs.
841      Node* var = ops::SourceOp("TestVariable",
842                                b.opts().WithName(strings::StrCat("var_", i)));
843      ops::BinaryOp("TestAssign", var, input,
844                    b.opts().WithName(strings::StrCat("assign_", i)));
845    }
846    for (int i = 10; i < 100; ++i) {
847      // Create a variable colocated with some existing variable, and
848      // an assignment colocated with a possibly-different variable.
849      Node* var = ops::SourceOp(
850          "TestVariable",
851          b.opts()
852              .WithName(strings::StrCat("var_", i))
853              .WithAttr("_class", {strings::StrCat("loc:@var_", i % 6)}));
854      ops::BinaryOp(
855          "TestAssign", var, input,
856          b.opts()
857              .WithName(strings::StrCat("assign_", i))
858              .WithAttr("_class", {strings::StrCat("loc:@assign_", i % 3)}));
859    }
860    TF_EXPECT_OK(BuildGraph(b, &g));
861  }
862
863  TF_EXPECT_OK(Place(&g));
864  for (int i = 0; i < 10; ++i) {
865    EXPECT_COLOCATED(g, strings::StrCat("var_", i),
866                     strings::StrCat("assign_", i));
867  }
868  for (int i = 10; i < 100; ++i) {
869    EXPECT_COLOCATED(g, strings::StrCat("var_", i),
870                     strings::StrCat("assign_", i));
871    EXPECT_COLOCATED(g, strings::StrCat("var_", i),
872                     strings::StrCat("var_", i % 6));
873    EXPECT_COLOCATED(g, strings::StrCat("assign_", i),
874                     strings::StrCat("assign_", i % 3));
875  }
876}
877
878// Test that placement fails when no devices are registered.
879TEST_F(PlacerTest, TestEmptyDeviceSet) {
880  Graph g(OpRegistry::Global());
881  {  // Scope for temporary variables used to construct g.
882    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
883    ops::SourceOp("TestInput", b.opts().WithName("in"));
884    TF_EXPECT_OK(BuildGraph(b, &g));
885  }
886
887  DeviceSet empty;
888
889  Status s = Place(&g, &empty);
890  EXPECT_TRUE(
891      StringPiece(s.error_message()).contains("No devices are registered"));
892}
893
894// Test that placement fails when the requested device forces an
895// indirect constraint to be violated.
896TEST_F(PlacerTest, TestHeterogeneousDeviceSetFailure) {
897  Graph g(OpRegistry::Global());
898  {  // Scope for temporary variables used to construct g.
899    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
900    Node* in = ops::SourceOp("TestInput", b.opts().WithName("in"));
901    Node* var = ops::SourceOp("VariableGPU", b.opts().WithName("var"));
902    ops::BinaryOp("TestAssign", var, in,
903                  b.opts().WithName("assign").WithDevice("/job:b/task:1"));
904    TF_EXPECT_OK(BuildGraph(b, &g));
905  }
906
907  DeviceSet heterogeneous;
908  std::unique_ptr<Device> gpu(
909      FakeDevice::MakeGPU("/job:b/replica:0/task:0/device:fakegpu:0"));
910  heterogeneous.AddDevice(gpu.get());
911  std::unique_ptr<Device> cpu(
912      FakeDevice::MakeCPU("/job:b/replica:0/task:1/device:fakecpu:0"));
913  heterogeneous.AddDevice(cpu.get());
914  Status s = Place(&g, &heterogeneous);
915  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
916  EXPECT_TRUE(StringPiece(s.error_message())
917                  .contains("colocated with a group of nodes that required "
918                            "incompatible device"));
919
920  // The error message should contain information that indicates which
921  // op types have which registered device types.
922  EXPECT_TRUE(StringPiece(s.error_message()).contains("VariableGPU: FakeGPU"))
923      << s;
924  EXPECT_TRUE(
925      StringPiece(s.error_message()).contains("TestAssign: FakeGPU FakeCPU"))
926      << s;
927}
928
929// Test that placement fails when an unknown device is requested.
930TEST_F(PlacerTest, TestUnknownDevice) {
931  Graph g(OpRegistry::Global());
932  {  // Scope for temporary variables used to construct g.
933    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
934    ops::SourceOp("TestInput", b.opts().WithName("in").WithDevice("/job:foo"));
935    TF_EXPECT_OK(BuildGraph(b, &g));
936  }
937
938  Status s = Place(&g);
939  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
940  EXPECT_TRUE(StringPiece(s.error_message()).contains("/job:foo"));
941}
942
943// Test that placement fails when the combination of partial
944// constraints leads to an unknown device.
945TEST_F(PlacerTest, TestUnknownMergedDevice) {
946  Graph g(OpRegistry::Global());
947  {  // Scope for temporary variables used to construct g.
948    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
949    ops::SourceOp("TestInput", b.opts().WithName("in").WithDevice("/job:foo"));
950    TF_EXPECT_OK(BuildGraph(b, &g));
951  }
952
953  Status s = Place(&g);
954  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
955  EXPECT_TRUE(StringPiece(s.error_message()).contains("/job:foo"));
956}
957
958// Test that placement fails when the previously-assigned device for a
959// node is unknown.
960TEST_F(PlacerTest, TestUnknownAssignedDevice) {
961  Graph g(OpRegistry::Global());
962  {  // Scope for temporary variables used to construct g.
963    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
964    ops::SourceOp("TestInput", b.opts().WithName("in"));
965    TF_EXPECT_OK(BuildGraph(b, &g));
966  }
967
968  GetNodeByName(g, "in")->set_assigned_device_name("/job:foo");
969
970  Status s = Place(&g);
971  EXPECT_EQ(error::INTERNAL, s.code());
972  EXPECT_TRUE(
973      StringPiece(s.error_message())
974          .contains("Assigned device '/job:foo' does not match any device"));
975}
976
977// Test that placement fails when an op with no registered kernels is
978// requested.
979TEST_F(PlacerTest, TestNoKernelsRegistered) {
980  Graph g(OpRegistry::Global());
981  {  // Scope for temporary variables used to construct g.
982    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
983    ops::SourceOp("VariableNoKernels", b.opts().WithName("var"));
984    TF_EXPECT_OK(BuildGraph(b, &g));
985  }
986
987  Status s = Place(&g);
988  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
989  EXPECT_TRUE(
990      StringPiece(s.error_message())
991          .contains(
992              "No OpKernel was registered to support Op 'VariableNoKernels'"));
993  EXPECT_TRUE(
994      StringPiece(s.error_message()).contains("<no registered kernels>"));
995}
996
997// Test that placement fails when a kernel is registered but no known
998// device supports it.
999TEST_F(PlacerTest, TestNoDevicesRegistered) {
1000  Graph g(OpRegistry::Global());
1001  {  // Scope for temporary variables used to construct g.
1002    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1003    ops::SourceOp("VariableGPU", b.opts().WithName("var"));
1004    TF_EXPECT_OK(BuildGraph(b, &g));
1005  }
1006
1007  DeviceSet cpu_only;
1008  std::unique_ptr<Device> cpu(
1009      FakeDevice::MakeCPU("/job:a/replica:0/task:0/device:fakecpu:0"));
1010  cpu_only.AddDevice(cpu.get());
1011
1012  Status s = Place(&g, &cpu_only);
1013  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1014  EXPECT_TRUE(StringPiece(s.error_message())
1015                  .contains("No OpKernel was registered to support "
1016                            "Op 'VariableGPU'"));
1017  EXPECT_TRUE(StringPiece(s.error_message()).contains("device='FakeGPU'"));
1018}
1019
1020// Test that placement fails when a requested device is malformed.
1021TEST_F(PlacerTest, TestMalformedDeviceSpecification) {
1022  Graph g(OpRegistry::Global());
1023  {  // Scope for temporary variables used to construct g.
1024    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1025    ops::SourceOp("TestInput", b.opts().WithName("in").WithDevice("/foo:bar"));
1026    TF_EXPECT_OK(BuildGraph(b, &g));
1027  }
1028
1029  Status s = Place(&g);
1030  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1031  EXPECT_TRUE(StringPiece(s.error_message())
1032                  .contains("Malformed device specification '/foo:bar'"));
1033}
1034
1035// Test that placement fails when a previously-assigned device is malformed.
1036TEST_F(PlacerTest, TestMalformedAssignedDevice) {
1037  Graph g(OpRegistry::Global());
1038  {  // Scope for temporary variables used to construct g.
1039    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1040    ops::SourceOp("TestInput", b.opts().WithName("in"));
1041    TF_EXPECT_OK(BuildGraph(b, &g));
1042  }
1043
1044  GetNodeByName(g, "in")->set_assigned_device_name("/foo:bar");
1045
1046  Status s = Place(&g);
1047  EXPECT_EQ(error::INTERNAL, s.code());
1048  EXPECT_TRUE(StringPiece(s.error_message())
1049                  .contains("Malformed assigned device '/foo:bar'"));
1050}
1051
1052// Test that placement fails when a device was previously assigned to
1053// a node, but it does not uniquely identify a particular device.
1054TEST_F(PlacerTest, TestNonUniqueAssignedDevice) {
1055  Graph g(OpRegistry::Global());
1056  {  // Scope for temporary variables used to construct g.
1057    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1058    ops::SourceOp("TestInput", b.opts().WithName("in"));
1059    TF_EXPECT_OK(BuildGraph(b, &g));
1060  }
1061
1062  GetNodeByName(g, "in")->set_assigned_device_name("/job:a");
1063
1064  Status s = Place(&g);
1065  EXPECT_EQ(error::INTERNAL, s.code());
1066  EXPECT_TRUE(
1067      StringPiece(s.error_message())
1068          .contains("Assigned device '/job:a' does not match any device"));
1069}
1070
1071// Test that ops request to be placed on non-existent devices will be relocated
1072// to existing device of the same type if allow_soft_placement is set.
1073TEST_F(PlacerTest, TestNonexistentGpuAllowSoftPlacement) {
1074  Graph g(OpRegistry::Global());
1075  {  // Scope for temporary variables used to construct g.
1076    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1077    ops::SourceOp("TestDevice",
1078                  b.opts().WithName("in").WithDevice("/device:fakegpu:11"));
1079    TF_EXPECT_OK(BuildGraph(b, &g));
1080  }
1081
1082  SessionOptions options;
1083  options.config.set_allow_soft_placement(true);
1084  TF_EXPECT_OK(Place(&g, &options));
1085  EXPECT_DEVICE_CONTAINS(g, "in", "/device:fakegpu:0");
1086}
1087
1088// Test that ops request to be placed on non-existent devices will fail if
1089// allow_soft_placement is not set.
1090TEST_F(PlacerTest, TestNonexistentGpuNoAllowSoftPlacement) {
1091  Graph g(OpRegistry::Global());
1092  {  // Scope for temporary variables used to construct g.
1093    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1094    ops::SourceOp("TestDevice",
1095                  b.opts().WithName("in").WithDevice("/device:fakegpu:11"));
1096    TF_EXPECT_OK(BuildGraph(b, &g));
1097  }
1098
1099  SessionOptions options;
1100  Status s = Place(&g, &options);
1101  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1102  EXPECT_TRUE(StringPiece(s.error_message()).contains("/device:fakegpu:11"));
1103}
1104
1105// Test that placement fails when a node requests an explicit device that is not
1106// supported by the registered kernels if allow_soft_placement is no set.
1107TEST_F(PlacerTest, TestUnsupportedDeviceNoAllowSoftPlacement) {
1108  Graph g(OpRegistry::Global());
1109  {  // Scope for temporary variables used to construct g.
1110    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1111    ops::SourceOp("VariableGPU",
1112                  b.opts().WithName("var").WithDevice("/device:fakecpu:0"));
1113    TF_EXPECT_OK(BuildGraph(b, &g));
1114  }
1115
1116  SessionOptions options;
1117  Status s = Place(&g, &options);
1118  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1119  EXPECT_TRUE(StringPiece(s.error_message()).contains("/device:fakecpu:0"));
1120  EXPECT_TRUE(
1121      StringPiece(s.error_message())
1122          .contains("no supported kernel for fakecpu devices is available"));
1123}
1124
1125// Test that placement fails when a node requests an explicit device that is not
1126// supported by the registered kernels if allow_soft_placement is no set.
1127TEST_F(PlacerTest, TestNonExistentDevice) {
1128  Graph g(OpRegistry::Global());
1129  {  // Scope for temporary variables used to construct g.
1130    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1131    ops::SourceOp("VariableGPU",
1132                  b.opts().WithName("var").WithDevice("/job:foo/replica:17"));
1133    TF_EXPECT_OK(BuildGraph(b, &g));
1134  }
1135
1136  SessionOptions options;
1137  Status s = Place(&g, &options);
1138  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1139  LOG(WARNING) << s.error_message();
1140  EXPECT_TRUE(StringPiece(s.error_message())
1141                  .contains("was explicitly assigned to /job:foo/replica:17 "
1142                            "but available devices"));
1143}
1144
1145TEST_F(PlacerTest, TestUnsupportedDeviceAllowSoftPlacement) {
1146  Graph g(OpRegistry::Global());
1147  {  // Scope for temporary variables used to construct g.
1148    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1149    ops::SourceOp("VariableGPU",
1150                  b.opts().WithName("var").WithDevice("/device:fakecpu:0"));
1151    TF_EXPECT_OK(BuildGraph(b, &g));
1152  }
1153
1154  SessionOptions options;
1155  options.config.set_allow_soft_placement(true);
1156  TF_EXPECT_OK(Place(&g, &options));
1157}
1158
1159// Test that a graph with device type and reference constraints on
1160// some of the ops will successfully assign nodes to the constrained
1161// device, and colocate nodes with reference connections.
1162TEST_F(PlacerTest, TestDeviceTypeConstraintsAllowSoftPlacement) {
1163  Graph g(OpRegistry::Global());
1164  {  // Scope for temporary variables used to construct g.
1165    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1166    // var_gpu has ref output and runs on GPU.
1167    // force_gpu takes var_gpu and requested CPU.
1168    // Verify that both are placed on GPU.
1169    Node* var_gpu = ops::SourceOp("VariableGPU", b.opts().WithName("var_gpu"));
1170    ops::UnaryOp(
1171        "TestDeviceEnforce", var_gpu,
1172        b.opts().WithName("force_gpu").WithDevice("/device:fakecpu:0"));
1173    // var_cpu has ref output and runs on CPU.
1174    // force_cpu takes var_cpu and requested GPU.
1175    // Verify that both are placed on CPU.
1176    Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu"));
1177    ops::UnaryOp(
1178        "TestDeviceEnforce", var_cpu,
1179        b.opts().WithName("force_cpu").WithDevice("/device:fakegpu:0"));
1180    TF_EXPECT_OK(BuildGraph(b, &g));
1181  }
1182
1183  SessionOptions options;
1184  options.config.set_allow_soft_placement(true);
1185  TF_EXPECT_OK(Place(&g, &options));
1186  EXPECT_DEVICE_TYPE(g, "var_gpu", "FakeGPU");
1187  EXPECT_DEVICE_TYPE(g, "force_gpu", "FakeGPU");
1188  EXPECT_COLOCATED(g, "var_gpu", "force_gpu");
1189  EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU");
1190  EXPECT_DEVICE_TYPE(g, "force_cpu", "FakeCPU");
1191  EXPECT_COLOCATED(g, "var_cpu", "force_cpu");
1192}
1193
1194// Test that placement fails when two nodes have a reference connection
1195// constraint, and each node requires a mutually incompatible device.
1196TEST_F(PlacerTest, TestUnsatisfiableConstraintWithReferenceConnections) {
1197  Graph g(OpRegistry::Global());
1198  {  // Scope for temporary variables used to construct g.
1199    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1200    Node* var = ops::SourceOp("VariableGPU", b.opts().WithName("var"));
1201    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
1202    ops::BinaryOp("AssignCPU", var, input, b.opts().WithName("assign"));
1203    TF_EXPECT_OK(BuildGraph(b, &g));
1204  }
1205
1206  Status s = Place(&g);
1207  EXPECT_EQ(error::INVALID_ARGUMENT, s.code());
1208  EXPECT_TRUE(StringPiece(s.error_message())
1209                  .contains("Cannot colocate nodes 'var' and 'assign'"));
1210}
1211
1212// Test that a generator node follows its consumers (where there are several
1213// consumer nodes on the same devices).
1214TEST_F(PlacerTest, TestGeneratorNodeFollowsConsumerNode) {
1215  Graph g(OpRegistry::Global());
1216  {  // Scope for temporary variables used to construct g.
1217    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1218
1219    // A variable is only on CPU
1220    Node* var1_cpu =
1221        ops::SourceOp("VariableCPU", b.opts().WithName("var1_cpu"));
1222    Node* var2_cpu =
1223        ops::SourceOp("VariableCPU", b.opts().WithName("var2_cpu"));
1224
1225    // The constant to be assigned can be on both GPU or CPU.
1226    //
1227    // Because of the heuristic, it gets placed on CPU to avoid a
1228    // copy.
1229    Node* input = ops::SourceOp("TestCPUGPUOutput", b.opts().WithName("in"));
1230
1231    // The assigns are bound to CPU by the reference edge.
1232    ops::BinaryOp("TestAssign", var1_cpu, input, b.opts().WithName("assign1"));
1233    ops::BinaryOp("TestAssign", var2_cpu, input, b.opts().WithName("assign2"));
1234
1235    TF_EXPECT_OK(BuildGraph(b, &g));
1236  }
1237
1238  TF_EXPECT_OK(Place(&g));
1239  EXPECT_COLOCATED(g, "var1_cpu", "in");
1240  EXPECT_COLOCATED(g, "assign1", "in");
1241  EXPECT_COLOCATED(g, "var2_cpu", "in");
1242  EXPECT_COLOCATED(g, "assign2", "in");
1243}
1244
1245// Test that a generator node does not follow its consumers (where there are
1246// several consumers on different devices).
1247TEST_F(PlacerTest, TestGeneratorNodeDoesntFollowNonColocatedConsumers) {
1248  Graph g(OpRegistry::Global());
1249  {  // Scope for temporary variables used to construct g.
1250    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
1251
1252    // A variable is only on CPU
1253    Node* var1_cpu =
1254        ops::SourceOp("VariableCPU", b.opts().WithName("var1_cpu"));
1255    Node* var2_cpu =
1256        ops::SourceOp("VariableCPU", b.opts().WithName("var2_cpu"));
1257
1258    // The constant to be assigned can be on both GPU or CPU.
1259    //
1260    // Because of the heuristic, it ought to be on the GPU (cannot be
1261    // co-located with both consumers, so goes to the 'standard' place)
1262    Node* input = ops::SourceOp("TestCPUGPUOutput", b.opts().WithName("in"));
1263
1264    // The assigns are bound to CPU by the reference edge.
1265    ops::BinaryOp("TestAssign", var1_cpu, input, b.opts().WithName("assign1"));
1266    ops::BinaryOp("TestAssign", var2_cpu, input, b.opts().WithName("assign2"));
1267
1268    TF_EXPECT_OK(BuildGraph(b, &g));
1269
1270    GetNodeByName(g, "var1_cpu")
1271        ->set_assigned_device_name("/job:a/replica:0/task:0/device:fakecpu:1");
1272
1273    GetNodeByName(g, "var2_cpu")
1274        ->set_assigned_device_name("/job:a/replica:0/task:0/device:fakecpu:2");
1275  }
1276
1277  TF_EXPECT_OK(Place(&g));
1278  EXPECT_COLOCATED(g, "assign1", "var1_cpu");
1279  EXPECT_COLOCATED(g, "assign2", "var2_cpu");
1280  EXPECT_DEVICE_TYPE(g, "in", "FakeGPU");
1281}
1282
1283}  // namespace
1284}  // namespace tensorflow
1285