diff options
4 files changed, 98 insertions, 27 deletions
diff --git a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc index fbac512785..3c0989c587 100644 --- a/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc +++ b/tensorflow/core/common_runtime/direct_session_with_tracking_alloc_test.cc @@ -102,7 +102,7 @@ TEST(DirectSessionWithTrackingAllocTest, CostModelTest) { if (node->name() == y->name()) { EXPECT_EQ(5, cm->AllocationId(node, 0)); } else { - EXPECT_EQ(7, cm->AllocationId(node, 0)); + EXPECT_EQ(6, cm->AllocationId(node, 0)); } } EXPECT_LE(0, cm->MaxExecutionTime(node)); diff --git a/tensorflow/core/common_runtime/simple_placer.cc b/tensorflow/core/common_runtime/simple_placer.cc index 99c74ce038..d3110cba04 100644 --- a/tensorflow/core/common_runtime/simple_placer.cc +++ b/tensorflow/core/common_runtime/simple_placer.cc @@ -716,8 +716,14 @@ Status SimplePlacer::Run() { if (!node->IsOp()) { continue; } - // Skip nodes that already have an assigned name. + + // The graph may have come pre-populated by the framework with assigned + // devices (e.g., for stateful placements), so the placer should not try to + // place nodes that are already placed. if (!node->assigned_device_name().empty()) { + // Although the device is already assigned, we run this function to + // possibly log pre-assigned placements. + AssignAndLog(node->assigned_device_name(), node); continue; } @@ -799,17 +805,17 @@ Status SimplePlacer::Run() { return Status::OK(); } -bool SimplePlacer::CanAssignToDevice(const string& candidate_device_name, - const std::vector<Device*> devices) const { +bool SimplePlacer::CanAssignToDevice( + const string& candidate_device_name, + const std::vector<Device*>& devices) const { if (!candidate_device_name.empty()) { - // Can we assign to the same device? Check by validating that - // the device type of 'candidate_device_name' is present - // in 'devices'. + // 'devices' lists the set of devices that the placer or the user has + // constrained the operation to. "candidate_device_name" must + // refer to a concrete Device that is in the list of 'devices'. const Device* other_device = devices_->FindDeviceByName(candidate_device_name); - if (std::any_of(devices.begin(), devices.end(), [other_device](Device* d) { - return d->device_type() == other_device->device_type(); - })) { + if (std::find(devices.begin(), devices.end(), other_device) != + devices.end()) { return true; } } @@ -822,10 +828,10 @@ void SimplePlacer::AssignAndLog(const string& assigned_device, node->set_assigned_device_name(assigned_device); // Log placement if log_device_placement is set. if (options_ && options_->config.log_device_placement()) { - printf("%s: (%s): %s\n", node->name().c_str(), - node->type_string().c_str(), + printf("%s: (%s): %s\n", node->name().c_str(), node->type_string().c_str(), node->assigned_device_name().c_str()); - LOG(INFO) << node->name() << ": " << "(" << node->type_string() << ")" + LOG(INFO) << node->name() << ": " + << "(" << node->type_string() << ")" << node->assigned_device_name(); } } diff --git a/tensorflow/core/common_runtime/simple_placer.h b/tensorflow/core/common_runtime/simple_placer.h index 59be859449..a041e96830 100644 --- a/tensorflow/core/common_runtime/simple_placer.h +++ b/tensorflow/core/common_runtime/simple_placer.h @@ -82,7 +82,7 @@ class SimplePlacer { // Returns true if the device type of 'candidate_device_name' is // found in 'devices'. bool CanAssignToDevice(const string& candidate_device_name, - const std::vector<Device*> devices) const; + const std::vector<Device*>& devices) const; // Assigns 'node's devices to 'assigned_device', and logs the // placement if the SessionOptions entry in 'options_' requests it. diff --git a/tensorflow/core/common_runtime/simple_placer_test.cc b/tensorflow/core/common_runtime/simple_placer_test.cc index bb13281cdb..06267d71ae 100644 --- a/tensorflow/core/common_runtime/simple_placer_test.cc +++ b/tensorflow/core/common_runtime/simple_placer_test.cc @@ -150,6 +150,9 @@ REGISTER_OP("TestCPUGPUOutput").Output("a: float"); REGISTER_KERNEL_BUILDER(Name("TestCPUGPUOutput").Device("FakeCPU"), DummyOp); REGISTER_KERNEL_BUILDER(Name("TestCPUGPUOutput").Device("FakeGPU"), DummyOp); +REGISTER_OP("TestGPUOutput").Output("a: float"); +REGISTER_KERNEL_BUILDER(Name("TestGPUOutput").Device("FakeGPU"), DummyOp); + REGISTER_OP("TestDevice").Output("a: float").Output("b: float"); REGISTER_KERNEL_BUILDER(Name("TestDevice").Device("FakeGPU"), DummyOp); @@ -251,11 +254,12 @@ class SimplePlacerTest : public ::testing::Test { GetNodeByName(g_, (name_b))->assigned_device_name()); \ } while (0) -#define EXPECT_DEVICE_TYPE(g, name, expected_device_type) \ - EXPECT_EQ(DeviceType(expected_device_type).type(), \ - devices_.FindDeviceByName( \ - GetNodeByName((g), (name))->assigned_device_name()) \ - ->attributes() \ +#define EXPECT_DEVICE_TYPE(g, name, expected_device_type) \ + EXPECT_EQ(DeviceType(expected_device_type).type(), \ + devices_ \ + .FindDeviceByName( \ + GetNodeByName((g), (name))->assigned_device_name()) \ + ->attributes() \ .device_type()) #define EXPECT_DEVICE_CONTAINS(g, name, device_substr) \ @@ -327,7 +331,7 @@ TEST_F(SimplePlacerTest, TestMetadataColocatedWithInput) { // Heuristic A implements "Island fusing": if a node only generates // an output and it has only one consumer, we place the node // with its consumer. -TEST_F(SimplePlacerTest, TestHeuristicA) { +TEST_F(SimplePlacerTest, TestHeuristicGeneratorFollowsSingleConsumer) { Graph g(OpRegistry::Global()); { // Scope for temporary variables used to construct g. GraphDefBuilder b(GraphDefBuilder::kFailImmediately); @@ -352,6 +356,65 @@ TEST_F(SimplePlacerTest, TestHeuristicA) { EXPECT_COLOCATED(g, "assign", "in"); } +TEST_F(SimplePlacerTest, TestIgnoreGeneratorHeuristicIfWrongDevice) { + Graph g(OpRegistry::Global()); + { // Scope for temporary variables used to construct g. + GraphDefBuilder b(GraphDefBuilder::kFailImmediately); + + // A variable is only on CPU + Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu")); + + // The constant to be assigned can only be on GPU. + // + // The heuristic to place the generator with its consumer does + // not apply since the consumer's device is not in the list + // of valid devices for the generator. + Node* input = ops::SourceOp("TestGPUOutput", b.opts().WithName("in")); + + // The assign is bound to CPU by the reference edge. + ops::BinaryOp("TestAssign", var_cpu, input, b.opts().WithName("assign")); + + TF_EXPECT_OK(BuildGraph(b, &g)); + } + + TF_EXPECT_OK(Place(&g)); + EXPECT_DEVICE_TYPE(g, "in", "FakeGPU"); + EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU"); + EXPECT_COLOCATED(g, "var_cpu", "assign"); +} + +TEST_F(SimplePlacerTest, TestIgnoreGeneratorHeuristicIfWrongPartialDevice) { + Graph g(OpRegistry::Global()); + { // Scope for temporary variables used to construct g. + GraphDefBuilder b(GraphDefBuilder::kFailImmediately); + + // A variable is only on CPU + Node* var_cpu = ops::SourceOp("VariableCPU", b.opts().WithName("var_cpu")); + + // The constant to be assigned can be on CPU or GPU, but is explicitly + // placed on CPU:1. + // + // The heuristic to place the generator with its consumer does + // not apply since the consumer's device is not in the list + // of valid devices for the generator. + Node* input = + ops::SourceOp("TestCPUGPUOutput", + b.opts().WithName("in").WithDevice("/device:fakecpu:1")); + + // The assign is bound to CPU by the reference edge. + ops::BinaryOp("TestAssign", var_cpu, input, b.opts().WithName("assign")); + + TF_EXPECT_OK(BuildGraph(b, &g)); + } + + TF_EXPECT_OK(Place(&g)); + EXPECT_DEVICE_TYPE(g, "in", "FakeCPU"); + EXPECT_DEVICE_CONTAINS(g, "in", "/device:fakecpu:1"); + EXPECT_DEVICE_TYPE(g, "var_cpu", "FakeCPU"); + EXPECT_COLOCATED(g, "var_cpu", "assign"); + EXPECT_DEVICE_CONTAINS(g, "var_cpu", "/device:fakecpu:0"); +} + // Test that a graph with partial device specifications on the ops // will successfully TEST_F(SimplePlacerTest, TestPartialSpec) { @@ -371,7 +434,7 @@ TEST_F(SimplePlacerTest, TestPartialSpec) { EXPECT_DEVICE_CONTAINS(g, "var", "/job:a"); } -// Test that a node with an assigned device is not relocated. +// Test that a node with a pre-assigned device is not relocated. TEST_F(SimplePlacerTest, TestAssignedDevicePreserved) { Graph g(OpRegistry::Global()); { // Scope for temporary variables used to construct g. @@ -666,9 +729,10 @@ TEST_F(SimplePlacerTest, TestMultipleColocationGroups) { Node* colocated_with_input = ops::UnaryOp( "TestRelu", input, b.opts().WithName("colocated_1").WithAttr("_class", {"loc:@in"})); - Node* colocated_with_input_and_other = ops::UnaryOp( - "TestRelu", input, b.opts().WithName("foo").WithAttr( - "_class", {"loc:@in", "loc:@colocated_1"})); + Node* colocated_with_input_and_other = + ops::UnaryOp("TestRelu", input, + b.opts().WithName("foo").WithAttr( + "_class", {"loc:@in", "loc:@colocated_1"})); CHECK(colocated_with_input); CHECK(colocated_with_input_and_other); TF_EXPECT_OK(BuildGraph(b, &g)); @@ -687,9 +751,10 @@ TEST_F(SimplePlacerTest, TestInvalidMultipleColocationGroups) { Node* colocated_with_input = ops::UnaryOp( "ReluCPU", input, b.opts().WithName("colocated_1").WithAttr("_class", {"loc:@in"})); - Node* colocated_with_input_and_other = ops::UnaryOp( - "ReluGPU", input, b.opts().WithName("foo").WithAttr( - "_class", {"loc:@in", "loc:@colocated_1"})); + Node* colocated_with_input_and_other = + ops::UnaryOp("ReluGPU", input, + b.opts().WithName("foo").WithAttr( + "_class", {"loc:@in", "loc:@colocated_1"})); CHECK(colocated_with_input); CHECK(colocated_with_input_and_other); TF_EXPECT_OK(BuildGraph(b, &g)); |