[MLIR] Fix operation clone
authorWilliam S. Moses <gh@wsmoses.com>
Sat, 26 Mar 2022 18:42:15 +0000 (14:42 -0400)
committerWilliam S. Moses <gh@wsmoses.com>
Fri, 15 Apr 2022 17:09:13 +0000 (13:09 -0400)
commited499ddcdaa6a47d21d98c12f29a94386cd4643a
tree0da689d83efc0d35399aecb622715d1c696c5fba
parentca2be81e34a6d87edb8e555dfac94ab68ee20f70
[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
mlir/include/mlir/IR/Operation.h
mlir/lib/IR/Operation.cpp
mlir/test/IR/test-clone.mlir [new file with mode: 0644]
mlir/test/lib/IR/CMakeLists.txt
mlir/test/lib/IR/TestClone.cpp [new file with mode: 0644]
mlir/tools/mlir-opt/mlir-opt.cpp