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

[Rafael Ubal]

rafaelubalmw wrote:

@sabauma @krzysz00: Thank you very much for your insightful feedback. The advantages of sticking to the collapse-expand pattern seemed a good enough reason to revert my lowering strategy and instead focus on enhancing it to support all `tosa.reshape` cases. The new version differs almost completely from my original PR, so I'd appreciate if you could take another look. Here's a brief description:

- I modified existing auxiliary functions to avoid double-checking for invariants guaranteed by the `tosa.reshape` op verifier, such as target shape compatibility. At conversion, these invariants are checked with assert's rather than pattern match failures.

- Previously unsupported cases were caused by failures to guarantee type consistency in the emitted `tensor.expand_shape` op (i.e., input and output shapes must be both static or both dynamic). The result type for this op is now manually constructed, and an additional `tensor.cast` op is emitted if this type differs from the `tosa.reshape` result type.

- An extended set of tests are intended to cover relevant conversion paths. Tests are named using patten `test_reshape_<rank>_{up|down|same}_{s2s|s2d|d2s|d2d}_{explicit|auto}[_empty][_identity]`, where:

  - `<rank>` is the input rank (e.g., 3d, 6d)
  - `{up|down|same}` indicates whether the reshape increases, decreases, or retains the input rank.
  - `{s2s|s2d|d2s|d2d}` indicates whether reshape converts a statically shaped input to a statically shaped result (`s2s`), a statically shaped input to a dynamically shaped result (`s2d`), etc.
  - `{explicit|auto}` is used to indicate that all values in the `new_shape` attribute are >=0 (`explicit`) or that a -1 placeholder value is used (`auto`).
  - `empty` is used to indicate that `new_shape` includes a component set to 0.
  - `identity` is used when the input and result shapes are the same.

Other considerations regarding @krzysz00, to the extent that it is still applicable:

- I created tensor types using `clone()`, as suggested.

- I avoided the use of `auto` for basic types. I agree it helps to spell these out, especially when it comes to integer sizes. I did leave all other uses of `auto` (`SmallVector`, `Location`, `Value`, ...) as it seems consistent with other occurrences in the repo.



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