[Mlir-commits] [mlir] Fixes in 'tosa.reshape' lowering and folder (PR #85798)

[Spenser Bauman]

sabauma wrote:

I think it is worth considering the downstream implications of changing the lowering from `tensor.collapse_shape`/`tensor.expand_shape` to `tensor.reshape`. Some existing optimizations rely on the fact that `tensor.expand_shape` and `tensor.collapse_shape` are easier to reason about. The `--linalg-fuse-elementwise-ops` pass and its associated rewrite patterns comes to mind as a good example.

For instance, the following example is able to produce a single fused `linalg.generic` with the existing lowering, but will no longer fuse with this change.

```mlir
// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(tosa-to-tensor,tosa-to-linalg,linalg-fuse-elementwise-ops))"

func.func @main(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<10x10xf32> {
  %add1 = tosa.add %arg0, %arg1 : (tensor<100xf32>, tensor<100xf32>) -> (tensor<100xf32>)

  %2 = "tosa.reshape"(%add1) {new_shape = array<i64: 10, 10>} : (tensor<100xf32>) -> tensor<10x10xf32>
  %3 = "tosa.reshape"(%arg1) {new_shape = array<i64: 10, 10>} : (tensor<100xf32>) -> tensor<10x10xf32>

  %add2 = tosa.add %2, %3 : (tensor<10x10xf32>, tensor<10x10xf32>) -> (tensor<10x10xf32>)
  return %add2 : tensor<10x10xf32>
}
```

There may be other fallout that other reviewers are aware of. Maybe @eric-k256 or @rsuderman could chime in with more examples, or know those who would know.

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