[Mlir-commits] [mlir] [mlir][linalg] Account for preserved outs in fusion bounds (PR #228467)

[Ziliang Zhang]

https://github.com/ziliangzl created https://github.com/llvm/llvm-project/pull/228467

When checking reduction fusion, include preserved producer outs in loop-bound coverage. This allows fusion when a missing loop dimension is defined by an out operand that remains on the fused op.

>From d5d1772dbf9f50e4579a3d31359377732f3124cb Mon Sep 17 00:00:00 2001
From: Ziliang Zhang <zzl.coding at gmail.com>
Date: Fri, 2 Oct 2026 12:20:11 +0800
Subject: [PATCH] [mlir][linalg] Account for preserved outs in fusion bounds

When checking reduction fusion, include preserved producer outs in loop-bound coverage. This allows fusion when a missing loop dimension is defined by an out operand that remains on the fused op.
---
 .../Linalg/Transforms/ElementwiseOpFusion.cpp | 15 +++++-
 .../Linalg/fusion-elementwise-ops.mlir        | 54 ++++++++++++++++---
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 8add883163fe95c..aba1ba484d2bbd5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -183,7 +183,7 @@ bool mlir::linalg::areElementwiseOpsFusable(OpOperand *fusedOperand) {
   // Ensure that the fusion does not remove size information required to
   // get the loop bounds. For non-reduction generics, this is trivially the
   // case due to the output operand. For reductions, we need to check that after
-  // the fusion, each loop dimension has at least one input that defines it.
+  // the fusion, each loop dimension has at least one operand that defines it.
   if ((consumer.getNumReductionLoops())) {
     BitVector coveredDims(consumer.getNumLoops(), false);
 
@@ -208,6 +208,19 @@ bool mlir::linalg::areElementwiseOpsFusable(OpOperand *fusedOperand) {
               operand, producerResultIndexMap, consumerIndexMap);
       addToCoveredDims(newIndexingMap);
     }
+
+    // Preserved producer outputs are also part of the fused op and may provide
+    // loop bounds that are not available from the producer inputs.
+    llvm::SmallDenseSet<int> preservedProducerResults =
+        getPreservedProducerResults(producer, consumer, fusedOperand);
+    for (const auto &operand : llvm::enumerate(producer.getDpsInitsMutable())) {
+      if (!preservedProducerResults.contains(operand.index()))
+        continue;
+      AffineMap newIndexingMap =
+          getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(
+              &operand.value(), producerResultIndexMap, consumerIndexMap);
+      addToCoveredDims(newIndexingMap);
+    }
     if (!coveredDims.all())
       return false;
   }
diff --git a/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir b/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
index 4f0d4bb0d8f5dc8..e0264bbb6b09aad 100644
--- a/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
+++ b/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
@@ -777,20 +777,22 @@ func.func @fuse_scalar_constant(%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, te
 
 // -----
 
-// Fusing the broadcast into a reduction would require to insert extra knowledge
-// about the size of the reduction dimension. As long, as this is not
-// implemented, we check that two linalg operations remain.
-// TODO: Support this case in element-wise fusion.
+// The producer output is preserved because its indexing map is needed to
+// compute the reduction dimension's bound after fusion.
 
 #map0 = affine_map<(d0, d1) -> ()>
 #map1 = affine_map<(d0, d1) -> (d0, d1)>
 #map2 = affine_map<(d0, d1) -> (d1, d0)>
 #map3 = affine_map<(d0, d1) -> (d0)>
 
-// CHECK-LABEL: @no_fusion_missing_reduction_shape
-// CHECK: linalg.generic
-// CHECK: linalg.generic
-func.func @no_fusion_missing_reduction_shape(%arg0: tensor<f32>, %arg1: index) -> tensor<?xf32> {
+// CHECK-LABEL: @fusion_reduction_shape_from_preserved_out
+// CHECK: %[[FUSED:.+]]:2 = linalg.generic
+// CHECK-SAME: ins(%{{.+}} : tensor<f32>)
+// CHECK-SAME: outs(%{{.+}}, %{{.+}} : tensor<?x?xf32>, tensor<?xf32>)
+// CHECK-NOT: linalg.generic
+// CHECK: return %[[FUSED]]#1
+func.func @fusion_reduction_shape_from_preserved_out(
+    %arg0: tensor<f32>, %arg1: index) -> tensor<?xf32> {
   %cst = arith.constant 0xFF800000 : f32
   %4 = tensor.empty(%arg1, %arg1) : tensor<?x?xf32>
   %5 = linalg.generic {
@@ -815,6 +817,42 @@ func.func @no_fusion_missing_reduction_shape(%arg0: tensor<f32>, %arg1: index) -
 
 // -----
 
+#map0 = affine_map<(d0, d1) -> (d0)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-LABEL: @fusion_reduction_shape_from_used_producer_out
+// CHECK-SAME: %[[ARG0:.+]]: tensor<4xf32>,
+// CHECK-SAME: %[[ARG1:.+]]: tensor<4x8xf32>,
+// CHECK-SAME: %[[ARG2:.+]]: tensor<4xf32>
+// CHECK: %[[FUSED:.+]]:2 = linalg.generic
+// CHECK-SAME: ins(%[[ARG0]] : tensor<4xf32>)
+// CHECK-SAME: outs(%[[ARG1]], %[[ARG2]] : tensor<4x8xf32>, tensor<4xf32>)
+// CHECK-NOT: linalg.generic
+// CHECK: return %[[FUSED]]#1
+func.func @fusion_reduction_shape_from_used_producer_out(
+    %arg0: tensor<4xf32>, %arg1: tensor<4x8xf32>,
+    %arg2: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = linalg.generic {
+      indexing_maps = [#map0, #map1],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%arg0 : tensor<4xf32>) outs(%arg1 : tensor<4x8xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %2 = arith.addf %in, %out : f32
+      linalg.yield %2 : f32
+  } -> tensor<4x8xf32>
+  %1 = linalg.generic {
+      indexing_maps = [#map1, #map0],
+      iterator_types = ["parallel", "reduction"]}
+      ins(%0 : tensor<4x8xf32>) outs(%arg2 : tensor<4xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %2 = arith.addf %in, %out : f32
+      linalg.yield %2 : f32
+  } -> tensor<4xf32>
+  return %1 : tensor<4xf32>
+}
+
+// -----
+
 func.func @fusion_different_axes(%arg0 : tensor<5000xi64>, %arg1 : tensor<5000xi32>) -> tensor<5000xi32> {
   %c1_i32 = arith.constant 1 : i32
   %0 = linalg.generic {