[MLIR] Vectorize tensor.extract on n-D tensor (n >= 2)
authorAndrzej Warzynski <andrzej.warzynski@arm.com>
Wed, 19 Oct 2022 17:04:31 +0000 (17:04 +0000)
committerAndrzej Warzynski <andrzej.warzynski@arm.com>
Mon, 12 Dec 2022 09:32:16 +0000 (09:32 +0000)
commitc181f21ac71f1d52ec3d11a381634bbcec3844a8
treedc388cd7b276ca618f709c7e35fc8816da7bed89
parentcb06b6ab0018d9aa07bf9fec0052a9f8c921a593
[MLIR] Vectorize tensor.extract on n-D tensor (n >= 2)

This patch implements the vectorization of tensor.extract for arbitrary
tensors. It basically extends https://reviews.llvm.org/D133786 by adding
support for n-D tensors (n >= 2). This is implemented by essentially
flattening the indices.

When benchmarking the vectorized code, we have observed that it is
slower than the scalar code. That's most likely due to sub-optimal (and,
in general slow) gather loads. More work is needed to identify an
implementation and/or a representation that would lead to better code.
In the meantime, the vectorization of n-D tensors (where n >= 2) has to
be explicitly enabled. This can be done either via:
  * transfer dialect's `vectorize_nd_extract` attribute,
  * dedicated bool argument in the `vectorize` method from
    "Vectorization.cpp".
The second option was added to control the new functionality through
means other than the transfer dialect.

Related discussion: https://github.com/iree-org/iree/issues/9198

Differential Revision: https://reviews.llvm.org/D137660
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization.mlir