[RFC] Add ONNX import support to torch-mlir

[qedawkins]

Motivation:

• Some customers and end users of torch-mlir would like to have a way to support ONNX as a “frontend” and lower down to dialects they care about (Linalg, MHLO, TOSA).
• Unify shape functions, and MHLO / TOSA exporters so customers don’t have to keep variations of these in multiple repositories.

Goals:

• Leverage as much of the onnx-mlir “frontend” without the backend kernels, shape functions and runtime and target the torch-mlir dialect / “torch-mlir backend contract”
• Retain the “EmitONNXBasic” export from onnx-mlir so exporting from torch-mlir can be used in onnx-mlir. (Discuss with onnx-mlir team if this is a possibility)
• Leverage OnnxtoOnnx decompositions from onnx-mlir (Discuss with onnx-mlir team if this is a possibility)
• Minimal dependencies in torch-mlir – add onnx submodule to externals/onnx.
• Request onnx-mlir team if they can dual license or LLVM license onnx-mlir code so we can share / importer code. If that is not a possibility add a dependency on a small onnx-mlir-exporter repo that can be pulled in at compile time. PoC: https://github.com/nod-ai/onnx-mlir-exporter
• Use torch-mlir shape functions etc once we are in Torch-MLIR dialect
• Use torch-mlir MHLO and TOSA exporters from Torch-MLIR dialect

Non-Goals:

Replicate all of onnx-mlir capabilities (Kernel dialect, Runtime etc)

Advantages:

Unified usage of torch-mlir shape library for both torch and ONNX paths.
Unified MHLO and TOSA export paths
ONNX – > Linalg is now a possibility to codegen ONNX graphs. (see onnx_add.py example)

Status / Proof of Concept:

An integration between onnx-mlir and torch-mlir is present in the following branch:
https://github.com/nod-ai/torch-mlir/tree/onnx-mlir

It incorporates onnx-mlir into the structure of torch-mlir to allow for natural access to the ONNX dialect. Without having to build all of the additional tools/features/runtime provided by onnx-mlir. We can take PyTorch modules, export to ONNX with torch.onnx.export, and subsequently load the ONNX module into the same frontend used for torch-mlir. This integration includes the ONNX-to-ONNX decompositions found in ONNX-MLIR.

E2E ONNX to Linalg Example

We have an end-to-end lowering PyTorch -> ONNX (torch.onnx.export) -> ONNX Dialect -> Torch Dialect -> Linalg -> Refbackend.
ONNX Add

Additional ONNX Lowerings

ONNX LogSoftmax

Resnet18

Looking forward to any feedback!

cc
@henrytwo @sjain-stanford @silvasean @ashay @stellaraccident @sstamenova @asaadaldien @stephenneuendorffer @fortianyou @ZihengJiang @byronyi @makslevental

[silvasean]

A few comments on this:

• If users are already using torch.onnx.export they can switch to torch_mlir.compile. So I would focus this effort on users that are producing ONNX but not via Torch
• torch.onnx.export already does shape inference/etc. so there is nothing to reuse from Torch-MLIR.
• it doesn't make sense for Torch-MLIR to depend on ONNX. If we want an ONNX path that reuses our "backend contract" -> {Linalg,TOSA,MHLO} lowerings then onnx-mlir can take a dependency on Torch-MLIR.

[qedawkins]

@silvasean Thank you for the feedback. Based on your suggestions I opened another RFC on the onnx-mlir side: onnx/onnx-mlir#1639

[silvasean]

This work seems to have gotten off the ground (see onnx/onnx-mlir#1639). Closing this issue.