[all-commits] [llvm/llvm-project] 27a713: [mlir][vector] Add scalable lowering for `transfer...

[Benjamin Maxwell via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 27a713f5b042bbcd88491c991877d0806aa66913
      https://github.com/llvm/llvm-project/commit/27a713f5b042bbcd88491c991877d0806aa66913
  Author: Benjamin Maxwell <benjamin.maxwell at arm.com>
  Date:   2024-08-12 (Mon, 12 Aug 2024)

  Changed paths:
    M mlir/include/mlir/Conversion/Passes.td
    M mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
    M mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
    M mlir/test/Conversion/VectorToSCF/tensor-transfer-ops.mlir
    M mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir

  Log Message:
  -----------
  [mlir][vector] Add scalable lowering for `transfer_write(transpose)` (#101353)

This specifically handles the case of a transpose from a vector type
like `vector<8x[4]xf32>` to `vector<[4]x8xf32>`. Such transposes occur
fairly frequently when scalably vectorizing `linalg.generic`s. There is
no direct lowering for these (as types like `vector<[4]x8xf32>` cannot
be represented in LLVM-IR). However, if the only use of the transpose is
a write, then it is possible to lower the `transfer_write(transpose)` as
a VLA loop.

Example:

```mlir
%transpose = vector.transpose %vec, [1, 0]
   : vector<4x[4]xf32> to vector<[4]x4xf32>
vector.transfer_write %transpose, %dest[%i, %j] {in_bounds = [true, true]}
   : vector<[4]x4xf32>,  memref<?x?xf32>
```

Becomes:

```mlir
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%c0 = arith.constant 0 : index
%0 = vector.extract %arg0[0] : vector<[4]xf32> from vector<4x[4]xf32>
%1 = vector.extract %arg0[1] : vector<[4]xf32> from vector<4x[4]xf32>
%2 = vector.extract %arg0[2] : vector<[4]xf32> from vector<4x[4]xf32>
%3 = vector.extract %arg0[3] : vector<[4]xf32> from vector<4x[4]xf32>
%vscale = vector.vscale
%c4_vscale = arith.muli %vscale, %c4 : index
scf.for %idx = %c0 to %c4_vscale step %c1 {
  %4 = vector.extract %0[%idx] : f32 from vector<[4]xf32>
  %5 = vector.extract %1[%idx] : f32 from vector<[4]xf32>
  %6 = vector.extract %2[%idx] : f32 from vector<[4]xf32>
  %7 = vector.extract %3[%idx] : f32 from vector<[4]xf32>
  %slice_i = affine.apply #map(%idx)[%i]
  %slice = vector.from_elements %4, %5, %6, %7 : vector<4xf32>
  vector.transfer_write %slice, %arg1[%slice_i, %j] {in_bounds = [true]}
    : vector<4xf32>, memref<?x?xf32>
}
```



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