Function BM_DynamicSlice

tensorflow/compiler/xla/tests/dynamic_ops_test.cc:734–806

}

void BM_DynamicSlice(int num_iters) {
  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();
  auto* transfer_manager =
      TransferManager::GetForPlatform(platform).ValueOrDie();
  int device_ordinal = client->default_device_ordinal();

  XlaBuilder builder("DynamicSlice");

  // Create input as a constant: shape [1, 2, 3, 4]
  auto input_literal = LiteralUtil::CreateR4(
      {{{{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}},
        {{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}});
  auto input = ConstantLiteral(&builder, input_literal);

  auto stream =
      client->mutable_backend()->BorrowStream(device_ordinal).ValueOrDie();

  // Create dynamic slice start indices as a parameter: shape [4]
  auto start_indices_shape = ShapeUtil::MakeShape(S32, {});
  std::vector<XlaOp> start_indices(4);
  std::vector<ScopedShapedBuffer> shaped_buffers;
  std::vector<const Shape*> host_shapes(4);
  for (int i = 0; i < 4; ++i) {
    start_indices[i] =
        Parameter(&builder, i, start_indices_shape, "start_indices");
    auto start_index_literal = LiteralUtil::CreateR0<int32>(i + 1);
    // Initialize and transfer parameter buffer.
    shaped_buffers.emplace_back(
        client->backend()
            .transfer_manager()
            ->AllocateScopedShapedBuffer(start_indices_shape, &allocator,
                                         /*device_ordinal=*/0)
            .ConsumeValueOrDie());
    host_shapes[i] = &shaped_buffers[i].on_host_shape();
    ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice(
        stream.get(), start_index_literal, shaped_buffers[i]));
  }

  // Add DynamicSlice op to the computation.
  DynamicSlice(input, start_indices, {1, 1, 1, 1});
  auto computation = builder.Build().ConsumeValueOrDie();

  TF_ASSERT_OK_AND_ASSIGN(
      auto executables,
      client->Compile(computation, host_shapes, ExecutableBuildOptions()));
  auto executable = std::move(executables[0]);

  // Run some warm-up executions.
  ExecutableRunOptions options;
  options.set_allocator(&allocator);
  const int kWarmups = 2;
  std::vector<const ShapedBuffer*> shaped_buffer_ptrs;