Function TEST

tensorflow/compiler/jit/partially_decluster_pass_test.cc:127–161

}

TEST(PartiallyDeclusterPassTest, ClusteredAndUnclustered) {
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  {
    GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
    Node* input =
        ops::SourceOp("FakeNullary", builder.opts().WithName("Input"));
    Node* clustered_producer =
        ops::BinaryOp("FakeBinary", input, input,
                      builder.opts().WithName("ClusteredProducer"));
    ops::BinaryOp("FakeBinary", clustered_producer, input,
                  builder.opts().WithName("UnclusteredConsumer"));
    Node* clustered_consumer =
        ops::BinaryOp("FakeBinary", {clustered_producer, 1}, input,
                      builder.opts().WithName("ClusteredConsumer"));
    clustered_producer->AddAttr(kXlaClusterAttr, "cluster_0");
    clustered_consumer->AddAttr(kXlaClusterAttr, "cluster_0");
    TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
  }

  TF_ASSERT_OK(PartiallyDecluster(&graph));
  std::vector<Node*> unclustered_consumer_inputs;
  ASSERT_TRUE(GetInputsForNode(*graph, "UnclusteredConsumer",
                               &unclustered_consumer_inputs));
  ASSERT_EQ(unclustered_consumer_inputs.size(), 2);
  EXPECT_EQ(unclustered_consumer_inputs[0]->name(),
            "ClusteredProducer/declustered");
  EXPECT_EQ(unclustered_consumer_inputs[1]->name(), "Input");

  std::vector<Node*> clustered_consumer_inputs;
  ASSERT_TRUE(GetInputsForNode(*graph, "ClusteredConsumer",
                               &clustered_consumer_inputs));
  ASSERT_EQ(clustered_consumer_inputs.size(), 2);
  EXPECT_EQ(clustered_consumer_inputs[0]->name(), "ClusteredProducer");
  EXPECT_EQ(clustered_consumer_inputs[1]->name(), "Input");
}

TEST(PartiallyDeclusterPassTest, DifferentClusters) {
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));