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