[Mlir-commits] [mlir] [mlir][linalg] Fix partial fuse by collapse (PR #136326)

[Ian Wood]

IanWood1 wrote:

The `expand_shape` patterns and the `collapse_shape` patterns share the same utilities to preform the fusion. There is a test for fusing `expand_shapes` that checks that reshapes can be partially fused (`fuse_only_one_reassociation`). That's the behavior I'm trying to replicate with this change. The bug causing the invalid IR was due `getCollapsableIterationSpaceDims` returning a re-association that was not equal to `reshapeOp.getReassociationIndices()` (this allows for the partial fusion). Using the test I added as an example, `getCollapsableIterationSpaceDims` returns a reassociation that partially matches the `collapse_shape` which means that the linalg op generated by `collapseOpIterationDims` isn't as collapsed as the `collapse_shape` it can't be used to replace the collapse shape op. Here is the output of running `-test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing`:

```mlir
module {
  func.func @partial_fuse_by_collapsing(%arg0: tensor<4x?x32x128x192xf16>, %arg1: tensor<4x128x192x?x32xf32>) -> tensor<512x192x?xf32> {
    %c1 = arith.constant 1 : index
    %dim = tensor.dim %arg0, %c1 : tensor<4x?x32x128x192xf16>
    %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4]] : tensor<4x?x32x128x192xf16> into tensor<4x?x128x192xf16>
    %collapsed_0 = tensor.collapse_shape %arg1 [[0], [1], [2], [3, 4]] : tensor<4x128x192x?x32xf32> into tensor<4x128x192x?xf32>
    %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%collapsed : tensor<4x?x128x192xf16>) outs(%collapsed_0 : tensor<4x128x192x?xf32>) {
    ^bb0(%in: f16, %out: f32):
      linalg.yield %out : f32
    } -> tensor<4x128x192x?xf32>
    %expanded = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] output_shape [4, 128, 192, %dim, 32] : tensor<4x128x192x?xf32> into tensor<4x128x192x?x32xf32>
    %collapsed_1 = tensor.collapse_shape %expanded [[0, 1], [2], [3, 4]] : tensor<4x128x192x?x32xf32> into tensor<512x192x?xf32>
    return %collapsed_1 : tensor<512x192x?xf32>
  }
}
```

The tensor.expand_shape -> tensor.collapse_shape can be folded into a single collapse.



https://github.com/llvm/llvm-project/pull/136326