[mlir][linalg] Refine `tensor.extract` vectorisation
authorAndrzej Warzynski <andrzej.warzynski@arm.com>
Wed, 12 Apr 2023 19:44:13 +0000 (12:44 -0700)
committerAndrzej Warzynski <andrzej.warzynski@gmail.com>
Fri, 21 Apr 2023 07:47:55 +0000 (08:47 +0100)
commita3ae3931d4e0c227a2d88fa827dbe536adda586a
tree4540e55c876ca02d7d183cc60ff2a695d2e2a0b0
parent735cb7b1f84a36049ea187fecdd281c6e93600bb
[mlir][linalg] Refine `tensor.extract` vectorisation

This patch updates the vectorisation of the extract Op so that the
permutation map for the transfer_read Op is defined explicitly by the
vectoriser (as opposed to being constructed implicitly by the
transfer_read builder).

This change is needed for cases where the rank of the source tensor is
lower than the rank of the output vector generated by the vectoriser:
```mlir
    %17 = vector.transfer_read %arg1[%14, %16], %cst_4 {in_bounds = [true, true]} : tensor<257x24xf32>, vector<1x1x4xf32>
```
In cases like this, the vectorize will create the following permutation map:
```
  (d0, d1) -> (0, d0, d1)
```

In other cases the behaviour remains unchanged.

Fixes https://github.com/openxla/iree/issues/13036. That's also where
the test case was extracted from.

Differential Revision: https://reviews.llvm.org/D148537
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization.mlir