[all-commits] [llvm/llvm-project] 6ba0d1: [mlir][linalg] Enable scalable vectorization of li...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/users/banach-space/linalg/unpack_vec_update
  Home:   https://github.com/llvm/llvm-project
  Commit: 6ba0d1791ab07bd0c505ab9684dff6a1f4732375
      https://github.com/llvm/llvm-project/commit/6ba0d1791ab07bd0c505ab9684dff6a1f4732375
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

  Changed paths:
    M mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
    M mlir/test/Dialect/Linalg/vectorization/linalg-ops.mlir

  Log Message:
  -----------
  [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

This patch updates `vectorizeAsTensorUnpackOp` to support scalable
vectorization by requiring user-specified vector sizes for both the
_read_ and _write_ operations involved in `linalg.unpack`. Detailed
rationale and an example are provided below.

Conceptually, `linalg.unpack` consists of the following high-level steps:
  1. _Read_ from the source tensor.
  2. Transpose the value read in step (1).
  3. _Write_ the value from step (2) into the destination tensor.

Currently, when vectorizing with user-provided vector sizes, only the
sizes for the _write_ operation (step 3) are required. Sizes for the
_read_ operation (step 1) are inferred from static shapes and inner tile
sizes.

This logic breaks when the input shapes or tile sizes are dynamic
(indeed, `vectorizeUnPackOpPrecondition` rejects such cases ATM and the
vectorization fails). This patch addresses the issue by requiring
explicit vector sizes for both the read and write sides, enabling
scalable vectorization in such cases.

Example:

```mlir
func.func @unpack(%in: tensor<1x1x8x?xf32>, %out: tensor<8x?xf32>) -> tensor<8x?xf32> {
  %vs = vector.vscale
  %c8 = arith.constant 8 : index
  %tile_size = arith.muli %vs, %c8 : index

  %unpack = linalg.unpack  %in
    inner_dims_pos = [0, 1]
    inner_tiles = [8, %tile_size]
    into %out : tensor<1x1x8x?xf32> -> tensor<8x?xf32>
  return %unpack : tensor<8x?xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    transform.structured.vectorize %0 vector_sizes [1, 1, 8, [8],  8, [8]] : !transform.any_op
    //                                              \         /    \    /
    //                                              read-sizes   write-sizes
    transform.yield
  }
}
```

Finally, this patch also extends `createReadOrMaskedRead` and
`createWriteOrMaskedWrite` to take scalable flags.


  Commit: 9dd68bb73d528ca5f5ce3a0627a43d5b9a237158
      https://github.com/llvm/llvm-project/commit/9dd68bb73d528ca5f5ce3a0627a43d5b9a237158
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

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

  Log Message:
  -----------
  fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Remove leftover code + comments


  Commit: c67523cfdbf4a8e11063ecdb1ea459a770a4be29
      https://github.com/llvm/llvm-project/commit/c67523cfdbf4a8e11063ecdb1ea459a770a4be29
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

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

  Log Message:
  -----------
  fixup! fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Fix pre-condition calculation


  Commit: 7229a6d6f89ddf64c832a4e93611c4c924c921f4
      https://github.com/llvm/llvm-project/commit/7229a6d6f89ddf64c832a4e93611c4c924c921f4
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

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

  Log Message:
  -----------
  fixup! fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Improve documentation + fix test after rebasing on top of
* https://github.com/llvm/llvm-project/pull/150602


  Commit: 78b335b06f64ccb809d02971dced1580c9aa2dad
      https://github.com/llvm/llvm-project/commit/78b335b06f64ccb809d02971dced1580c9aa2dad
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

  Changed paths:
    M mlir/test/Dialect/Linalg/vectorization/linalg-ops.mlir

  Log Message:
  -----------
  fixup! fixup! fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Remove unintended test change


  Commit: 724cc9a5b6a0fc5cc9a9aa81938e05370f0c5d4e
      https://github.com/llvm/llvm-project/commit/724cc9a5b6a0fc5cc9a9aa81938e05370f0c5d4e
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

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

  Log Message:
  -----------
  fixup! fixup! fixup! fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Remove TODO


  Commit: 1d9f73778a74103bdfddb9112d2fcb7af8025981
      https://github.com/llvm/llvm-project/commit/1d9f73778a74103bdfddb9112d2fcb7af8025981
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-07-28 (Mon, 28 Jul 2025)

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

  Log Message:
  -----------
  fixup! fixup! fixup! fixup! fixup! [mlir][linalg] Enable scalable vectorization of linalg.unpack (WIP)

Fix comment


Compare: https://github.com/llvm/llvm-project/compare/cae3fb470c2c...1d9f73778a74

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