Inside Apple's Neural Engine: Reverse-Engineering ANE for Neural Network Training

The Apple Neural Engine (ANE) is a marvel of specialized silicon. Integrated into every Apple Silicon chip since the A11, it's a high-performance, power-efficient NPU designed to accelerate machine learning inference. With the M4 chip achieving 15.8 TFLOPS of FP16 performance, the ANE is a powerhouse. Yet, for all its might, Apple has kept it on a tight leash. Officially, its capabilities are only accessible for inference through the high-level CoreML framework. Training neural networks? That's strictly off-limits, reserved for the GPU and CPU.

Or is it?

The open-source project ANE Training challenges this fundamental assumption. Through meticulous reverse-engineering of private Apple frameworks, it demonstrates a complete, end-to-end training implementation, encompassing dynamic program generation, zero-copy memory I/O, and custom kernel development for both forward and backward passes.

Core Technical Architecture

At its heart, this project is a feat of reverse engineering. The entire toolchain is built upon undocumented, private APIs within macOS. Accessing them required peeling back the layers of Apple's frameworks to understand how CoreML communicates with the ANE driver stack.

Private API Reverse Engineering

The official path to the ANE is through CoreML, which abstracts away the hardware details. To bypass this, ANE Training interacts directly with the underlying private frameworks: AppleNeuralEngine.framework and ANECompiler.framework. The key private classes and their roles:

• _ANECompilerService: Compiles models into ANE-executable binary programs.
• _ANEInMemoryModel / _ANEInMemoryModelDescriptor: Allows creating ANE models directly from memory buffers, bypassing the need for a .mlmodelc file on disk. This is the central mechanism enabling dynamic network generation.
• _ANEClient: Communicates with the ANE hardware daemon aned.
• _ANERequest: Encapsulates a specific computation request, containing input/output IOSurfaces and the model to execute.

Since these are private APIs, the project employs dynamic runtime resolution rather than compile-time linking:

// Dynamic API resolution via objc_msgSend
typedef id (*objc_msgSend_t)(id, SEL, ...);
objc_msgSend_t msgSend = (objc_msgSend_t)objc_msgSend;

// Dynamically get the _ANEClient class
Class ANEClient = NSClassFromString(@"_ANEClient");
id client = msgSend(ANEClient, sel_getUid("sharedClient"));

// Dynamically call methods on private class instances
SEL evalSelector = sel_getUid("evaluateWithQoS:options:request:error:");
msgSend(client, evalSelector, qos, options, request, &error);

The core functions built on this foundation:

1. ane_init(): Uses dlopen to load AppleNeuralEngine.framework at runtime, then resolves private class handles via NSClassFromString.
2. ane_compile(): Takes MIL text and weight blob, compiles via _ANEInMemoryModelDescriptor and _ANECompilerService. Includes 100ms retry on load failure for ANE resource reclamation.
3. ane_eval(): Creates an _ANERequest with bound IOSurface buffers, submits to ANE via evaluateWithQoS.

MIL (Model Intermediate Language)

MIL is the language through which ANE Training communicates with the ANE compiler. It's a text-based intermediate representation describing neural network computation graphs. The project generates MIL program text dynamically at runtime.

Alongside the MIL text, the compiler requires a binary "weight blob" containing the model's parameters. Reverse engineering revealed its format:

• A 64-byte global header (buf[0]=0x01, buf[4]=0x02)
• Per-weight chunks, each with a 64-byte header containing magic number 0xDEADBEEF, data size, and data offset
• Weight data stored in FP16 format (F32 converted via (_Float16) truncation)

The MIL generator supports the essential operations needed for transformers: convolution (for linear layers), matmul (for attention), softmax, and element-wise operations.

IOSurface Tensor I/O

Data movement between CPU and ANE is a critical performance factor. ANE Training uses IOSurface — Apple's framework for zero-copy memory sharing between processes and hardware accelerators. An IOSurface is a buffer in system memory simultaneously mapped into both CPU and ANE address spaces, eliminating costly data copies.

Working with shared memory requires synchronization:

• IOSurfaceLock(): CPU gains exclusive access for reading/writing
• IOSurfaceUnlock(): CPU releases the lock, signaling the buffer is ready for ANE

A crucial discovery was the ANE's preferred tensor layout: channels-first format [1, Channels, 1, Spatial]. The project maintains this layout throughout all CPU-side computations, avoiding expensive transpose operations at every ANE boundary.

Training Implementation

A complete transformer training step is decomposed into multiple ANE kernels and CPU tasks.

Forward Pass (6 ANE Kernels)

The forward pass uses six distinct ANE kernels:

This decomposition is deliberate — forward kernels "tap" intermediate values by concatenating Q, K, V, and attention scores into output IOSurfaces, making them directly available to backward kernels without CPU intervention.

Operations like RMSNorm are fused directly into these kernels as MIL operations (reduce_sum + pow + mul), reducing kernel launch overhead and memory bandwidth requirements.

Backward Pass (Hybrid CPU/ANE)

The backward pass is a hybrid computation:

ANE handles the compute-intensive transposed matrix multiplications (dx = W^T @ dy) — all input gradient computations via kFFNBwd, kSdpaBwd1/2, and kQKVb.

CPU handles the remaining operations:

