From 3955f7b76e83ba7ab8a667520b658a69f02557b6 Mon Sep 17 00:00:00 2001
From: Niklas Vangerow <nikv@google.com>
Date: Mon, 3 Feb 2025 11:27:53 -0800
Subject: [PATCH] Fix includes in `HloRunnerInterface` implementations.

PiperOrigin-RevId: 722732996
---
 xla/service/BUILD                   | 15 ++++++++++++++-
 xla/service/hlo_runner.cc           | 26 ++++++++++++++++++++++++++
 xla/service/hlo_runner_interface.cc |  1 -
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/xla/service/BUILD b/xla/service/BUILD
index 8c9432b61d578b..6820e011903e06 100644
--- a/xla/service/BUILD
+++ b/xla/service/BUILD
@@ -4649,7 +4649,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
-        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -4664,7 +4663,12 @@ cc_library(
         ":executable",
         ":hlo_module_util",
         ":hlo_runner_interface",
+        ":maybe_owning_device_memory",
+        ":shaped_buffer",
         ":transfer_manager",
+        "//xla:executable_run_options",
+        "//xla:literal",
+        "//xla:shape_tree",
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:types",
@@ -4674,10 +4678,19 @@ cc_library(
         "//xla/hlo/ir:hlo_module_group",
         "//xla/service/gpu:gpu_executable_run_options",
         "//xla/stream_executor:device_description",
+        "//xla/stream_executor:device_memory",
         "//xla/stream_executor:device_memory_allocator",
+        "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor:stream_executor_memory_allocator",
+        "//xla/tsl/platform:env",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:logging",
+        "//xla/tsl/platform:status",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
     ],
diff --git a/xla/service/hlo_runner.cc b/xla/service/hlo_runner.cc
index 0f87c13a12caf0..7edf48951642d1 100644
--- a/xla/service/hlo_runner.cc
+++ b/xla/service/hlo_runner.cc
@@ -15,24 +15,50 @@ limitations under the License.
 
 #include "xla/service/hlo_runner.h"
 
+#include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 #include <utility>
 #include <variant>
 #include <vector>
 
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
 #include "unsupported/Eigen/CXX11/Tensor"
+#include "xla/executable_run_options.h"
+#include "xla/hlo/ir/hlo_input_output_alias_config.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_module_group.h"
+#include "xla/literal.h"
+#include "xla/service/backend.h"
+#include "xla/service/computation_placer.h"
 #include "xla/service/executable.h"
 #include "xla/service/gpu/gpu_executable_run_options.h"
 #include "xla/service/hlo_module_util.h"
 #include "xla/service/hlo_runner_interface.h"
+#include "xla/service/maybe_owning_device_memory.h"
+#include "xla/service/service_executable_run_options.h"
+#include "xla/service/shaped_buffer.h"
 #include "xla/service/transfer_manager.h"
 #include "xla/shape.h"
+#include "xla/shape_tree.h"
 #include "xla/shape_util.h"
+#include "xla/status_macros.h"
 #include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/device_memory_allocator.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_memory_allocator.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/logging.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
 
 namespace xla {
 
diff --git a/xla/service/hlo_runner_interface.cc b/xla/service/hlo_runner_interface.cc
index 510ccdba6e4453..82b24855161eaa 100644
--- a/xla/service/hlo_runner_interface.cc
+++ b/xla/service/hlo_runner_interface.cc
@@ -31,7 +31,6 @@ limitations under the License.
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
-#include "tsl/platform/statusor.h"
 
 namespace xla {