OxiONNX

Pure Rust ONNX Inference Engine -- Zero C/C++ Dependencies

OxiONNX is a high-performance ONNX inference engine written in pure Rust. It supports 189 ONNX operators, GPU acceleration via wgpu, SIMD optimization, and runs in the browser on wasm32-unknown-unknown (CPU inference via wasm-bindgen; build with -C target-feature=+simd128 to enable matrixmultiply's SIMD sgemm kernel) as well as every native target -- GPU acceleration ships for native targets today; an async WebGPU path for wasm32 now compiles and is exercised by native tests, but is not yet wired into the browser bindings (see Feature Flags).

159,672 lines of Rust | 3,595 tests | 0 clippy warnings

Live Demo

swap.cooljapan.tech is a face-swap application running on OxiONNX -- every forward pass in the pipeline goes through this crate, with no C/C++ dependency anywhere in the stack.

Source face and target image are read straight out of a canvas and stay on the machine running the engine. The result panel reports the wall-clock time OxiONNX spent in the engine -- 621 ms for the 1600 x 914 frame above, at the 640 px detector input setting. Below is that output at full resolution:

Above: synthetic image, generated by the demo. Not a photograph.

Synthetic media notice. Every image this demo produces is synthetic, and the demo labels it as such. Please keep that label attached wherever such output is reshared.

Features

• Pure Rust -- Zero C/C++/Fortran dependencies. Safe, portable, auditable.
• 189 ONNX operators -- Math, NN, Conv, Shape, Indexing, Comparison, RNN, Attention, ML; real-world detection models run, including YOLOv8 and YOLO11 (opset 11+)
• GPU dispatch is device- and shape-aware -- every oxionnx-gpu entry point decides whether a node is worth dispatching from GpuTuning, derived once from the adapter's own wgpu::AdapterInfo, rather than from compile-time constants. Three consequences worth knowing before turning the gpu feature on: a total FLOP count is not sufficient (a skinny [1,25088] x [25088,512] GEMM clears any FLOP floor and measured 1.54x slower than the CPU kernel on an RTX A4000, because it moves a 51.4 MB B to do 25.7 MFLOP -- so the gate also tests arithmetic intensity, 2mkn/(mk+kn+mn)); the same shape is 0.43x, a 2.3x win, once that B is served from the weight-residency cache, so the gate takes residency as an input rather than applying one rule to both; and a software adapter (Mesa lavapipe, SwiftShader, Direct3D WARP -- what a headless container gets from mesa-vulkan-drivers) declines every size, because it is the same CPU running one invocation per shader thread without rayon or matrixmultiply. On Linux, wgpu also needs the Vulkan loader package (libvulkan1 / vulkan-loader / vulkan-icd-loader), which GPU driver packages do not pull in; GpuContext::try_new_diagnosed() names it instead of returning a bare None
• GPU acceleration -- wgpu compute shaders for MatMul, Gemm (incl. transposed-B gemm_nt), Conv2D (direct implicit-GEMM kernel with no im2col materialization, ~692 GFLOP/s at InSwapper-128's 128x128 decoder layer on M3), Softmax, ReLU, PRelu, Resize, broadcast elementwise, etc. Session-lifetime weight residency caches Conv's W/B and Gemm's B/C initializers on the device after their first upload (Session::gpu_resident_bytes(); e.g. InSwapper-128's 502.7 MB/frame of convolution weights now cross the host<->device bus once per session instead of once per dispatch), and run-scoped activation residency lets one GPU node's output stay device-resident for the next GPU consumer to bind in place instead of a read-back plus re-upload at every node boundary (toggle via Session::activation_residency_enabled()/set_activation_residency()). A measured two-tier size gate (gpu_min_transfer_elements) declines to dispatch memory-bound elementwise ops whose operands must still cross the bus -- several measured strictly slower on the GPU than the rayon CPU kernel by more than an order of magnitude (up to ~36x for Relu) -- so turning on the gpu feature does not by itself guarantee a speedup on every op; residency and this size gate are what keep OpPlacement::Auto from trading a CPU win for a GPU loss
• SIMD optimization -- NEON (aarch64) and AVX2 (x86_64) for element-wise ops
• Multi-dtype -- f32, f16, bf16, i8, i32, i64 with automatic type promotion
• INT8 quantization -- Quantized MatMul with per-channel scale/zero-point
• Mixed precision -- f16 activations with f32 accumulation
• Graph optimization -- Constant folding, operator fusion, CSE, dead code elimination
• Memory efficiency -- Arena allocator, buffer pooling, strided tensor views
• Streaming inference -- Token-by-token generation for autoregressive models: session.generate(prompt, GenerationConfig) returns a TokenStream iterator that runs one forward pass per next(), feeds the model's present.* key/value outputs back in as the next step's past.* inputs, and stops on EOS, on a token cap, or on cancellation. Greedy (argmax) selection only -- temperature / top-k / top-p / beam search are deliberately out of scope; set emit_logits and sample outside the crate. No tokenizer: token ids in, token ids out
• Async execution -- Non-blocking inference via Arc::clone(&session).run_async(inputs), which starts the model on a std::thread immediately and returns a RunFuture. Executor-agnostic (no async-runtime dependency at all): .await it under tokio/async-std/smol, or drive it with the crate's own dependency-free block_on. spawn_run() returns a blocking RunHandle for callers with no executor. The receiver is Arc<Self> because the worker thread outlives the call; thread-per-inference is the right tool for one long inference, not for many small concurrent ones. Native targets only -- wasm32-unknown-unknown cannot spawn OS threads, so this whole module is compiled out there (a compile-time error on that target rather than the runtime panic it used to be); call Session::run synchronously instead
• Cancellation -- SessionBuilder::with_session_cancellation(token) makes every operator the model uses check a CancellationToken before it runs, so run() unwinds with OnnxError::Cancelled at the first node boundary after token.cancel() -- on the sequential path, the rayon parallel path, and inside If/Loop/Scan bodies. The token is session-scoped: cancelling stops every run in flight on that session. For per-request cancellation of a generation, use GenerationConfig::with_cancellation, which is checked between decode steps. Nodes claimed by a GPU execution provider are dispatched before the registry and are not cancellation points
• Control flow -- If/Loop/Scan operators with nested subgraph execution
• ONNX local functions -- FunctionProto bodies are inlined into the graph at load time, in both the eager and streaming parsers, so models built from reusable function definitions execute like any other graph
• Rank-generic convolution -- Conv / ConvTranspose support 1D/2D/3D spatial ranks through a shared N-D im2col path (pads uses the ONNX [begin_0..begin_r-1, end_0..end_r-1] layout, which for r == 2 is exactly the classic [top, left, bottom, right] array); the 2D case keeps its dedicated fast path
• Rank-0 (scalar) tensor support -- Tensor/TensorView represent a true ONNX scalar (shape: vec![], one element) instead of silently promoting every rank-reducing result to shape [1]; Det and the loss ops (NegativeLogLikelihoodLoss/SoftmaxCrossEntropyLoss with reduction=mean|sum) emit it end-to-end today, confirmed through shape resolution and output-slot allocation -- most other rank-reducing ops (e.g. an all-axes Reduce*) still promote to [1], a known tracked gap
• Opset-aware execution -- Softmax/LogSoftmax/Hardmax branch on the model's declared ai.onnx opset (parsed from opset_import) instead of hardcoding opset-13+ semantics, so a pre-13 model gets the spec's default-axis-1-and-flatten-to-2D contract rather than the post-13 per-axis one
• Einsum ellipsis and broadcasting -- the equation parser handles numpy-compatible ... tokens (e.g. ...ij,...jk->...ik for broadcast batched matmul) with numpy's right-aligned broadcasting rule, and a label shared across operands broadcasts when one side's extent is 1; large contractions lower to matrixmultiply::sgemm via greedy pairwise decomposition instead of a scalar loop nest
• Model encryption -- AES-GCM encrypted model files, keyed with CSPRNG-derived nonces
• WebAssembly -- wasm32-unknown-unknown CPU inference via wasm-bindgen (wasm feature); load a model from bytes and run it with no native code path involved. -C target-feature=+simd128 routes matrixmultiply's MatMul/Conv sgemm kernel through its v128 SIMD path (measured ~3-4.3x over the scalar fallback). Session::run_async/spawn_run (thread-per-inference) are unavailable on this target -- wasm32-unknown-unknown cannot spawn OS threads -- call Session::run synchronously instead; GPU acceleration (gpu feature) now builds for wasm32 too, with a working async execution path -- Session::enable_gpu_async/run_gpu_async, GpuContext::try_new_async acquiring a wgpu::Backends::BROWSER_WEBGPU adapter, kernels awaiting a real map_async read-back -- proven on native targets so far; it is not yet wired into these wasm-bindgen bindings (WasmSession still runs every GPU-eligible node on the CPU) or exercised in an actual browser
• no_std -- Core types work without std (alloc only)
• Session caching -- session.save_optimized(path) writes the post-optimization graph (nodes, rewritten weight table, value-info, model metadata, nested subgraphs) in a version-tagged, length-prefixed pure-Rust binary format; Session::load_optimized(path) / SessionBuilder::load_optimized(path) rebuilds it at OptLevel::None, so constant folding, CSE, fusion and dead-node elimination do not run again (the test suite proves this by counting operator executions during load: exactly zero). The encoding is deterministic, so a cache file can be content-hashed; a truncated, foreign or wrong-version file is always a typed OnnxError::Parse. Runtime settings (threads, providers, profiling, memory pool) are not cached -- they come from the builder that loads it
• Native dtype dispatch -- run_typed() path executes 40+ operators natively (no f32 round-trip) via TypedOpContext; MatMul and Gemm natively handle F32/F16/BF16/I8→I32/I32 dtypes, and Conv/ConvTranspose, Attention/MultiHeadAttention, and LSTM/GRU also have dedicated native F32/F16/BF16 typed kernels (cast-compute-cast with an f32 accumulator) beyond the original pilot set
• DirectML backend -- Windows D3D12 execution provider (directml feature) with CPU fallback on other platforms; opt-in (OXIONNX_DIRECTML=1 / .with_directml(true)), compile- and lint-verified for Windows and proven on Linux against a CPU oracle, but not yet executed on GPU hardware
• Zero-copy output reuse -- Operators write into pre-allocated output slots via execute_into_slots; a large subset of hot operators (elementwise, activations, normalization, reduce, pooling, shape, indexing, attention, conv, RNN) have hand-coded zero-copy slot-write kernels that avoid the intermediate copy and preserve pointer identity across inference runs with IoBinding. Operators without a hand-coded kernel fall back to a correct copy-based default (execute() then copy_from_slice)
• Graph introspection -- Enumerate a model's compute nodes (op type, inputs, outputs, attributes) via Session::nodes() / NodeInfo

Status

Crate	Status	Tests (all-features)*
oxionnx (root)	Alpha	1,113 passing
oxionnx-core	Stable	69 passing
oxionnx-ops	Alpha	1,325 passing
oxionnx-proto	Stable	134 passing
oxionnx-gpu	Alpha	272 passing
oxionnx-cuda	Partial	397 passing -- MatMul/Gemm (batched, with session-lifetime buffer-pool/weight/PTX-module caching), Conv (direct dispatch to oxicuda-dnn's Conv1x1/DepthwiseConv/ImplicitGemmConv), 16 unary activations, Add/Sub/Mul/Div (incl. channel and scalar broadcast), PRelu, BatchNorm/OxiInstanceNorm, ReduceSum/Max/Mean, Softmax, MaxPool/AveragePool, Resize, Pad, Slice, Concat, and zero-cost Reshape/Squeeze/Unsqueeze/Flatten -- 40 ops total via is_supported_op, plus opt-in CUDA Graph capture for MatMul/Gemm. Most of these tests need no GPU; the on-device gpu-tests suite skips gracefully on a host with no CUDA device, so this row's count is not evidence the device paths ran
oxionnx-directml	Implemented (opt-in; GPU path not yet hardware-verified)	242 tests, all Linux-executed or cross-target type-checked. Dual backend — DirectML operators + HLSL/D3D12 compute fallback — routing 15 ops: MatMul, Gemm, Add, Sub, Mul, Div, Relu, Sigmoid, Tanh, Softmax, ReduceSum, ReduceMean, ReduceMax, ReduceMin, Conv; kernels compile/lint-verified for Windows and proven on Linux vs a CPU oracle, but not yet run on GPU hardware
oxionnx-coreml	Alpha (opt-in; predict path not yet verified against a real model)	8 passing on non-Apple hosts (compiles to a stub); 43 passing + 8 skipped on macOS/iOS/tvOS/visionOS. The 43 cover the surrounding logic -- tensor/MLMultiArray layout conversion, f16 up-conversion, metadata handling. The 8 that skip are exactly the end-to-end ones (test_load_arcface, test_predict_arcface_returns_512_dim_embedding, the two compute-plan tests, test_model_metadata_returns_ok_map): they need a real .mlpackage, so predict/predict_raw/predict_features, ANE engagement and the compute plan are not exercised by the passing count. Apple frameworks are reached through objc2/objc2-core-ml, behind the off-by-default coreml feature

* Per-crate figures are all-features passing counts from this release (the oxionnx-coreml row uses its macOS figure, 43, since that is the host this run used) and sum exactly to the workspace total below: 1,113 + 69 + 1,325 + 134 + 272 + 397 + 242 + 43 = 3,595 passing.

Total: 3,595 tests passing with all features (cargo nextest run --workspace --all-features, run on macOS; the exact count drifts between runs as the workspace evolves under active development), 3,289 with default features (19 pre-filtered/ignored tests either way). 0 clippy warnings on the host target (cargo clippy --all-features --all-targets -- -D warnings). Platform-gated suites still run only on their target OS — this run's macOS host exercises oxionnx-coreml's Apple-only paths, included above (8 of its tests still skip, needing a real .mlpackage), but not oxionnx-directml's Windows FFI tests, and oxionnx-cuda's on-device tests skip for want of an NVIDIA device — so no single machine runs the entire cross-platform set, and a green total is not by itself evidence every hardware path ran. 190,392 lines of Rust (159,672 excluding blanks/comments).

Quick Start

use oxionnx::{Session, Tensor};
use std::collections::HashMap;

// Load model
let session = Session::from_file("model.onnx".as_ref())?;

// Prepare input
let mut inputs = HashMap::new();
inputs.insert("input", Tensor::new(vec![1.0, 2.0, 3.0], vec![1, 3]));

// Run inference
let outputs = session.run(&inputs)?;
println!("{:?}", outputs);

Session Builder

use oxionnx::{Session, OptLevel};

let session = Session::builder()
    .with_optimization_level(OptLevel::All)
    .with_memory_pool(true)
    .with_parallel_execution(true)
    .with_profiling()
    .load("model.onnx".as_ref())?;

Supported Operators

OxiONNX implements 189 ONNX operators (plus 15 aliases: short-forms like LayerNorm/RMSNorm/Silu/CeLU, plus the 11-name ai.onnx.ml.* domain) -- 204 op-type strings resolve through the registry in total.

Category	Count	Examples
Math	47	MatMul, Gemm, Add, Mul, Pow, Sqrt, Reduce* (incl. L1/L2/LogSum/LogSumExp/SumSquare), Trig, ArgMax/Min, CumSum, TopK, BitShift, VariadicMin/Max/Mean/Sum, Det
Neural Network	35	Relu, Sigmoid, Softmax, LayerNorm, BatchNorm, GELU, SiLU, Mish, GroupNorm, InstanceNorm, RmsNorm, Hardmax, Shrink, NegativeLogLikelihoodLoss, SoftmaxCrossEntropyLoss
Convolution / Pool	15	Conv, ConvTranspose (rank-generic: 1D/2D/3D), MaxPool, AveragePool, GlobalAvgPool, GlobalMaxPool, Pad, Resize, LRN, LpPool, GlobalLpPool, MaxUnpool, MaxRoiPool, Upsample, Col2Im
Shape	15	Reshape, Transpose, Concat, Slice, Split, Flatten, Tile, DepthToSpace, SpaceToDepth, ReverseSequence, Size, Expand, Squeeze, Unsqueeze, CenterCropPad
Indexing / Quant	16	Gather, GatherElements, GatherND, Scatter, ScatterND, Where, OneHot, Compress, Unique, QuantizeLinear, DequantizeLinear, QLinearConv, QLinearMatMul, MatMulInteger, ConvInteger, DynamicQuantizeLinear
Comparison / Logic	26	Equal, Greater, Less, And, Or, Not, Xor, Bitwise* (And/Or/Xor/Not), IsInf, IsNaN, NonZero, Cast, CastLike, Constant, Einsum, ConstantOfShape, EyeLike, Trilu, Identity, Shape, NonMaxSuppression
RNN / Attention	8	RNN, LSTM, GRU, Attention, MultiHeadAttention, RotaryEmbedding, GridSample, RoiAlign
DSP	7	DFT, STFT, HannWindow, HammingWindow, BlackmanWindow, MelWeightMatrix, Bernoulli
Control Flow	3	If, Loop, Scan
ONNX-ML	11	LinearClassifier, LinearRegressor, TreeEnsembleClassifier/Regressor, SVMClassifier/Regressor, Normalizer, Scaler, LabelEncoder, TfIdfVectorizer, StringNormalizer
Random / Generator	5	RandomNormal, RandomUniform, RandomNormalLike, RandomUniformLike, Multinomial

Feature Flags

Feature	Description
gpu	GPU acceleration via wgpu
simd	SIMD-accelerated element-wise ops
encryption	AES-GCM model encryption
cuda	CUDA GPU acceleration via OxiCUDA
mmap	Memory-mapped weight loading
wasm	WebAssembly browser bindings -- CPU inference on wasm32-unknown-unknown via wasm-bindgen. Combining this with gpu now builds: Session's Send + Sync requirement is wasm32-exempt and an async WebGPU path exists (enable_gpu_async/run_gpu_async, BROWSER_WEBGPU adapter, map_async read-back), proven on native targets only -- not yet wired into these bindings or run in a browser
wasm-threads	Re-enables rayon's intra-operator parallelism on wasm32. Needs a nightly -Z build-std browser build with +atomics,+bulk-memory, a cross-origin-isolated page, and a wasm-bindgen-rayon initThreadPool() call before the first operator. Inert on every non-wasm32 target
ndarray	ndarray interop for Tensor conversion
directml	DirectML GPU acceleration (Windows, via D3D12)
coreml	CoreML execution provider (Apple Silicon: macOS/iOS/tvOS/visionOS)

Architecture

oxionnx (root)           -- Session, optimizer, execution engine
  oxionnx-core           -- Tensor, DType, Graph, Operator trait, OnnxError
  oxionnx-ops            -- 189 operator implementations
  oxionnx-proto          -- Pure Rust ONNX protobuf parser
  oxionnx-gpu            -- wgpu compute backend (optional)
  oxionnx-cuda           -- CUDA dispatch layer via OxiCUDA (optional)
  oxionnx-directml       -- DirectML execution provider for Windows D3D12 (optional)
  oxionnx-coreml         -- CoreML execution provider for macOS/iOS/tvOS/visionOS (optional)

Performance

OxiONNX is a pure Rust implementation with no C/C++ BLAS dependency. Run cargo bench --bench performance to measure on your hardware.

Operator Microbenchmarks

benches/performance.rs covers MatMul (512², 1024², 2048², via matrixmultiply$), \text{Conv2D} (64\text{ch} 56 \times 56 3 \times 3, \text{im2col} + \text{matmul}), \text{Softmax} \text{and} \text{LayerNorm} \text{at} $[1, 128, 768], GELU over 100K elements (SIMD-accelerated under --features simd), and broadcast Add. Numbers are hardware-specific and deliberately not reproduced here -- run the bench to get yours.

End-to-End Model Workloads

Workload	Description	Notes
ResNet-50 backbone	Conv(3→64, 7×7) → BN → ReLU → MaxPool → 4 residual blocks	batch=1, 224×224 input
BERT attention	Q/K/V projections → scaled dot-product attention → output proj	seq=128, hidden=768
Transformer block	LayerNorm → Attention → FFN(GELU) → Residual	Stacked 4-layer encoder
Optimization passes	Session load with/without graph optimization	20-layer graph with dead code

Performance Characteristics

• Pure Rust, zero C/BLAS: All computation uses matrixmultiply (pure Rust BLAS-like) and hand-written kernels
• SIMD: Optional NEON (aarch64) and AVX2 (x86_64) acceleration for element-wise ops via --features simd
• Graph optimization: Constant folding, operator fusion, CSE, and dead code elimination reduce runtime overhead
• Memory pooling: Buffer reuse across inference calls reduces allocation pressure
• Parallelism: Rayon-based parallel execution of independent graph branches

Comparison note: OxiONNX prioritizes portability and safety: pure Rust with zero C/C++/Fortran dependencies, built on memory-safe pure-Rust crates. It is not unsafe-free — the ops crate confines unsafe to a few documented sites: the call into matrixmultiply (itself pure Rust) that backs MatMul/Conv in the default build, plus optional SIMD intrinsics compiled only under --features simd. In the default (non-simd) build oxionnx-ops is #![deny(unsafe_code)], with those matrixmultiply call sites the only explicitly-allowed exceptions. For absolute peak throughput, C++ runtimes like onnxruntime (with MKL/cuDNN) will be faster on operations dominated by BLAS. OxiONNX targets use cases where pure Rust, WebAssembly compatibility, and zero native dependencies are more important than raw FLOPS.

License

Apache-2.0

Author

COOLJAPAN OU (Team Kitasan)