Function TEST_F

tensorflow/core/grappler/optimizers/loop_optimizer_test.cc:81–146

};

TEST_F(LoopOptimizerTest, Basic) {
  GraphDef graph;
  AddSimpleNode("In", "Identity", {}, &graph);
  AddEnterNode("InvariantEnter", "while/while_context", true, 1, {"In"},
               &graph);
  AddSimpleNode("InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"},
                &graph);
  AddSimpleNode("VariantAdd", "Add", {"InvariantAdd", "Identity"}, &graph);
  AddEnterNode("VariantEnter", "while/while_context", false, 1, {"In"}, &graph);
  AddSimpleNode("Merge", "Merge", {"VariantEnter", "NextIteration"}, &graph);
  AddSimpleNode("Less/y", "Const", {"^Identity"}, &graph);
  AddSimpleNode("Less", "Less", {"VariantAdd", "Less/y"}, &graph);
  AddSimpleNode("LoopCond", "LoopCond", {"Less"}, &graph);
  AddSimpleNode("Switch", "Switch", {"Merge", "LoopCond"}, &graph);
  AddSimpleNode("Identity", "Identity", {"Switch:1"}, &graph);
  AddSimpleNode("NextIteration", "NextIteration", {"VariantAdd"}, &graph);
  AddSimpleNode("Exit", "Exit", {"Switch"}, &graph);
  AddSimpleNode("Out", "Identity", {"Exit"}, &graph);

  GrapplerItem item;
  item.graph = graph;

  LoopOptimizer optimizer;
  EnableOnlyLoopInvariantNodeMotion(&optimizer);
  GraphDef output;
  TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));

  {  // Original graph.
    Status status;
    utils::GraphView view(&graph, &status);
    TF_ASSERT_OK(status);
    FrameView frames;
    TF_EXPECT_OK(frames.InferFromGraphView(view));

    EXPECT_EQ(frames.num_frames(), 1);
    const auto* invariant_add_node = view.GetNode("InvariantAdd");
    ASSERT_NE(invariant_add_node, nullptr);
    const auto* invariant_add_node_def = invariant_add_node->node();
    ASSERT_EQ(frames.Frames(*invariant_add_node_def).size(), 1);
    EXPECT_EQ(frames.Frames(*invariant_add_node_def).back(), 0);
    const auto* variant_add_node = view.GetNode("VariantAdd");
    ASSERT_NE(variant_add_node, nullptr);
    const auto* variant_add_node_def = variant_add_node->node();
    ASSERT_EQ(frames.Frames(*variant_add_node_def).size(), 1);
    EXPECT_EQ(frames.Frames(*variant_add_node_def).back(), 0);
  }

  {  // Optimized graph.
    Status status;
    utils::GraphView view(&output, &status);
    TF_ASSERT_OK(status);
    FrameView frames;
    TF_EXPECT_OK(frames.InferFromGraphView(view));

    EXPECT_EQ(frames.num_frames(), 1);
    const auto* invariant_add_node = view.GetNode("InvariantAdd");
    ASSERT_NE(invariant_add_node, nullptr);
    const auto* invariant_add_node_def = invariant_add_node->node();