[mlir][Linalg] Add pattern for folding reshape by collapsing.
authorMahesh Ravishankar <ravishankarm@google.com>
Mon, 7 Feb 2022 03:43:52 +0000 (03:43 +0000)
committerMahesh Ravishankar <ravishankarm@google.com>
Wed, 16 Feb 2022 03:15:20 +0000 (03:15 +0000)
commit2c58cde003eb7b5aaae3eb5ac94f25a52f151df2
treecd756d5412df5c1f67011ad1c8b5d94ed5b74c59
parent2e2f3158c604adb8401a2a44a03f58d4b6f1c7f9
[mlir][Linalg] Add pattern for folding reshape by collapsing.

Fusion of `linalg.generic` with
`tensor.expand_shape/tensor.collapse_shape` currently handles fusion
with reshape by expanding the dimensionality of the `linalg.generic`
operation. This helps fuse elementwise operations better since they
are fused at the highest dimensionality while keeping all indexing
maps involved projected permutations. The intent of these is to push
the reshape to the boundaries of functions.

The presence of named ops (or other ops across which the reshape
cannot be propagated) stops the propagation to the edges of the
function. At this stage, the converse patterns that fold the reshapes
with generic ops by collapsing the dimensions of the generic op can
push the reshape towards edges. In particular it helps the case where
reshapes exist in between named ops and generic ops.

`linalg.named_op` -> `tensor.expand_shape` -> `linalg.generic`

Pushing the reshape down will help fusion of `linalg.named_op` ->
`linalg.generic` using tile + fuse transformations.

This pattern is intended to replace the following patterns

1) FoldReshapeByLinearization : These patterns create indexing maps
that are not projected permutations that affect future
transformations. They are only useful for folding unit-dimensions.
2) PushReshapeByExpansion : This pattern has the same functionality
but has some restrictions
    a) It tries to avoid creating new reshapes that limits its
    applicability. The pattern added here can achieve the same
    functionality through use of the `controlFn` that allows clients
    of the pattern freedom to make this decision.
    b) It does not work for ops with indexing semantics.

These patterns will be deprecated in a future patch.

Differential Revision: https://reviews.llvm.org/D119365
mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir [new file with mode: 0644]
mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp