[mlir][linalg] Fix vectorisation of tensor.extract with dynamic shapes
authorAndrzej Warzynski <andrzej.warzynski@arm.com>
Sat, 15 Jul 2023 16:21:34 +0000 (16:21 +0000)
committerAndrzej Warzynski <andrzej.warzynski@arm.com>
Mon, 17 Jul 2023 17:28:17 +0000 (17:28 +0000)
commit84d96947ef227e408d27c9db0ee08622d0b20996
treeb58136cfdcbd2633d89bfed3820f966943e26fcc
parent5784c47806cf276e8fcc1311201f437cb8169261
[mlir][linalg] Fix vectorisation of tensor.extract with dynamic shapes

The Linalg vectoriser incorrectly recognises the following
`tensor.extract` as contiguous:
```
func.func @example(%in: tensor<123x321xf32>, %arg1: tensor<1x?x8xf32>) -> tensor<1x?x8xf32> {
  %c0 = arith.constant 1 : index
  %2 = linalg.generic {
    indexing_maps = [#map1],
    iterator_types = ["parallel", "parallel", "parallel"]
  } outs(%arg1 : tensor<1x?x8xf32>)
  {
  ^bb0(%arg3: f32):
    %idx_0 = linalg.index 0 : index
    %idx_1 = linalg.index 1 : index
    %idx = arith.addi %idx_0, %idx_1 : index
    %7 = tensor.extract %in[%c0, %idx] : tensor<123x321xf32>
    linalg.yield %7 : f32
  } -> tensor<1x?x8xf32>
  return %2 : tensor<1x?x8xf32>
}
```

However, the following index Op corresponds to the dynamic dimension
in the iteration space:
```
    %idx_1 = linalg.index 1 : index
```
The vectoriser should assume that:
  * this index Op _is not_ loop invariant,
  * the resulting memory access is a gather load
This is what this patch fixes.

Differential Revision: https://reviews.llvm.org/D155373
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir