Function BM_WhileLoop

tensorflow/compiler/xla/tests/while_test.cc:1262–1337

}

void BM_WhileLoop(int num_iters) {
  // Benchmark a simple kernel to measure while loop overheads.
  tensorflow::testing::StopTiming();

  se::Platform* platform = PlatformUtil::GetDefaultPlatform().ValueOrDie();
  auto executors = PlatformUtil::GetStreamExecutors(platform).ValueOrDie();
  se::StreamExecutorMemoryAllocator allocator(platform, executors);
  LocalClient* client =
      ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie();

  const int64 seq_len = 100;
  Shape loop_state_shape = ShapeUtil::MakeTupleShape(
      {ShapeUtil::MakeShape(S32, {}),
       ShapeUtil::MakeShape(F32, {seq_len, 1024, 1024})});

  // Create while condition computation with 'loop_limit'.
  const int32 loop_limit = 100;
  XlaComputation condition;
  {
    XlaBuilder builder("condition");
    auto prev = Parameter(&builder, 0, loop_state_shape, "prev");
    auto iteration = GetTupleElement(prev, 0);
    Lt(iteration, ConstantR0<int32>(&builder, loop_limit));
    condition = builder.Build().ConsumeValueOrDie();
  }

  // Create while body computation with unit loop increment.
  XlaComputation body;
  {
    XlaBuilder builder("body");
    auto prev = Parameter(&builder, 0, loop_state_shape, "prev");
    // TupleElement 0
    auto iteration = GetTupleElement(prev, 0);
    auto out0 = Add(iteration, ConstantR0<int32>(&builder, 1));
    // TupleElement 1
    auto input = GetTupleElement(prev, 1);
    // Update.
    auto one = ConstantR0<float>(&builder, 1.0);
    auto update = Broadcast(one, {1, 1024, 1024});
    // Starts = iteration * 2;
    auto zero = ConstantR0<int32>(&builder, 0);
    // UpdateSlice.
    auto out1 = DynamicUpdateSlice(input, update, {zero, zero, zero});
    Tuple(&builder, {out0, out1});
    body = builder.Build().ConsumeValueOrDie();
  }

  // Create a While instruction.
  XlaBuilder builder("while");
  auto zero = ConstantR0<float>(&builder, 0.0);
  auto input = Broadcast(zero, {seq_len, 1024, 1024});
  auto init = Tuple(&builder, {ConstantR0<int32>(&builder, 0), input});
  While(condition, body, init);
  auto computation = builder.Build().ConsumeValueOrDie();

  TF_ASSERT_OK_AND_ASSIGN(
      auto executables,
      client->Compile(computation, {}, ExecutableBuildOptions()));