[ANN] AppleNeuralEngine.jl: Julia on the Apple Neural Engine

Package Announcements
,
1

Every Apple Silicon Mac has a Neural Engine (ANE), an fp16 accelerator that Apple exposes only through Core ML. AppleNeuralEngine.jl drives it directly from Julia through ObjectiveC.jl, as Metal.jl does for the GPU, with no native code to build. You write ordinary Julia, and it is traced into an ANE program:

using AppleNeuralEngine
silu(z) = z / (1 + exp(-z))
function ffn(x)  # LLaMA-style feed-forward block: RMSNorm, SwiGLU, residual
    n = x ./ sqrt.(sum(abs2, x; dims=2) ./ size(x, 2) .+ 1f-5) .* g'
    return x .+ (silu.(n * W1') .* (n * W3')) * W2'
end
x   = randn(Float16, 256, 512)
ane = compile(ffn, x)   # trace + compile
y   = ane(x)            # runs on the Neural Engine

Results on an M2 Max

16-core ANE, macOS 27.0.1, Julia 1.13.3, 8 BLAS threads. The same Julia function runs on each device; ANE times include copying inputs and outputs.

workload ANE fp16 ms GFLOP/s Accelerate fp32 ms GFLOP/s OpenBLAS fp32 ms GFLOP/s rel. error
FFN block D=512 H=1536 T=256 0.261 4629 2.646 457 3.866 312 6.5e-04
x*W 256×1024 · 1024×1024 0.148 3628 0.210 2553 0.728 737 3.3e-04
x*W 1024×1024 · 1024×1024 0.366 5859 0.791 2715 2.839 756 3.3e-04
x*W 2048×2048 · 2048×2048 1.774 9682 7.223 2378 24.062 714 3.6e-04

ANE times include copying inputs and outputs. The ANE computes in fp16 and the CPU in fp32, so each runs at its native precision.

Reproduce it and post your numbers

git clone https://github.com/pitsianis/AppleNeuralEngine.jl
cd AppleNeuralEngine.jl
julia --project=benchmark -e 'using Pkg; Pkg.instantiate()'
julia --project=benchmark benchmark/run.jl     # prints the table above for your Mac
julia --project=benchmark tools/probe_ops.jl   # checks every op on your Mac

The benchmark prints a Markdown table with your chip and macOS version: please paste it in this thread. So far everything has been tested only with Julia 1.13.3 on a M2 Max.

What runs on the ANE

On the M2 Max, 129 of 167 probed MIL operations compile and give correct results. Works: elementwise arithmetic, exp/log/sqrt/tanh/erf, activations (relu, gelu, silu, softmax, …), reductions, convolutions (strided, depthwise, int8 weights), pooling, matmul/linear/einsum, layer/batch/instance norm, scaled dot-product attention, reshape/transpose/slice/concat/split/tile, upsampling. Rejected by the ANE compiler: sin/cos/tan, pad, gather/scatter, cumsum, topk/argsort/argmax, LSTM/GRU/RNN, random ops, fp32 compute. Full table: docs/src/op_support.md.

PRs are very welcome if you get a rejected op to work, or if tools/probe_ops.jl gives a different table on a newer M chip or macOS version. The probe rewrites docs/src/op_support.md, so the PR is that file.

Caveats: it uses a private, undocumented API that any macOS update can break. It computes in fp16 only, and each call costs about 0.1–0.2 ms, so it pays off for matrix–matrix work. It is not registered yet.

Credits: this builds on the reverse engineering of the private ANE runtime in maderix/ANE, and on ObjectiveC.jl and Metal.jl. Most of the code was written together with Claude, Anthropic’s AI assistant, in Claude Code.

Repository: GitHub - pitsianis/AppleNeuralEngine.jl: Julia on the Apple Neural Engine: trace Julia functions, build MIL graphs, run them on the ANE · GitHub

Edit: Corrected a bug that underestimated the performance of Accelerate and OpenBLAS.

2 Likes