Skip to content
Commit 84d96947 authored by Andrzej Warzynski's avatar Andrzej Warzynski
Browse files

[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
parent 5784c478
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment