Function TEST_F

tensorflow/core/graph/graph_partition_test.cc:338–377

};

TEST_F(FuseRecvTest, FuseRecvNormal) {
  string worker_device = "/job:worker/replica:0/task:0/cpu:0";
  string ps_device = "/job:ps/replica:0/task:0/cpu:0";

  auto w1 = FloatInput(in_.WithOpName("W1"), worker_device);
  auto w2 = Scatter1to2(in_.WithOpName("W2"), worker_device, w1);
  auto w3 = Scatter1to2(in_.WithOpName("W3"), worker_device, w1);
  ASSERT_EQ(w2.size(), 2);
  ASSERT_EQ(w3.size(), 2);
  auto p1 = FakeIdentity(in_.WithOpName("P1"), ps_device, w2[0]);
  auto p2 = FakeIdentity(in_.WithOpName("P2"), ps_device, w2[1]);
  auto p3 = FakeIdentity(in_.WithOpName("P3"), ps_device, w3[0]);
  auto w4 = FakeIdentity(in_.WithOpName("W4"), worker_device, w3[1]);

  std::shared_ptr<Graph> g = ConstructGraph();

  PartitionOptions popts;
  popts.node_to_loc = SplitByWorker;
  popts.new_name = [&g](const string& prefix) { return g->NewName(prefix); };
  popts.get_incarnation = [](const string& name) {
    return (name[0] - 'A') + 100;
  };
  std::unordered_map<string, GraphDef> partitions;
  PartitionWithTensorFuse(popts, g.get(), &partitions);

  ASSERT_EQ(partitions.size(), 2);
  for (auto p : partitions) {
    int fuse_recv_count = 0;
    for (int i = 0; i < p.second.node_size(); ++i) {
      if (p.second.node(i).op() == "_FuseRecv") {
        ++fuse_recv_count;
      }
    }
    if (p.first.find("ps") != std::string::npos) {
      ASSERT_EQ(fuse_recv_count, 1);
    } else {
      ASSERT_EQ(fuse_recv_count, 0);
    }
  }
}

TEST_F(FuseRecvTest, FuseRecvNormal2) {
  string worker_device = "/job:worker/replica:0/task:0/cpu:0";