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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Kunal Dubey  (xakep8)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/216564.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp (+6) 
- (modified) mlir/test/Dialect/Affine/loop-fusion-sibling.mlir (+17) 


``````````diff
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
+}

``````````

</details>


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