[all-commits] [llvm/llvm-project] ed499d: [MLIR] Fix operation clone

[William Moses via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: ed499ddcdaa6a47d21d98c12f29a94386cd4643a
      https://github.com/llvm/llvm-project/commit/ed499ddcdaa6a47d21d98c12f29a94386cd4643a
  Author: William S. Moses <gh at wsmoses.com>
  Date:   2022-04-15 (Fri, 15 Apr 2022)

  Changed paths:
    M mlir/include/mlir/IR/Operation.h
    M mlir/lib/IR/Operation.cpp
    A mlir/test/IR/test-clone.mlir
    M mlir/test/lib/IR/CMakeLists.txt
    A mlir/test/lib/IR/TestClone.cpp
    M mlir/tools/mlir-opt/mlir-opt.cpp

  Log Message:
  -----------
  [MLIR] Fix operation clone

Operation clone is currently faulty.

Suppose you have a block like as follows:

```
(%x0 : i32) {
   %x1 = f(%x0)
   return %x1
}
```

The test case we have is that we want to "unroll" this, in which we want to change this to compute `f(f(x0))` instead of just `f(x0)`. We do so by making a copy of the body at the end of the block and set the uses of the argument in the copy operations with the value returned from the original block.
This is implemented as follows:
1) map to the block arguments to the returned value (`map[x0] = x1`).
2) clone the body

Now for this small example, this works as intended and we get the following.

```
(%x0 : i32) {
   %x1 = f(%x0)
   %x2 = f(%x1)
   return %x2
}
```

This is because the current logic to clone `x1 = f(x0)` first looks up the arguments in the map (which finds `x0` maps to `x1` from the initialization), and then sets the map of the result to the cloned result (`map[x1] = x2`).

However, this fails if `x0` is not an argument to the op, but instead used inside the region, like below.

```
(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   return %x1
}
```

This is because cloning an op currently first looks up the args (none), sets the map of the result (`map[%x1] = %x2`), and then clones the regions. This results in the following, which is clearly illegal:

```
(%x0 : i32) {
   %x1 = f() {
      yield %x0
   }
   %x2 = f() {
      yield %x2
   }
   return %x2
}
```

Diving deeper, this is partially due to the ordering (how this PR fixes it), as well as how region cloning works. Namely it will first clone with the mapping, and then it will remap all operands. Since the ordering above now has a map of `x0 -> x1` and `x1 -> x2`, we end up with the incorrect behavior here.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D122531