Rust EP spike — proving a Rust rewrite of the MLX execution provider

This is a vertical-slice spike (not the full EP). It de-risks a full Rust rewrite of onnxruntime-mlx by proving the two boundaries that were the only real unknowns, end-to-end, against the existing test harness.

What it proves

1. The ORT plugin-EP C ABI can be implemented entirely from Rust. src/factory.rs + src/ep.rs fill the OrtEpFactory, OrtEp, and OrtNodeComputeInfo C vtables with extern "C" functions, using the "embed the ORT struct as the first field" pattern so a *OrtEpFactory handed to ORT is pointer-identical to our *MlxEpFactory (repr(C), offset 0). Ownership crosses the C boundary via Box::into_raw / Box::from_raw, mirroring the C++ new/Release.

2. mlx-c can be bound DIRECTLY (no mlx-rs crate) and driven from Rust. build.rs runs bindgen over the mlx-c headers (mlx/c/mlx.h) to get a 1:1 mlx_* binding, and compute_add runs the op through mlx_array_new_data → mlx_add → mlx_array_eval → mlx_array_data_float32. We do NOT link libonnxruntime — ORT is reached purely through the OrtApi function-pointer table passed to CreateEpFactories.

Scope: claims Add (fp32) only. The oracle is the repo's own pytest suite: tests/ops/test_mlx_ops.py::test_binary_fp32[Add-...] (compares the EP output against ORT's CPU EP, tolerance-gated).

Results

[rust-mlx-ep] GetSupportedDevices: bound to GPU device
[rust-mlx-ep] GetCapability: claimed 1 Add node(s) of 1
[rust-mlx-ep] Add computed via mlx-c (6 elems)
1 passed

• Correctness: test_binary_fp32[Add] passes (MLX output == ORT CPU).
• Memory safety: 500 back-to-back sessions under macOS leaks → 0 leaks / 0 bytes, now enforced on every CI run (the Leak check gate runs three multi-op stress loops under leaks --atExit). The spike caught a real per-session mlx_stream leak (499 leaks / 15968 bytes) that a 3-line impl Drop for MlxEp fixed — the exact RAII win that motivates the rewrite (the C++ EP has hit this class of bug repeatedly: teardown UAF, the MRR MTLBuffer leak, manual ctx.Keep).

Build & run

export ORT_INCLUDE_DIR=<onnx-genai>/target/*/build/onnx-genai-ort-sys-*/out/ort-prebuilt/include
cargo build --release            # -> target/release/libonnxruntime_mlx_ep.dylib
# needs: brew install mlx-c mlx

ORT_LIB=<...ort-prebuilt/lib>
DYLD_LIBRARY_PATH=$ORT_LIB \
  ONNXRUNTIME_MLX_EP_LIB=$PWD/target/release/libonnxruntime_mlx_ep.dylib \
  python -m pytest ../tests/ops/test_mlx_ops.py -k "binary_fp32 and Add" -q -s

Observability tracing (src/trace.rs)

Env-gated GPU tracing into the pure-Rust onnx-runtime-tracer (Chrome/Perfetto JSON). It ships the feasible slice of docs/METAL_TRACING.md: because mlx-c only exposes the Xcode mlx_metal_start_capture and MLX fuses a whole subgraph into ONE hidden, synchronous mlx_eval, per-op MTLCommandBuffer gpuStartTime and per-kernel counters (design §4/§6) are unreachable — so the mlx_eval wall time (GPU-inclusive, since eval blocks) is the granularity.

# Write a Chrome/Perfetto trace (loads in https://ui.perfetto.dev):
ONNX_GENAI_MLX_TRACE=/tmp/mlx_trace.json DYLD_LIBRARY_PATH=$ORT_LIB \
  ONNXRUNTIME_MLX_EP_LIB=$PWD/target/release/libonnxruntime_mlx_ep.dylib \
  python -m pytest ../tests/ops/test_mlx_ops.py -k "binary_fp32 and Add" -q

• Spans: mlx.subgraph (cat ep, whole Compute) → nested mlx.eval (cat gpu, the synchronous eval = GPU-inclusive time); one <op_type> (cat op) span per node at graph-build time, plus (when tracing is on) a rich per-op detail span carrying input/output shapes, dtype, element count and byte size so every op has resource context even without fine mode.
• GPU counters (Chrome "C" phase, own Perfetto tracks): mlx.gpu_mem_bytes (MTLDevice.currentAllocatedSize), mlx.gpu_mem_pct (÷ recommendedMaxWorkingSetSize), and mlx.gpu_util_pct — GPU active-residency % read from the private IOReport framework (the GPUPH "GPU Performance States" channel, the same signal macmon/powermetrics use; no sudo). IOReport is resolved by dlopen/dlsym at runtime, so a missing/ABI-changed framework just disables the util counter — never a crash.
• os_signpost intervals around the same subgraph/eval regions for an Instruments Metal System Trace (ONNX_GENAI_MLX_SIGNPOST=1 forces them on).
• Events are stamped with the real pid (std::process::id()) so they merge into onnx-genai's Perfetto timeline. Written on EP teardown; the collector accumulates across sessions, so each teardown rewrites the full cumulative trace.
• Cost when off (env unset): a single relaxed atomic load + early return per entry point — no signpost log, no device handle, no IOReport sampler, no allocation, and the single fused mlx_eval is left untouched.

