[Mlir-commits] [mlir] [MLIR][linalg] Erase effectful producer after fusion if it's no longer used (PR #170036)

[Artemiy Bulavin]

https://github.com/abulavin created https://github.com/llvm/llvm-project/pull/170036

This PR fixes a subtle bug in the `linalg-fuse-elementwise-ops` pass where some producers are not being erased after fusion.

The pass relies on the DCE done by the greedy pattern rewriter in order to erase any trivially dead producer ops that no longer have uses after fusion.
However if the producer's region contains ops with memory effects, the producer itself cannot be erased because `linalg.generic` ops have the `RecursiveMemoryEffects` trait. This makes the producer not trivially dead even when it has no uses. This causes miscompiles as the body of the 'old' producer is essentially repeated in the IR.

The fix is to erase the producer within the pattern if it no longer has uses instead of solely relying on DCE.

>From cc3d40671af427f860b1f2ce0564c43e8963468e Mon Sep 17 00:00:00 2001
From: Artemiy Bulavin <artemiyb at graphcore.ai>
Date: Sun, 30 Nov 2025 14:30:26 +0000
Subject: [PATCH] Remove effectful producer after fusion if no longer used

---
 .../Linalg/Transforms/ElementwiseOpFusion.cpp |  6 ++++
 .../Linalg/fusion-elementwise-ops.mlir        | 36 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 05fc7cbbb90af..b65e5dd35ff3c 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -493,6 +493,12 @@ class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
         });
       }
       rewriter.eraseOp(genericOp);
+      // If after fusion, the producer no longer has uses, erase it. Usually the
+      // greedy pattern driver takes care of this, however if the producer
+      // contains ops with memory effects it won't be considered trivially dead.
+      if (producer->use_empty())
+        rewriter.eraseOp(producer);
+
       return success();
     }
     return failure();
diff --git a/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir b/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
index bc55c12c02f29..dedb95db4f4b6 100644
--- a/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
+++ b/mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir
@@ -1079,4 +1079,38 @@ module {
 // CHECK-NOT:     linalg.generic
 // CHECK:         tensor.expand_shape
 // CHECK:         linalg.generic {{.*}}, iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]}
-// CHECK-SAME:     ins(%[[ARG0]], %[[FUSED]]#1 : tensor<1x1x2x1xf32>, tensor<4x1x1x1xf32>)
\ No newline at end of file
+// CHECK-SAME:     ins(%[[ARG0]], %[[FUSED]]#1 : tensor<1x1x2x1xf32>, tensor<4x1x1x1xf32>)
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+
+module {
+  func.func @remove_effectful_producer_after_fusion_if_no_uses(%arg0: bf16, %arg1: memref<2xbf16>, %arg2: memref<2xbf16>) {
+    %c0 = arith.constant 0 : index
+    %0 = tensor.empty() : tensor<2xbf16>
+    %1 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%0 : tensor<2xbf16>) {
+    ^bb0(%out: bf16):
+      %2 = memref.atomic_rmw addf %arg0, %arg1[%c0] : (bf16, memref<2xbf16>) -> bf16
+      linalg.yield %2 : bf16
+    } -> tensor<2xbf16>
+  
+    linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} ins(%1 : tensor<2xbf16>) {
+    ^bb0(%in: bf16):
+      memref.store %in, %arg2[%c0] : memref<2xbf16>
+      linalg.yield
+    }
+    return
+  }
+}
+
+// CHECK-LABEL:   func.func @remove_effectful_producer_after_fusion_if_no_uses
+// CHECK-SAME:      %[[ARG0:.*]]: bf16,
+// CHECK-SAME:      %[[ARG1:.*]]: memref<2xbf16>,
+// CHECK-SAME:      %[[ARG2:.*]]: memref<2xbf16>)
+// CHECK-NOT:       memref.atomic_rmw addf %[[ARG0]], %[[ARG1]]
+// CHECK:           linalg.generic
+// CHECK:             %[[ATOMIC_RMW:.*]] = memref.atomic_rmw addf %[[ARG0]], %[[ARG1]]
+// CHECK:             memref.store %[[ATOMIC_RMW]], %[[ARG2]]
+// CHECK:             linalg.yield %[[ATOMIC_RMW]]
+// CHECK:           return