[all-commits] [llvm/llvm-project] a63853: [mlir] Broadcast scalars when vectorising tensor.e...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: a63853e6acc594d33b2fd96036e804d9e4a53f76
      https://github.com/llvm/llvm-project/commit/a63853e6acc594d33b2fd96036e804d9e4a53f76
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2023-01-12 (Thu, 12 Jan 2023)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/test/Dialect/Linalg/vectorization.mlir

  Log Message:
  -----------
  [mlir] Broadcast scalars when vectorising tensor.extract

When vectorizing tensor.extract embedded within linalg.generic, the
default option is to rewrite it as vector.gather. When doing so, we need
to make sure that the corresponding indices are vectorized accordingly.
However, the Linalg vectorizer will not vectorize constants like in the
following example. This is fixed by simply broadcasting %c0 and %c1.

```
  func.func @example(%arg0: tensor<3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
    %c0 = arith.constant 1 : index
    %c1 = arith.constant 2 : index
    %1 = linalg.generic {
      (...)
    } outs(...) {
    ^bb0(...):
      %2 = tensor.extract %arg0[%c0, %c1] : tensor<3x3xf32>
      linalg.yield %2 : f32
    } -> tensor<1x1x3xf32>
    return %1 : tensor<1x1x3xf32>
  }
```

This patch makes sure that in the case above (and other similar cases),
the vectorizer broadcasts %c0 and %c1.

Differential Revision: https://reviews.llvm.org/D140781