Quick tour

onnx-light is a rewrite of onnx that keeps the same format and the same API but drops protobuf along the way. It is protobuf free, written in C++ only, and parallelized.

Protobuf

Since the early days of ONNX, protobuf has been a recurring pain point:

• frequent conflicts between the installed version and the one used by ONNX,

• no way to switch from Python to C++ without going through serialization,

• complex compilation since a header has to be generated first,

• an API that cannot easily be extended,

• a hard 2 GB message-size limitation.

So the question is simple: can we have protobuf without protobuf?

First question

The prototype onnx/onnx#7208 demonstrated that a replacement keeping the same API and the same format, but without protobuf, is possible. The next question was whether to do a smooth rewriting or a complete rewriting from scratch. I chose to rewrite ONNX rather than modify the existing project: two ONNX implementations are better than one. Users can implement a switch and check both at the same time rather than going through major breaking changes in the original code base.

No protobuf, but also more

Removing protobuf brings several benefits:

• removes the 2 GB limitation,

• enables parallelization,

• allows custom loading strategies (read only part of a model without loading the entire file, control alignment, …),

• adds native C++ support for weight files,

• minimizes dependencies and security issues.

On top of that:

• no serialization layer between Python and C++,

• the same C++ and Python APIs,

• no auto-generated Markdown files in the repository, so pull requests no longer modify 200 files at once.

Migrating existing code only requires changing the import:

import onnx_light.onnx as onnx

Other requirements

A few other things change too:

• backend tests are no longer stored as files; they are implemented as C++ functions,

• backend tests are generated from a shared random test suite,

• expected outputs are produced by C++ kernels,

• kernels can be parallelized except for reductions, to ensure deterministic results,

• flexible C++ linking: users can link only the reading / writing / serialization components, or include schemas, kernels and/or backend tests as needed.

Symbolic shape inference

Shapes use real symbolic values instead of placeholders such as unk__0, unk__1, …

Core pieces: take what you need

The C++ code ships as small libraries so downstream projects link only what they need:

• parse and serialize, no schema needed (1.5 Mb),

• existing C++ ONNX library, no change (5 Mb),

• light schema (2 Mb, not loaded as a static variable),

• graph manipulations (0.6 Mb),

• new shape inference (1.5 Mb),

• C++ kernels (4.8 Mb),

• 2000+ C++ backend tests (7 Mb).

Parsing and serializing options

• data can be aligned when loading a tensor,

• parallelization is possible,

• onnxruntime FlatBuffers format is supported,

• more secure: recursion is limited,

• other scenarios can be implemented, such as a callback to allocate directly on GPU or any other memory.

Kernels

Kernels are kept simple, with a dispatch mechanism to support all types.

Backend tests

They are very similar to the existing ones, except they do not write any file.

Shape inference tests

New tests are dedicated to shape inference: value_info is filled with the expected values.

Symbolic expressions

Expressions only use symbols, with no function calls, to keep them short. There are two divisions:

• // rounds the result,

• /: is used when the exact result is known to be an integer.

For example, Reshape(x[N], [-1, 2]) tells us N is even, so we write N /: 2 and not N // 2: 2 * (N /: 2) == N while 2 * (N // 2) may differ.

Next steps

• implement a switch in ir-py to check it is compatible with onnx_light.onnx,

• do the same for onnxruntime and validate the C++ API,

• support more loading / saving scenarios,

• expose standard graph transformations (dead-end removal, …).