[Mlir-commits] [mlir] [tosa] : Enhance EqualizeRanks to handle dynamic dimensions. (PR #168564)

[Sayan Saha]

sahas3 wrote:

>  How about adding the test case you constructed somewhere here? https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

Thanks for the suggestion, but can you please elaborate more on this? 

The `EqualizeRank` function will operate on `tensor<?x?x?x?xf32>` and `tensor<5xf32>` to expand the latter to `tensor<1x1x1x5xf32>` when the `tosa.add` op is being constructed in the legalization example I have. However, it's not possible to create a `tosa.add` or other such tosa ops with broadcasting semantics with operands of types `tensor<?x?x?x?xf32>` and `tensor<5xf32>` that will trigger the code-path.

The only use of `EqualizeRanks` in `llvm-project` is in `TosaDecomposeDepthWise`, `TosaDecomposeTransposeConv` and `TosaMakeBroadcastable`. The first two passes don't handle dynamic dimensions so the fix in this change cannot be triggered -- I think it's possible to enhance these passes to handle dynamic batch and then maybe the change in this code can be triggered. I can look into that separately. 

For `TosaMakeBroadcastable`, it seems the pass may even be no-op now since it's expected that tosa ops will already have it's operands broadcasted during construction time, otherwise the `verify` method will error.


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