[mlir][Linalg] Reimplement hoisting on tensors as a subset-based transformation
authorNicolas Vasilache <nicolas.vasilache@gmail.com>
Wed, 22 Feb 2023 13:24:25 +0000 (05:24 -0800)
committerNicolas Vasilache <nicolas.vasilache@gmail.com>
Mon, 27 Feb 2023 16:15:44 +0000 (08:15 -0800)
commit4521b113978d9ddaaae038e3cdd9d8902e2392f9
tree262827760ad7fe0f5bd2d365f34b920e4c82e817
parent779d54fd75c19a6d4c27b0a8b115d2e5ccd9768d
[mlir][Linalg] Reimplement hoisting on tensors as a subset-based transformation

This revision significantly rewrites hoisting on tensors.
Previously, `vector.transfer_read/write` and `tensor.extract/insert_slice` would
be clumped together when looking for candidate pairs.
This would significantly increase the complexity of the logic and would not apply
independently to `tensor.extract/insert_slice`.

The new implementation decouples the cases and starts to cast the problem
as a generic matching subset extract/insert, which will be future proof when
other such operation pairs are introduced.

Lastly, the implementation makes the distinction clear between `vector.transfer_read/write` for
which we allow bypasses of the disjoint subsets from `tensor.extract/insert_slice` for which we
do not yet allow it.

This can be extended in the future and unified once we have subset disjunction implemented more generally.

The algorithm can be rewritten to be less of a fixed point with interspersed canonicalizations.
As a consequence, the test explicitly adds a canonicalization to clean up the IR and verify we end up in the same state.

That extra canonicalization exhibited that one of the uses in one of the tests was dead, so we fix the appropriate test.

Differential Revision: https://reviews.llvm.org/D144656
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
mlir/lib/Dialect/Linalg/Transforms/SubsetHoisting.cpp [new file with mode: 0644]
mlir/lib/Dialect/Utils/StaticValueUtils.cpp
mlir/test/Dialect/Linalg/hoisting.mlir