[Mlir-commits] [mlir] [mlir][SCF]-Fix loop coalescing with iteration arguements (PR #105488)

[llvmlistbot at llvm.org]

================
@@ -864,6 +864,18 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
 
     Operation *innerTerminator = innerLoop.getBody()->getTerminator();
     auto yieldedVals = llvm::to_vector(innerTerminator->getOperands());
+    assert(llvm::equal(outerLoop.getRegionIterArgs(), innerLoop.getInitArgs()));
+    for (Value &yieldedVal : yieldedVals) {
+      // The yielded value may be an iteration argument of the inner loop
+      // which is about to be inlined.
+      auto iter = llvm::find(innerLoop.getRegionIterArgs(), yieldedVal);
+      if (iter != innerLoop.getRegionIterArgs().end()) {
+        unsigned iterArgIndex = iter - innerLoop.getRegionIterArgs().begin();
+        // `outerLoop` iter args identical to the `innerLoop` init args.
----------------
MaheshRavishankar wrote:

I think this is better to just return the corresponding `iter_args` operand of the inner loop?

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