diff --git a/xla/backends/interpreter/BUILD b/xla/backends/interpreter/BUILD
index 85af2b52aac00..030bf8d6ea589 100644
--- a/xla/backends/interpreter/BUILD
+++ b/xla/backends/interpreter/BUILD
@@ -39,10 +39,10 @@ cc_library(
"//xla/hlo/ir:hlo",
"//xla/hlo/ir:hlo_module_group",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:cholesky_expander",
- "//xla/hlo/transforms:dynamic_index_splitter",
- "//xla/hlo/transforms:eigh_expander",
- "//xla/hlo/transforms:qr_expander",
+ "//xla/hlo/transforms/expanders:dynamic_index_splitter",
+ "//xla/hlo/transforms/expanders:cholesky_expander",
+ "//xla/hlo/transforms/expanders:eigh_expander",
+ "//xla/hlo/transforms/expanders:qr_expander",
"//xla/service:batchnorm_expander",
"//xla/service:compiler",
"//xla/service:computation_placer",
diff --git a/xla/hlo/tools/hlo_opt/BUILD b/xla/hlo/tools/hlo_opt/BUILD
index 490019fa0e96a..e2a2abe738973 100644
--- a/xla/hlo/tools/hlo_opt/BUILD
+++ b/xla/hlo/tools/hlo_opt/BUILD
@@ -34,28 +34,28 @@ cc_library(
"//xla/hlo/analysis:indexed_array_analysis",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:cholesky_expander",
- "//xla/hlo/transforms:comparison_expander",
"//xla/hlo/transforms:convert_memory_placement_to_internal_annotations",
- "//xla/hlo/transforms:convolution_4d_expander",
- "//xla/hlo/transforms:convolution_pred_expander",
- "//xla/hlo/transforms:dot_decomposer",
- "//xla/hlo/transforms:dynamic_index_splitter",
- "//xla/hlo/transforms:eigh_expander",
- "//xla/hlo/transforms:logistic_expander",
"//xla/hlo/transforms:operand_upcaster",
- "//xla/hlo/transforms:optimization_barrier_expander",
- "//xla/hlo/transforms:qr_expander",
- "//xla/hlo/transforms:real_imag_expander",
- "//xla/hlo/transforms:reduce_decomposer",
- "//xla/hlo/transforms:reshape_decomposer",
- "//xla/hlo/transforms:rng_expander",
- "//xla/hlo/transforms:stable_sort_expander",
- "//xla/hlo/transforms:stochastic_convert_decomposer",
"//xla/hlo/transforms:while_loop_trip_count_annotator",
"//xla/hlo/transforms/collectives:all_gather_broadcast_reorder",
"//xla/hlo/transforms/collectives:all_reduce_contiguous",
"//xla/hlo/transforms/collectives:collective_quantizer",
+ "//xla/hlo/transforms/expanders:cholesky_expander",
+ "//xla/hlo/transforms/expanders:comparison_expander",
+ "//xla/hlo/transforms/expanders:convolution_4d_expander",
+ "//xla/hlo/transforms/expanders:convolution_pred_expander",
+ "//xla/hlo/transforms/expanders:dot_decomposer",
+ "//xla/hlo/transforms/expanders:dynamic_index_splitter",
+ "//xla/hlo/transforms/expanders:eigh_expander",
+ "//xla/hlo/transforms/expanders:logistic_expander",
+ "//xla/hlo/transforms/expanders:optimization_barrier_expander",
+ "//xla/hlo/transforms/expanders:qr_expander",
+ "//xla/hlo/transforms/expanders:real_imag_expander",
+ "//xla/hlo/transforms/expanders:reduce_decomposer",
+ "//xla/hlo/transforms/expanders:reshape_decomposer",
+ "//xla/hlo/transforms/expanders:rng_expander",
+ "//xla/hlo/transforms/expanders:stable_sort_expander",
+ "//xla/hlo/transforms/expanders:stochastic_convert_decomposer",
"//xla/hlo/transforms/simplifiers:all_reduce_folder",
"//xla/hlo/transforms/simplifiers:batch_dot_simplification",
"//xla/hlo/transforms/simplifiers:broadcast_canonicalizer",
diff --git a/xla/hlo/transforms/BUILD b/xla/hlo/transforms/BUILD
index 9530c89d5494e..7ae0a25510cc6 100644
--- a/xla/hlo/transforms/BUILD
+++ b/xla/hlo/transforms/BUILD
@@ -71,331 +71,6 @@ xla_cc_test(
],
)
-cc_library(
- name = "op_expander_pass",
- srcs = ["expanders/op_expander_pass.cc"],
- hdrs = ["expanders/op_expander_pass.h"],
- deps = [
- "//xla:util",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- "@tsl//tsl/platform:statusor",
- ],
-)
-
-cc_library(
- name = "optimization_barrier_expander",
- srcs = ["expanders/optimization_barrier_expander.cc"],
- hdrs = ["expanders/optimization_barrier_expander.h"],
- deps = [
- ":op_expander_pass",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
-cc_library(
- name = "comparison_expander",
- srcs = ["expanders/comparison_expander.cc"],
- hdrs = ["expanders/comparison_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:comparison_util",
- "//xla:literal_util",
- "//xla:shape_util",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/log",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_absl//absl/types:span",
- ],
-)
-
-cc_library(
- name = "cholesky_expander",
- srcs = ["expanders/cholesky_expander.cc"],
- hdrs = ["expanders/cholesky_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal",
- "//xla:shape_util",
- "//xla:status_macros",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder/lib:arithmetic",
- "//xla/hlo/builder/lib:constants",
- "//xla/hlo/builder/lib:loops",
- "//xla/hlo/builder/lib:math",
- "//xla/hlo/builder/lib:matrix",
- "//xla/hlo/builder/lib:slicing",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:str_format",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_absl//absl/types:span",
- "@tsl//tsl/platform:errors",
- ],
-)
-
-cc_library(
- name = "qr_expander",
- srcs = ["expanders/qr_expander.cc"],
- hdrs = ["expanders/qr_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal",
- "//xla:shape_util",
- "//xla:status_macros",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder/lib:arithmetic",
- "//xla/hlo/builder/lib:constants",
- "//xla/hlo/builder/lib:loops",
- "//xla/hlo/builder/lib:math",
- "//xla/hlo/builder/lib:matrix",
- "//xla/hlo/builder/lib:qr",
- "//xla/hlo/builder/lib:slicing",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:str_format",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_absl//absl/types:span",
- "@tsl//tsl/platform:errors",
- ],
-)
-
-cc_library(
- name = "real_imag_expander",
- srcs = ["expanders/real_imag_expander.cc"],
- hdrs = ["expanders/real_imag_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal_util",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
-xla_cc_test(
- name = "real_imag_expander_test",
- size = "small",
- srcs = ["expanders/real_imag_expander_test.cc"],
- deps = [
- ":real_imag_expander",
- "//xla:literal",
- "//xla:shape_util",
- "//xla:types",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:pattern_matcher_gmock",
- "//xla/hlo/testlib:test",
- "//xla/hlo/utils:hlo_matchers",
- "//xla/service:hlo_creation_utils",
- "//xla/service:pattern_matcher",
- "//xla/tsl/lib/core:status_test_util",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
-cc_library(
- name = "eigh_expander",
- srcs = ["expanders/eigh_expander.cc"],
- hdrs = ["expanders/eigh_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal_util",
- "//xla:shape_util",
- "//xla:status_macros",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder/lib:arithmetic",
- "//xla/hlo/builder/lib:comparators",
- "//xla/hlo/builder/lib:constants",
- "//xla/hlo/builder/lib:loops",
- "//xla/hlo/builder/lib:math",
- "//xla/hlo/builder/lib:matrix",
- "//xla/hlo/builder/lib:slicing",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/status",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- "@com_google_absl//absl/strings:str_format",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_absl//absl/types:span",
- "@tsl//tsl/platform:errors",
- ],
-)
-
-cc_library(
- name = "convolution_4d_expander",
- srcs = ["expanders/convolution_4d_expander.cc"],
- hdrs = ["expanders/convolution_4d_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:shape_util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- ],
-)
-
-xla_cc_test(
- name = "convolution_4d_expander_test",
- srcs = ["expanders/convolution_4d_expander_test.cc"],
- deps = [
- "convolution_4d_expander",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:test",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:statusor",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
-cc_library(
- name = "convolution_pred_expander",
- srcs = ["expanders/convolution_pred_expander.cc"],
- hdrs = ["expanders/convolution_pred_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:shape_util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/service:hlo_creation_utils",
- "//xla/service:pattern_matcher",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:inlined_vector",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- ],
-)
-
-xla_cc_test(
- name = "convolution_pred_expander_test",
- srcs = ["expanders/convolution_pred_expander_test.cc"],
- deps = [
- ":convolution_pred_expander",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:pattern_matcher_gmock",
- "//xla/service:pattern_matcher",
- "@com_google_googletest//:gtest_main",
- "@tsl//tsl/platform:statusor",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
-cc_library(
- name = "logistic_expander",
- srcs = ["expanders/logistic_expander.cc"],
- hdrs = ["expanders/logistic_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:shape_util",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- "@tsl//tsl/platform:logging",
- ],
-)
-
-xla_cc_test(
- name = "logistic_expander_test",
- srcs = ["expanders/logistic_expander_test.cc"],
- deps = [
- ":logistic_expander",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/parser:hlo_parser",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:pattern_matcher_gmock",
- "//xla/hlo/testlib:test",
- "//xla/service:dynamic_padder",
- "//xla/service:pattern_matcher",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:statusor",
- "@tsl//tsl/platform:test_main", # fixdeps: keep
- ],
-)
-
-cc_library(
- name = "bitcast_dtypes_expander",
- srcs = ["expanders/bitcast_dtypes_expander.cc"],
- hdrs = ["expanders/bitcast_dtypes_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal_util",
- "//xla:shape_util",
- "//xla:status_macros",
- "//xla:types",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder:xla_computation",
- "//xla/hlo/builder/lib:arithmetic",
- "//xla/hlo/builder/lib:broadcast",
- "//xla/hlo/builder/lib:constants",
- "//xla/hlo/ir:hlo",
- "//xla/service:hlo_creation_utils",
- "//xla/service:hlo_module_config",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- "@com_google_absl//absl/strings:str_format",
- "@tsl//tsl/platform:logging",
- "@tsl//tsl/platform:statusor",
- ],
-)
-
-xla_cc_test(
- name = "bitcast_dtypes_expander_test",
- srcs = ["expanders/bitcast_dtypes_expander_test.cc"],
- deps = [
- ":bitcast_dtypes_expander",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:filecheck",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/utils:hlo_matchers",
- "//xla/tsl/lib/core:status_test_util",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_googletest//:gtest_main",
- "@tsl//tsl/platform:statusor",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
cc_library(
name = "while_loop_trip_count_annotator",
srcs = ["while_loop_trip_count_annotator.cc"],
@@ -482,109 +157,6 @@ xla_cc_test(
],
)
-cc_library(
- name = "dot_decomposer",
- srcs = ["expanders/dot_decomposer.cc"],
- hdrs = ["expanders/dot_decomposer.h"],
- deps = [
- "//xla:shape_util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "//xla/service:shape_inference",
- "@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:errors",
- "@tsl//tsl/platform:logging",
- "@tsl//tsl/platform:statusor",
- ],
-)
-
-xla_cc_test(
- name = "dot_decomposer_test",
- srcs = ["expanders/dot_decomposer_test.cc"],
- deps = [
- ":dot_decomposer",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:pattern_matcher_gmock",
- "//xla/hlo/utils:hlo_matchers",
- "//xla/service:pattern_matcher",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:statusor",
- "@tsl//tsl/platform:test",
- "@tsl//tsl/platform:test_main", # fixdeps: keep
- ],
-)
-
-cc_library(
- name = "reshape_decomposer",
- srcs = ["expanders/reshape_decomposer.cc"],
- hdrs = ["expanders/reshape_decomposer.h"],
- deps = [
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/log",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
-cc_library(
- name = "reduce_decomposer",
- srcs = ["expanders/reduce_decomposer.cc"],
- hdrs = ["expanders/reduce_decomposer.h"],
- deps = [
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/status",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
-xla_cc_test(
- name = "reduce_decomposer_test",
- srcs = ["expanders/reduce_decomposer_test.cc"],
- deps = [
- ":reduce_decomposer",
- "//xla/hlo/parser:hlo_parser",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:test",
- "//xla/hlo/testlib:test_helpers",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
-xla_cc_test(
- name = "reshape_decomposer_test",
- srcs = ["expanders/reshape_decomposer_test.cc"],
- deps = [
- ":reshape_decomposer",
- "//xla/hlo/parser:hlo_parser",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:test",
- "//xla/hlo/testlib:test_helpers",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
cc_library(
name = "memory_space_propagation",
srcs = ["memory_space_propagation.cc"],
@@ -616,41 +188,6 @@ xla_cc_test(
],
)
-cc_library(
- name = "stable_sort_expander",
- srcs = ["expanders/stable_sort_expander.cc"],
- hdrs = ["expanders/stable_sort_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
-xla_cc_test(
- name = "stable_sort_expander_test",
- srcs = ["expanders/stable_sort_expander_test.cc"],
- deps = [
- ":stable_sort_expander",
- "//xla/hlo/parser:hlo_parser",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:pattern_matcher_gmock",
- "//xla/hlo/testlib:test",
- "//xla/hlo/transforms/simplifiers:algebraic_simplifier",
- "//xla/hlo/utils:hlo_matchers",
- "//xla/service:pattern_matcher",
- "//xla/tsl/lib/core:status_test_util",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
cc_library(
name = "convert_memory_placement_to_internal_annotations",
srcs = ["convert_memory_placement_to_internal_annotations.cc"],
@@ -903,91 +440,16 @@ xla_cc_test(
],
)
-cc_library(
- name = "dynamic_index_splitter",
- srcs = ["expanders/dynamic_index_splitter.cc"],
- hdrs = ["expanders/dynamic_index_splitter.h"],
- deps = [
- "//xla:shape_util",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/container:inlined_vector",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings",
- "@com_google_absl//absl/types:span",
- ],
-)
-
-xla_cc_test(
- name = "dynamic_index_splitter_test",
- srcs = ["expanders/dynamic_index_splitter_test.cc"],
- deps = [
- ":dynamic_index_splitter",
- "//xla:xla_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/testlib:test",
- "//xla/hlo/testlib:test_helpers",
- "//xla/hlo/utils:hlo_matchers",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
-cc_library(
- name = "rng_expander",
- srcs = ["expanders/rng_expander.cc"],
- hdrs = ["expanders/rng_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:literal_util",
- "//xla:shape_util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder/lib:prng",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/algorithm:container",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/log",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- "@com_google_absl//absl/synchronization",
- ],
-)
-
-cc_library(
- name = "rng_bit_generator_expander",
- srcs = ["expanders/rng_bit_generator_expander.cc"],
- hdrs = ["expanders/rng_bit_generator_expander.h"],
- deps = [
- ":op_expander_pass",
- "//xla:shape_util",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/builder:xla_builder",
- "//xla/hlo/builder/lib:prng",
- "//xla/hlo/ir:hlo",
- "//xla/service:hlo_creation_utils",
- "@com_google_absl//absl/container:flat_hash_map",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- ],
-)
-
cc_library(
name = "operand_upcaster",
srcs = ["operand_upcaster.cc"],
hdrs = ["operand_upcaster.h"],
deps = [
- ":op_expander_pass",
"//xla:shape_util",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"//xla/service:hlo_creation_utils",
"//xla/service:shape_inference",
"@com_google_absl//absl/algorithm:container",
@@ -1015,44 +477,6 @@ xla_cc_test(
],
)
-cc_library(
- name = "stochastic_convert_decomposer",
- srcs = ["expanders/stochastic_convert_decomposer.cc"],
- hdrs = ["expanders/stochastic_convert_decomposer.h"],
- deps = [
- "//xla:shape_util",
- "//xla:util",
- "//xla:xla_data_proto_cc",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/pass:hlo_pass",
- "//xla/service:hlo_creation_utils",
- "//xla/service:shape_inference",
- "@com_google_absl//absl/container:flat_hash_set",
- "@com_google_absl//absl/log",
- "@com_google_absl//absl/log:check",
- "@com_google_absl//absl/status",
- "@com_google_absl//absl/status:statusor",
- "@com_google_absl//absl/strings:string_view",
- "@tsl//tsl/platform:errors",
- "@tsl//tsl/platform:statusor",
- ],
-)
-
-xla_cc_test(
- name = "stochastic_convert_decomposer_test",
- srcs = ["expanders/stochastic_convert_decomposer_test.cc"],
- deps = [
- ":stochastic_convert_decomposer",
- "//xla/hlo/ir:hlo",
- "//xla/hlo/parser:hlo_parser",
- "//xla/hlo/testlib:hlo_hardware_independent_test_base",
- "//xla/hlo/utils:hlo_matchers",
- "@com_google_absl//absl/status",
- "@com_google_googletest//:gtest",
- "@tsl//tsl/platform:test_main",
- ],
-)
-
cc_library(
name = "sharding_format_picker",
testonly = True,
diff --git a/xla/hlo/transforms/collectives/BUILD b/xla/hlo/transforms/collectives/BUILD
index b7f8aba244bdf..2e4274f677f40 100644
--- a/xla/hlo/transforms/collectives/BUILD
+++ b/xla/hlo/transforms/collectives/BUILD
@@ -442,7 +442,7 @@ cc_library(
deps = [
"//xla:shape_util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"//xla/service:hlo_creation_utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
diff --git a/xla/hlo/transforms/expanders/BUILD b/xla/hlo/transforms/expanders/BUILD
new file mode 100644
index 0000000000000..a911e6de36054
--- /dev/null
+++ b/xla/hlo/transforms/expanders/BUILD
@@ -0,0 +1,595 @@
+# Description:
+# Implementation of XLA’s HLO expander transformations.
+
+load("//xla:xla.bzl", "xla_cc_test")
+
+package(
+ # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+ default_visibility = [":friends"],
+ licenses = ["notice"],
+)
+
+package_group(
+ name = "friends",
+ includes = [
+ "//xla:friends",
+ ],
+)
+
+cc_library(
+ name = "op_expander_pass",
+ srcs = ["op_expander_pass.cc"],
+ hdrs = ["op_expander_pass.h"],
+ deps = [
+ "//xla:util",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ "@tsl//tsl/platform:statusor",
+ ],
+)
+
+cc_library(
+ name = "optimization_barrier_expander",
+ srcs = ["optimization_barrier_expander.cc"],
+ hdrs = ["optimization_barrier_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+cc_library(
+ name = "comparison_expander",
+ srcs = ["comparison_expander.cc"],
+ hdrs = ["comparison_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:comparison_util",
+ "//xla:literal_util",
+ "//xla:shape_util",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/log",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_absl//absl/types:span",
+ ],
+)
+
+cc_library(
+ name = "cholesky_expander",
+ srcs = ["cholesky_expander.cc"],
+ hdrs = ["cholesky_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal",
+ "//xla:shape_util",
+ "//xla:status_macros",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder/lib:arithmetic",
+ "//xla/hlo/builder/lib:constants",
+ "//xla/hlo/builder/lib:loops",
+ "//xla/hlo/builder/lib:math",
+ "//xla/hlo/builder/lib:matrix",
+ "//xla/hlo/builder/lib:slicing",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:str_format",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_absl//absl/types:span",
+ "@tsl//tsl/platform:errors",
+ ],
+)
+
+cc_library(
+ name = "qr_expander",
+ srcs = ["qr_expander.cc"],
+ hdrs = ["qr_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal",
+ "//xla:shape_util",
+ "//xla:status_macros",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder/lib:arithmetic",
+ "//xla/hlo/builder/lib:constants",
+ "//xla/hlo/builder/lib:loops",
+ "//xla/hlo/builder/lib:math",
+ "//xla/hlo/builder/lib:matrix",
+ "//xla/hlo/builder/lib:qr",
+ "//xla/hlo/builder/lib:slicing",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:str_format",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_absl//absl/types:span",
+ "@tsl//tsl/platform:errors",
+ ],
+)
+
+cc_library(
+ name = "real_imag_expander",
+ srcs = ["real_imag_expander.cc"],
+ hdrs = ["real_imag_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal_util",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+xla_cc_test(
+ name = "real_imag_expander_test",
+ size = "small",
+ srcs = ["real_imag_expander_test.cc"],
+ deps = [
+ ":real_imag_expander",
+ "//xla:literal",
+ "//xla:shape_util",
+ "//xla:types",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:pattern_matcher_gmock",
+ "//xla/hlo/testlib:test",
+ "//xla/hlo/utils:hlo_matchers",
+ "//xla/service:hlo_creation_utils",
+ "//xla/service:pattern_matcher",
+ "//xla/service:pattern_matcher_gmock",
+ "//xla/tsl/lib/core:status_test_util",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "eigh_expander",
+ srcs = ["eigh_expander.cc"],
+ hdrs = ["eigh_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal_util",
+ "//xla:shape_util",
+ "//xla:status_macros",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder/lib:arithmetic",
+ "//xla/hlo/builder/lib:comparators",
+ "//xla/hlo/builder/lib:constants",
+ "//xla/hlo/builder/lib:loops",
+ "//xla/hlo/builder/lib:math",
+ "//xla/hlo/builder/lib:matrix",
+ "//xla/hlo/builder/lib:slicing",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/strings:str_format",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_absl//absl/types:span",
+ "@tsl//tsl/platform:errors",
+ ],
+)
+
+cc_library(
+ name = "convolution_4d_expander",
+ srcs = ["convolution_4d_expander.cc"],
+ hdrs = ["convolution_4d_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:shape_util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ ],
+)
+
+xla_cc_test(
+ name = "convolution_4d_expander_test",
+ srcs = ["convolution_4d_expander_test.cc"],
+ deps = [
+ "convolution_4d_expander",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:test",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:statusor",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "convolution_pred_expander",
+ srcs = ["convolution_pred_expander.cc"],
+ hdrs = ["convolution_pred_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:shape_util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/service:hlo_creation_utils",
+ "//xla/service:pattern_matcher",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:inlined_vector",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ ],
+)
+
+xla_cc_test(
+ name = "convolution_pred_expander_test",
+ srcs = ["convolution_pred_expander_test.cc"],
+ deps = [
+ ":convolution_pred_expander",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:pattern_matcher_gmock",
+ "//xla/service:pattern_matcher",
+ "@com_google_googletest//:gtest_main",
+ "@tsl//tsl/platform:statusor",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "logistic_expander",
+ srcs = ["logistic_expander.cc"],
+ hdrs = ["logistic_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:shape_util",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ "@tsl//tsl/platform:logging",
+ ],
+)
+
+xla_cc_test(
+ name = "logistic_expander_test",
+ srcs = ["logistic_expander_test.cc"],
+ deps = [
+ ":logistic_expander",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/parser:hlo_parser",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:pattern_matcher_gmock",
+ "//xla/hlo/testlib:test",
+ "//xla/service:dynamic_padder",
+ "//xla/service:pattern_matcher",
+ "//xla/service:pattern_matcher_gmock",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:statusor",
+ "@tsl//tsl/platform:test_main", # fixdeps: keep
+ ],
+)
+
+cc_library(
+ name = "bitcast_dtypes_expander",
+ srcs = ["bitcast_dtypes_expander.cc"],
+ hdrs = ["bitcast_dtypes_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal_util",
+ "//xla:shape_util",
+ "//xla:status_macros",
+ "//xla:types",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder:xla_computation",
+ "//xla/hlo/builder/lib:arithmetic",
+ "//xla/hlo/builder/lib:broadcast",
+ "//xla/hlo/builder/lib:constants",
+ "//xla/hlo/ir:hlo",
+ "//xla/service:hlo_creation_utils",
+ "//xla/service:hlo_module_config",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/strings:str_format",
+ "@tsl//tsl/platform:logging",
+ "@tsl//tsl/platform:statusor",
+ ],
+)
+
+xla_cc_test(
+ name = "bitcast_dtypes_expander_test",
+ srcs = ["bitcast_dtypes_expander_test.cc"],
+ deps = [
+ ":bitcast_dtypes_expander",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:filecheck",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/utils:hlo_matchers",
+ "//xla/tsl/lib/core:status_test_util",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_googletest//:gtest_main",
+ "@tsl//tsl/platform:statusor",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "dot_decomposer",
+ srcs = ["dot_decomposer.cc"],
+ hdrs = ["dot_decomposer.h"],
+ deps = [
+ "//xla:shape_util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "//xla/service:shape_inference",
+ "@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:errors",
+ "@tsl//tsl/platform:logging",
+ "@tsl//tsl/platform:statusor",
+ ],
+)
+
+xla_cc_test(
+ name = "dot_decomposer_test",
+ srcs = ["dot_decomposer_test.cc"],
+ deps = [
+ ":dot_decomposer",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:pattern_matcher_gmock",
+ "//xla/hlo/utils:hlo_matchers",
+ "//xla/service:pattern_matcher",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:statusor",
+ "@tsl//tsl/platform:test",
+ "@tsl//tsl/platform:test_main", # fixdeps: keep
+ ],
+)
+
+cc_library(
+ name = "reshape_decomposer",
+ srcs = ["reshape_decomposer.cc"],
+ hdrs = ["reshape_decomposer.h"],
+ deps = [
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/log",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+cc_library(
+ name = "reduce_decomposer",
+ srcs = ["reduce_decomposer.cc"],
+ hdrs = ["reduce_decomposer.h"],
+ deps = [
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+xla_cc_test(
+ name = "reduce_decomposer_test",
+ srcs = ["reduce_decomposer_test.cc"],
+ deps = [
+ ":reduce_decomposer",
+ "//xla/hlo/parser:hlo_parser",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:test",
+ "//xla/hlo/testlib:test_helpers",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+xla_cc_test(
+ name = "reshape_decomposer_test",
+ srcs = ["reshape_decomposer_test.cc"],
+ deps = [
+ ":reshape_decomposer",
+ "//xla/hlo/parser:hlo_parser",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:test",
+ "//xla/hlo/testlib:test_helpers",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "stable_sort_expander",
+ srcs = ["stable_sort_expander.cc"],
+ hdrs = ["stable_sort_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+xla_cc_test(
+ name = "stable_sort_expander_test",
+ srcs = ["stable_sort_expander_test.cc"],
+ deps = [
+ ":stable_sort_expander",
+ "//xla/hlo/parser:hlo_parser",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:pattern_matcher_gmock",
+ "//xla/hlo/testlib:test",
+ "//xla/hlo/transforms/simplifiers:algebraic_simplifier",
+ "//xla/hlo/utils:hlo_matchers",
+ "//xla/service:pattern_matcher",
+ "//xla/tsl/lib/core:status_test_util",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "dynamic_index_splitter",
+ srcs = ["dynamic_index_splitter.cc"],
+ hdrs = ["dynamic_index_splitter.h"],
+ deps = [
+ "//xla:shape_util",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/container:inlined_vector",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings",
+ "@com_google_absl//absl/types:span",
+ ],
+)
+
+xla_cc_test(
+ name = "dynamic_index_splitter_test",
+ srcs = ["dynamic_index_splitter_test.cc"],
+ deps = [
+ ":dynamic_index_splitter",
+ "//xla:xla_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/testlib:test",
+ "//xla/hlo/testlib:test_helpers",
+ "//xla/hlo/utils:hlo_matchers",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
+
+cc_library(
+ name = "rng_expander",
+ srcs = ["rng_expander.cc"],
+ hdrs = ["rng_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:literal_util",
+ "//xla:shape_util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder/lib:prng",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/algorithm:container",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/log",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_absl//absl/synchronization",
+ ],
+)
+
+cc_library(
+ name = "rng_bit_generator_expander",
+ srcs = ["rng_bit_generator_expander.cc"],
+ hdrs = ["rng_bit_generator_expander.h"],
+ deps = [
+ ":op_expander_pass",
+ "//xla:shape_util",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/builder:xla_builder",
+ "//xla/hlo/builder/lib:prng",
+ "//xla/hlo/ir:hlo",
+ "//xla/service:hlo_creation_utils",
+ "@com_google_absl//absl/container:flat_hash_map",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ ],
+)
+
+cc_library(
+ name = "stochastic_convert_decomposer",
+ srcs = ["stochastic_convert_decomposer.cc"],
+ hdrs = ["stochastic_convert_decomposer.h"],
+ deps = [
+ "//xla:shape_util",
+ "//xla:util",
+ "//xla:xla_data_proto_cc",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/pass:hlo_pass",
+ "//xla/service:hlo_creation_utils",
+ "//xla/service:shape_inference",
+ "@com_google_absl//absl/container:flat_hash_set",
+ "@com_google_absl//absl/log",
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/status",
+ "@com_google_absl//absl/status:statusor",
+ "@com_google_absl//absl/strings:string_view",
+ "@tsl//tsl/platform:errors",
+ "@tsl//tsl/platform:statusor",
+ ],
+)
+
+xla_cc_test(
+ name = "stochastic_convert_decomposer_test",
+ srcs = ["stochastic_convert_decomposer_test.cc"],
+ deps = [
+ ":stochastic_convert_decomposer",
+ "//xla/hlo/ir:hlo",
+ "//xla/hlo/parser:hlo_parser",
+ "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+ "//xla/hlo/utils:hlo_matchers",
+ "@com_google_absl//absl/status",
+ "@com_google_googletest//:gtest",
+ "@tsl//tsl/platform:test_main",
+ ],
+)
diff --git a/xla/hlo/transforms/simplifiers/BUILD b/xla/hlo/transforms/simplifiers/BUILD
index e19f32f2cbbb7..3a5892cb14c95 100644
--- a/xla/hlo/transforms/simplifiers/BUILD
+++ b/xla/hlo/transforms/simplifiers/BUILD
@@ -1352,7 +1352,7 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"//xla/service:shape_inference",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
@@ -1384,7 +1384,7 @@ cc_library(
"//xla:shape_util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@tsl//tsl/platform:errors",
@@ -1454,7 +1454,7 @@ cc_library(
"//xla:permutation_util",
"//xla:shape_util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"//xla/service:gather_scatter_utils",
"//xla/service:hlo_creation_utils",
"@com_google_absl//absl/algorithm:container",
diff --git a/xla/pjrt/interpreter/BUILD b/xla/pjrt/interpreter/BUILD
index cbaf62d362445..5f9c1d597303a 100644
--- a/xla/pjrt/interpreter/BUILD
+++ b/xla/pjrt/interpreter/BUILD
@@ -22,10 +22,10 @@ cc_library(
"//xla/hlo/evaluator:hlo_evaluator",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:cholesky_expander",
- "//xla/hlo/transforms:dynamic_index_splitter",
- "//xla/hlo/transforms:eigh_expander",
- "//xla/hlo/transforms:qr_expander",
+ "//xla/hlo/transforms/expanders:cholesky_expander",
+ "//xla/hlo/transforms/expanders:dynamic_index_splitter",
+ "//xla/hlo/transforms/expanders:eigh_expander",
+ "//xla/hlo/transforms/expanders:qr_expander",
"//xla/pjrt:layout_mode",
"//xla/pjrt:mlir_to_hlo",
"//xla/pjrt:pjrt_client",
diff --git a/xla/service/BUILD b/xla/service/BUILD
index a497df7bfad6c..5ee5d58c939cd 100644
--- a/xla/service/BUILD
+++ b/xla/service/BUILD
@@ -2090,8 +2090,8 @@ cc_library(
cc_library(
name = "op_expander_pass",
hdrs = ["op_expander_pass.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:op_expander_pass instead.",
- deps = ["//xla/hlo/transforms:op_expander_pass"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:op_expander_pass instead.",
+ deps = ["//xla/hlo/transforms/expanders:op_expander_pass"],
)
cc_library(
@@ -2105,7 +2105,7 @@ cc_library(
"//xla:literal_util",
"//xla:util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status:statusor",
@@ -2116,15 +2116,15 @@ cc_library(
cc_library(
name = "optimization_barrier_expander",
hdrs = ["optimization_barrier_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:optimization_barrier_expander instead.",
- deps = ["//xla/hlo/transforms:optimization_barrier_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:optimization_barrier_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:optimization_barrier_expander"],
)
cc_library(
name = "comparison_expander",
hdrs = ["comparison_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:comparison_expander instead.",
- deps = ["//xla/hlo/transforms:comparison_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:comparison_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:comparison_expander"],
)
cc_library(
@@ -2158,7 +2158,7 @@ cc_library(
"//xla:literal_util",
"//xla:shape_util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
@@ -2180,7 +2180,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
@@ -2241,7 +2241,7 @@ cc_library(
"//xla/hlo/builder/lib:matrix",
"//xla/hlo/builder/lib:slicing",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
@@ -2274,42 +2274,43 @@ xla_cc_test(
cc_library(
name = "cholesky_expander",
hdrs = ["cholesky_expander.h"],
- deps = ["//xla/hlo/transforms:cholesky_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:cholesky_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:cholesky_expander"],
)
cc_library(
name = "qr_expander",
hdrs = ["qr_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:qr_expander instead.",
- deps = ["//xla/hlo/transforms:qr_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:qr_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:qr_expander"],
)
cc_library(
name = "real_imag_expander",
hdrs = ["real_imag_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:real_imag_expander instead.",
- deps = ["//xla/hlo/transforms:real_imag_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:real_imag_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:real_imag_expander"],
)
cc_library(
name = "eigh_expander",
hdrs = ["eigh_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:eigh_expander instead.",
- deps = ["//xla/hlo/transforms:eigh_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:eigh_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:eigh_expander"],
)
cc_library(
name = "convolution_4d_expander",
hdrs = ["convolution_4d_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:convolution_4d_expander instead.",
- deps = ["//xla/hlo/transforms:convolution_4d_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:convolution_4d_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:convolution_4d_expander"],
)
cc_library(
name = "convolution_pred_expander",
hdrs = ["convolution_pred_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:convolution_pred_expander instead.",
- deps = ["//xla/hlo/transforms:convolution_pred_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:convolution_pred_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:convolution_pred_expander"],
)
xla_test(
@@ -2373,8 +2374,8 @@ cc_library(
cc_library(
name = "logistic_expander",
hdrs = ["logistic_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:logistic_expander instead.",
- deps = ["//xla/hlo/transforms:logistic_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:logistic_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:logistic_expander"],
)
cc_library(
@@ -2428,8 +2429,8 @@ cc_library(
cc_library(
name = "bitcast_dtypes_expander",
hdrs = ["bitcast_dtypes_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:bitcast_dtypes_expander instead.",
- deps = ["//xla/hlo/transforms:bitcast_dtypes_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:bitcast_dtypes_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:bitcast_dtypes_expander"],
)
cc_library(
@@ -3004,8 +3005,8 @@ cc_library(
cc_library(
name = "dot_decomposer",
hdrs = ["dot_decomposer.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:dot_decomposer instead.",
- deps = ["//xla/hlo/transforms:dot_decomposer"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:dot_decomposer instead.",
+ deps = ["//xla/hlo/transforms/expanders:dot_decomposer"],
)
cc_library(
@@ -3037,7 +3038,7 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
],
@@ -3098,15 +3099,15 @@ cc_library(
cc_library(
name = "reshape_decomposer",
hdrs = ["reshape_decomposer.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:reshape_decomposer instead.",
- deps = ["//xla/hlo/transforms:reshape_decomposer"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:reshape_decomposer instead.",
+ deps = ["//xla/hlo/transforms/expanders:reshape_decomposer"],
)
cc_library(
name = "reduce_decomposer",
hdrs = ["reduce_decomposer.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:reduce_decomposer instead.",
- deps = ["//xla/hlo/transforms:reduce_decomposer"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:reduce_decomposer instead.",
+ deps = ["//xla/hlo/transforms/expanders:reduce_decomposer"],
)
cc_library(
@@ -4721,8 +4722,8 @@ cc_library(
cc_library(
name = "stable_sort_expander",
hdrs = ["stable_sort_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:stable_sort_expander instead.",
- deps = ["//xla/hlo/transforms:stable_sort_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:stable_sort_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:stable_sort_expander"],
)
cc_library(
@@ -5269,8 +5270,8 @@ cc_library(
cc_library(
name = "dynamic_index_splitter",
hdrs = ["dynamic_index_splitter.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:dynamic_index_splitter instead.",
- deps = ["//xla/hlo/transforms:dynamic_index_splitter"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:dynamic_index_splitter instead.",
+ deps = ["//xla/hlo/transforms/expanders:dynamic_index_splitter"],
)
xla_cc_test(
@@ -5440,15 +5441,15 @@ cc_library(
cc_library(
name = "rng_expander",
hdrs = ["rng_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:rng_expander instead.",
- deps = ["//xla/hlo/transforms:rng_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:rng_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:rng_expander"],
)
cc_library(
name = "rng_bit_generator_expander",
hdrs = ["rng_bit_generator_expander.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:rng_bit_generator_expander instead.",
- deps = ["//xla/hlo/transforms:rng_bit_generator_expander"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:rng_bit_generator_expander instead.",
+ deps = ["//xla/hlo/transforms/expanders:rng_bit_generator_expander"],
)
cc_library(
@@ -5876,7 +5877,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/types:span",
@@ -5906,7 +5907,7 @@ cc_library(
":call_inliner",
"//xla:literal_util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
],
)
@@ -5986,7 +5987,7 @@ cc_library(
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
@@ -6007,8 +6008,8 @@ cc_library(
cc_library(
name = "stochastic_convert_decomposer",
hdrs = ["stochastic_convert_decomposer.h"],
- deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms:stochastic_convert_decomposer instead.",
- deps = ["//xla/hlo/transforms:stochastic_convert_decomposer"],
+ deprecation = "This library is deprecated. Use //third_party/tensorflow/compiler/xla/hlo/transforms/expanders:stochastic_convert_decomposer instead.",
+ deps = ["//xla/hlo/transforms/expanders:stochastic_convert_decomposer"],
)
cc_library(
diff --git a/xla/service/cpu/BUILD b/xla/service/cpu/BUILD
index 3c37459fe7423..4b0386f6690b1 100644
--- a/xla/service/cpu/BUILD
+++ b/xla/service/cpu/BUILD
@@ -255,23 +255,23 @@ cc_library(
"//xla/hlo/ir:hlo_module_group",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:bitcast_dtypes_expander",
- "//xla/hlo/transforms:cholesky_expander",
- "//xla/hlo/transforms:comparison_expander",
- "//xla/hlo/transforms:dot_decomposer",
- "//xla/hlo/transforms:dynamic_index_splitter",
- "//xla/hlo/transforms:eigh_expander",
"//xla/hlo/transforms:literal_canonicalizer",
- "//xla/hlo/transforms:logistic_expander",
"//xla/hlo/transforms:operand_upcaster",
- "//xla/hlo/transforms:optimization_barrier_expander",
- "//xla/hlo/transforms:qr_expander",
- "//xla/hlo/transforms:reduce_decomposer",
- "//xla/hlo/transforms:reshape_decomposer",
- "//xla/hlo/transforms:rng_bit_generator_expander",
- "//xla/hlo/transforms:rng_expander",
- "//xla/hlo/transforms:stochastic_convert_decomposer",
"//xla/hlo/transforms:while_loop_trip_count_annotator",
+ "//xla/hlo/transforms/expanders:bitcast_dtypes_expander",
+ "//xla/hlo/transforms/expanders:cholesky_expander",
+ "//xla/hlo/transforms/expanders:comparison_expander",
+ "//xla/hlo/transforms/expanders:dot_decomposer",
+ "//xla/hlo/transforms/expanders:dynamic_index_splitter",
+ "//xla/hlo/transforms/expanders:eigh_expander",
+ "//xla/hlo/transforms/expanders:logistic_expander",
+ "//xla/hlo/transforms/expanders:optimization_barrier_expander",
+ "//xla/hlo/transforms/expanders:qr_expander",
+ "//xla/hlo/transforms/expanders:reduce_decomposer",
+ "//xla/hlo/transforms/expanders:reshape_decomposer",
+ "//xla/hlo/transforms/expanders:rng_bit_generator_expander",
+ "//xla/hlo/transforms/expanders:rng_expander",
+ "//xla/hlo/transforms/expanders:stochastic_convert_decomposer",
"//xla/hlo/transforms/simplifiers:algebraic_simplifier",
"//xla/hlo/transforms/simplifiers:batch_dot_simplification",
"//xla/hlo/transforms/simplifiers:broadcast_canonicalizer",
diff --git a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD
index 4a010cf3b6c03..c5ef6942c8060 100644
--- a/xla/service/gpu/BUILD
+++ b/xla/service/gpu/BUILD
@@ -1434,6 +1434,11 @@ cc_library(
"//xla/hlo/ir:hlo_module_group",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
+ "//xla/hlo/transforms:convert_memory_placement_to_internal_annotations",
+ "//xla/hlo/transforms:host_offload_legalize",
+ "//xla/hlo/transforms:host_offloader",
+ "//xla/hlo/transforms:operand_upcaster",
+ "//xla/hlo/transforms:while_loop_trip_count_annotator",
"//xla/hlo/transforms/collectives:all_gather_broadcast_reorder",
"//xla/hlo/transforms/collectives:all_gather_combiner",
"//xla/hlo/transforms/collectives:all_reduce_combiner",
@@ -1442,6 +1447,23 @@ cc_library(
"//xla/hlo/transforms/collectives:collective_quantizer",
"//xla/hlo/transforms/collectives:collectives_schedule_linearizer",
"//xla/hlo/transforms/collectives:convert_async_collectives_to_sync",
+ "//xla/hlo/transforms/expanders:bitcast_dtypes_expander",
+ "//xla/hlo/transforms/expanders:comparison_expander",
+ "//xla/hlo/transforms/expanders:convolution_4d_expander",
+ "//xla/hlo/transforms/expanders:convolution_pred_expander",
+ "//xla/hlo/transforms/expanders:dot_decomposer",
+ "//xla/hlo/transforms/expanders:dynamic_index_splitter",
+ "//xla/hlo/transforms/expanders:eigh_expander",
+ "//xla/hlo/transforms/expanders:logistic_expander",
+ "//xla/hlo/transforms/expanders:optimization_barrier_expander",
+ "//xla/hlo/transforms/expanders:qr_expander",
+ "//xla/hlo/transforms/expanders:real_imag_expander",
+ "//xla/hlo/transforms/expanders:reduce_decomposer",
+ "//xla/hlo/transforms/expanders:reshape_decomposer",
+ "//xla/hlo/transforms/expanders:rng_bit_generator_expander",
+ "//xla/hlo/transforms/expanders:rng_expander",
+ "//xla/hlo/transforms/expanders:stable_sort_expander",
+ "//xla/hlo/transforms/expanders:stochastic_convert_decomposer",
"//xla/hlo/transforms/simplifiers:algebraic_simplifier",
"//xla/hlo/transforms/simplifiers:all_reduce_folder",
"//xla/hlo/transforms/simplifiers:broadcast_canonicalizer",
@@ -1467,28 +1489,6 @@ cc_library(
"//xla/hlo/transforms/simplifiers:sub_byte_normalization",
"//xla/hlo/transforms/simplifiers:tuple_simplifier",
"//xla/hlo/transforms/simplifiers:zero_sized_hlo_elimination",
- "//xla/hlo/transforms:bitcast_dtypes_expander",
- "//xla/hlo/transforms:comparison_expander",
- "//xla/hlo/transforms:convert_memory_placement_to_internal_annotations",
- "//xla/hlo/transforms:convolution_4d_expander",
- "//xla/hlo/transforms:convolution_pred_expander",
- "//xla/hlo/transforms:dot_decomposer",
- "//xla/hlo/transforms:dynamic_index_splitter",
- "//xla/hlo/transforms:eigh_expander",
- "//xla/hlo/transforms:host_offload_legalize",
- "//xla/hlo/transforms:host_offloader",
- "//xla/hlo/transforms:logistic_expander",
- "//xla/hlo/transforms:operand_upcaster",
- "//xla/hlo/transforms:optimization_barrier_expander",
- "//xla/hlo/transforms:qr_expander",
- "//xla/hlo/transforms:real_imag_expander",
- "//xla/hlo/transforms:reduce_decomposer",
- "//xla/hlo/transforms:reshape_decomposer",
- "//xla/hlo/transforms:rng_bit_generator_expander",
- "//xla/hlo/transforms:rng_expander",
- "//xla/hlo/transforms:stable_sort_expander",
- "//xla/hlo/transforms:stochastic_convert_decomposer",
- "//xla/hlo/transforms:while_loop_trip_count_annotator",
"//xla/hlo/translate/hlo_to_mhlo:hlo_utils",
"//xla/hlo/translate/mhlo_to_hlo:location_exporter",
"//xla/hlo/utils:hlo_query",
diff --git a/xla/service/gpu/transforms/BUILD b/xla/service/gpu/transforms/BUILD
index 1f7752a8cd938..11d5b3860f712 100644
--- a/xla/service/gpu/transforms/BUILD
+++ b/xla/service/gpu/transforms/BUILD
@@ -371,7 +371,7 @@ cc_library(
hdrs = ["all_gather_dynamic_slice_simplifier.h"],
deps = [
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"//xla/service:collective_opt_utils",
"//xla/service:hlo_module_config",
"@com_google_absl//absl/status:statusor",
@@ -1030,7 +1030,7 @@ xla_test(
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/parser:hlo_parser",
- "//xla/hlo/transforms:reshape_decomposer",
+ "//xla/hlo/transforms/expanders:reshape_decomposer",
"//xla/hlo/transforms/simplifiers:algebraic_simplifier",
"//xla/hlo/transforms/simplifiers:hlo_dce",
"//xla/service:computation_layout",
@@ -1450,7 +1450,7 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@tsl//tsl/platform:errors",
@@ -1481,7 +1481,7 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:op_expander_pass",
+ "//xla/hlo/transforms/expanders:op_expander_pass",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@tsl//tsl/platform:errors",
diff --git a/xla/service/spmd/BUILD b/xla/service/spmd/BUILD
index 07781581968e1..407b6280d4d02 100644
--- a/xla/service/spmd/BUILD
+++ b/xla/service/spmd/BUILD
@@ -274,7 +274,7 @@ xla_cc_test(
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:rng_expander",
+ "//xla/hlo/transforms/expanders:rng_expander",
"//xla/hlo/utils:hlo_matchers",
"//xla/service:hlo_module_config",
"//xla/service:hlo_verifier",
diff --git a/xla/tests/BUILD b/xla/tests/BUILD
index b98fef95d0f51..fe6d37ccbac96 100644
--- a/xla/tests/BUILD
+++ b/xla/tests/BUILD
@@ -2222,8 +2222,8 @@ xla_test(
"//xla:literal",
"//xla:literal_util",
"//xla/hlo/ir:hlo",
- "//xla/hlo/transforms:rng_bit_generator_expander",
- "//xla/hlo/transforms:rng_expander",
+ "//xla/hlo/transforms/expanders:rng_bit_generator_expander",
+ "//xla/hlo/transforms/expanders:rng_expander",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
diff --git a/xla/tools/BUILD b/xla/tools/BUILD
index 7ebb2f4130677..5b025edb98823 100644
--- a/xla/tools/BUILD
+++ b/xla/tools/BUILD
@@ -228,9 +228,9 @@ cc_library(
deps = [
"//xla:xla_data_proto_cc",
"//xla/hlo/pass:hlo_pass_pipeline",
- "//xla/hlo/transforms:cholesky_expander",
- "//xla/hlo/transforms:rng_bit_generator_expander",
- "//xla/hlo/transforms:rng_expander",
+ "//xla/hlo/transforms/expanders:cholesky_expander",
+ "//xla/hlo/transforms/expanders:rng_bit_generator_expander",
+ "//xla/hlo/transforms/expanders:rng_expander",
"//xla/service:batchnorm_expander",
"//xla/service:hlo_proto_cc",
"//xla/service:hlo_verifier",
diff --git a/xla/tools/hlo_opt/BUILD b/xla/tools/hlo_opt/BUILD
index 43f2218f573e8..ccf3947d68143 100644
--- a/xla/tools/hlo_opt/BUILD
+++ b/xla/tools/hlo_opt/BUILD
@@ -34,7 +34,7 @@ cc_library(
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/tools/hlo_opt:opt_lib",
- "//xla/hlo/transforms:bitcast_dtypes_expander",
+ "//xla/hlo/transforms/expanders:bitcast_dtypes_expander",
"//xla/service:all_reduce_simplifier",
"//xla/service:all_to_all_decomposer",
"//xla/service:batched_gather_scatter_normalizer",
@@ -157,7 +157,7 @@ cc_library(
"//xla/backends/cpu/codegen:target_machine_features",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
- "//xla/hlo/transforms:rng_bit_generator_expander",
+ "//xla/hlo/transforms/expanders:rng_bit_generator_expander",
"//xla/hlo/transforms/simplifiers:algebraic_simplifier",
"//xla/hlo/transforms/simplifiers:reduce_window_rewriter",
"//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",