[Mlir-commits] [mlir] [mlir][tensor] Loosen restrictions on folding dynamic reshapes (PR #137963)

[Ian Wood]

IanWood1 wrote:

I'm not sure if this is related to the failure seen here but `ComposeExpandOfCollapseOp` needs to be updated/fixed. It uses `getReassociationIndicesForCollapse` but doesn't correctly bail on dynamic cases. For example

```
func.func public @test(%arg0 : tensor<?x8x128x?xf16>, %arg1: index) -> tensor<?x128x?xf16> {
  %collapse = tensor.collapse_shape %arg0 [[0, 1, 2, 3]] : tensor<?x8x128x?xf16> into tensor<?xf16>
  %expanded_19 = tensor.expand_shape %collapse [[0, 1, 2]] output_shape [%arg1, 8, %arg1] : tensor<?xf16> into tensor<?x128x?xf16>
  return %expanded_19 : tensor<?x128x?xf16>
}
```

gets folded into

```
func.func public @test(%arg0: tensor<?x8x128x?xf16>, %arg1: index) -> tensor<?x128x?xf16> {
  %collapsed = tensor.collapse_shape %arg0 [[0, 1], [2], [3]] : tensor<?x8x128x?xf16> into tensor<?x128x?xf16>
  return %collapsed : tensor<?x128x?xf16>
}
```

This isn't valid. `expand(collapse)` can reshape to in any valid output shape but a single `collapse` can only collapse contiguous dims.

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