[all-commits] [llvm/llvm-project] c02d07: [mlir][vector] Add pattern to drop unit dim from e...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: c02d07fdf007afc6b928cda0342751889cc2604b
      https://github.com/llvm/llvm-project/commit/c02d07fdf007afc6b928cda0342751889cc2604b
  Author: Andrzej Warzyński <andrzej.warzynski at arm.com>
  Date:   2023-12-13 (Wed, 13 Dec 2023)

  Changed paths:
    M mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
    M mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
    M mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
    M mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
    M mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

  Log Message:
  -----------
  [mlir][vector] Add pattern to drop unit dim from elementwise(a, b)) (#74817)

For vectors with either leading or trailing unit dim, replaces:

    elementwise(a, b)

with:

    sc_a = shape_cast(a)
    sc_b = shape_cast(b)
    res = elementwise(sc_a, sc_b)
    return shape_cast(res)

The newly inserted shape_cast Ops fold (before elementwise Op) and then
restore (after elementwise Op) the unit dim. Vectors `a` and `b` are
required to be rank > 1.

Example:
```mlir
  %mul = arith.mulf %B_row, %A_row : vector<1x[4]xf32>
  %cast = vector.shape_cast %mul : vector<1x[4]xf32> to vector<[4]xf32>
```

gets converted to:

```mlir
  %B_row_sc = vector.shape_cast %B_row : vector<1x[4]xf32> to vector<[4]xf32>
  %A_row_sc = vector.shape_cast %A_row : vector<1x[4]xf32> to vector<[4]xf32>
  %mul = arith.mulf %B_row_sc, %A_row_sc : vector<[4]xf32>
  %mul_sc = vector.shape_cast %mul : vector<[4]xf32> to vector<1x[4]xf32>
  %cast = vector.shape_cast %mul_sc : vector<1x[4]xf32> to vector<[4]xf32>
```

In practice, the bottom 2 shape_cast(s) will be folded away.