ONNX support

Caution - under development

Support for a broad set of ONNX operators and data types is an active investment area. See the ONNX Op Support tracking issue for the latest status.

Overview

Machine learning models using the Open Neural Network Exchange (ONNX) format can be deployed using the IREE compiler and runtime:

graph LR
  accTitle: ONNX to runtime deployment workflow overview
  accDescr {
    Programs start as ONNX protobufs.
    Programs are imported into MLIR using iree-import-onnx.
    The IREE compiler uses the imported MLIR.
    Compiled programs are used by the runtime.
  }

  A["ONNX\n(protobuf)"]
  B["MLIR\n(torch-mlir)"]
  C[IREE compiler]
  D[Runtime deployment]

  A -- iree-import-onnx --> B
  B --> C
  C --> D

Prerequisites

1. Install ONNX:

   python -m pip install onnx
   

2. Install IREE packages, either by building from source or from pip:

   Stable releases Nightly releases

   Stable release packages are published to PyPI.

   python -m pip install \
     iree-base-compiler[onnx] \
     iree-base-runtime
   

   Nightly pre-releases are published on GitHub releases.

   python -m pip install \
     --find-links https://iree.dev/pip-release-links.html \
     --upgrade \
     --pre \
     iree-base-compiler[onnx] \
     iree-base-runtime
   

Quickstart

1. Start with a .onnx protobuf file, such as a model from https://github.com/onnx/models.

2. Convert the .onnx file into MLIR using the iree-import-onnx tool:

   iree-import-onnx [model.onnx] -o [model.mlir]
   

   This tool produces a MLIR file with the help of the torch-mlir project.

3. Once imported, the standard set of tools and APIs available for any of IREE's deployment configurations and API bindings can be used:

   iree-compile \
     model.mlir \
     --iree-hal-target-backends=llvm-cpu \
     -o model_cpu.vmfb
   
   iree-run-module \
     --module=model_cpu.vmfb \
     --device=local-task \
     --function=... \
     --input=... \
     ...
   

Samples

Code samples
Generated op tests	iree-test-suites onnx_ops
Public model tests	iree-test-suites onnx_models
Curated op and model tests	SHARK-TestSuite e2eshark/onnx and alt_e2eshark/onnx_tests
Importer tests	torch-mlir test/python/onnx_importer

Troubleshooting

Failed to legalize operation that was explicitly marked illegal

If you see an error compiling a converted .mlir file like this:

$ iree-compile model.mlir --iree-hal-target-backends=llvm-cpu -o model.vmfb

model.mlir:507:12: error: failed to legalize operation 'torch.operator' that was explicitly marked illegal
    %503 = torch.operator "onnx.Identity"(%arg0) : (!torch.vtensor<[?],si64>) -> !torch.vtensor<[?],si64>
           ^

There are several possible scenarios:

1. The operator is not implemented, or the implementation is missing a case. Search for a matching issue in one of these places:
   • https://github.com/llvm/torch-mlir/issues
   • https://github.com/nod-ai/SHARK-ModelDev/issues

2. The operator is implemented but only for a more recent ONNX version. You can try upgrading your .onnx file using the ONNX Version Converter:

   convert_onnx_model.py

   import onnx
   original_model = onnx.load_model("model.onnx")
   converted_model = onnx.version_converter.convert_version(original_model, 17)
   onnx.save(converted_model, "model_17.onnx")
   

   and then attempting the convert -> compile again:

   iree-import-onnx model_17.onnx -o model_17.mlir
   iree-compile model_17.mlir ...