diff --git a/xla/hlo/utils/BUILD b/xla/hlo/utils/BUILD index d17c81d526b1e..5547a0fe6a278 100644 --- a/xla/hlo/utils/BUILD +++ b/xla/hlo/utils/BUILD @@ -5,6 +5,7 @@ load( "//xla:xla.bzl", "xla_cc_test", ) +load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable") load("//xla/tsl/platform:rules_cc.bzl", "cc_library") package( @@ -194,3 +195,36 @@ xla_cc_test( "@tsl//tsl/platform:statusor", ], ) + +cc_library( + name = "hlo_traversal", + srcs = ["hlo_traversal.cc"], + hdrs = ["hlo_traversal.h"], + compatible_with = get_compatible_with_portable(), + deps = [ + "//xla:shape_util", + "//xla/hlo/ir:hlo", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/types:span", + ], +) + +xla_cc_test( + name = "hlo_traversal_test", + srcs = ["hlo_traversal_test.cc"], + deps = [ + ":hlo_traversal", + "//xla/hlo/ir:hlo", + "//xla/service:pattern_matcher", + "//xla/service:pattern_matcher_gmock", + "//xla/tests:hlo_test_base", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/xla/service/gpu/hlo_traversal.cc b/xla/hlo/utils/hlo_traversal.cc similarity index 99% rename from xla/service/gpu/hlo_traversal.cc rename to xla/hlo/utils/hlo_traversal.cc index 30f0517d352aa..cf05f92b8b194 100644 --- a/xla/service/gpu/hlo_traversal.cc +++ b/xla/hlo/utils/hlo_traversal.cc @@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include #include @@ -36,7 +36,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" namespace xla { -namespace gpu { namespace { template @@ -686,5 +685,4 @@ std::vector HloFindUseChain(HloInstructionAdaptor parent, return result; } -} // namespace gpu } // namespace xla diff --git a/xla/service/gpu/hlo_traversal.h b/xla/hlo/utils/hlo_traversal.h similarity index 98% rename from xla/service/gpu/hlo_traversal.h rename to xla/hlo/utils/hlo_traversal.h index 7f42775948485..71a8fe5c725ce 100644 --- a/xla/service/gpu/hlo_traversal.h +++ b/xla/hlo/utils/hlo_traversal.h @@ -12,8 +12,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_SERVICE_GPU_HLO_TRAVERSAL_H_ -#define XLA_SERVICE_GPU_HLO_TRAVERSAL_H_ +#ifndef XLA_HLO_UTILS_HLO_TRAVERSAL_H_ +#define XLA_HLO_UTILS_HLO_TRAVERSAL_H_ #include #include @@ -30,7 +30,6 @@ limitations under the License. #include "xla/shape.h" namespace xla { -namespace gpu { class HloFusionAdaptor; @@ -225,7 +224,6 @@ bool HloAnyOf(const HloFusionAdaptor& fusion, Pred&& pred) { std::vector HloFindUseChain(HloInstructionAdaptor parent, HloInstructionAdaptor root); -} // namespace gpu } // namespace xla -#endif // XLA_SERVICE_GPU_HLO_TRAVERSAL_H_ +#endif // XLA_HLO_UTILS_HLO_TRAVERSAL_H_ diff --git a/xla/service/gpu/hlo_traversal_test.cc b/xla/hlo/utils/hlo_traversal_test.cc similarity index 99% rename from xla/service/gpu/hlo_traversal_test.cc rename to xla/hlo/utils/hlo_traversal_test.cc index fcab8f47a3100..00aa9b8c3064c 100644 --- a/xla/service/gpu/hlo_traversal_test.cc +++ b/xla/hlo/utils/hlo_traversal_test.cc @@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include #include @@ -29,7 +29,6 @@ limitations under the License. #include "xla/tests/hlo_test_base.h" namespace xla { -namespace gpu { namespace { namespace m = ::xla::match; @@ -707,5 +706,4 @@ TEST_F(HloTraversalTest, HloFindUseChain) { } } // namespace -} // namespace gpu } // namespace xla diff --git a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD index c4bb8f8d6a2bc..488f54b927f8f 100644 --- a/xla/service/gpu/BUILD +++ b/xla/service/gpu/BUILD @@ -637,7 +637,6 @@ cc_library( hdrs = ["ir_emission_utils.h"], compatible_with = get_compatible_with_portable(), deps = [ - ":hlo_traversal", ":target_util", "//xla:literal", "//xla:shape_util", @@ -646,6 +645,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:backend_config", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:buffer_assignment", "//xla/service/llvm_ir:llvm_type_conversion_util", "//xla/service/llvm_ir:llvm_util", @@ -672,7 +672,6 @@ xla_cc_test( srcs = ["ir_emission_utils_test.cc"], deps = [ ":backend_configs_cc", - ":hlo_traversal", ":ir_emission_utils", "//xla:literal", "//xla:literal_util", @@ -680,6 +679,7 @@ xla_cc_test( "//xla:types", "//xla/hlo/ir:backend_config", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:buffer_assignment", "//xla/tests:hlo_test_base", "//xla/tests:xla_internal_test_main", # fixdeps: keep @@ -2431,11 +2431,11 @@ cc_library( compatible_with = get_compatible_with_portable(), deps = [ ":backend_configs_cc", - ":hlo_traversal", ":ir_emission_utils", ":reduction_utils", "//xla:shape_util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/stream_executor:device_description", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:inlined_vector", @@ -2455,9 +2455,9 @@ xla_cc_test( ":backend_configs_cc", ":gpu_device_info_for_tests", ":hlo_fusion_analysis", - ":hlo_traversal", ":ir_emission_utils", "//xla:protobuf_util", + "//xla/hlo/utils:hlo_traversal", "//xla/stream_executor:device_description", "//xla/stream_executor:device_description_proto_cc", "//xla/tests:hlo_test_base", @@ -2574,7 +2574,6 @@ cc_library( deps = [ ":backend_configs_cc", ":hlo_fusion_analysis", - ":hlo_traversal", ":ir_emission_utils", ":launch_dimensions", ":reduction_utils", @@ -2583,6 +2582,7 @@ cc_library( "//xla:util", "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:instruction_fusion", "//xla/stream_executor:device_description", "@com_google_absl//absl/algorithm:container", @@ -2944,39 +2944,6 @@ cc_library( ], ) -cc_library( - name = "hlo_traversal", - srcs = ["hlo_traversal.cc"], - hdrs = ["hlo_traversal.h"], - compatible_with = get_compatible_with_portable(), - deps = [ - "//xla:shape_util", - "//xla/hlo/ir:hlo", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/container:inlined_vector", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/types:span", - ], -) - -xla_cc_test( - name = "hlo_traversal_test", - srcs = ["hlo_traversal_test.cc"], - deps = [ - ":hlo_traversal", - "//xla/hlo/ir:hlo", - "//xla/service:pattern_matcher", - "//xla/service:pattern_matcher_gmock", - "//xla/tests:hlo_test_base", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - ], -) - xla_test( name = "determinism_test", srcs = ["determinism_test.cc"], diff --git a/xla/service/gpu/autotuning/BUILD b/xla/service/gpu/autotuning/BUILD index 78468abbe2c9b..9ac46f74649b5 100644 --- a/xla/service/gpu/autotuning/BUILD +++ b/xla/service/gpu/autotuning/BUILD @@ -129,6 +129,7 @@ cc_library( "//xla/hlo/pass:hlo_pass_pipeline", "//xla/hlo/transforms:float_normalization", "//xla/hlo/utils:hlo_query", + "//xla/hlo/utils:hlo_traversal", "//xla/pjrt/distributed:key_value_store_interface", "//xla/service:algorithm_util", "//xla/service:call_inliner", @@ -141,7 +142,6 @@ cc_library( "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:buffer_comparator", "//xla/service/gpu:gpu_float_support", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:matmul_indexing_utils", "//xla/service/gpu:matmul_utils", diff --git a/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc b/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc index 8fafd3fb82faf..d4829237f9063 100644 --- a/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc +++ b/xla/service/gpu/autotuning/gemm_fusion_autotuner.cc @@ -53,6 +53,7 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_pipeline.h" #include "xla/hlo/transforms/simplifiers/float_normalization.h" #include "xla/hlo/utils/hlo_query.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/pjrt/distributed/key_value_store_interface.h" #include "xla/primitive_util.h" #include "xla/service/algorithm_util.h" @@ -65,7 +66,6 @@ limitations under the License. #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/buffer_comparator.h" #include "xla/service/gpu/gpu_float_support.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/kernels/custom_kernel.h" #include "xla/service/gpu/kernels/custom_kernel_fusion.h" diff --git a/xla/service/gpu/fusions/BUILD b/xla/service/gpu/fusions/BUILD index dd9cdfc85d317..d547af3bbfd12 100644 --- a/xla/service/gpu/fusions/BUILD +++ b/xla/service/gpu/fusions/BUILD @@ -17,9 +17,9 @@ cc_library( "//xla:shape_util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions/mlir:computation_partitioner", @@ -44,9 +44,9 @@ cc_library( ":fusion_emitter", "//xla:shape_util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:buffer_assignment", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emitter_context", "//xla/service/gpu/runtime:copy_thunk", "//xla/service/gpu/runtime:thunk", @@ -72,6 +72,7 @@ cc_library( "//xla/ffi:attribute_map", "//xla/ffi:ffi_api", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:buffer_assignment", "//xla/service:custom_call_status", "//xla/service:custom_call_target_registry", @@ -80,7 +81,6 @@ cc_library( "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:cublas_cudnn", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:ir_emitter_context", "//xla/service/gpu:kernel_arguments", @@ -222,10 +222,10 @@ cc_library( ":triton", "//xla:shape_util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:buffer_assignment", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/status:statusor", @@ -242,9 +242,9 @@ cc_library( "//xla:status_macros", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions/ir:xla_gpu", "//xla/service/gpu/fusions/mlir:computation_partitioner", @@ -338,9 +338,9 @@ cc_library( "//xla:shape_util", "//xla:status_macros", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:ir_emitter_context", "//xla/service/gpu:kernel_arguments", @@ -466,9 +466,9 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_query", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:reduction_utils", @@ -499,8 +499,8 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:reduction_utils", @@ -577,8 +577,8 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions/ir:xla_gpu", "//xla/service/gpu/fusions/mlir:computation_partitioner", diff --git a/xla/service/gpu/fusions/copy.cc b/xla/service/gpu/fusions/copy.cc index c4fd6bbb1ed71..71a1d3381db23 100644 --- a/xla/service/gpu/fusions/copy.cc +++ b/xla/service/gpu/fusions/copy.cc @@ -21,9 +21,9 @@ limitations under the License. #include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/fusions/fusion_emitter.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emitter_context.h" #include "xla/service/gpu/runtime/copy_thunk.h" #include "xla/service/gpu/runtime/thunk.h" diff --git a/xla/service/gpu/fusions/custom.cc b/xla/service/gpu/fusions/custom.cc index e416674a0e70b..ba827d77d444a 100644 --- a/xla/service/gpu/fusions/custom.cc +++ b/xla/service/gpu/fusions/custom.cc @@ -42,6 +42,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/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/service/buffer_assignment.h" #include "xla/service/custom_call_status.h" @@ -50,7 +51,6 @@ limitations under the License. #include "xla/service/gpu/cublas_cudnn.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/ir_emitter_context.h" #include "xla/service/gpu/kernel_arguments.h" diff --git a/xla/service/gpu/fusions/fusions.cc b/xla/service/gpu/fusions/fusions.cc index 8f610f3d315c9..05da9a663c9e9 100644 --- a/xla/service/gpu/fusions/fusions.cc +++ b/xla/service/gpu/fusions/fusions.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/strings/match.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/layout_util.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusions/concatenate_mlir.h" @@ -37,7 +38,6 @@ limitations under the License. #include "xla/service/gpu/fusions/transpose_mlir.h" #include "xla/service/gpu/fusions/triton.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/shape.h" #include "xla/shape_util.h" diff --git a/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h b/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h index ab1a71f1c6c8a..2f3a7fbe6173b 100644 --- a/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h +++ b/xla/service/gpu/fusions/in_place_dynamic_update_slice_mlir.h @@ -24,11 +24,11 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/mlir/computation_partitioner.h" #include "xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h" #include "xla/service/gpu/gpu_fusible.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/indexing_map.h" diff --git a/xla/service/gpu/fusions/input_slices_mlir.cc b/xla/service/gpu/fusions/input_slices_mlir.cc index ae6730de098ff..d7c90eb6a17ca 100644 --- a/xla/service/gpu/fusions/input_slices_mlir.cc +++ b/xla/service/gpu/fusions/input_slices_mlir.cc @@ -38,10 +38,10 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/ir/xla_gpu_ops.h" #include "xla/service/gpu/fusions/mlir/computation_partitioner.h" #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" diff --git a/xla/service/gpu/fusions/loop_mlir.cc b/xla/service/gpu/fusions/loop_mlir.cc index f7f760afcf223..4e50755fdd7fe 100644 --- a/xla/service/gpu/fusions/loop_mlir.cc +++ b/xla/service/gpu/fusions/loop_mlir.cc @@ -35,11 +35,11 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/ir/xla_gpu_ops.h" #include "xla/service/gpu/fusions/mlir/computation_partitioner.h" #include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" diff --git a/xla/service/gpu/fusions/mlir/BUILD b/xla/service/gpu/fusions/mlir/BUILD index 85d86a0cf03b9..5cc1e97a49119 100644 --- a/xla/service/gpu/fusions/mlir/BUILD +++ b/xla/service/gpu/fusions/mlir/BUILD @@ -69,10 +69,10 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/translate/hlo_to_mhlo:hlo_utils", + "//xla/hlo/utils:hlo_traversal", "//xla/mlir_hlo", "//xla/mlir_hlo:map_mhlo_to_scalar_op", "//xla/service:algorithm_util", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu/fusions/ir:xla_gpu", "//xla/service/gpu/model:indexing_analysis", "//xla/stream_executor:device_description", diff --git a/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h b/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h index 34fabee1ee86e..95003ff80218f 100644 --- a/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h +++ b/xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h @@ -31,8 +31,8 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/mlir/computation_partitioner.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/stream_executor/device_description.h" diff --git a/xla/service/gpu/fusions/reduction_base.cc b/xla/service/gpu/fusions/reduction_base.cc index 0d49478abf7af..574afd95fddaa 100644 --- a/xla/service/gpu/fusions/reduction_base.cc +++ b/xla/service/gpu/fusions/reduction_base.cc @@ -32,11 +32,11 @@ limitations under the License. #include "mlir/IR/AffineExpr.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/utils/hlo_query.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/primitive_util.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/gpu_fusible.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/service/gpu/reduction_utils.h" diff --git a/xla/service/gpu/fusions/reduction_mlir.cc b/xla/service/gpu/fusions/reduction_mlir.cc index 2a4ee9168fa61..1e24d857029dd 100644 --- a/xla/service/gpu/fusions/reduction_mlir.cc +++ b/xla/service/gpu/fusions/reduction_mlir.cc @@ -48,6 +48,7 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/fusions/ir/xla_gpu_ops.h" #include "xla/service/gpu/fusions/mlir/computation_partitioner.h" @@ -55,7 +56,6 @@ limitations under the License. #include "xla/service/gpu/fusions/mlir/type_util.h" #include "xla/service/gpu/fusions/reduction_base.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/indexing_analysis.h" diff --git a/xla/service/gpu/fusions/triton.cc b/xla/service/gpu/fusions/triton.cc index e8a2c421c2db8..22e809f3f40b4 100644 --- a/xla/service/gpu/fusions/triton.cc +++ b/xla/service/gpu/fusions/triton.cc @@ -37,11 +37,11 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/fusions/triton/triton_fusion_emitter.h" #include "xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/ir_emitter_context.h" #include "xla/service/gpu/kernel_arguments.h" diff --git a/xla/service/gpu/fusions/triton/BUILD b/xla/service/gpu/fusions/triton/BUILD index 55498f755e559..7a88cadc8475e 100644 --- a/xla/service/gpu/fusions/triton/BUILD +++ b/xla/service/gpu/fusions/triton/BUILD @@ -81,12 +81,12 @@ cc_library( "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/translate/hlo_to_mhlo:hlo_function_importer", + "//xla/hlo/utils:hlo_traversal", "//xla/mlir_hlo", "//xla/service:dump", "//xla/service:hlo_module_config", "//xla/service:instruction_fusion", "//xla/service/gpu:backend_configs_cc", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:matmul_utils", @@ -177,11 +177,11 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_query", + "//xla/hlo/utils:hlo_traversal", "//xla/mlir_hlo", "//xla/mlir_hlo:map_mhlo_to_scalar_op", "//xla/mlir_hlo:transformation_helpers", "//xla/service:algorithm_util", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:matmul_indexing_utils", @@ -230,8 +230,8 @@ cc_library( deps = [ "//xla:autotuning_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_module_config", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu:matmul_utils", "//xla/service/gpu:triton_fusion_analysis", @@ -258,7 +258,7 @@ xla_cc_test( "//xla:literal", "//xla:literal_util", "//xla/hlo/ir:hlo", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu/model:tiled_hlo_instruction_or_computation", "@com_google_googletest//:gtest_main", "@llvm-project//mlir:IR", @@ -642,8 +642,8 @@ xla_cc_test( deps = [ ":triton_fusion_emitter", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_device_info_for_tests", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu/model:symbolic_tile_analysis", "//xla/service/gpu/model:tiled_hlo_instruction_or_computation", "//xla/service/gpu/model:triton_emitter_constraints", diff --git a/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.cc b/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.cc index 55fa20c42ca90..62fdc9ff20b22 100644 --- a/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.cc +++ b/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.cc @@ -59,6 +59,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/utils/hlo_query.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" #include "xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h" @@ -67,7 +68,6 @@ limitations under the License. #include "xla/service/algorithm_util.h" #include "xla/service/gpu/fusions/triton/emitter_helpers.h" #include "xla/service/gpu/fusions/triton/xla_triton_ops.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/matmul_indexing_utils.h" diff --git a/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h b/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h index 75802e04e46fa..540f511ec0306 100644 --- a/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h +++ b/xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h @@ -21,7 +21,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "mlir/IR/Builders.h" #include "xla/hlo/ir/hlo_instructions.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/matmul_utils.h" #include "xla/service/gpu/model/tiled_hlo_computation.h" diff --git a/xla/service/gpu/fusions/triton/triton_fusion_emitter_mem_utils_test.cc b/xla/service/gpu/fusions/triton/triton_fusion_emitter_mem_utils_test.cc index 543a66db68a86..e570cb8a8bb7b 100644 --- a/xla/service/gpu/fusions/triton/triton_fusion_emitter_mem_utils_test.cc +++ b/xla/service/gpu/fusions/triton/triton_fusion_emitter_mem_utils_test.cc @@ -43,9 +43,9 @@ limitations under the License. #include "mlir/IR/ValueRange.h" #include "mlir/Support/LLVM.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/triton/triton_fusion_emitter.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/symbolic_tile_analysis.h" #include "xla/service/gpu/model/tiled_hlo_computation.h" #include "xla/service/gpu/model/tiled_hlo_instruction.h" diff --git a/xla/service/gpu/fusions/triton/triton_fusion_emitter_stub_test.cc b/xla/service/gpu/fusions/triton/triton_fusion_emitter_stub_test.cc index 0accfa716fae7..48c14f18c8ab9 100644 --- a/xla/service/gpu/fusions/triton/triton_fusion_emitter_stub_test.cc +++ b/xla/service/gpu/fusions/triton/triton_fusion_emitter_stub_test.cc @@ -20,11 +20,11 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "mlir/Pass/PassManager.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/gpu/fusions/triton/triton_fusion_emitter.h" #include "xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/tiled_hlo_instruction.h" #include "tsl/platform/test.h" diff --git a/xla/service/gpu/gpu_fusible.cc b/xla/service/gpu/gpu_fusible.cc index 0972e86a61b11..073f4e76753c7 100644 --- a/xla/service/gpu/gpu_fusible.cc +++ b/xla/service/gpu/gpu_fusible.cc @@ -34,10 +34,10 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/permutation_util.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/reduction_utils.h" diff --git a/xla/service/gpu/gpu_fusible.h b/xla/service/gpu/gpu_fusible.h index 08f2c0f4ac4a6..cae7ff9f2a411 100644 --- a/xla/service/gpu/gpu_fusible.h +++ b/xla/service/gpu/gpu_fusible.h @@ -27,8 +27,8 @@ limitations under the License. #include "absl/synchronization/mutex.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" diff --git a/xla/service/gpu/hlo_fusion_analysis.cc b/xla/service/gpu/hlo_fusion_analysis.cc index a6e87bb662421..9f81f459fe784 100644 --- a/xla/service/gpu/hlo_fusion_analysis.cc +++ b/xla/service/gpu/hlo_fusion_analysis.cc @@ -33,9 +33,9 @@ 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/utils/hlo_traversal.h" #include "xla/primitive_util.h" #include "xla/service/gpu/backend_configs.pb.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/reduction_utils.h" #include "xla/shape.h" diff --git a/xla/service/gpu/hlo_fusion_analysis.h b/xla/service/gpu/hlo_fusion_analysis.h index fa311e0bd8954..22aa7fb55a889 100644 --- a/xla/service/gpu/hlo_fusion_analysis.h +++ b/xla/service/gpu/hlo_fusion_analysis.h @@ -25,8 +25,8 @@ limitations under the License. #include "absl/log/check.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/stream_executor/device_description.h" diff --git a/xla/service/gpu/hlo_fusion_analysis_test.cc b/xla/service/gpu/hlo_fusion_analysis_test.cc index 26b86b6c1f7d5..d053d660ce424 100644 --- a/xla/service/gpu/hlo_fusion_analysis_test.cc +++ b/xla/service/gpu/hlo_fusion_analysis_test.cc @@ -15,10 +15,10 @@ limitations under the License. #include "xla/service/gpu/hlo_fusion_analysis.h" #include +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/protobuf_util.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/stream_executor/device_description.h" #include "xla/stream_executor/device_description.pb.h" diff --git a/xla/service/gpu/ir_emission_utils.cc b/xla/service/gpu/ir_emission_utils.cc index eccb56ae1259f..c72c05042ec73 100644 --- a/xla/service/gpu/ir_emission_utils.cc +++ b/xla/service/gpu/ir_emission_utils.cc @@ -50,10 +50,10 @@ 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/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/target_util.h" #include "xla/service/llvm_ir/llvm_type_conversion_util.h" #include "xla/service/llvm_ir/llvm_util.h" diff --git a/xla/service/gpu/ir_emission_utils.h b/xla/service/gpu/ir_emission_utils.h index a38bcef4be5f4..3432cec278bee 100644 --- a/xla/service/gpu/ir_emission_utils.h +++ b/xla/service/gpu/ir_emission_utils.h @@ -33,9 +33,9 @@ limitations under the License. #include "xla/hlo/ir/backend_config.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/util.h" diff --git a/xla/service/gpu/ir_emission_utils_test.cc b/xla/service/gpu/ir_emission_utils_test.cc index 458c723c043d1..bdc976675efdc 100644 --- a/xla/service/gpu/ir_emission_utils_test.cc +++ b/xla/service/gpu/ir_emission_utils_test.cc @@ -25,11 +25,11 @@ limitations under the License. #include "xla/hlo/ir/backend_config.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/backend_configs.pb.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/shape_util.h" #include "xla/tests/hlo_test_base.h" #include "xla/types.h" diff --git a/xla/service/gpu/model/BUILD b/xla/service/gpu/model/BUILD index c4de389c88380..27c668b03cd1d 100644 --- a/xla/service/gpu/model/BUILD +++ b/xla/service/gpu/model/BUILD @@ -198,9 +198,9 @@ cc_library( "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions", "//xla/service/gpu/fusions:fusion_emitter", @@ -307,11 +307,11 @@ cc_library( "//xla:util", "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions", "//xla/service/gpu/fusions:fusion_emitter", @@ -359,11 +359,11 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", "//xla/service:instruction_fusion", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/service/gpu/fusions:triton", @@ -396,10 +396,10 @@ xla_cc_test( "//xla:shape_util", "//xla:test_helpers", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_device_info_for_tests", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:launch_dimensions", "//xla/stream_executor:device_description", @@ -460,7 +460,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/transforms:gather_simplifier", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:matmul_indexing_utils", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", @@ -552,7 +552,7 @@ xla_cc_test( ":indexing_analysis", ":indexing_test_utils", "//xla/hlo/ir:hlo", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/tests:xla_internal_test_main", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", @@ -619,7 +619,7 @@ xla_cc_test( ":symbolic_tile", ":symbolic_tiled_hlo_instruction", "//xla/hlo/ir:hlo", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/tests:hlo_test_base", "//xla/tests:verified_hlo_module", "@com_google_googletest//:gtest_main", @@ -698,9 +698,9 @@ cc_library( "//xla:shape_util", "//xla:status_macros", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:instruction_fusion", "//xla/service:name_uniquer", - "//xla/service/gpu:hlo_traversal", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -731,8 +731,8 @@ xla_cc_test( "//xla:util", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", + "//xla/hlo/utils:hlo_traversal", "//xla/service:instruction_fusion", - "//xla/service/gpu:hlo_traversal", "//xla/tests:hlo_test_base", "//xla/tests:verified_hlo_module", "//xla/tsl/lib/core:status_test_util", @@ -760,7 +760,7 @@ cc_library( ":symbolic_tiled_hlo_instruction", "//xla:shape_util", "//xla/hlo/ir:hlo", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/stream_executor:device_description", "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", @@ -779,9 +779,9 @@ xla_cc_test( ":triton_emitter_constraints", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", + "//xla/hlo/utils:hlo_traversal", "//xla/service:instruction_fusion", "//xla/service/gpu:gpu_device_info_for_tests", - "//xla/service/gpu:hlo_traversal", "//xla/stream_executor:device_description", "//xla/tests:hlo_test_base", "//xla/tests:verified_hlo_module", @@ -805,9 +805,9 @@ cc_library( "//xla:shape_util", "//xla:util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu/fusions:fusion_emitter", "//xla/stream_executor:device_description", "@com_google_absl//absl/container:flat_hash_map", @@ -829,10 +829,10 @@ xla_cc_test( ":tiled_hlo_instruction_or_computation", "//xla:shape_util", "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_module_config", "//xla/service/gpu:gpu_device_info_for_tests", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu/fusions", "//xla/service/gpu/fusions:fusion_emitter", "//xla/stream_executor:device_description", diff --git a/xla/service/gpu/model/coalescing_analysis.cc b/xla/service/gpu/model/coalescing_analysis.cc index bb6ce98e33697..65c7efc39c36e 100644 --- a/xla/service/gpu/model/coalescing_analysis.cc +++ b/xla/service/gpu/model/coalescing_analysis.cc @@ -34,11 +34,11 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/layout.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/gpu_fusible.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/affine_map_evaluator.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" diff --git a/xla/service/gpu/model/coalescing_analysis.h b/xla/service/gpu/model/coalescing_analysis.h index ac579d030bab9..6681c19c68784 100644 --- a/xla/service/gpu/model/coalescing_analysis.h +++ b/xla/service/gpu/model/coalescing_analysis.h @@ -20,9 +20,9 @@ limitations under the License. #include "absl/types/span.h" #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/tiled_hlo_instruction.h" #include "xla/stream_executor/device_description.h" diff --git a/xla/service/gpu/model/coalescing_analysis_test.cc b/xla/service/gpu/model/coalescing_analysis_test.cc index 0a8d93c5074ec..928914835a1f5 100644 --- a/xla/service/gpu/model/coalescing_analysis_test.cc +++ b/xla/service/gpu/model/coalescing_analysis_test.cc @@ -26,11 +26,11 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/fusions/fusions.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/symbolic_tile_analysis.h" #include "xla/service/gpu/model/tiled_hlo_computation.h" #include "xla/service/gpu/model/tiled_hlo_instruction.h" diff --git a/xla/service/gpu/model/gpu_indexing_performance_model.cc b/xla/service/gpu/model/gpu_indexing_performance_model.cc index a72d0d1adef78..0d4ad97e7bee5 100644 --- a/xla/service/gpu/model/gpu_indexing_performance_model.cc +++ b/xla/service/gpu/model/gpu_indexing_performance_model.cc @@ -32,10 +32,10 @@ limitations under the License. #include "llvm/Support/MathExtras.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusions/triton.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/coalescing_analysis.h" diff --git a/xla/service/gpu/model/gpu_indexing_performance_model.h b/xla/service/gpu/model/gpu_indexing_performance_model.h index a8e7fed4314a3..e293fbbdf14df 100644 --- a/xla/service/gpu/model/gpu_indexing_performance_model.h +++ b/xla/service/gpu/model/gpu_indexing_performance_model.h @@ -25,8 +25,8 @@ limitations under the License. #include "absl/types/span.h" #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h" diff --git a/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/xla/service/gpu/model/gpu_indexing_performance_model_test.cc index 5d1321cc13873..a092e2bb9bc72 100644 --- a/xla/service/gpu/model/gpu_indexing_performance_model_test.cc +++ b/xla/service/gpu/model/gpu_indexing_performance_model_test.cc @@ -28,9 +28,9 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" diff --git a/xla/service/gpu/model/gpu_performance_model_base.cc b/xla/service/gpu/model/gpu_performance_model_base.cc index a760bf384f346..8cbf262c6f8b3 100644 --- a/xla/service/gpu/model/gpu_performance_model_base.cc +++ b/xla/service/gpu/model/gpu_performance_model_base.cc @@ -28,12 +28,12 @@ limitations under the License. #include "absl/time/time.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusions/fusion_emitter.h" #include "xla/service/gpu/fusions/fusions.h" #include "xla/service/gpu/fusions/triton.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h" #include "xla/shape_util.h" diff --git a/xla/service/gpu/model/indexing_analysis.cc b/xla/service/gpu/model/indexing_analysis.cc index e0a433211e7e9..2ea954bbdf32e 100644 --- a/xla/service/gpu/model/indexing_analysis.cc +++ b/xla/service/gpu/model/indexing_analysis.cc @@ -45,9 +45,9 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/transforms/simplifiers/gather_simplifier.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/layout.h" #include "xla/permutation_util.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/matmul_indexing_utils.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/shape.h" diff --git a/xla/service/gpu/model/indexing_analysis.h b/xla/service/gpu/model/indexing_analysis.h index e05a598cc2a32..695680f1de215 100644 --- a/xla/service/gpu/model/indexing_analysis.h +++ b/xla/service/gpu/model/indexing_analysis.h @@ -30,7 +30,7 @@ limitations under the License. #include "mlir/IR/AffineMap.h" #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/shape.h" diff --git a/xla/service/gpu/model/indexing_analysis_test.cc b/xla/service/gpu/model/indexing_analysis_test.cc index cb4aefc17e642..57f2f1794399d 100644 --- a/xla/service/gpu/model/indexing_analysis_test.cc +++ b/xla/service/gpu/model/indexing_analysis_test.cc @@ -19,7 +19,7 @@ limitations under the License. #include #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/indexing_map_serialization.h" #include "xla/service/gpu/model/indexing_test_utils.h" #include "tsl/platform/test.h" diff --git a/xla/service/gpu/model/symbolic_tile_analysis.cc b/xla/service/gpu/model/symbolic_tile_analysis.cc index 27cdf6e1adf72..6c8ae9e51c2ba 100644 --- a/xla/service/gpu/model/symbolic_tile_analysis.cc +++ b/xla/service/gpu/model/symbolic_tile_analysis.cc @@ -49,7 +49,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/service/gpu/model/indexing_map_serialization.h" diff --git a/xla/service/gpu/model/symbolic_tile_analysis.h b/xla/service/gpu/model/symbolic_tile_analysis.h index e8e0cea2fef4b..a82e56dd79c31 100644 --- a/xla/service/gpu/model/symbolic_tile_analysis.h +++ b/xla/service/gpu/model/symbolic_tile_analysis.h @@ -30,7 +30,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/symbolic_tile.h" #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h" #include "xla/service/gpu/model/tiled_hlo_computation.h" diff --git a/xla/service/gpu/model/symbolic_tile_analysis_test.cc b/xla/service/gpu/model/symbolic_tile_analysis_test.cc index 500def83da5c3..b408d96cff4ff 100644 --- a/xla/service/gpu/model/symbolic_tile_analysis_test.cc +++ b/xla/service/gpu/model/symbolic_tile_analysis_test.cc @@ -33,7 +33,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/verified_hlo_module.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/indexing_test_utils.h" #include "xla/service/gpu/model/symbolic_tile.h" #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h" diff --git a/xla/service/gpu/model/symbolic_tiled_hlo_instruction_test.cc b/xla/service/gpu/model/symbolic_tiled_hlo_instruction_test.cc index 21a270cfd7866..d8f0324417c80 100644 --- a/xla/service/gpu/model/symbolic_tiled_hlo_instruction_test.cc +++ b/xla/service/gpu/model/symbolic_tiled_hlo_instruction_test.cc @@ -23,7 +23,7 @@ limitations under the License. #include #include "mlir/IR/MLIRContext.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" #include "xla/service/gpu/model/symbolic_tile.h" diff --git a/xla/service/gpu/model/triton_emitter_constraints.cc b/xla/service/gpu/model/triton_emitter_constraints.cc index 6beffb69badad..d11b0390d69dd 100644 --- a/xla/service/gpu/model/triton_emitter_constraints.cc +++ b/xla/service/gpu/model/triton_emitter_constraints.cc @@ -32,7 +32,7 @@ limitations under the License. #include "mlir/IR/AffineMap.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/affine_map_evaluator.h" #include "xla/service/gpu/model/indexing_analysis.h" #include "xla/service/gpu/model/indexing_map.h" diff --git a/xla/service/gpu/model/triton_emitter_constraints.h b/xla/service/gpu/model/triton_emitter_constraints.h index 3c36ba14d24f8..c8f1356c398d6 100644 --- a/xla/service/gpu/model/triton_emitter_constraints.h +++ b/xla/service/gpu/model/triton_emitter_constraints.h @@ -22,7 +22,7 @@ limitations under the License. #include "absl/types/span.h" #include "llvm/ADT/SmallVector.h" #include "mlir/IR/AffineMap.h" -#include "xla/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/model/symbolic_tile.h" #include "xla/service/gpu/model/symbolic_tile_analysis.h" #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h" diff --git a/xla/service/gpu/model/triton_emitter_constraints_test.cc b/xla/service/gpu/model/triton_emitter_constraints_test.cc index c4dc4a6117346..c3b7e65150778 100644 --- a/xla/service/gpu/model/triton_emitter_constraints_test.cc +++ b/xla/service/gpu/model/triton_emitter_constraints_test.cc @@ -27,8 +27,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/model/symbolic_tile_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" diff --git a/xla/service/gpu/transforms/BUILD b/xla/service/gpu/transforms/BUILD index a521a94d61970..ceda991ef7706 100644 --- a/xla/service/gpu/transforms/BUILD +++ b/xla/service/gpu/transforms/BUILD @@ -379,10 +379,10 @@ cc_library( deps = [ "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", "//xla/service:instruction_fusion", "//xla/service/gpu:backend_configs_cc", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu/fusions/triton:triton_support", "//xla/service/gpu/model:fusion_analysis_cache", @@ -760,8 +760,8 @@ cc_library( deps = [ "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:gpu_fusible", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:reduction_utils", "//xla/stream_executor:device_description", @@ -1587,12 +1587,12 @@ cc_library( "//xla/ffi:ffi_api", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/hlo/utils:hlo_traversal", "//xla/service:custom_call_target_registry", "//xla/service:pattern_matcher", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:cublas_cudnn", "//xla/service/gpu:gpu_constants", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu/kernels:custom_fusion_library", "@com_google_absl//absl/algorithm:container", @@ -2334,6 +2334,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/hlo/utils:hlo_traversal", "//xla/service:dump", "//xla/service:hlo_cost_analysis", "//xla/service:hlo_graph_dumper", @@ -2343,7 +2344,6 @@ cc_library( "//xla/service/gpu:fusion_process_dump_proto_cc", "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:hlo_fusion_analysis", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu/fusions/triton:triton_support", "//xla/service/gpu/model:fusion_analysis_cache", @@ -2595,7 +2595,7 @@ cc_library( deps = [ "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", - "//xla/service/gpu:hlo_traversal", + "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_set", @@ -2779,11 +2779,11 @@ cc_library( "//xla/hlo/pass:hlo_pass", "//xla/hlo/pass:hlo_pass_pipeline", "//xla/hlo/utils:hlo_query", + "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", "//xla/service:instruction_fusion", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:fusion_pipeline", - "//xla/service/gpu:hlo_traversal", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu/fusions/triton:triton_support", "//xla/service/gpu/model:fusion_analysis_cache", diff --git a/xla/service/gpu/transforms/copy_fusion.cc b/xla/service/gpu/transforms/copy_fusion.cc index 9912316f5e945..c7e6278ffd877 100644 --- a/xla/service/gpu/transforms/copy_fusion.cc +++ b/xla/service/gpu/transforms/copy_fusion.cc @@ -24,8 +24,8 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/gpu_fusible.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/reduction_utils.h" #include "tsl/platform/errors.h" diff --git a/xla/service/gpu/transforms/dynamic_slice_fusion_rewriter.cc b/xla/service/gpu/transforms/dynamic_slice_fusion_rewriter.cc index ece36dbc81fb1..3fc07a0ac980f 100644 --- a/xla/service/gpu/transforms/dynamic_slice_fusion_rewriter.cc +++ b/xla/service/gpu/transforms/dynamic_slice_fusion_rewriter.cc @@ -38,11 +38,11 @@ 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/utils/hlo_traversal.h" #include "xla/service/custom_call_target_registry.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/cublas_cudnn.h" #include "xla/service/gpu/gpu_constants.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/pattern_matcher.h" #include "xla/shape.h" diff --git a/xla/service/gpu/transforms/fusion_block_level_rewriter.cc b/xla/service/gpu/transforms/fusion_block_level_rewriter.cc index 809183953d306..57a54a0436290 100644 --- a/xla/service/gpu/transforms/fusion_block_level_rewriter.cc +++ b/xla/service/gpu/transforms/fusion_block_level_rewriter.cc @@ -32,9 +32,9 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusions/triton/triton_support.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_indexing_performance_model.h" diff --git a/xla/service/gpu/transforms/priority_fusion.cc b/xla/service/gpu/transforms/priority_fusion.cc index b33a1073f662e..d1de6ae252c1d 100644 --- a/xla/service/gpu/transforms/priority_fusion.cc +++ b/xla/service/gpu/transforms/priority_fusion.cc @@ -42,6 +42,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/dump.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusion_deduplication_cache.h" @@ -49,7 +50,6 @@ limitations under the License. #include "xla/service/gpu/fusions/triton/triton_support.h" #include "xla/service/gpu/gpu_fusible.h" #include "xla/service/gpu/hlo_fusion_analysis.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h" diff --git a/xla/service/gpu/transforms/rename_fusions.cc b/xla/service/gpu/transforms/rename_fusions.cc index 9ab62f68664eb..29f3edf968fb3 100644 --- a/xla/service/gpu/transforms/rename_fusions.cc +++ b/xla/service/gpu/transforms/rename_fusions.cc @@ -29,7 +29,7 @@ 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/service/gpu/hlo_traversal.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" namespace xla { diff --git a/xla/service/gpu/transforms/softmax_rewriter_triton.cc b/xla/service/gpu/transforms/softmax_rewriter_triton.cc index 3350124e64c1e..c9c536c5752e7 100644 --- a/xla/service/gpu/transforms/softmax_rewriter_triton.cc +++ b/xla/service/gpu/transforms/softmax_rewriter_triton.cc @@ -40,11 +40,11 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_fix.h" #include "xla/hlo/pass/hlo_pass_pipeline.h" #include "xla/hlo/utils/hlo_query.h" +#include "xla/hlo/utils/hlo_traversal.h" #include "xla/layout_util.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/fusion_pipeline.h" #include "xla/service/gpu/fusions/triton/triton_support.h" -#include "xla/service/gpu/hlo_traversal.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"