Seeing inside the fused mlx_eval

By default a whole fused subgraph evaluates as ONE opaque mlx.eval blob (MLX hides its per-kernel Metal command buffers). Two opt-in modes break that open:

• Fine-grained per-op GPU timing — ONNX_GENAI_MLX_TRACE_FINE=1 (implies tracing on). After each node's handler binds its outputs, they are mlx_array_eval'd individually and timed, emitting a gpu.op span per op with the op's GPU-inclusive wall time + shape/dtype/bytes Args. The single blob becomes per-op bars showing which op dominates. This BREAKS fusion (materialising per node defeats MLX's lazy graph) so it is slower and strictly a debug tool — the normal path keeps the single fused eval.

  ONNX_GENAI_MLX_TRACE=/tmp/fine.json ONNX_GENAI_MLX_TRACE_FINE=1 ... python ...
  

• Xcode/Instruments GPU capture — ONNX_GENAI_MLX_GPU_CAPTURE=<path.gputrace> (or =1 for a default path) wraps the first eval only in mlx_metal_start_capture … stop_capture, producing a .gputrace bundle with full per-kernel GPU timing / occupancy / memory-bandwidth. Requires MTL_CAPTURE_ENABLED=1 exported before the process starts (the capture layer is inserted at device creation); without it the mode logs a clear message and skips rather than aborting.

  MTL_CAPTURE_ENABLED=1 ONNX_GENAI_MLX_GPU_CAPTURE=/tmp/cap.gputrace ... python ...
  open /tmp/cap.gputrace   # in Xcode
  

• Slowest-ops summary — whenever tracing is on, EP teardown logs a compact top-10 (op_type → total µs, %, call count) to stderr and as an mlx.slowest_ops trace-metadata event, so an agent sees e.g. "Add = 80% of GPU time" without parsing the JSON. Times are GPU-inclusive in fine mode, otherwise build-time (noted in the summary).

Full-port plan (what this unlocks)

The two boundaries are proven; the rest is mechanical, guarded by the language-agnostic pytest suite (ONNX models vs ORT-CPU reference):

