[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
Loading
Please sign in to comment