[Mlir-commits] [mlir] [mlir][affine] Avoid fusion when the sibling loop has result usage (PR #216564)

[Kunal Dubey]

https://github.com/xakep8 created https://github.com/llvm/llvm-project/pull/216564

Sibling loop fusion copies sibling loop to the destination loop and removes the sibling without checking if the sibiling has loop results being used at other places. Added check for this to reject fusion if the sibling has usage.

Added test for the same.

Fixes #216464

>From 7963e687a35414d6c9e2784483e04f1b5a426e06 Mon Sep 17 00:00:00 2001
From: Kunal Dubey <xakep8 at protonmail.com>
Date: Sun, 16 Aug 2026 17:47:57 +0530
Subject: [PATCH] [mlir][affine] Avoid fusion when the sibling loop has result
 usage

Sibling loop fusion copies sibling loop to the destination loop and
removes the sibling without checking if the sibiling has loop results
being used at other places. Added check for this to reject fusion if the
sibling has usage.

Added test for the same.
---
 .../Dialect/Affine/Transforms/LoopFusion.cpp    |  6 ++++++
 .../Dialect/Affine/loop-fusion-sibling.mlir     | 17 +++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..8adbb97563e6f 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -1369,6 +1369,12 @@ struct GreedyFusion {
       // TODO: Remove restrict to single load op restriction.
       if (sibNode->getLoadOpCount(memref) != 1)
         return false;
+
+      // Sibling fusion removes the sibling node after cloning it into the
+      // destination. Do not fuse siblings whose results still have uses.
+      if (!sibNode->op->use_empty())
+        return false;
+
       // Skip if there exists a path of dependent edges between
       // 'sibNode' and 'dstNode'.
       if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 937c855b86b50..6c3b1db56d0e9 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -21,3 +21,20 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
   // CHECK-NOT: affine.for
   return
 }
+
+// CHECK-LABEL: func @sibling_with_used_loop_result
+func.func @sibling_with_used_loop_result(%m: memref<4xi32>, %n: memref<4xi32>,
+                                         %init: i32) -> i32 {
+  // CHECK: %[[RESULT:.*]] = affine.for {{.*}} iter_args
+  %a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
+    %v = affine.load %m[%i] : memref<4xi32>
+    %t = arith.addi %x, %v : i32
+    affine.yield %t : i32
+  }
+  affine.for %i = 0 to 4 {
+    %v = affine.load %m[%i] : memref<4xi32>
+    affine.store %v, %n[%i] : memref<4xi32>
+  }
+  // CHECK: return %[[RESULT]]
+  return %a : i32
+}