[Mlir-commits] [mlir] [mlir][affine] Use iter argument replace init when delete loop in the coalesceLoops function (PR #169514)

[lonely eagle]

https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/169514

>From 96b1865ee7acaffda55fc831dbbcea69c8f1eb8f Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 25 Nov 2025 16:20:33 +0000
Subject: [PATCH 1/2] fix affine-loop-coalescing.

---
 mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp   |  5 +++
 mlir/test/Dialect/Affine/loop-coalescing.mlir | 31 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 4743941deff3f..fe245940e69af 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1711,6 +1711,11 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
   outermost.getBody()->getOperations().splice(
       Block::iterator(secondOutermostLoop.getOperation()),
       innermost.getBody()->getOperations());
+  for (auto [iter, init] : llvm::zip(secondOutermostLoop.getRegionIterArgs(),
+                                     secondOutermostLoop.getInits())) {
+    iter.replaceAllUsesWith(init);
+    iter.dropAllUses();
+  }
   secondOutermostLoop.erase();
   return success();
 }
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index 3be14eaf5c326..f78c210e0ad56 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -416,3 +416,34 @@ func.func @test_loops_do_not_get_coalesced() {
 // CHECK-NEXT: }
 // CHECK-NEXT: }
 // CHECK-NEXT: return
+
+// -----
+
+// CHECK-LABEL: func @inner_loop_has_iter_args
+func.func @inner_loop_has_iter_args() {
+  %c17 = arith.constant 17 : index
+  %c79 = arith.constant 79 : index
+  %memref = gpu.alloc  (%c79) : memref<?xi64>
+  affine.for %arg0 = 0 to 79 {
+    %0 = affine.for %arg1 = 0 to 64 iter_args(%arg2 = %memref) -> (memref<?xi64>) {
+      %1 = arith.remui %arg1, %c17 : index
+      %2 = arith.index_cast %arg1 : index to i64
+      memref.store %2, %arg2[%1] : memref<?xi64>
+      affine.yield %arg2 : memref<?xi64>
+    }
+  }
+  return
+}
+
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 17 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 79 : index
+// CHECK: %[[ALLOC_0:.*]] = gpu.alloc  (%[[CONSTANT_1]]) : memref<?xi64>
+// CHECK: %[[APPLY_0:.*]] = affine.apply affine_map<() -> (79)>()
+// CHECK: %[[APPLY_1:.*]] = affine.apply affine_map<() -> (64)>()
+// CHECK: %[[APPLY_2:.*]] = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%[[APPLY_0]]){{\[}}%[[APPLY_1]]]
+// CHECK:   affine.for %[[IV:.*]] = 0 to %[[APPLY_2]] {
+// CHECK:     %[[APPLY_3:.*]] = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%[[IV]]){{\[}}%[[APPLY_1]]]
+// CHECK:     %[[REMUI_0:.*]] = arith.remui %[[APPLY_3]], %[[CONSTANT_0]] : index
+// CHECK:     %[[INDEX_CAST_0:.*]] = arith.index_cast %[[APPLY_3]] : index to i64
+// CHECK:      memref.store %[[INDEX_CAST_0]], %[[ALLOC_0]]{{\[}}%[[REMUI_0]]] : memref<?xi64>
+// CHECK:   }

>From 6643cf045e840e33d28e959b0e301b0536912052 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 30 Nov 2025 16:03:07 +0000
Subject: [PATCH 2/2] use zip_equal

---
 mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index fe245940e69af..8f1249e3afaf0 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1711,8 +1711,9 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
   outermost.getBody()->getOperations().splice(
       Block::iterator(secondOutermostLoop.getOperation()),
       innermost.getBody()->getOperations());
-  for (auto [iter, init] : llvm::zip(secondOutermostLoop.getRegionIterArgs(),
-                                     secondOutermostLoop.getInits())) {
+  for (auto [iter, init] :
+       llvm::zip_equal(secondOutermostLoop.getRegionIterArgs(),
+                       secondOutermostLoop.getInits())) {
     iter.replaceAllUsesWith(init);
     iter.dropAllUses();
   }