From 675237fd0af29df7ebffd9dd2a2f721cd542475b Mon Sep 17 00:00:00 2001 From: Eugene Zhulenev Date: Wed, 12 Jun 2024 18:48:44 -0700 Subject: [PATCH] [xla] NFC: Fix hlo_memory_scheduler warnings PiperOrigin-RevId: 642810364 --- third_party/xla/xla/service/BUILD | 18 ++++++++++++++++++ .../xla/xla/service/hlo_memory_scheduler.cc | 16 +++++++++++----- .../xla/xla/service/hlo_memory_scheduler.h | 7 ++++++- .../xla/service/hlo_memory_scheduler_test.cc | 15 +++++++++++++-- 4 files changed, 48 insertions(+), 8 deletions(-) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 7eaf85827fe5e3..f752786dd374f0 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -1985,6 +1985,7 @@ cc_library( hdrs = ["hlo_memory_scheduler.h"], local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]), deps = [ + ":buffer_value", ":hlo_alias_analysis", ":hlo_pass", ":logical_buffer", @@ -1998,9 +1999,15 @@ cc_library( "//xla/service/heap_simulator", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@local_tsl//tsl/lib/gtl:map_util", "@local_tsl//tsl/platform:errors", "@local_tsl//tsl/platform:logging", + "@local_tsl//tsl/platform:numbers", + "@local_tsl//tsl/platform:statusor", "@local_tsl//tsl/profiler/lib:scoped_annotation", ], ) @@ -2009,9 +2016,14 @@ xla_cc_test( name = "hlo_memory_scheduler_test", srcs = ["hlo_memory_scheduler_test.cc"], deps = [ + ":buffer_value", + ":hlo_alias_analysis", ":hlo_dce", ":hlo_memory_scheduler", ":hlo_ordering", + ":hlo_value", + ":logical_buffer", + "//xla:literal_util", "//xla:shape_util", "//xla:types", "//xla:xla_data_proto_cc", @@ -2021,7 +2033,13 @@ xla_cc_test( "//xla/tests:xla_internal_test_main", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/status", + "@com_google_absl//absl/types:span", "@local_tsl//tsl/lib/core:status_test_util", + "@local_tsl//tsl/platform:errors", + "@local_tsl//tsl/platform:logging", + "@local_tsl//tsl/platform:statusor", + "@local_tsl//tsl/platform:test", ], ) diff --git a/third_party/xla/xla/service/hlo_memory_scheduler.cc b/third_party/xla/xla/service/hlo_memory_scheduler.cc index 912ac546fb51d7..163030c35e3444 100644 --- a/third_party/xla/xla/service/hlo_memory_scheduler.cc +++ b/third_party/xla/xla/service/hlo_memory_scheduler.cc @@ -16,28 +16,34 @@ limitations under the License. #include "xla/service/hlo_memory_scheduler.h" #include +#include +#include #include #include -#include +#include #include #include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" +#include "xla/service/buffer_value.h" #include "xla/service/heap_simulator/heap_simulator.h" +#include "xla/service/hlo_alias_analysis.h" +#include "xla/service/logical_buffer.h" #include "xla/service/tuple_points_to_analysis.h" #include "xla/shape_util.h" -#include "xla/status_macros.h" -#include "xla/statusor.h" -#include "xla/types.h" #include "xla/util.h" -#include "tsl/lib/gtl/map_util.h" #include "tsl/platform/errors.h" #include "tsl/platform/logging.h" +#include "tsl/platform/numbers.h" +#include "tsl/platform/statusor.h" #include "tsl/profiler/lib/scoped_annotation.h" namespace xla { diff --git a/third_party/xla/xla/service/hlo_memory_scheduler.h b/third_party/xla/xla/service/hlo_memory_scheduler.h index 24ca085a9a40b9..1b1281fd5eb62a 100644 --- a/third_party/xla/xla/service/hlo_memory_scheduler.h +++ b/third_party/xla/xla/service/hlo_memory_scheduler.h @@ -16,7 +16,13 @@ limitations under the License. #ifndef XLA_SERVICE_HLO_MEMORY_SCHEDULER_H_ #define XLA_SERVICE_HLO_MEMORY_SCHEDULER_H_ +#include +#include + #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_schedule.h" @@ -24,7 +30,6 @@ limitations under the License. #include "xla/service/hlo_pass_interface.h" #include "xla/service/logical_buffer.h" #include "xla/service/tuple_points_to_analysis.h" -#include "xla/statusor.h" namespace xla { diff --git a/third_party/xla/xla/service/hlo_memory_scheduler_test.cc b/third_party/xla/xla/service/hlo_memory_scheduler_test.cc index 17a10ae1ea4d44..fb2794d08ed185 100644 --- a/third_party/xla/xla/service/hlo_memory_scheduler_test.cc +++ b/third_party/xla/xla/service/hlo_memory_scheduler_test.cc @@ -15,23 +15,34 @@ limitations under the License. #include "xla/service/hlo_memory_scheduler.h" +#include #include #include +#include #include "absl/algorithm/container.h" #include "absl/container/flat_hash_map.h" +#include "absl/types/span.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/buffer_value.h" #include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_dce.h" +#include "xla/service/hlo_alias_analysis.h" #include "xla/service/hlo_ordering.h" +#include "xla/service/hlo_value.h" +#include "xla/service/logical_buffer.h" +#include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tests/hlo_test_base.h" -#include "xla/types.h" #include "xla/xla_data.pb.h" #include "tsl/lib/core/status_test_util.h" +#include "tsl/platform/logging.h" +#include "tsl/platform/statusor.h" +#include "tsl/platform/test.h" namespace xla { namespace {