simple_placement_test.cc

// Test that a graph with no constraints will successfully assign nodes to the
// "best available" device (i.e. prefer GPU over CPU).
TEST_F(SimplePlacerTest, TestNoConstraints) {
  Graph g(OpRegistry::Global());
  {  // Scope for temporary variables used to construct g.
    GraphDefBuilder b(GraphDefBuilder::kFailImmediately);
    Node* input = ops::SourceOp("TestInput", b.opts().WithName("in"));
    ops::UnaryOp("TestRelu", ops::NodeOut(input, 0), b.opts().WithName("n1"));
    ops::UnaryOp("TestRelu", ops::NodeOut(input, 1), b.opts().WithName("n2"));
    EXPECT_OK(BuildGraph(b, &g));
  }

  EXPECT_OK(Place(&g));
  EXPECT_DEVICE_TYPE(g, "in", DEVICE_CPU);
  EXPECT_DEVICE_TYPE(g, "n1", DEVICE_GPU);
  EXPECT_DEVICE_TYPE(g, "n2", DEVICE_GPU);
}