[mlir] Structured transforms: introduce op splitting
authorAlex Zinenko <zinenko@google.com>
Thu, 7 Jul 2022 11:10:40 +0000 (13:10 +0200)
committerAlex Zinenko <zinenko@google.com>
Thu, 7 Jul 2022 11:19:44 +0000 (13:19 +0200)
commitff6e5508d686395dfb5f26085fadeae174847d52
treef93bcdf027b330c13736c5ebb6b58df80aa3543c
parent1d9086bf054c2e734940620d02d4451156b424e6
[mlir] Structured transforms: introduce op splitting

Introduce a new transformation on structured ops that splits the iteration
space into two parts along the specified dimension. The index at which the
splitting happens may be static or dynamic. This transformation can be seen as
a rudimentary form of index-set splitting that only supports the splitting
along hyperplanes parallel to the iteration space hyperplanes, and is therefore
decomposable into per-dimension application.

It is a key low-level transformation that enables independent scheduling for
different parts of the iteration space of the same op, which hasn't been
possible previously. It may be used to implement, e.g., multi-sized tiling. In
future, peeling can be implemented as a combination of split-off amount
computation and splitting.

The transformation is conceptually close to tiling in its separation of the
iteration and data spaces, but cannot be currently implemented on top of
TilingInterface as the latter does not properly support `linalg.index`
offsetting.

Note that the transformation intentionally bypasses folding of
`tensor.extract_slice` operations when creating them as this folding was found
to prevent repeated splitting of the same operation because due to internal
assumptions about extract/insert_slice combination in dialect utilities.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D129090
16 files changed:
mlir/include/mlir-c/BuiltinTypes.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Bindings/Python/IRTypes.cpp
mlir/lib/CAPI/IR/BuiltinTypes.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/lib/Dialect/Linalg/Transforms/Split.cpp [new file with mode: 0644]
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
mlir/python/mlir/dialects/_structured_transform_ops_ext.py
mlir/test/Dialect/Linalg/transform-op-split.mlir [new file with mode: 0644]
mlir/test/python/dialects/transform_structured_ext.py