From 96abe88eb21c8396a2f25bab0700c1fe91a0a22d Mon Sep 17 00:00:00 2001 From: Sandeep Dasgupta Date: Sun, 20 Oct 2024 17:05:33 -0700 Subject: [PATCH] [HLO Componentization] Create hlo/analysis sub-component (Phase II). This CL takes care of 1. Migrating external projects dependencies from xla/service --> xla/hlo/analysis Phase I takes care of 1. Migrating xla/service/ --> xla/hlo/analysis/ 2. Setting up build aliases in xla/service ensuring external dependencies are still satisfied. PiperOrigin-RevId: 687936257 --- xla/hlo/experimental/auto_sharding/BUILD | 6 +++--- xla/hlo/experimental/auto_sharding/auto_sharding.cc | 2 +- xla/hlo/experimental/auto_sharding/auto_sharding.h | 2 +- xla/hlo/experimental/auto_sharding/auto_sharding_option.h | 2 +- xla/hlo/experimental/auto_sharding/auto_sharding_test.cc | 2 +- xla/hlo/experimental/auto_sharding/auto_sharding_util.cc | 2 +- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/xla/hlo/experimental/auto_sharding/BUILD b/xla/hlo/experimental/auto_sharding/BUILD index 5b559373ee327..1340031a1cdc7 100644 --- a/xla/hlo/experimental/auto_sharding/BUILD +++ b/xla/hlo/experimental/auto_sharding/BUILD @@ -49,6 +49,7 @@ cc_library( "//xla:status_macros", "//xla:util", "//xla:xla_data_proto_cc", + "//xla/hlo/analysis:hlo_alias_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/hlo/transforms:hlo_constant_splitter", @@ -62,7 +63,6 @@ cc_library( "//xla/service:computation_layout", "//xla/service:dot_as_convolution_util", "//xla/service:dump", - "//xla/service:hlo_alias_analysis", "//xla/service:hlo_buffer", "//xla/service:hlo_cost_analysis", "//xla/service:hlo_value", @@ -295,6 +295,7 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", + "//xla/hlo/analysis:while_loop_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/ir:ptrvec", "//xla/hlo/ir:tile_assignment", @@ -302,7 +303,6 @@ cc_library( "//xla/service:call_graph", "//xla/service:computation_layout", "//xla/service:sharding_propagation", - "//xla/service:while_loop_analysis", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", @@ -382,6 +382,7 @@ xla_cc_test( ":auto_sharding_option", ":auto_sharding_strategy", ":auto_sharding_util", + "//xla/hlo/analysis:hlo_alias_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", "//xla/hlo/testlib:verified_hlo_module", @@ -389,7 +390,6 @@ xla_cc_test( "//xla/hlo/utils:hlo_live_range", "//xla/hlo/utils:hlo_matchers", "//xla/service:buffer_value", - "//xla/service:hlo_alias_analysis", "//xla/service:hlo_value", "//xla/tests:hlo_test_base", "//xla/tests:xla_internal_test_main", diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding.cc b/xla/hlo/experimental/auto_sharding/auto_sharding.cc index 169acb6c8ec2d..21f49fc8498e6 100644 --- a/xla/hlo/experimental/auto_sharding/auto_sharding.cc +++ b/xla/hlo/experimental/auto_sharding/auto_sharding.cc @@ -48,6 +48,7 @@ limitations under the License. #include "absl/time/clock.h" #include "absl/time/time.h" #include "absl/types/span.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_cost_graph.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_device_mesh.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_memory.h" @@ -78,7 +79,6 @@ limitations under the License. #include "xla/service/call_graph.h" #include "xla/service/computation_layout.h" #include "xla/service/dump.h" -#include "xla/service/hlo_alias_analysis.h" #include "xla/service/hlo_buffer.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/service/hlo_value.h" diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding.h b/xla/hlo/experimental/auto_sharding/auto_sharding.h index 639c6143077aa..40cbc6295b0bc 100644 --- a/xla/hlo/experimental/auto_sharding/auto_sharding.h +++ b/xla/hlo/experimental/auto_sharding/auto_sharding.h @@ -32,6 +32,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_cost_graph.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_device_mesh.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_option.h" @@ -44,7 +45,6 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/hlo/utils/hlo_live_range.h" #include "xla/service/call_graph.h" -#include "xla/service/hlo_alias_analysis.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/shape.h" diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding_option.h b/xla/hlo/experimental/auto_sharding/auto_sharding_option.h index 02ca89b334531..33b6f4385ee20 100644 --- a/xla/hlo/experimental/auto_sharding/auto_sharding_option.h +++ b/xla/hlo/experimental/auto_sharding/auto_sharding_option.h @@ -165,7 +165,7 @@ struct AutoShardingOption { // Static estimate for iteration count of a while loop, used in the cost // model. This estimate is used when we cannot infer an upper bound on the // number of iterations in the loop (as implemented in - // third_party/tensorflow/compiler/xla/service/while_loop_analysis.h) + // third_party/tensorflow/compiler/xla/hlo/analysis/while_loop_analysis.h) int64_t loop_iteration_count_estimate = 100; // Allows the conversion of aliases to followers if their pairwise strategy diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding_test.cc b/xla/hlo/experimental/auto_sharding/auto_sharding_test.cc index 93e718dd62fd7..7cb4d8cc11e15 100644 --- a/xla/hlo/experimental/auto_sharding/auto_sharding_test.cc +++ b/xla/hlo/experimental/auto_sharding/auto_sharding_test.cc @@ -31,6 +31,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_cost_graph.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_device_mesh.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_option.h" @@ -48,7 +49,6 @@ limitations under the License. #include "xla/hlo/utils/hlo_live_range.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/service/buffer_value.h" -#include "xla/service/hlo_alias_analysis.h" #include "xla/service/hlo_value.h" #include "xla/tests/hlo_test_base.h" #include "xla/tsl/lib/core/status_test_util.h" diff --git a/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc b/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc index 6971a134973a9..ccdbbe4c0ff13 100644 --- a/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc +++ b/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc @@ -41,6 +41,7 @@ limitations under the License. #include "absl/types/span.h" #include "json/json.h" #include "xla/array.h" +#include "xla/hlo/analysis/while_loop_analysis.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_device_mesh.h" #include "xla/hlo/experimental/auto_sharding/auto_sharding_strategy.h" #include "xla/hlo/ir/hlo_computation.h" @@ -55,7 +56,6 @@ limitations under the License. #include "xla/service/call_graph.h" #include "xla/service/computation_layout.h" #include "xla/service/sharding_propagation.h" -#include "xla/service/while_loop_analysis.h" #include "xla/shape.h" #include "xla/shape_layout.h" #include "xla/shape_tree.h"