[mlir][SCF] Add utility method to add new yield values to a loop.
authorMahesh Ravishankar <ravishankarm@google.com>
Fri, 6 May 2022 21:44:26 +0000 (21:44 +0000)
committerMahesh Ravishankar <ravishankarm@google.com>
Tue, 10 May 2022 18:44:11 +0000 (18:44 +0000)
commit567fd523bf538523f58779e5af9d20c3e48838a2
tree2b9aa5c97c006bfec293bfeeb09f7f847cbe897d
parent82c5e302f9e63a3491b5e40aa33771f355791598
[mlir][SCF] Add utility method to add new yield values to a loop.

The current implementation of `cloneWithNewYields` has a few issues
- It clones the loop body of the original loop to create a new
  loop. This is very expensive.
- It performs `erase` operations which are incompatible when this
  method is called from within a pattern rewrite. All erases need to
  go through `PatternRewriter`.

To address these a new utility method `replaceLoopWithNewYields` is added
which
- moves the operations from the original loop into the new loop.
- replaces all uses of the original loop with the corresponding
  results of the new loop
- use a call back to allow caller to generate the new yield values.
- the original loop is modified to just yield the basic block
  arguments corresponding to the iter_args of the loop. This
  represents a no-op loop. The loop itself is dead (since all its uses
  are replaced), but is not removed. The caller is expected to erase
  the op. Consequently, this method can be called from within a
  `matchAndRewrite` method of a `PatternRewriter`.

The `cloneWithNewYields` could be replaces with
`replaceLoopWithNewYields`, but that seems to trigger a failure during
walks, potentially due to the operations being moved. That is left as
a TODO.

Differential Revision: https://reviews.llvm.org/D125147
mlir/include/mlir/Dialect/SCF/Utils/Utils.h
mlir/lib/Dialect/SCF/Utils/Utils.cpp
mlir/test/Transforms/scf-loop-utils.mlir
mlir/test/Transforms/scf-replace-with-new-yields.mlir [new file with mode: 0644]
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp