[all-commits] [llvm/llvm-project] 66fed3: [mlir][vector] Update `castAwayContractionLeadingO...

[Kojo Acquah via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 66fed33db014bd705433e4b4f1ea766a8d71cadf
      https://github.com/llvm/llvm-project/commit/66fed33db014bd705433e4b4f1ea766a8d71cadf
  Author: Kojo Acquah <KoolJBlack at users.noreply.github.com>
  Date:   2024-04-03 (Wed, 03 Apr 2024)

  Changed paths:
    M mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
    M mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir

  Log Message:
  -----------
  [mlir][vector] Update `castAwayContractionLeadingOneDim` to omit transposes solely on leading unit dims.  (#85694)

Updates `castAwayContractionLeadingOneDim` to check for leading unit
dimensions before inserting `vector.transpose` ops.

Currently `castAwayContractionLeadingOneDim` removes all leading unit
dims based on the accumulator and transpose any subsequent operands to
match the accumulator indexing. This does not take into account if the
transpose is strictly necessary, for instance when given this
vector-matrix contract:
```mlir
  %result = vector.contract {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %lhs, %rhs, %acc : vector<1x1x8xi32>, vector<1x8x8xi32> into vector<1x8xi32>
```
Passing this through `castAwayContractionLeadingOneDim` pattern produces
the following:
```mlir
    %0 = vector.transpose %arg0, [1, 0, 2] : vector<1x1x8xi32> to vector<1x1x8xi32>
    %1 = vector.extract %0[0] : vector<1x8xi32> from vector<1x1x8xi32>
    %2 = vector.extract %arg2[0] : vector<8xi32> from vector<1x8xi32>
    %3 = vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d1)>], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %1, %arg1, %2 : vector<1x8xi32>, vector<1x8x8xi32> into vector<8xi32>
    %4 = vector.broadcast %3 : vector<8xi32> to vector<1x8xi32>
```
The `vector.transpose` introduced does not affect the underlying data
layout (effectively a no op), but it cannot be folded automatically.
This change avoids inserting transposes when only leading unit
dimensions are involved.

Fixes #85691



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