Add explicit includes to fix Kokoro compile issues.

PiperOrigin-RevId: 677916620
openxla · Sep 28, 2024 · e911ef4 · e911ef4
1 parent 683725f
commit e911ef4
Show file tree

Hide file tree

Showing 14 changed files with 158 additions and 34 deletions.
diff --git a/xla/hlo/utils/BUILD b/xla/hlo/utils/BUILD
@@ -46,16 +46,17 @@ xla_cc_test(
     srcs = ["hlo_live_range_test.cc"],
     deps = [
         ":hlo_live_range",
-        "//xla:literal",
-        "//xla:status_macros",
+        "//xla:comparison_util",
+        "//xla:literal_util",
+        "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/service:hlo_alias_analysis",
-        "//xla/service:hlo_ordering",
         "//xla/service:hlo_value",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:statusor",
     ],
 )

diff --git a/xla/hlo/utils/hlo_live_range_test.cc b/xla/hlo/utils/hlo_live_range_test.cc
@@ -14,19 +14,26 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/hlo/utils/hlo_live_range.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include <gtest/gtest.h>
 #include "absl/container/flat_hash_map.h"
+#include "xla/comparison_util.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_schedule.h"
+#include "xla/literal_util.h"
 #include "xla/service/hlo_alias_analysis.h"
 #include "xla/service/hlo_value.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "tsl/platform/statusor.h"

