OxiONNX

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

OxiONNX is a high-performance ONNX inference engine written in pure Rust. It supports 165 ONNX operators, GPU acceleration via wgpu, SIMD optimization, and runs on any platform including WebAssembly.

67,129 lines of Rust | 1,264 tests | 0 clippy warnings

Features

Pure Rust -- Zero C/C++/Fortran dependencies. Safe, portable, auditable.
165 ONNX operators -- Math, NN, Conv, Shape, Indexing, Comparison, RNN, Attention, ML; real-world detection models run, including YOLOv8 and YOLO11 (opset 11+)
GPU acceleration -- wgpu compute shaders for MatMul, Softmax, ReLU, etc.
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
Async execution -- Non-blocking inference via run_async()
Control flow -- If/Loop/Scan operators with nested subgraph execution
Model encryption -- AES-GCM encrypted model files
WebAssembly -- Run in the browser via wasm-bindgen
no_std -- Core types work without std (alloc only)
Session caching -- Save/load pre-optimized graphs to skip re-optimization
Native dtype dispatch -- run_typed() path executes 40+ operators natively (no f32 round-trip) via TypedOpContext; MatMul natively handles F32/F16/BF16/I8→I32/I32 dtypes
DirectML backend -- Windows D3D12 execution provider (directml feature) with CPU fallback on other platforms
Zero-copy output reuse -- All 121 operators support pre-allocated output slot reuse via executeintoslots; 105 operators have hand-coded zero-copy kernels (normalization, reduce, shape, attention, control-flow, and all elementwise/pool ops) — no memcpy, pointer-identity across inference runs with IoBinding
Graph introspection -- Enumerate a model's compute nodes (op type, inputs, outputs, attributes) via Session::nodes() / NodeInfo

Status

| Crate | Status | Tests | |-------|--------|-------| | oxionnx (root) | Alpha | 536 passing | | oxionnx-core | Stable | 36 passing | | oxionnx-ops | Alpha | 624 passing | | oxionnx-proto | Stable | 41 passing | | oxionnx-gpu | Alpha | 17 passing | | oxionnx-cuda | Partial | 4 passing (GEMM/elementwise/softmax via OxiCUDA; Conv stubbed) | | oxionnx-directml | Planned | 4 passing (Windows scaffold; HLSL shaders defined but not yet bound) | | oxionnx-coreml | Partial | 2 passing (CoreML session bridge; macOS only) |

Total: 1,264 tests passing, 0 clippy warnings, 67,129 SLoC

Quick Start

use oxionnx::{Session, Tensor};
use std::collections::HashMap;
// Load model let session = Session::fromfile("model.onnx".asref())?;
// 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()     .withoptimizationlevel(OptLevel::All)     .withmemorypool(true)     .withparallelexecution(true)     .with_profiling()     .load("model.onnx".as_ref())?;

Supported Operators

OxiONNX implements 165 ONNX operators (plus 21 aliases including the ai.onnx.ml.* domain)

| Category | Count | Examples | |----------|-------|---------| | Math | 46 | MatMul, Gemm, Add, Mul, Pow, Sqrt, Reduce* (incl. L1/L2/LogSum/LogSumExp/SumSquare), Trig, ArgMax/Min, CumSum, TopK, BitShift, VariadicMin/Max/Mean/Sum | | Neural Network | 33 | Relu, Sigmoid, Softmax, LayerNorm, BatchNorm, GELU, SiLU, Mish, GroupNorm, InstanceNorm, RmsNorm, Hardmax, Shrink | | Convolution / Pool | 8 | Conv, ConvTranspose, MaxPool, AveragePool, GlobalAvgPool, GlobalMaxPool, Pad, Resize | | Shape | 14 | Reshape, Transpose, Concat, Slice, Split, Flatten, Tile, DepthToSpace, SpaceToDepth, ReverseSequence, Size, Expand, Squeeze, Unsqueeze | | Indexing / Quant | 11 | Gather, GatherElements, GatherND, Scatter, ScatterND, Where, OneHot, Compress, Unique, QuantizeLinear, DequantizeLinear | | Comparison / Logic | 25 | Equal, Greater, Less, And, Or, Not, Xor, Bitwise* (And/Or/Xor/Not), IsInf, IsNaN, NonZero, Cast, Constant, Einsum, ConstantOfShape, EyeLike, Trilu, Identity, Shape, NonMaxSuppression | | RNN / Attention | 7 | 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 |

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 | | ndarray | ndarray interop for Tensor conversion | | directml | DirectML GPU acceleration (Windows, via D3D12) |

Architecture

oxionnx (root)           -- Session, optimizer, execution engine
  oxionnx-core           -- Tensor, DType, Graph, Operator trait, OnnxError
  oxionnx-ops            -- 165 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 (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

| Operation | Size | Implementation | Notes | |-----------|------|----------------|-------| | MatMul | 512×512 | matrixmultiply crate | Run cargo bench to measure | | MatMul | 1024×1024 | matrixmultiply crate | Run cargo bench to measure | | MatMul | 2048×2048 | matrixmultiply crate | Run cargo bench to measure | | Conv2D | 64ch, 56×56, 3×3 | im2col + matmul | Run cargo bench to measure | | Softmax | [1, 128, 768] | Numerically stable (log-sum-exp) | Run cargo bench to measure | | LayerNorm | [1, 128, 768] | Fused mean/var + scale/bias | Run cargo bench to measure | | GELU | 100K elements | SIMD-accelerated (with simd feature) | Run cargo bench to measure | | Add (broadcast) | [1, 128, 768] + [768] | Auto-broadcast | Run cargo bench to measure |

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