[mlir] Bubble up tensor.extract_slice above linalg operation
authorOkwan Kwon <okkwon@gmail.com>
Fri, 18 Mar 2022 16:48:34 +0000 (16:48 +0000)
committerOkwan Kwon <okkwon@gmail.com>
Thu, 31 Mar 2022 16:48:38 +0000 (16:48 +0000)
commit65bdeddb1e5c0d27be0397379131b2d712c7a227
tree676c9d819f1890a91073dba41ba659fb0e6850a4
parent19054163e11a6632b4973c936e5aa93ec742c866
[mlir] Bubble up tensor.extract_slice above linalg operation

Bubble up extract_slice above Linalg operation.

A sequence of operations

    %0 = linalg.<op> ... arg0, arg1, ...
    %1 = tensor.extract_slice %0 ...

can be replaced with

    %0 = tensor.extract_slice %arg0
    %1 = tensor.extract_slice %arg1
    %2 = linalg.<op> ... %0, %1, ...

This results in the reduce computation of the linalg operation.

The implementation uses the tiling utility functions. One difference
from the tiling process is that we don't need to insert the checking
code for the out-of-bound accesses. The use of the slice itself
represents that the code writer is sure about the boundary condition.
To avoid adding the boundary condtion check code, `omitPartialTileCheck`
is introduced for the tiling utility functions.

Differential Revision: https://reviews.llvm.org/D122437
13 files changed:
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/include/mlir/IR/AffineMap.h
mlir/lib/Dialect/Linalg/Transforms/BubbleUpExtractSlice.cpp [new file with mode: 0644]
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
mlir/lib/IR/AffineMap.cpp
mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir [new file with mode: 0644]
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp