[XLA:GPU][Mlir-based emitters] Split MlirReductionFusion into MlirRow…

…ReductionFusion and MlirColumnReductionFusion. PiperOrigin-RevId: 642556476
tensorflow · Jun 12, 2024 · 5445fac · 5445fac
1 parent c53d3a5
commit 5445fac
Show file tree

Hide file tree

Showing 5 changed files with 917 additions and 625 deletions.
diff --git a/third_party/xla/xla/service/gpu/fusions/BUILD b/third_party/xla/xla/service/gpu/fusions/BUILD
@@ -839,6 +839,7 @@ cc_library(
         "//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",
         "//xla/service/gpu/fusions/mlir:computation_partitioner",
         "//xla/service/gpu/fusions/mlir:elemental_hlo_to_mlir",
@@ -849,6 +850,7 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@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",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",

diff --git a/third_party/xla/xla/service/gpu/fusions/fusions.cc b/third_party/xla/xla/service/gpu/fusions/fusions.cc
@@ -176,7 +176,7 @@ std::unique_ptr<FusionInterface> GetFusionEmitter(const FusionInfo& fusion_info,
     }
     case HloFusionAnalysis::EmitterFusionKind::kReduction:
       if (check_mlir_emitters(/*required_level=*/4)) {
-        return std::make_unique<MlirReductionFusion>(analysis);
+        return CreateMlirReductionFusion(analysis);
       }
       return std::make_unique<ReductionFusion>(analysis);
     case HloFusionAnalysis::EmitterFusionKind::kScatter: {