diff --git a/xla/service/BUILD b/xla/service/BUILD
@@ -630,8 +630,11 @@ xla_cc_test(
     srcs = ["collective_pipeliner_test.cc"],
     deps = [
         ":collective_pipeliner",
+        ":hlo_module_config",
         ":hlo_parser",
+        ":hlo_verifier",
         ":host_memory_offload_annotations_hdr",
+        "//xla:test_helpers",
         "//xla:util",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass_pipeline",
@@ -1062,12 +1065,16 @@ xla_cc_test(
     name = "hlo_dfs_reachability_test",
     srcs = ["hlo_dfs_reachability_test.cc"],
     deps = [
+        ":computation_placer_hdr",
+        ":hlo_module_config",
+        "//xla:literal_util",
+        "//xla:shape_util",
         "//xla:test",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_dfs_reachability",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
-        "@com_google_absl//absl/random",
+        "@tsl//tsl/platform:status",
         "@tsl//tsl/platform:test_benchmark",
     ],
 )
@@ -1077,13 +1084,17 @@ xla_cc_test(
     srcs = ["hlo_reachability_test.cc"],
     deps = [
         ":computation_placer",
+        ":hlo_module_config",
+        "//xla:literal_util",
+        "//xla:shape_util",
         "//xla:test",
         "//xla:test_helpers",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_reachability",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "@com_google_absl//absl/random",
+        "@tsl//tsl/platform:status",
         "@tsl//tsl/platform:test_benchmark",
     ],
 )
@@ -2023,24 +2034,36 @@ xla_cc_test(
         ":cpu_plugin",
         ":flatten_call_graph",
         ":hlo_alias_analysis",
+        ":hlo_buffer",
         ":hlo_dce",
         ":hlo_memory_scheduler",
         ":hlo_ordering",
         ":hlo_parser",
         ":hlo_proto_cc",
         ":hlo_proto_util",
+        ":hlo_value",
+        ":logical_buffer",
+        "//xla:comparison_util",
         "//xla:literal",
+        "//xla:literal_util",
         "//xla:shape_util",
         "//xla:test",
         "//xla:test_helpers",
         "//xla:types",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/service/memory_space_assignment",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@tsl//tsl/platform:status",
         "@tsl//tsl/platform:statusor",
     ],
 )
@@ -3888,10 +3911,13 @@ xla_cc_test(
     deps = [
         ":defuser",
         "//xla:literal",
+        "//xla:literal_util",
         "//xla:shape_util",
+        "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_matchers",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest_main",
     ],
 )
 
@@ -5030,10 +5056,13 @@ xla_cc_test(
     deps = [
         ":flatten_call_graph",
         ":hlo_alias_analysis",
+        ":hlo_buffer",
         ":hlo_graph_dumper",
         ":hlo_ordering",
+        ":hlo_value",
         ":instruction_fusion",
         "//xla:literal",
+        "//xla:literal_util",
         "//xla:shape_util",
         "//xla:test",
         "//xla:test_helpers",
@@ -5043,7 +5072,10 @@ xla_cc_test(
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings:string_view",
         "@tsl//tsl/platform:logging",
+        "@tsl//tsl/platform:statusor",
         "@tsl//tsl/platform:test",
     ],
 )
@@ -5487,6 +5519,7 @@ cc_library(
     testonly = 1,
     hdrs = ["hlo_rematerialization_test_utils.h"],
     deps = [
+        "//xla:literal_util",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
@@ -5502,6 +5535,7 @@ xla_cc_test(
         ":hlo_rematerialization_test_utils",
         "//xla/hlo/ir:hlo",
         "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest_main",
     ],
 )
 

diff --git a/xla/service/buffer_assignment_test.cc b/xla/service/buffer_assignment_test.cc
@@ -15,40 +15,49 @@ limitations under the License.
 
 #include "xla/service/buffer_assignment.h"
 
+#include <cstdint>
 #include <memory>
 #include <optional>
-#include <set>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/comparison_util.h"
 #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_schedule.h"
 #include "xla/literal.h"
+#include "xla/literal_util.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/call_graph.h"
 #include "xla/service/copy_insertion.h"
 #include "xla/service/flatten_call_graph.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_alias_analysis.h"
-#include "xla/service/hlo_dce.h"
+#include "xla/service/hlo_buffer.h"
 #include "xla/service/hlo_memory_scheduler.h"
 #include "xla/service/hlo_ordering.h"
 #include "xla/service/hlo_parser.h"
-#include "xla/service/hlo_proto_util.h"
+#include "xla/service/hlo_value.h"
+#include "xla/service/logical_buffer.h"
+#include "xla/service/memory_space_assignment/memory_space_assignment.h"
 #include "xla/shape_util.h"
 #include "xla/test.h"
 #include "xla/test_helpers.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/lib/core/status_test_util.h"
-#include "xla/types.h"
 #include "xla/xla_data.pb.h"
+#include "tsl/platform/status.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {

diff --git a/xla/service/collective_pipeliner_test.cc b/xla/service/collective_pipeliner_test.cc
@@ -38,8 +38,11 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/hlo/utils/hlo_matchers.h"
+#include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_parser.h"
+#include "xla/service/hlo_verifier.h"
 #include "xla/service/host_memory_offload_annotations.h"
+#include "xla/test_helpers.h"
 #include "xla/tests/filecheck.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/util.h"

diff --git a/xla/service/defuser_test.cc b/xla/service/defuser_test.cc
@@ -15,8 +15,16 @@ limitations under the License.
 
 #include "xla/service/defuser.h"
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
 

diff --git a/xla/service/hlo_alias_analysis_test.cc b/xla/service/hlo_alias_analysis_test.cc
@@ -15,22 +15,30 @@ limitations under the License.
 
 #include "xla/service/hlo_alias_analysis.h"
 
-#include <map>
 #include <memory>
-
+#include <set>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/hlo/utils/hlo_matchers.h"
-#include "xla/literal.h"
+#include "xla/literal_util.h"
 #include "xla/service/flatten_call_graph.h"
+#include "xla/service/hlo_buffer.h"
 #include "xla/service/hlo_ordering.h"
-#include "xla/service/instruction_fusion.h"
+#include "xla/service/hlo_value.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/test.h"
 #include "xla/test_helpers.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 
 namespace xla {

diff --git a/xla/service/hlo_dfs_reachability_test.cc b/xla/service/hlo_dfs_reachability_test.cc
@@ -16,14 +16,21 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_dfs_reachability.h"
 
 #include <cstddef>
-#include <set>
+#include <memory>
+#include <string>
 #include <string_view>
 
-#include "absl/random/random.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/literal_util.h"
+#include "xla/service/computation_placer.h"
+#include "xla/service/hlo_module_config.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
 #include "xla/test.h"
 #include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/status.h"
 #include "tsl/platform/test_benchmark.h"
 
 namespace xla {

diff --git a/xla/service/hlo_reachability_test.cc b/xla/service/hlo_reachability_test.cc
@@ -17,14 +17,21 @@ limitations under the License.
 
 #include <memory>
 #include <set>
+#include <string>
 #include <string_view>
 
 #include "absl/random/random.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/literal_util.h"
 #include "xla/service/computation_placer.h"
+#include "xla/service/hlo_module_config.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
 #include "xla/test.h"
 #include "xla/test_helpers.h"
 #include "xla/tests/hlo_test_base.h"
+#include "tsl/platform/status.h"
 #include "tsl/platform/test_benchmark.h"
 
 namespace xla {

diff --git a/xla/service/hlo_rematerialization_test_utils.h b/xla/service/hlo_rematerialization_test_utils.h
@@ -18,12 +18,15 @@ limitations under the License.
 #ifndef XLA_SERVICE_HLO_REMATERIALIZATION_TEST_UTILS_H_
 #define XLA_SERVICE_HLO_REMATERIALIZATION_TEST_UTILS_H_
 
+#include <cstdint>
 #include <memory>
 #include <string>
 
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/literal_util.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/xla_data.pb.h"

diff --git a/xla/service/hlo_rematerialization_test_utils_test.cc b/xla/service/hlo_rematerialization_test_utils_test.cc
@@ -17,6 +17,8 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"