Skip to content
Commit c181f21a authored by Andrzej Warzynski's avatar Andrzej Warzynski
Browse files

[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
parent cb06b6ab
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment