[all-commits] [llvm/llvm-project] 84d969: [mlir][linalg] Fix vectorisation of tensor.extract...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 84d96947ef227e408d27c9db0ee08622d0b20996
      https://github.com/llvm/llvm-project/commit/84d96947ef227e408d27c9db0ee08622d0b20996
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2023-07-17 (Mon, 17 Jul 2023)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir

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