[all-commits] [llvm/llvm-project] c181f2: [MLIR] Vectorize tensor.extract on n-D tensor (n >...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: c181f21ac71f1d52ec3d11a381634bbcec3844a8
      https://github.com/llvm/llvm-project/commit/c181f21ac71f1d52ec3d11a381634bbcec3844a8
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2022-12-12 (Mon, 12 Dec 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
    M mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    M mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/test/Dialect/Linalg/vectorization.mlir

  Log Message:
  -----------
  [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