[all-commits] [llvm/llvm-project] 8868c0: [mlir][linalg] Relax tensor.extract vectorization ...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 8868c02cda875d1efe1646affa01656ef268ffed
      https://github.com/llvm/llvm-project/commit/8868c02cda875d1efe1646affa01656ef268ffed
  Author: Andrzej Warzyński <andrzej.warzynski at arm.com>
  Date:   2024-08-06 (Tue, 06 Aug 2024)

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

  Log Message:
  -----------
  [mlir][linalg] Relax tensor.extract vectorization (#99299)

Simplifies the vectorization of tensor.extract so that:
* all cases that read into a genuinely multi-dim vector (*) are
  considered a gather load,
* all other cases are considered as potential contiguous loads.

This change means that the following extraction from a "column" tensor
will be correctly identified as a scalar load followed by a broadcast (rather
than a gather load).

```mlir
func.func @vectorize_scalar_broadcast_column_tensor(%in: tensor<1x1x4xi32>) -> tensor<1x1x4xi32> {
  %c4 = arith.constant 4 : index
  %c0 = arith.constant 0 : index
  %cst = arith.constant dense<[...]> : tensor<15x1xi32>

  %out = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]}
    outs(%in : tensor<1x1x4xi32>) {

  ^bb0(%out: i32):
    %idx_0 = linalg.index 0 : index
    %extracted = tensor.extract %cst[%idx_0, %c0] : tensor<15x1xi32>
    linalg.yield %extracted : i32
  } -> tensor<1x1x4xi32>

  return %out:tensor<1x1x4xi32>
}
```

(*) `vector<1x4x1xf32>` is considered as 1D vector in this context.



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications