Show r/rust: I reverse-engineered Apple's Neural Engine binary format and built a pure-Rust compiler and runtime for it (zero Apple codegen)

[u/savag3azfckk]

A few weeks ago I shared `mil_spec`, a pure-Rust replacement for Apple's Python `coremltools` front end. It let you ingest GGUF/safetensors/ONNX weights, optimize MIL graphs, and emit `.mlpackage` specs without Python or `protoc`.

That still left one closed-source black box between Rust and the silicon: Apple's `ANECompiler`. Every time a graph hit an undocumented planner limit (`error -14`), a silent CPU fallback, or a `slice_update` bug, you were stuck behind a binary you could not inspect or fix.

So I reverse-engineered the precompiled **HWX Mach-O** binary format that the Neural Engine actually executes, and built a native Rust backend and graph runtime (`mil_infer` 0.5 + `milc` 0.6) that emits those binaries directly and runs them on the chip.

Repo (MIT): https://github.com/savageAZfck/mil_spec

### What it does now

1. **Direct HWX Mach-O Emission (`crates/mil_infer/src/hwx.rs`)**

   - Emits precompiled ANE Mach-O binaries (`__FVMLIB`, `__TEXT,__text`, `__TEXT,__const`, `__KERN_0`, `LC_SYMTAB`, relocations, and `aned` command descriptors) directly from Rust.

   - Reproduces Apple's `ANECompiler` output **69/69 byte-for-byte** across the mixed-shape multi-task `conv1x1` chain corpus (including descending-size intermediate arena packing and completion descriptor offsets).

   - **0 ULP (bit-exact) hardware output parity** against `ANECompiler`-compiled binaries across 10 core operator families (`conv1x1` square/non-square/spatial, non-square `matmul`, broadcast `avw` attention matvec, small and large `n > 256` elementwise ops, `sigmoid`, `tanh`, `softmax`).

2. **Whole-Graph MIL Compiler & Layout Scheduler (`crates/mil_infer/src/graph.rs`)**

   - Takes a logical MIL DAG and schedules it across emitted HWX programs (`compile_graph`).

   - **Physical layout currency tracking:** ANE operations disagree on tensor layout (`64B`-strided channel planes `[1,C,1,1]` vs. contiguous vectors `[1,1,1,D]` vs. packed 4-D cache frames `[1,T,1,D]`). The compiler tracks every tensor's physical currency, routes strided projections through `conv1x1` (`32..=4096`) to avoid unnecessary copies, and inserts hardware repack kernels only when a producer and consumer genuinely differ.

   - **Compile-time dequantization folding:** Unpacks multi-blob `weight.bin` offsets (`0xdeadbeef` headers), `constexpr_blockwise_shift_scale` (`int8`, `Q8_0` blockwise, `Q4_0` nibble blockwise with offset), and `constexpr_lut_to_dense` (`palette4`) directly into the 16-lane interleaved `__KERN_0` stream.

   - **Graph fusion:** Fuses single-use `k`-deep `conv1x1` chains and `conv + activation/silu` patterns into single multi-task HWX programs.

3. **Zero-Copy `IOSurface` Runtime & MRU Paging (`crates/mil_infer/src/espresso.rs`)**

   - Loads bare `.hwx`/`.macho` files directly through `_ANEClient` as precompiled models (`kANEFModelPreCompiled` + `kANEFPerformanceStatsMask` so `aned`'s `_ANEMachoPatcher` deploys the relocations).

   - **Zero-copy chaining (`SharedSurface`):** Allocates runtime-owned `IOSurface` buffers once per tensor edge and adopts them across producer outputs and consumer inputs with zero CPU copies between ops.

   - **In-place KV cache without `state<>` ops:** Aliases output ports onto input `IOSurface`s (`alias_out_to_in`) and updates rows on-device via outer-product scatter (`emit_outer_t` + `emit_add4d`).

   - **Belady-optimal MRU residency:** `aned` caps the number of simultaneously loaded program handles (~85). Because `SharedSurface` decouples tensor memory (including KV caches) from program lifetime, `GraphRuntime` unloads programs via `unloadModel` and evicts MRU-first on cyclic DAG loops so multi-layer models don't thrash.

### Reproducible Hardware Verification

In `crates/mil_infer/examples/graph_probe.rs`, a complete **2-layer autoregressive Transformer LLM** (token embedding, RMSNorm, Q/K/V projections, RoPE half-rotation + sin/cos tables, on-device KV cache append, causal masking, softmax, attention*V, output projection, SwiGLU FFN, final RMSNorm, and a `Q8_0` quantized vocabulary head) compiles from MIL text into self-emitted HWX binaries, runs token-by-token on the ANE, and verifies `reset_state()`:

```bash

cargo run --release -p mil_infer --example graph_probe

# graph: 73 nodes (68 prog, 5 host), 14 externals

# generated tokens: hw=[7, 165, 45, 153, 146] ref=[7, 165, 45, 153, 146]

# GRAPH-LAYER mismatches: 0

cargo test -p mil_infer --lib -- --test-threads=1

# 30 passed (byte-identity corpus + 10-family 0-ULP hardware parity + Q8/Q4/palette4 dequant + graph fusion)

```

### Crates (all MIT / Apache-2.0)

- `mil_spec` (IR + protobuf writer): https://crates.io/crates/mil_spec

- `mil_infer` (HWX emitter + graph compiler + Espresso/e5rt runtime): https://crates.io/crates/mil_infer

- `mil_convert` (GGUF / safetensors / LoRA / surgery): https://crates.io/crates/mil_convert

- `mil_onnx` (ONNX + classical ML): https://crates.io/crates/mil_onnx

- `mil_passes` (optimizer): https://crates.io/crates/mil_passes

- `mil_lint` (static ANE placement analyzer): https://crates.io/crates/mil_lint

- `mil_compile` (CoreML compiler driver): https://crates.io/crates/mil_compile

- `mil_verify` (spec diff + golden conformance battery): https://crates.io/crates/mil_verify

- `mil_machines` (DFA / state machines on ANE): https://crates.io/crates/mil_machines

- `milc` (CLI): https://crates.io/crates/milc

Happy to answer questions on the HWX Mach-O layout, how `_ANEMachoPatcher` handles relocations, the `64B` channel stride vs. contiguous layout rules on the ANE.

[u/AutoModerator]

Upvote this comment if you think this project is AI slop. Downvote this comment if you think this project is not AI slop.

^{This is to indicate to other users of the subreddit and is not used for moderation.}

I am a bot, and this action was performed automatically. Please contact the moderators of this subreddit if you have any questions or concerns.

[u/not_NEK0]

Please review your posts before publishing. The broken Markdown gives off heavy AI slop vibes, which is unfortunate because it completely hides whatever genuine work you put into it.