Function TEST

tensorflow/compiler/jit/cluster_scoping_pass_test.cc:85–127

}

TEST(XlaCompilationTest, StagePipelinePreserved) {
  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  {
    // Graph:
    //       b
    //       |
    //       v
    // a -> add0 (ClusterX) -> relu0 (ClusterX) -> stage
    //
    //             b
    //             |
    //             v
    // unstage -> add1 (ClusterY) -> relu1 (ClusterY)
    GraphDefBuilder builder(GraphDefBuilder::kFailImmediately);
    Node* a = ops::SourceOp("Const", builder.opts()
                                         .WithName("a")
                                         .WithAttr("dtype", DT_FLOAT)
                                         .WithAttr("value", Tensor()));
    Node* b = ops::SourceOp("Const", builder.opts()
                                         .WithName("b")
                                         .WithAttr("dtype", DT_FLOAT)
                                         .WithAttr("value", Tensor()));
    Node* unstage = ops::SourceOp(
        "Unstage",
        builder.opts().WithName("unstage").WithAttr("dtypes", {DT_FLOAT}));

    Node* add0 = ops::BinaryOp("Add", a, b, builder.opts().WithName("add0"));
    Node* add1 =
        ops::BinaryOp("Add", unstage, b, builder.opts().WithName("add1"));
    Node* relu0 = ops::UnaryOp("Relu", add0, builder.opts().WithName("relu0"));
    ops::UnaryOp("Relu", add1, builder.opts().WithName("relu1"));
    BuildStageNode(builder, "stage", {DT_FLOAT}, {relu0});

    TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
  }

  TF_ASSERT_OK(ClusterScoping(&graph));

  auto scopes = GetXlaInternalScopes(*graph);
  EXPECT_NE(scopes["add0"], scopes["add1"]);
  EXPECT_EQ(scopes["add0"], scopes["relu0"]);
  EXPECT_EQ(scopes["add1"], scopes["relu1"]);
}

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