[Mlir-commits] [mlir] [mlir][affine] Do not coalesce loops that carry a value (PR #216494)

[Alessandro Potenza]

https://github.com/alepot55 commented:

Disclosure first: I have #216853 open against the `scf.for` path of this same pass, so I have an interest here. That is also why I went looking at this one.

**The fix does what it says.** I applied the `LoopUtils.cpp` hunk on `11e915f2b75d` and built it. Your reproducer, before:

```
$ mlir-opt --affine-loop-coalescing affine.mlir
mlir-opt: mlir/include/mlir/IR/UseDefLists.h:198:
Assertion `use_empty() && "Cannot destroy a value that still has uses!"' failed.
```

after: the nest comes out untouched, no crash. A two-loop nest whose outer loop carries a value and whose inner one is invariant is left alone on both, so nothing that used to coalesce stops coalescing in the cases I tried.

**This branch is in conflict and nothing in the thread says so.** `mergeable` is `false` and `mergeable_state` is `dirty`. The conflict is in `mlir/test/Dialect/Affine/loop-coalescing.mlir`: #216903 appended `@inner_loop_yields_induction_var` at the end of that file after this branch was cut, and your two tests land in the same place. I hit exactly the same conflict on my own PR in the same file this morning, and it is invisible unless you look at `mergeable_state`, which is very likely why this has sat since 15 August with no reviewer.

**The `scf.for` path has a sibling bug that this does not cover.** `--affine-loop-coalescing` dispatches to two different functions: `mlir::affine::coalesceLoops` in `Affine/Utils/LoopUtils.cpp`, which you are fixing, and `mlir::coalesceLoops` in `SCF/Utils/Utils.cpp` for `scf.for` nests. With your hunk applied:

```mlir
func.func @scf_reads_outer(%init: i64, %lb: index, %ub: index, %st: index) -> i64 {
  %0 = scf.for %i = %lb to %ub step %st iter_args(%outer = %init) -> (i64) {
    %1 = scf.for %j = %lb to %ub step %st iter_args(%inner = %outer) -> (i64) {
      %2 = arith.addi %inner, %outer : i64
      scf.yield %2 : i64
    }
    scf.yield %1 : i64
  }
  return %0 : i64
}
```

still coalesces to a body containing `arith.addi %arg5, %arg5`: the two iteration arguments have become one, though the outer is fixed for a whole run of the inner loop. That is #216853, and it is not something you need to fix here; I mention it because the two are complementary rather than overlapping, and because a reviewer looking at one will want to know about the other.

**One question that is worth a maintainer's opinion**, and I do not have a preference: the two guards will use different criteria for the same transformation. Yours is that every loop but the outermost must yield its iteration arguments unchanged. Mine is that an outer iteration argument may only be used as the inner loop's init operand. Yours is the more precise of the two, and mine declines some shapes that are probably sound, which I say in my own description. Whether the two dialect paths of one pass should end up with different legality rules is not obvious to me.

**No reviewers are requested on this.** Going by the history of `Affine/Utils/LoopUtils.cpp`, @bondhugula has by far the most commits in it, and @kuhar approved #169514, which is the change that introduced the code you are correcting.

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