[all-commits] [llvm/llvm-project] 2f6b43: [mlir][linalg] Update vectorization logic for lina...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 2f6b433f2898f6d431dfefd393c2c7777c740418
      https://github.com/llvm/llvm-project/commit/2f6b433f2898f6d431dfefd393c2c7777c740418
  Author: Andrzej Warzyński <andrzej.warzynski at arm.com>
  Date:   2025-09-18 (Thu, 18 Sep 2025)

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

  Log Message:
  -----------
  [mlir][linalg] Update vectorization logic for linalg.pack (#149156) (#158926)

NOTE: See #149156 for a smilar change for `linalg.unpack`

This PR makes sure that we don't generate unnecessary `tensor.empty`
when vectorizing `linalg.pack`.

To better visualize the changes implemented here, consider this IR:
```mlir
func.func @example(
    %src: tensor<64x4xf32>,
    %dest: tensor<2x4x16x2xf32>) -> tensor<2x4x16x2xf32> {

  %pack = linalg.pack %src
    outer_dims_perm = [1, 0]
    inner_dims_pos = [0, 1]
    inner_tiles = [16, 2]
    into %dest : tensor<64x4xf32> -> tensor<2x4x16x2xf32>

  return %pack : tensor<2x4x16x2xf32>
}
```

Below is the output after vectorization, BEFORE and AFTER this PR.

BEFORE (note `tensor.empty` and the fact that `%arg1` is not used):
```mlir
  func.func @example(%arg0: tensor<64x4xf32>, %arg1: tensor<2x4x16x2xf32>) -> tensor<2x4x16x2xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %c0 = arith.constant 0 : index
    %0 = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : tensor<64x4xf32>, vector<64x4xf32>
    %1 = vector.shape_cast %0 : vector<64x4xf32> to vector<4x16x2x2xf32>
    %2 = vector.transpose %1, [2, 0, 1, 3] : vector<4x16x2x2xf32> to vector<2x4x16x2xf32>
    %3 = tensor.empty() : tensor<2x4x16x2xf32>
    %c0_0 = arith.constant 0 : index
    %4 = vector.transfer_write %2, %3[%c0_0, %c0_0, %c0_0, %c0_0] {in_bounds = [true, true, true, true]} : vector<2x4x16x2xf32>, tensor<2x4x16x2xf32>
    return %4 : tensor<2x4x16x2xf32>
  }
```

AFTER (note that `%arg1` is correctly used):
```mlir
func.func @example(%arg0: tensor<64x4xf32>, %arg1: tensor<2x4x16x2xf32>) -> tensor<2x4x16x2xf32> {
  %cst = arith.constant 0.000000e+00 : f32
  %c0 = arith.constant 0 : index
  %0 = vector.transfer_read %arg0[%c0, %c0], %cst {in_bounds = [true, true]} : tensor<64x4xf32>, vector<64x4xf32>
  %1 = vector.shape_cast %0 : vector<64x4xf32> to vector<4x16x2x2xf32>
  %2 = vector.transpose %1, [2, 0, 1, 3] : vector<4x16x2x2xf32> to vector<2x4x16x2xf32>
  %c0_0 = arith.constant 0 : index
  %3 = vector.transfer_write %2, %arg1[%c0_0, %c0_0, %c0_0, %c0_0] {in_bounds = [true, true, true, true]} : vector<2x4x16x2xf32>, tensor<2x4x16x2xf32>
  return %3 : tensor<2x4x16x2xf32>
}
```

ADDITIONAL CHANGES:
  * Adds missing `CHECK-LABEL` in tests.
  * Capitalize LIT test variables names.



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