[mlir] Fix loop unrolling: properly replace the arguments of the epilogue loop.
authorgrosul1 <grosul@google.com>
Thu, 12 May 2022 01:44:13 +0000 (01:44 +0000)
committerMogball <jeffniu22@gmail.com>
Thu, 12 May 2022 01:54:39 +0000 (01:54 +0000)
commita4b227c28aba3487ffbfde12fb59fdb69a6b5bfe
tree615aa1c636211ca6f03310287b6cd681d26629b8
parent24532d05f8da67058c41ea4ad885c96ee918aa57
[mlir] Fix loop unrolling: properly replace the arguments of the epilogue loop.

Using "replaceUsesOfWith" is incorrect because the same initializer value may appear multiple times.

For example, if the epilogue is needed when this loop is unrolled
```
%x:2 = scf.for ... iter_args(%arg1 = %c1, %arg2 = %c1) {
  ...
}
```
then both epilogue's arguments will be incorrectly renamed to use the same result index (note #1 in both cases):
```
%x_unrolled:2 = scf.for ... iter_args(%arg1 = %c1, %arg2 = %c1) {
  ...
}
%x_epilogue:2 = scf.for ... iter_args(%arg1 = %x_unrolled#1, %arg2 = %x_unrolled#1) {
  ...
}
```
mlir/lib/Dialect/SCF/Utils/Utils.cpp
mlir/test/Dialect/SCF/loop-unroll.mlir