Skip to content

Commit

Permalink
[HLO Componentization] Create pass sub-component
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 674391572
  • Loading branch information
sdasgup3 authored and Google-ML-Automation committed Sep 18, 2024
1 parent 409ecf9 commit 590c0c6
Show file tree
Hide file tree
Showing 18 changed files with 37 additions and 37 deletions.
30 changes: 15 additions & 15 deletions xla/service/gpu/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -1305,12 +1305,12 @@ cc_library(
hdrs = ["fusion_pipeline.h"],
deps = [
"//xla:xla_proto_cc",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:cpu_gpu_shape_verifier",
"//xla/service:hlo_cost_analysis",
"//xla/service:hlo_cse",
"//xla/service:hlo_dce",
"//xla/service:hlo_pass",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_verifier",
"//xla/service:layout_assignment",
"//xla/service/gpu/model:gpu_hlo_cost_analysis",
Expand All @@ -1333,11 +1333,11 @@ cc_library(
deps = [
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:copy_insertion",
"//xla/service:cpu_gpu_shape_verifier",
"//xla/service:hlo_dataflow_analysis",
"//xla/service:hlo_dce",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_verifier",
"//xla/service:layout_assignment",
"//xla/service:loop_schedule_linearizer",
Expand Down Expand Up @@ -1407,6 +1407,8 @@ cc_library(
"@llvm-project//mlir:Support",
"//xla/hlo/ir:hlo",
"//xla/hlo/ir:hlo_module_group",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/pjrt/distributed:key_value_store_interface",
"//xla/service/gpu/autotuning:autotuner_util",
"//xla/service/gpu/autotuning:custom_kernel_fusion_autotuner",
Expand Down Expand Up @@ -1515,8 +1517,6 @@ cc_library(
"//xla/service:hlo_dce",
"//xla/service:hlo_module_config",
"//xla/service:hlo_ordering",
"//xla/service:hlo_pass",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_proto_cc",
"//xla/service:hlo_rematerialization",
"//xla/service:hlo_verifier",
Expand Down Expand Up @@ -1735,6 +1735,8 @@ cc_library(
"//xla:xla_data_proto_cc",
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/pjrt/distributed:key_value_store_interface",
"//xla/service:algebraic_simplifier",
"//xla/service:call_inliner",
Expand All @@ -1748,8 +1750,6 @@ cc_library(
"//xla/service:hlo_dataflow_analysis",
"//xla/service:hlo_dce",
"//xla/service:hlo_module_config",
"//xla/service:hlo_pass",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_verifier",
"//xla/service:reshape_mover",
"//xla/service:tuple_simplifier",
Expand Down Expand Up @@ -1981,6 +1981,8 @@ cc_library(
"//xla:util",
"//xla:xla_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:algebraic_simplifier",
"//xla/service:call_inliner",
"//xla/service:convert_mover",
Expand All @@ -1989,8 +1991,6 @@ cc_library(
"//xla/service:float_support",
"//xla/service:hlo_constant_folding",
"//xla/service:hlo_module_config",
"//xla/service:hlo_pass",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_verifier",
"//xla/service:reshape_mover",
"//xla/service:tuple_simplifier",
Expand Down Expand Up @@ -2078,11 +2078,11 @@ cc_library(
"//xla:shape_util",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/hlo/utils:hlo_query",
"//xla/service:buffer_value",
"//xla/service:collective_ops_utils",
"//xla/service:hlo_memory_scheduler",
"//xla/service:hlo_pass_pipeline",
"//xla/service:latency_hiding_scheduler",
"//xla/service:p2p_schedule_preparation",
"//xla/service:profile_guided_latency_estimator",
Expand Down Expand Up @@ -2149,10 +2149,10 @@ cc_library(
deps = [
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:collective_ops_utils",
"//xla/service:collective_pipeliner",
"//xla/service:hlo_parser",
"//xla/service:hlo_pass_pipeline",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
Expand All @@ -2168,9 +2168,9 @@ xla_cc_test(
":gpu_p2p_pipeliner",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:hlo_module_config",
"//xla/service:hlo_parser",
"//xla/service:hlo_pass_pipeline",
"//xla/service:hlo_verifier",
"//xla/tests:filecheck",
"//xla/tests:hlo_test_base",
Expand All @@ -2188,15 +2188,15 @@ cc_library(
deps = [
":runtime_intrinsics",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/hlo/transforms:hlo_constant_splitter",
"//xla/service:algebraic_simplifier",
"//xla/service:conditional_simplifier",
"//xla/service:gather_expander",
"//xla/service:hlo_constant_folding",
"//xla/service:hlo_dce",
"//xla/service:hlo_module_config",
"//xla/service:hlo_pass",
"//xla/service:hlo_pass_pipeline",
"//xla/service:reshape_mover",
"//xla/service:scatter_expander",
"//xla/service:sharding_propagation",
Expand Down Expand Up @@ -2225,10 +2225,10 @@ xla_cc_test(
"//xla:util",
"//xla/client:executable_build_options",
"//xla/hlo/ir:hlo",
"//xla/hlo/pass:hlo_pass_pipeline",
"//xla/service:algebraic_simplifier",
"//xla/service:hlo_module_config",
"//xla/service:hlo_parser",
"//xla/service:hlo_pass_pipeline",
"//xla/stream_executor:device_description",
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
Expand Down
4 changes: 2 additions & 2 deletions xla/service/gpu/amdgpu_compiler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ 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_fix.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/call_inliner.h"
#include "xla/service/convert_mover.h"
Expand All @@ -49,8 +51,6 @@ limitations under the License.
#include "xla/service/gpu/transforms/triangular_solve_rewriter.h"
#include "xla/service/hlo_constant_folding.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_fix.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_verifier.h"
#include "xla/service/reshape_mover.h"
#include "xla/service/tuple_simplifier.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/amdgpu_compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,10 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "llvm/IR/Module.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/gpu/autotuning/autotuner_util.h"
#include "xla/service/gpu/gpu_compiler.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_memory_allocator.h"
#include "xla/stream_executor/dnn.h"
Expand Down
4 changes: 2 additions & 2 deletions xla/service/gpu/fusion_pipeline.cc
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ limitations under the License.
#include <memory>
#include <utility>

#include "xla/hlo/pass/hlo_pass_fix.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/cpu_gpu_shape_verifier.h"
#include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
#include "xla/service/gpu/transforms/fusion_merger.h"
Expand All @@ -30,8 +32,6 @@ limitations under the License.
#include "xla/service/hlo_cost_analysis.h"
#include "xla/service/hlo_cse.h"
#include "xla/service/hlo_dce.h"
#include "xla/service/hlo_pass_fix.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_verifier.h"
#include "xla/service/layout_assignment.h"
#include "xla/stream_executor/device_description.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/fusion_pipeline.h
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@ limitations under the License.
#ifndef XLA_SERVICE_GPU_FUSION_PIPELINE_H_
#define XLA_SERVICE_GPU_FUSION_PIPELINE_H_

#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/hlo_cost_analysis.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/stream_executor/device_description.h"
#include "xla/xla.pb.h"
#include "tsl/platform/threadpool.h"
Expand Down
4 changes: 2 additions & 2 deletions xla/service/gpu/gpu_compiler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ limitations under the License.
#include "xla/hlo/ir/hlo_module_group.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/hlo/ir/hlo_schedule.h"
#include "xla/hlo/pass/hlo_pass_fix.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/maybe_owning.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/all_gather_broadcast_reorder.h"
Expand Down Expand Up @@ -193,8 +195,6 @@ limitations under the License.
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_dce.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_fix.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_rematerialization.h"
#include "xla/service/hlo_verifier.h"
#include "xla/service/host_memory_transfer_asyncifier.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ limitations under the License.
#include "xla/autotune_results.pb.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/pjrt/distributed/key_value_store_interface.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/compiler.h"
Expand All @@ -41,7 +42,6 @@ limitations under the License.
#include "xla/service/hlo_cost_analysis.h"
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/llvm_compiler.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_description.pb.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_hlo_schedule.cc
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ limitations under the License.
#include "xla/hlo/ir/hlo_instructions.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/hlo/ir/hlo_schedule.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/hlo/utils/hlo_query.h"
#include "xla/service/buffer_value.h"
#include "xla/service/collective_ops_utils.h"
Expand All @@ -51,7 +52,6 @@ limitations under the License.
#include "xla/service/gpu/transforms/schedule_postprocessing.h"
#include "xla/service/gpu/transforms/scheduling_instruction_annotator.h"
#include "xla/service/hlo_memory_scheduler.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/latency_hiding_scheduler.h"
#include "xla/service/p2p_schedule_preparation.h"
#include "xla/service/profile_guided_latency_estimator.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_p2p_pipeliner.cc
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,10 @@ limitations under the License.
#include "absl/strings/str_join.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/collective_ops_utils.h"
#include "xla/service/collective_pipeliner.h"
#include "xla/service/hlo_parser.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/util.h"

namespace xla {
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_p2p_pipeliner.h
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ limitations under the License.
#ifndef XLA_SERVICE_GPU_GPU_P2P_PIPELINER_H_
#define XLA_SERVICE_GPU_GPU_P2P_PIPELINER_H_

#include "xla/service/hlo_pass_pipeline.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"

namespace xla {
namespace gpu {
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_p2p_pipeliner_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,9 @@ limitations under the License.
#include "xla/hlo/ir/hlo_instructions.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/hlo_module_config.h"
#include "xla/service/hlo_parser.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_verifier.h"
#include "xla/tests/filecheck.h"
#include "xla/tests/hlo_test_base.h"
Expand Down
4 changes: 2 additions & 2 deletions xla/service/gpu/gpu_spmd_pipeline.cc
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ limitations under the License.
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/ir/hlo_schedule.h"
#include "xla/hlo/pass/hlo_pass_fix.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/hlo/transforms/hlo_constant_splitter.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/conditional_simplifier.h"
Expand All @@ -31,8 +33,6 @@ limitations under the License.
#include "xla/service/hlo_constant_folding.h"
#include "xla/service/hlo_dce.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_fix.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/reshape_mover.h"
#include "xla/service/scatter_expander.h"
#include "xla/service/sharding_propagation.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_spmd_pipeline.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ limitations under the License.

#include "absl/functional/function_ref.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/hlo_pass_pipeline.h"

namespace xla {
namespace gpu {
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/gpu_spmd_pipeline_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,10 @@ limitations under the License.
#include "absl/log/log.h"
#include "xla/client/executable_build_options.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/algebraic_simplifier.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/stream_executor/device_description.h"
#include "xla/tests/hlo_test_base.h"
Expand Down
4 changes: 2 additions & 2 deletions xla/service/gpu/nvptx_compiler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ limitations under the License.
#include "llvm/Support/raw_ostream.h"
#include "xla/hlo/ir/hlo_computation.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/hlo/pass/hlo_pass_fix.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/pjrt/distributed/key_value_store_interface.h"
#include "xla/service/algebraic_simplifier.h"
#include "xla/service/call_inliner.h"
Expand Down Expand Up @@ -86,8 +88,6 @@ limitations under the License.
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_dce.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_fix.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_verifier.h"
#include "xla/service/llvm_ir/llvm_util.h"
#include "xla/service/reshape_mover.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/nvptx_compiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,13 @@ limitations under the License.
#include "llvm/IR/Module.h"
#include "xla/autotune_results.pb.h"
#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/pjrt/distributed/key_value_store_interface.h"
#include "xla/service/gpu/autotuning/autotuner_util.h"
#include "xla/service/gpu/gpu_compiler.h"
#include "xla/service/gpu/ir_emission_utils.h"
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_module_config.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/stream_executor/cuda/ptx_linking_method.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_memory_allocator.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/prepare_hlo_for_ir_emitting_pipeline.cc
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ limitations under the License.
#include <utility>

#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/copy_insertion.h"
#include "xla/service/cpu_gpu_shape_verifier.h"
#include "xla/service/gpu/transforms/alias_passthrough_params.h"
Expand All @@ -28,7 +29,6 @@ limitations under the License.
#include "xla/service/gpu/transforms/sanitize_constant_names.h"
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_dce.h"
#include "xla/service/hlo_pass_pipeline.h"
#include "xla/service/hlo_verifier.h"
#include "xla/service/layout_assignment.h"
#include "xla/service/loop_schedule_linearizer.h"
Expand Down
2 changes: 1 addition & 1 deletion xla/service/gpu/prepare_hlo_for_ir_emitting_pipeline.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@ limitations under the License.
#define XLA_SERVICE_GPU_PREPARE_HLO_FOR_IR_EMITTING_PIPELINE_H_

#include "xla/hlo/ir/hlo_module.h"
#include "xla/hlo/pass/hlo_pass_pipeline.h"
#include "xla/service/hlo_dataflow_analysis.h"
#include "xla/service/hlo_pass_pipeline.h"

namespace xla {
namespace gpu {
Expand Down

0 comments on commit 590c0c6

Please sign in to comment.