ANE — Training Neural Networks on Apple Neural Engine

GitHub: https://github.com/maderix/ANE

🔍 Overview

ANE is a research project that demonstrates, for the first time publicly, that full neural network training (forward + backward pass) can be executed directly on Apple's Neural Engine (ANE) — the dedicated AI accelerator built into Apple Silicon chips (M1/M2/M4 series).

Apple officially restricts the ANE to inference-only use through the CoreML framework. This project reverse-engineers Apple's private APIs (_ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor) to bypass that restriction and run custom compute graphs — including backpropagation — natively on the ANE hardware.

It is not a production framework. It is a proof of concept and a benchmark reference, proving that the limitation is software-imposed, not a hardware constraint.

🎯 Why This Matters

Apple Silicon chips (especially M4) contain an ANE rated at 15.8 TFLOPS — a massive amount of dedicated AI compute that Apple locks to inference. Every time you run CoreML, you get inference only. Training always falls back to the GPU or CPU.

This project answers the question: "Can you train on the ANE at all?"

The answer is yes — and this repo is the proof.

⚙️ How It Works

The project implements a complete transformer layer training loop by:

MIL Program Generation — Constructs Apple's Model Intermediate Language (MIL) programs at runtime in Objective-C, defining convolutions (for linear layers), matrix multiplications (for attention), softmax, and element-wise ops.
In-Memory Compilation — Uses _ANEInMemoryModelDescriptor to compile MIL text + weight blobs directly into ANE programs without needing to write .mlmodelc files to disk.
IOSurface I/O — Passes input/output tensors via IOSurface shared memory in [1, channels, 1, spatial] fp16 format — the format ANE hardware natively expects.
Weight Embedding — Weights are baked into ANE programs as BLOBFILE constants and recompiled each batch when weights are updated.
Gradient Flow — Forward "taps" expose intermediate activations needed for backward passes. Backward kernels compute input gradients (dx) on ANE; weight gradients (dW) are computed on CPU via cblas_sgemm.

🧱 Architecture: 6 ANE Kernels Per Training Step

Kernel	Function
`kFwdAttn`	RMSNorm + QKV projection + SDPA + output projection
`kFwdFFN`	RMSNorm + SwiGLU FFN (W1, W3, SiLU, W2)
`kFFNBwd`	FFN backward (W2ᵀ + SiLU_bwd + W1ᵀ + W3ᵀ)
`kSdpaBwd1`	Woᵀ + SDPA backward part 1 (dV, probs, dp)
`kSdpaBwd2`	SDPA backward part 2 (softmax grad, dQ, dK)
`kQKVb`	QKV backward (Wqᵀ + Wkᵀ + Wvᵀ → dx)

CPU handles: RMSNorm backward, residual connections, loss computation, dW gradient accumulation (cblas_sgemm), and Adam optimizer updates.

📊 Benchmark Results (M4 Mac, single transformer layer, dim=768, seq=512)

Optimization	ms/step	ANE Utilization
Baseline (vDSP transpose)	33.5	3.1%
Channel-first layout	20.3	5.2%
vDSP vectorized RMSNorm	14.2	7.4%
GCD async cblas overlap	11.4	9.2%
ANE RMSNorm fusion	11.4	9.2%
Wo^T fusion (7→6 kernels)	11.4	9.2%
Deferred cblas wait	9.3	11.2%

Best result: 9.3 ms/step, sustaining 1.78 TFLOPS (11.2% of the M4 ANE's 15.8 TFLOPS peak).

🔑 Key Optimizations

Channel-first CPU layout — Matches ANE's native [1,C,1,S] IOSurface format, eliminating all transpose overhead.
vDSP vectorized RMSNorm — 10× faster than naive implementation (6.7ms → 0.7ms).
GCD async cblas overlap — Weight gradient sgemm operations run in parallel with ANE evaluations on a serial dispatch queue.
Deferred cblas wait — The wait for dW sgemms is pushed into the next step's forward pass for maximum compute overlap.
ANE RMSNorm fusion — RMSNorm folded directly into forward ANE kernels as MIL ops.
Forward taps — Q, K, V, attention scores, and hidden states are exposed via concat outputs, avoiding CPU recompute during backward passes.
exec() restart — Bypasses the ~119 ANE compile limit per process using checkpoint/resume.

📁 Repository Structure

├── api_exploration.m       # Initial ANE API discovery experiments
├── inmem_basic.m           # In-memory MIL compilation proof-of-concept
├── inmem_bench.m           # ANE dispatch latency benchmarks
├── inmem_peak.m            # Peak TFLOPS measurement (2048×2048 matmul)
├── sram_bench.m            # ANE SRAM bandwidth probing
├── sram_probe.m            # SRAM size/layout exploration
└── training/
    ├── ane_runtime.h       # ANE private API wrapper (compile, eval, IOSurface)
    ├── ane_mil_gen.h       # MIL program generation helpers
    ├── model.h             # Model weight initialization and blob builders
    ├── forward.h           # Forward pass MIL generators
    ├── backward.h          # Backward pass MIL generators
    ├── train.m             # Minimal training loop (early prototype)
    ├── tiny_train.m        # 2-layer tiny model training
    ├── train_large.m       # Main: single-layer dim=768 training (optimized)
    ├── test_*.m            # Unit tests for individual kernels
    └── Makefile

🛠️ Building & Running

Requirements: macOS 15+ on Apple Silicon (tested on M4).

# Build
xcrun clang -O2 -framework Foundation -framework IOSurface \
  -framework CoreML -framework Accelerate -ldl -lobjc \
  -o train_large training/train_large.m

# Run
./train_large

No external dependencies — uses only system frameworks plus private ANE APIs resolved at runtime via objc_msgSend.

⚠️ Known Limitations

Limitation	Detail
Low ANE utilization	~11.2% of peak; many element-wise ops still fall back to CPU
~119 compile limit	ANE compiler leaks resources; worked around via `exec()` restart
Single transformer layer	Multi-layer pipeline scheduling not yet implemented
Synthetic data only	Real tokenized data support is work-in-progress
SDPA causal masking	ANE ignores `attn_mask` in SDPA; workaround via manual decomposition
Private APIs	Uses undocumented APIs that may break with any macOS update

📚 Related Research Articles

🏷️ Tech Stack

Language: Objective-C, C
Platform: macOS 15+ / Apple Silicon (M1/M2/M4)
Frameworks: Foundation, IOSurface, CoreML, Accelerate
Private APIs: _ANEClient, _ANECompiler, _ANEInMemoryModelDescriptor
Format: MIL (Model Intermediate Language), fp16

⚖️ Legal Disclaimer

This project uses Apple's private, undocumented APIs. These APIs carry no stability guarantee and may break with any macOS update. The project is independent research under fair use and interoperability provisions (Sega v. Accolade, 1992; DMCA §1201(f)). No Apple proprietary code or binaries are included. Not affiliated with or endorsed by Apple Inc.