Skip to content

Commit

Permalink
[HLO Componentization] Create hlo/transforms/ sub-component (Phase II).
Browse files Browse the repository at this point in the history 
This CL takes care of
1. Migrating external projects dependencies from  xla/service --> xla/hlo/transforms

Phase I takes care of
1. Migrating xla/service/<passes> --> xla/hlo/transforms/<passes>
2. Setting up build aliases in xla/service ensuring external dependencies are still satisfied.

PiperOrigin-RevId: 687508139
  • Loading branch information
@sdasgup3 @Google-ML-Automation
sdasgup3 authored and Google-ML-Automation committed Oct 19, 2024
1 parent 63eefd0 commit 33ff534
Show file tree
Hide file tree
Showing 33 changed files with 92 additions and 98 deletions.
98 changes: 47 additions & 51 deletions xla/service/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -362,8 +362,8 @@ cc_library(
"//xla/hlo/ir:hlo_instruction_utils",
"//xla/hlo/parser:hlo_parser",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/utils:hlo_query",
"//xla/service:hlo_dce",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -587,8 +587,8 @@ xla_cc_test(
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
"//xla/hlo/transforms:hlo_constant_splitter",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/utils:hlo_matchers",
"//xla/service:hlo_dce",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"@com_google_absl//absl/log",
Expand Down Expand Up @@ -898,13 +898,13 @@ cc_library(
hdrs = ["call_inliner.h"],
deps = [
":call_graph",
":hlo_dce",
":hlo_domain_isolator",
"//xla:status_macros",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:hlo_dce",
"//xla/service/spmd/shardy:constants",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -1163,7 +1163,7 @@ xla_cc_test(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/service:async_collective_creator",
"//xla/hlo/transforms:async_collective_creator",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"@com_google_absl//absl/algorithm:container",
Expand Down Expand Up @@ -1695,9 +1695,9 @@ xla_cc_test(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
"//xla/service:flatten_call_graph",
"//xla/service:hlo_dce",
"//xla/service:hlo_memory_scheduler",
"//xla/hlo/transforms:flatten_call_graph",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:hlo_memory_scheduler",
"//xla/service/memory_space_assignment",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
Expand Down Expand Up @@ -1808,8 +1808,8 @@ xla_cc_test(
"//xla:types",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/service:hlo_dce",
"//xla/service:hlo_memory_scheduler",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:hlo_memory_scheduler",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
Expand All @@ -1829,8 +1829,6 @@ xla_cc_test(
"//xla:test_helpers",
"//xla:types",
"//xla/hlo/ir:hlo",
"//xla/service:hlo_dce",
"//xla/service:hlo_memory_scheduler",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
Expand Down Expand Up @@ -1917,7 +1915,7 @@ cc_library(
"//xla/hlo/ir:hlo",
"//xla/hlo/ir:hlo_reachability",
"//xla/hlo/pass:hlo_pass",
"//xla/service:hlo_dce",
"//xla/hlo/transforms:hlo_dce",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status:statusor",
Expand Down Expand Up @@ -2043,7 +2041,7 @@ cc_library(
"//xla:literal_util",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
Expand Down Expand Up @@ -2094,7 +2092,7 @@ cc_library(
":while_util",
"//xla:literal_util",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/status:statusor",
],
Expand All @@ -2106,7 +2104,6 @@ cc_library(
hdrs = ["scatter_determinism_expander.h"],
deps = [
":hlo_creation_utils",
":op_expander_pass",
":scatter_utils",
"//xla:array",
"//xla:array2d",
Expand All @@ -2116,6 +2113,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings:str_format",
"@tsl//tsl/platform:logging",
Expand Down Expand Up @@ -2165,7 +2163,6 @@ cc_library(
deps = [
":hlo_creation_utils",
":hlo_module_config",
":op_expander_pass",
"//xla:shape_util",
"//xla:util",
"//xla/hlo/builder:xla_builder",
Expand All @@ -2175,6 +2172,7 @@ cc_library(
"//xla/hlo/builder/lib:matrix",
"//xla/hlo/builder/lib:slicing",
"//xla/hlo/ir:hlo",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
Expand Down Expand Up @@ -2606,8 +2604,8 @@ cc_library(
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:inlined_vector",
Expand Down Expand Up @@ -2715,7 +2713,7 @@ cc_library(
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:tuple_simplifier",
"//xla/hlo/transforms:tuple_simplifier",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status:statusor",
Expand Down Expand Up @@ -2758,7 +2756,7 @@ cc_library(
"//xla/hlo/evaluator:hlo_evaluator",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:algebraic_simplifier",
"//xla/hlo/transforms:algebraic_simplifier",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -2811,9 +2809,9 @@ cc_library(
"//xla/hlo/evaluator:hlo_evaluator",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:flatten_call_graph",
"//xla/hlo/transforms:tuple_simplifier",
"//xla/hlo/utils:hlo_query",
"//xla/service:flatten_call_graph",
"//xla/service:tuple_simplifier",
"@com_google_absl//absl/algorithm",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -2873,8 +2871,8 @@ cc_library(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/utils:hlo_query",
"//xla/service:hlo_dce",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand All @@ -2900,9 +2898,9 @@ xla_cc_test(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"//xla/hlo/utils:hlo_matchers",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
Expand Down Expand Up @@ -2962,7 +2960,7 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
],
Expand Down Expand Up @@ -3118,7 +3116,7 @@ cc_library(
"//xla/hlo/builder:xla_builder",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:hlo_dce",
"//xla/hlo/transforms:hlo_dce",
"//xla/tsl/lib/monitoring:gauge",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -3156,11 +3154,11 @@ xla_test(
"//xla/hlo/builder:xla_builder",
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
"//xla/hlo/transforms:algebraic_simplifier",
"//xla/hlo/transforms:dynamic_dimension_simplifier",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"//xla/hlo/utils:hlo_matchers",
"//xla/service:algebraic_simplifier",
"//xla/service:dynamic_dimension_simplifier",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/tests:client_library_test_base",
"//xla/tests:hlo_test_base",
"//xla/tests:literal_test_util",
Expand Down Expand Up @@ -3458,8 +3456,8 @@ xla_cc_test(
"//xla:xla_data_proto_cc",
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/transforms:hlo_memory_scheduler",
"//xla/hlo/utils:hlo_matchers",
"//xla/service:hlo_memory_scheduler",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
Expand Down Expand Up @@ -3550,8 +3548,6 @@ cc_library(
],
)

# Deprecated, use
# //third_party/tensorflow/compiler/xla/hlo/analysis:hlo_dataflow_analysis instead.
cc_library(
name = "hlo_dataflow_analysis",
hdrs = ["hlo_dataflow_analysis.h"],
Expand Down Expand Up @@ -3688,8 +3684,8 @@ cc_library(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -3722,11 +3718,9 @@ cc_library(
":hlo_alias_analysis",
":hlo_buffer",
":hlo_dataflow_analysis",
":hlo_dce",
":hlo_graph_dumper",
":hlo_ordering",
":hlo_value",
":tuple_simplifier",
"//xla:frontend_attributes",
"//xla:shape_tree",
"//xla:shape_util",
Expand All @@ -3735,6 +3729,8 @@ cc_library(
"//xla/hlo/ir:hlo",
"//xla/hlo/ir:hlo_reachability",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"//xla/hlo/utils:hlo_query",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
Expand Down Expand Up @@ -3852,8 +3848,8 @@ cc_library(
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:errors",
Expand Down Expand Up @@ -4003,7 +3999,7 @@ xla_cc_test(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
"//xla/service:algebraic_simplifier",
"//xla/hlo/transforms:algebraic_simplifier",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
Expand Down Expand Up @@ -4849,8 +4845,8 @@ cc_library(
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:hlo_dce",
"//xla/service:tuple_simplifier",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -4893,7 +4889,7 @@ cc_library(
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/service:hlo_dce",
"//xla/hlo/transforms:hlo_dce",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
Expand Down Expand Up @@ -5405,12 +5401,12 @@ xla_cc_test(
name = "topk_rewriter_test",
srcs = ["topk_rewriter_test.cc"],
deps = [
":hlo_dce",
":pattern_matcher",
":pattern_matcher_gmock",
":topk_rewriter",
":tuple_simplifier",
"//xla/hlo/ir:hlo",
"//xla/hlo/transforms:hlo_dce",
"//xla/hlo/transforms:tuple_simplifier",
"//xla/hlo/utils:hlo_matchers",
"//xla/tests:hlo_test_base",
"//xla/tests:literal_test_util",
Expand Down Expand Up @@ -5709,7 +5705,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/types:span",
Expand Down Expand Up @@ -5739,7 +5735,7 @@ cc_library(
":call_inliner",
"//xla:literal_util",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
],
)

Expand Down Expand Up @@ -5812,7 +5808,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/service:op_expander_pass",
"//xla/hlo/transforms:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
Expand Down Expand Up @@ -6204,11 +6200,11 @@ cc_library(
hdrs = ["while_loop_pipeline_unroller.h"],
deps = [
":call_inliner",
":flatten_call_graph",
":hlo_dce",
":while_util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/transforms:flatten_call_graph",
"//xla/hlo/transforms:hlo_dce",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log:check",
Expand Down
2 changes: 1 addition & 1 deletion xla/service/all_to_all_decomposer.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/service/op_expander_pass.h"
#include "xla/hlo/transforms/expanders/op_expander_pass.h"

namespace xla {

Expand Down
2 changes: 1 addition & 1 deletion xla/service/batched_gather_scatter_normalizer.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/service/op_expander_pass.h"
#include "xla/hlo/transforms/expanders/op_expander_pass.h"

namespace xla {

Expand Down
Loading

0 comments on commit 33ff534

Please sign in to comment.