[Mlir-commits] [mlir] [mlir][tosa] Convert tosa.transpose_conv2d to linalg.generic directly (PR #79824)

[Rafael Ubal]

rafaelubalmw wrote:

Hi, Hsiangkai. Thank you for your contribution. A first look at the proposed lowering led me to the following questions and suggestions:

- Does this code work for any combination of input arguments and attributes?

- What would it take to have full support for dynamic dimensions? There is technically no requirement for `tosa.transpose_conv2d` to use static dimensions for weight or bias, which makes the corresponding error message rather misleading. We (at MathWorks) have recently been working on comprehensive lowerings for the `tfl` (Tensorflow Lite) dialect counterparts of 2D/3D transpose convs. We're aware it gets significantly trickier to support dynamic dims, but it is possible to emit generic code that is also efficient when only static dimensions are used (by relying on op folders).

- Has the correctness of this lowering been validated with any sort of runtime framework? We can help with that, but it would be desirable to first make sure that this lowering is complete according to the TOSA op spec.

- The test coverage looks minimal. An extensive set of unit tests should verify that the lowering strategy works correctly for a reasonable set of input argument/attribute combinations. But again, since FileCheck tests are very sensitive to the exact sequence of emitted ops, it'd desirable to attain full spec support before addressing this.

- A thorough verification of valid combinations of input arguments can save lots of headaches to end users (e.g., verify that `out_shape` is consistent with input/weight/stride/pad values when known, verify that `out_shape` matches output shape when static, verify that the number of input channels matches in input/weights when known, etc.). Failing to verify these will likely lead to obscure error messages when the `linalg.generic` op is emitted. Maybe some of these verifications belong in the op verifier instead. My personal preference here is including most lightweight verifications in the op verifier, and then adding `assert`s in the lowering with a comment indicating that the op verifier enforces a certain invariant.
	
- Is this intended to act as a replacement to the `TosaDecomposeTranposeConv` pass? If so, is there still a point in keeping that pass?

- It is not clear to me that it is preferable to lower directly to `linalg.generic` versus the decomposition approach based on pads + reverses + conv2d. Our approach when working on TFL lowerings was the latter, with the benefit that the resulting `tosa.conv2d`/`linalg.conv2d_xxx` ops may be mapped to device-specific intrinsics for better performance. Maybe allowing the user to optionally run the decomposition pass ahead of time would be a good tradeoff here.


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