[HLO Componentization] Create pass sub-component

PiperOrigin-RevId: 675905538
openxla · Sep 18, 2024 · 768b667 · 768b667
1 parent 8893550
commit 768b667
Show file tree

Hide file tree

Showing 22 changed files with 25 additions and 25 deletions.
diff --git a/xla/backends/interpreter/BUILD b/xla/backends/interpreter/BUILD
@@ -38,6 +38,7 @@ cc_library(
         "//xla/hlo/evaluator:hlo_evaluator",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_module_group",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/service:batchnorm_expander",
         "//xla/service:cholesky_expander",
         "//xla/service:compiler",
@@ -49,7 +50,6 @@ cc_library(
         "//xla/service:executable",
         "//xla/service:hlo_cost_analysis",
         "//xla/service:hlo_module_config",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service:layout_assignment",
         "//xla/service:qr_expander",
         "//xla/service:topk_rewriter",

diff --git a/xla/backends/interpreter/compiler.cc b/xla/backends/interpreter/compiler.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_module_group.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/literal.h"
 #include "xla/service/batchnorm_expander.h"
 #include "xla/service/cholesky_expander.h"
@@ -41,7 +42,6 @@ limitations under the License.
 #include "xla/service/eigh_expander.h"
 #include "xla/service/executable.h"
 #include "xla/service/hlo_cost_analysis.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/service/layout_assignment.h"
 #include "xla/service/qr_expander.h"
 #include "xla/service/topk_rewriter.h"

diff --git a/xla/hlo/transforms/BUILD b/xla/hlo/transforms/BUILD
@@ -23,7 +23,7 @@ cc_library(
     hdrs = ["hlo_constant_splitter.h"],
     deps = [
         "//xla/hlo/ir:hlo",
-        "//xla/service:hlo_pass",
+        "//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",

diff --git a/xla/hlo/transforms/hlo_constant_splitter.h b/xla/hlo/transforms/hlo_constant_splitter.h
@@ -21,7 +21,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/service/hlo_pass_interface.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 
 namespace xla {
 

diff --git a/xla/python/BUILD b/xla/python/BUILD
@@ -611,13 +611,13 @@ cc_library(
         "//xla:util",
         "//xla/client:xla_computation",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/pjrt:mlir_to_hlo",
         "//xla/pjrt/c:pjrt_c_api_custom_partitioner_extension_hdrs",
         "//xla/pjrt/c:pjrt_c_api_hdrs",
         "//xla/pjrt/c:pjrt_c_api_helpers",
         "//xla/service:call_inliner",
         "//xla/service:custom_call_sharding_helper",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service/spmd:spmd_partitioner",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -1160,6 +1160,7 @@ cc_library(
         "//xla/ffi/api:c_api",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_module_group",
+        "//xla/hlo/pass:hlo_pass",
         "//xla/pjrt:exceptions",
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt:status_casters",
@@ -1171,7 +1172,6 @@ cc_library(
         "//xla/service:hlo_graph_dumper",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_parser",
-        "//xla/service:hlo_pass",
         "//xla/service:hlo_proto_cc",
         "//xla/service:name_uniquer",
         "//xla/service:tuple_simplifier",

diff --git a/xla/python/custom_partition_callback.cc b/xla/python/custom_partition_callback.cc
@@ -40,12 +40,12 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/pjrt/c/pjrt_c_api_custom_partitioner_extension.h"
 #include "xla/pjrt/c/pjrt_c_api_helpers.h"
 #include "xla/pjrt/mlir_to_hlo.h"
 #include "xla/service/call_inliner.h"
 #include "xla/service/custom_call_sharding_helper.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/service/spmd/spmd_partitioner_util.h"
 #include "xla/util.h"
 

diff --git a/xla/python/xla_compiler.cc b/xla/python/xla_compiler.cc
@@ -55,6 +55,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_module_group.h"
 #include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
 #include "xla/literal.h"
@@ -75,7 +76,6 @@ limitations under the License.
 #include "xla/service/hlo_graph_dumper.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_parser.h"
-#include "xla/service/hlo_pass_interface.h"
 #include "xla/service/name_uniquer.h"
 #include "xla/service/tuple_simplifier.h"
 #include "xla/shape.h"

diff --git a/xla/service/gpu/fusions/triton/BUILD b/xla/service/gpu/fusions/triton/BUILD
@@ -294,9 +294,9 @@ cc_library(
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/hlo/utils:hlo_query",
         "//xla/service:float_normalization",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:gpu_device_info_for_tests",
         "//xla/service/gpu:gpu_float_support",

diff --git a/xla/service/gpu/fusions/triton/triton_test_utils.cc b/xla/service/gpu/fusions/triton/triton_test_utils.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/hlo/utils/hlo_query.h"
 #include "xla/primitive_util.h"
 #include "xla/service/float_normalization.h"
@@ -47,7 +48,6 @@ limitations under the License.
 #include "xla/service/gpu/gpu_float_support.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/model/tiled_hlo_computation.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/status_macros.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/filecheck.h"

diff --git a/xla/service/gpu/model/BUILD b/xla/service/gpu/model/BUILD
@@ -107,9 +107,9 @@ cc_library(
         ":gpu_performance_model",
         ":gpu_performance_model_base",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
         "//xla/service:hlo_cost_analysis",
         "//xla/service:hlo_graph_dumper",
-        "//xla/service:hlo_pass",
         "//xla/stream_executor:device_description",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status:statusor",

diff --git a/xla/service/gpu/model/gpu_cost_model_stats_collection.h b/xla/service/gpu/model/gpu_cost_model_stats_collection.h
@@ -21,9 +21,9 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
 #include "xla/service/hlo_cost_analysis.h"
-#include "xla/service/hlo_pass_interface.h"
 #include "xla/stream_executor/device_description.h"
 
 namespace xla {

diff --git a/xla/service/gpu/tests/BUILD b/xla/service/gpu/tests/BUILD
@@ -498,8 +498,8 @@ xla_test(
     deps = [
         ":gpu_codegen_test",
         "//xla:shape_util",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/service:hlo_cost_analysis",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service/gpu:gpu_device_info_for_tests",
         "//xla/service/gpu/transforms:fusion_merger",
         "//xla/service/gpu/transforms:instruction_fusion",

diff --git a/xla/service/gpu/tests/gpu_fusion_pipeline_test.cc b/xla/service/gpu/tests/gpu_fusion_pipeline_test.cc
@@ -18,13 +18,13 @@ limitations under the License.
 #include <utility>
 
 #include "absl/strings/string_view.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/service/gpu/tests/gpu_codegen_test.h"
 #include "xla/service/gpu/transforms/fusion_merger.h"
 #include "xla/service/gpu/transforms/instruction_fusion.h"
 #include "xla/service/gpu/transforms/multi_output_fusion.h"
 #include "xla/service/hlo_cost_analysis.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_description.h"

diff --git a/xla/tests/BUILD b/xla/tests/BUILD
@@ -2388,10 +2388,10 @@ xla_test(
         "//xla:error_spec",
         "//xla:util",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/service:collective_pipeliner",
         "//xla/service:hlo_dce",
         "//xla/service:hlo_parser",
-        "//xla/service:hlo_pass_pipeline",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",

diff --git a/xla/tests/collective_pipeliner_execution_test.cc b/xla/tests/collective_pipeliner_execution_test.cc
@@ -26,10 +26,10 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/service/collective_pipeliner.h"
 #include "xla/service/hlo_dce.h"
 #include "xla/service/hlo_parser.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/util.h"
 

diff --git a/xla/tools/BUILD b/xla/tools/BUILD
@@ -328,7 +328,7 @@ cc_library(
         ":hlo_module_loader",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
-        "//xla/service:hlo_pass_pipeline",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/tsl/util:command_line_flags",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
@@ -345,9 +345,9 @@ cc_library(
     hdrs = ["hlo_expand.h"],
     deps = [
         "//xla:xla_data_proto_cc",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/service:batchnorm_expander",
         "//xla/service:cholesky_expander",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service:hlo_proto_cc",
         "//xla/service:hlo_verifier",
         "//xla/service:rng_bit_generator_expander",
@@ -684,10 +684,10 @@ cc_library(
         "//xla:shape_util",
         "//xla:util",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
         "//xla/service:call_graph",
         "//xla/service:collective_ops_utils",
         "//xla/service:hlo_dce",
-        "//xla/service:hlo_pass",
         "//xla/service:tuple_util",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",

diff --git a/xla/tools/hlo_control_flow_flattening.h b/xla/tools/hlo_control_flow_flattening.h
@@ -26,8 +26,8 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
 #include "xla/service/call_graph.h"
-#include "xla/service/hlo_pass_interface.h"
 
 namespace xla {
 

diff --git a/xla/tools/hlo_expand.cc b/xla/tools/hlo_expand.cc
@@ -17,10 +17,10 @@ limitations under the License.
 
 #include <vector>
 
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/service/batchnorm_expander.h"
 #include "xla/service/cholesky_expander.h"
 #include "xla/service/hlo.pb.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/service/hlo_verifier.h"
 #include "xla/service/rng_bit_generator_expander.h"
 #include "xla/service/rng_expander.h"

diff --git a/xla/tools/hlo_expand.h b/xla/tools/hlo_expand.h
@@ -19,7 +19,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
-#include "xla/service/hlo_pass_pipeline.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/tsl/util/command_line_flags.h"
 
 namespace xla {

diff --git a/xla/tools/hlo_expand_main.cc b/xla/tools/hlo_expand_main.cc
@@ -22,7 +22,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/service/hlo_pass_pipeline.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/tools/hlo_expand.h"
 #include "xla/tools/hlo_module_loader.h"
 #include "xla/tsl/util/command_line_flags.h"

diff --git a/xla/tools/multihost_hlo_runner/BUILD b/xla/tools/multihost_hlo_runner/BUILD
@@ -137,6 +137,7 @@ cc_library(
         "//xla/client:executable_build_options",
         "//xla/client:xla_computation",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass_pipeline",
         "//xla/pjrt:host_memory_spaces",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_compiler",
@@ -147,7 +148,6 @@ cc_library(
         "//xla/service:computation_placer_hdr",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_parser",
-        "//xla/service:hlo_pass_pipeline",
         "//xla/service:hlo_proto_cc",
         "//xla/tests:test_utils",
         "//xla/tools:hlo_control_flow_flattening",

diff --git a/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_input_output_alias_config.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/layout.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
@@ -57,7 +58,6 @@ limitations under the License.
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_parser.h"
-#include "xla/service/hlo_pass_pipeline.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/tests/test_utils.h"