oxionnx

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

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

Feature Flags

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), Conv2D (64ch 56×56 3×3, im2col + matmul), Softmax and LayerNorm 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

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)