1. mlx-c-sys crate — bindgen over all mlx-c headers (the C++ EP uses 181 mlx_* symbols, incl. fast_scaled_dot_product_attention, fast_rope, fast_rms_norm, quantized_matmul, compile), plus safe RAII wrappers (Array/Stream/VectorArray with Drop).
2. ort-ep-sys — bindgen over the ORT EP C ABI (reuse/extend onnx-genai's onnx-genai-ort-sys).
3. Engine + registry — port TranslationContext/NodeDesc and the (domain,op,[min,max]opset) registry; add the ~24 op modules in waves, each validated against its pytest module.
4. DataTransfer + allocator — the unified-memory memcpy transfer + a Metal-buffer allocator (the C++ has ~5 raw Metal calls; use metal-rs). The spike keeps I/O on the CPU allocator, which was sufficient to prove the boundaries; the GPU-memory path is a known-simple follow-up.
5. pyo3 packaging — abi3 + free-threaded (abi3t) wheels, replacing nanobind.

Update: foundation + wave-1 (engine generalized)

The single-Add spike has been generalized into a real engine + registry that ports the first wave of ops. This is no longer a single hardcoded op.

Module structure

• src/mlx.rs — RAII layer. Safe Stream, Array, VectorArray wrappers over sys::mlx, each with impl Drop calling the matching mlx_*_free. All ownership of mlx refs lives here; op handlers never free manually. This is where the 0-leak result comes from. Raw bindgen stays in sys::mlx.
• src/engine.rs — Engine core. NodeDesc (op_type/domain/since_version + int/float/array/string attrs + input/output tensor names), the Plan (one per fused subgraph), and TranslationContext which owns a name -> mlx_array environment plus an arena: Vec<Array> (freed at run end) and a persistent cache for constants. Provides resolve/bind/keep and the eager execute/finish_boundary/copy_out. mlx_dtype_from_onnx maps ONNX element types to mlx_dtype (fp32/fp16/bf16/int32/int64/… ) for the copy path.
• src/registry.rs — Registry. (domain, op_type) -> { handler, claim predicate }, an OnceLock singleton wired by register_builtin_ops. claim and translate are the single source of truth so claimed == translatable. NodeView is the claim-time FFI wrapper (reads inputs/outputs/attrs off the OrtNode before compile). Includes claim helpers (is_mlx_float, suffix_broadcast, …).
• src/ops/elementwise.rs, src/ops/math.rs — Wave-1 handlers. Each op is handler + claim predicate + registry entry.
• src/ep.rs — Generalized boundary. GetCapability claims nodes via the registry then groups them into maximal convex connected subgraphs with a faithful port of BuildConvexClusters (union-find + reachability bitsets, prevents the cycles ORT rejects). Compile extracts each node's NodeDesc (attrs via Node_GetAttributes/ReadOpAttr, tensor names via Node_GetInputs/GetOutputs) and builds one Plan per subgraph. Compute (RunPlan port) resolves subgraph inputs from the KernelContext, runs each node's handler in topo order, does a single mlx_eval at the boundary, then writes each output via KernelContext_GetOutput + a unified-memory memcpy.

Wave-1 ops (all pass through the Rust EP)

Add, Sub, Mul, Div, Neg, Abs, Sqrt, Exp, Log, Relu, Sigmoid, Tanh (+ Softmax last-axis and Cast). fp32 required; the copy path also handles fp16/bf16/int32/int64.

Normalization + attention ops (ops/norm.rs, ops/attention.rs)

The transformer decode path. All claim + translate through the same registry and run on MLX (fp32/fp16/bf16), verified against ORT CPU by tests/ops:

• Normalization — RMSNormalization, LayerNormalization, SimplifiedLayerNormalization, SkipLayerNormalization, SkipSimplifiedLayerNormalization, GroupNormalization, BatchNormalization (inference form), LpNormalization. The last-axis forms use mlx_fast_rms_norm / mlx_fast_layer_norm; the rest compose mean/var/rsqrt.
• Attention — GroupQueryAttention (in-op RoPE + KV-cache append + causal SDPA, multi-output attn/present_key/present_value), Attention (ai.onnx opset 23 & 24, 3D/4D, optional attn_mask + past/present KV), MultiHeadAttention (com.microsoft, optional projection bias), RotaryEmbedding (ai.onnx opset 23 & com.microsoft, gather / offset / absent position_ids, rotate-half + interleaved, partial rotation). SDPA maps onto mlx_fast_scaled_dot_product_attention.
• Leak check — stress_norm_attn.py under MallocStackLogging=1 leaks --atExit → 0 leaks / 0 total leaked bytes (exercises the fast-norm / fast-SDPA / RoPE / multi-output present-K/V paths).

Edge cases intentionally left on CPU (claim returns false), matching the C++ EP: attention softcap, the qk_matmul_output extra output, the opset-24 nonpad_kv_seqlen input, and the is_causal + explicit attn_mask combination (MLX fast SDPA cannot mix a causal mode with an array mask); GQA smooth_softmax / qk_output; MHA packed-QKV and every masked / past-KV form (they imply an interior optional gap the subgraph builder cannot consume); norm Mean/InvStdDev extra outputs. The compiled-decode fast-path (dynamic cos/sin slice, rotate-half matmul) is next-wave — the eager single-mlx_eval path is implemented here.

Results

• cargo build --release — clean.
• pytest tests/ops/test_mlx_ops.py -k "binary_fp32 or sigmoid_fp32 or softmax_fp32" — 5 passed, all through MLX (verified by [rust-mlx-ep] GetCapability/Compute stderr). Full test_mlx_ops.py — 31 passed; math ops (Relu/Tanh/Neg/Abs/Sqrt/Exp/Log/Div) — 21 passed, 18 through MLX.
• Leak check — 500-session Add stress loop under MallocStackLogging=1 leaks --atExit → 0 leaks / 0 total leaked bytes (rust/stress_add.py).

What the next wave needs from the engine

• Initializer / constant handling — Src::Initializer + the constant-flag cache path exist but are only lightly exercised; reductions/norm/matmul/quant need weights resolved once and kept in the persistent Plan.cache.
• Reading constant host bytes (a RawHost accessor) — ops like Slice/Trilu/OneHot/Reshape read integer/shape operands on the host, not as mlx arrays.
• Multi-output nodes — the plan/output binding currently assumes the common 1-output case; TopK/Split/attention need N outputs bound and copied.
• Subgraph / control-flow attrs — SubgraphDesc (If/Scan/Loop bodies) is not yet ported; the convex-cluster singleton special-case for control-flow ops was intentionally omitted in wave-1.
• Reductions & shape/data-movement helpers — axis normalization, keepdims, and gather/concat/reshape wiring in TranslationContext.
• Compiled-decode fast-path (mlx_compile) — omitted; a later perf item.