[all-commits] [llvm/llvm-project] a2acb2: [mlir][linalg] Fix vectorization of tensor.extract...

[Andrzej Warzyński via All-commits]

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

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

  Log Message:
  -----------
  [mlir][linalg] Fix vectorization of tensor.extract (#118105)

The example below demonstrates a "scalar read followed by a broadcast"
pattern for `tensor.extract`:

```mlir
 #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func.func @scalar_broadcast(
    %init : tensor<1x1x3xi32>,
    %src: tensor<1x3x2x4xi32>,
    %idx :index) -> tensor<1x1x3xi32> {

   %c0 = arith.constant 0 :index

   %res = linalg.generic {
     indexing_maps = [#map],
     iterator_types = ["parallel", "parallel", "parallel"]}
     outs(%init : tensor<1x1x3xi32>) {
     ^bb0(%out: i32):
       %val = tensor.extract %src[%idx, %idx, %idx, %idx] : tensor<1x3x2x4xi32>
       linalg.yield %val : i32
   } -> tensor<1x1x3xi32>

   return %res : tensor<1x1x3xi32>
}
```

The default masking path within the Linalg vectorizer, which assumes an
identity masking map, is not suitable here. Indeed:

* identity != broadcast.

This patch ensures masking is handled in the `vectorizeTensorExtract`
hook, which has the necessary context for proper handling.

Fixes #116197



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