[all-commits] [llvm/llvm-project] 62e503: Reapply "[mlir][linalg] Relax tensor.extract vecto...

[Andrzej Warzyński via All-commits]

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

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

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

[This reverts commit 6662523d6b2ca0198141c94ee80ebbb41601df9f]

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
is 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):
    %8 = linalg.index 0 : index
    %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>
}
```

Overview of the delta compared to the original submission (#99299):
  * removed an assert representing a condition that is being relaxed
    here,
  * added a test (reading from a column tensor) based on a repro from
    @hanhanW.

(*) `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