Skip to content

ONNX supportlink

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.

Overviewlink

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

Prerequisiteslink

  1. Install ONNX:

    python -m pip install onnx
    
  2. Install IREE packages, either by building from source or from pip:

    Stable release packages are published to PyPI.

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

    Nightly releases are published on GitHub releases.

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

Quickstartlink

  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 \
      model_cpu.vmfb \
      --device=local-task \
      --entry_function=... \
      --input=... \
      ...
    

Sampleslink

Code samples
Curated op and model tests SHARK-TestSuite e2eshark/onnx
Generated op tests SHARK-TestSuite iree_tests/onnx
Importer tests torch-mlir test/python/onnx_importer

Troubleshootinglink

Failed to legalize operation that was explicitly marked illegallink

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:
  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 ...