[mlir] make `fuse_into_containing_op` preserve the containing op handle
authorAlex Zinenko <zinenko@google.com>
Fri, 26 May 2023 11:26:52 +0000 (11:26 +0000)
committerAlex Zinenko <zinenko@google.com>
Fri, 26 May 2023 16:01:40 +0000 (16:01 +0000)
commit5310be521db2aa8c05a1c1adb7e108fc2c7c9ddc
tree4c2f9e1f4d2ff694d1b46a7cfda4c26897007073
parent44f6e8625e634a1c26805a752132c74aeb99c7d5
[mlir] make `fuse_into_containing_op` preserve the containing op handle

This partially undoes the intent of https://reviews.llvm.org/D151418 by
cheating its way to keep the "containing op" (aka loop) handle read-only
in fusion. It is crucial to do so for composability of tiling and
fusion. Specfically, after the "containing op" handle started being
consumed, it became impossible to perform additional tiling after fusion
except tiling the last-fused op:

  %tiled1, %loop1 = tile %op
  %producer1, %loop2 = fuse %producer into %loop1
  // invalid, because %tiled1 is invalidated by consuming %loop1
  // that points to its parent
  tile %tiled1

or

  %tiled1, %loop1 = tile %op
  %tiled2, %loop2 = tile %tiled1
  %p2 = fuse %producer into %loop1
  // invalid, because %loop2 is invalidated by consuming %loop1
  // that points to its parent
  fuse %p2 into %loop2

The approach here makes creative use of the state extension mechanism to
update the payload operation associted with the operand handle. Further
investigation is necessary to understand if is consistent with the
overall execution model of the transform dialect, but it is crucial to
restore composability ASAP.

Reviewed By: springerm, nicolasvasilache

Differential Revision: https://reviews.llvm.org/D151555
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/test/Dialect/Linalg/transform-ops.mlir