• RMSNorm backward: Chain rule via vDSP-vectorized implementation
• Residual connections: Simple gradient additions
• Loss calculation: Numerically stable softmax + log-prob with max subtraction
• Gradient accumulation (dW): dW += dy @ x^T via Accelerate's cblas_sgemm
• Adam optimizer: Full optimizer step (m/v states, bias correction) in FP32
• Gradient clipping: Global L2 norm computation, scaled when exceeding max_norm

This is a pragmatic heterogeneous computing design: ANE handles parallel, compute-intensive tensor operations while CPU manages serial, control-flow-heavy, or ANE-unsupported operations.

Key Optimizations

The project employs several system-level optimizations to squeeze out performance:

1. Channel-First CPU Layout

By maintaining [1, C, 1, S] data layout on the CPU side throughout, the project eliminates all transpose operations. For a 110M model, each transpose costs several milliseconds — a significant overhead when accumulated across every training step.

2. Vectorized RMSNorm (vDSP)

The initial C-loop RMSNorm backward took 6.7ms. Rewriting with Accelerate's vDSP (leveraging SIMD units) dropped it to 0.7ms — a nearly 10x speedup.

3. Asynchronous cblas Overlap with GCD

Gradient accumulation (cblas_sgemm) is dispatched to a background serial GCD queue, running in parallel with ANE kernel execution. While the ANE processes kFwdAttn, the CPU simultaneously computes the previous step's dW gradients.

4. Deferred cblas Wait

The main thread doesn't immediately wait for async sgemm completion. The wait is deferred until the next forward pass actually needs the result, maximizing CPU/ANE parallelism.

5. ANE Kernel Fusion

RMSNorm is fused into kFwdAttn and kFwdFFN as MIL operations. The backward pass fuses Wo^T into kSdpaBwd1, reducing the kernel count from 7 to 6.

6. exec() Restart

The ANE compiler service fails after approximately 119 compilations per process — likely a resource leak in the aned daemon. The workaround: save a binary checkpoint, then use execl() to restart the process and resume from the checkpoint. To the user, it appears as a single continuous training run.

Optimization History

Dynamic Weight Pipeline

The initial "static" pipeline bakes weights into compiled ANE programs, requiring full recompilation after every weight update. Profiling revealed that 76% of training time was spent on compilation alone.

The "dynamic" pipeline solves this by treating weights as model inputs rather than static data. Kernels are compiled once at startup, and weights are packed into IOSurface input tensors alongside activations:

Input tensor layout (sdpaFwd example):
[1, DIM, 1, SEQ + 4*DIM] fp32
  [0:SEQ]           = xnorm (activation)
  [SEQ:SEQ+DIM]     = Wq (weight matrix)
  [SEQ+DIM:SEQ+2D]  = Wk
  [SEQ+2D:SEQ+3D]   = Wv
  [SEQ+3D:SEQ+4D]   = Wo

The ANE kernel slices out the weights it needs at runtime. Weight updates require only IOSurfaceLock → memcpy → IOSurfaceUnlock — negligible cost compared to recompilation.

SRAM and Performance Analysis

The ANE contains fast on-chip SRAM, estimated at approximately ~16MB. When a kernel's working set fits within SRAM, performance is excellent. Exceeding this threshold causes a dramatic "performance cliff" as data spills to system DRAM.

Peak performance was measured via chained convolution benchmarks, confirming the M4 ANE's 15.8 TFLOPS FP16 capability. However, actual transformer training utilization is only 5-9% (1-2 TFLOPS). This indicates the bottleneck lies not in ANE compute, but in CPU fallbacks, CPU/ANE synchronization overhead, and kernel scheduling latency.

Stories110M Results

End-to-end training was benchmarked on a 109M parameter transformer (12 layers, dim=768, 12 heads, seq=256):

The M3 Ultra outperforming the M4 is notable — the bottleneck is the surrounding CPU orchestration and DRAM bandwidth, not ANE raw compute power.

Engineering Observations

The implementation reflects several careful engineering choices:

• Pure Objective-C: ~23,600 lines of code, zero external dependencies
• System frameworks only: Foundation, IOSurface, CoreML, Accelerate — maximum portability
• Bridge API: C-style bridge for Python ctypes bindings, extending accessibility to the ML community
• Comprehensive testing: 13 test files covering the full stack from low-level API to end-to-end training
• Binary checkpointing: Custom format storing weights, Adam states, and cumulative training time for interrupt/resume support

Limitations and Future Work

The project reveals several boundary conditions:

• SDPA causal masking: The ANE appears to ignore attn_mask in SDPA operations, requiring decomposition into separate Q@K^T → mask+softmax → scores@V operations
• ~119 compile limit: The exec() restart is a workaround; a principled fix requires understanding the resource leak within Apple's compiler service
• Low utilization: 5-9% utilization indicates significant optimization headroom — future work should focus on moving gradient accumulation and optimizer steps to ANE

Feasibility and Hardware Constraints

ANE Training demonstrates that training on the Apple Neural Engine is feasible. Current performance bottlenecks stem from software limitations — private API constraints, incomplete hardware feature exposure, and CPU/ANE coordination overhead — not fundamental hardware deficiencies.

If Apple were to expose lower-level ANE programming interfaces or provide official training support in CoreML, Apple Silicon devices could become powerful and energy-efficient on-device AI training platforms.