[all-commits] [llvm/llvm-project] 5ed5d7: [mlir][ArmSME] Lower multi-tile stores to a single...

[Benjamin Maxwell via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 5ed5d723db58f7556782427444642d0571cdf649
      https://github.com/llvm/llvm-project/commit/5ed5d723db58f7556782427444642d0571cdf649
  Author: Benjamin Maxwell <benjamin.maxwell at arm.com>
  Date:   2024-06-25 (Tue, 25 Jun 2024)

  Changed paths:
    M mlir/lib/Dialect/ArmSME/Transforms/VectorLegalization.cpp
    M mlir/test/Dialect/ArmSME/vector-legalization.mlir
    M mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/multi-tile-matmul.mlir

  Log Message:
  -----------
  [mlir][ArmSME] Lower multi-tile stores to a single loop (#96187)

This adds a new pattern that can legalize a multi-tile transfer_write as
a single store loop. This is done as part of type decomposition as at
this level we know each tile write is disjoint, but that information is
lost after decomposition (without analysis to reconstruct it).

Example (pseudo-MLIR):

```
vector.transfer_write %vector, %dest[%y, %x], %mask
  : vector<[16]x[8]xi16>, memref<?x?xi16>
```
Is rewritten to:
```
scf.for %slice_idx = %c0 to %c8_vscale step %c1 {
  %upper_slice_mask = vector.extract %mask[%slice_idx] ─┐
    : vector<[8]xi1> from vector<[16]x[8]xi1>           |
  %upper_slice = vector.extract %upper_tile[%slice_idx] |- Store upper tile
    : vector<[8]xi16> from vector<[8]x[8]xi16>          |
  vector.transfer_write %upper_slice,                   |
    %dest[%slice_idx + %y, %x], %upper_slice_mask       |
    : vector<[8]xi16>, memref<?x?xi16>                  ┘
  %lower_slice_idx = %slice_idx + %c8_vscale                 ─┐
  %lower_slice_mask = vector.extract %mask[%lower_slice_idx]  |
    : vector<[8]xi1> from vector<[16]x[8]xi1>                 |
  %lower_slice = vector.extract %lower_tile[%slice_idx]       |- Store lower
    : vector<[8]xi16> from vector<[8]x[8]xi16>                |  tile
  vector.transfer_write %lower_slice,                         |
    %dest[%lower_slice_idx + %y, %x], %lower_slice_mask       |
    : vector<[8]xi16>, memref<?x?xi16>                        ┘
}
```



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