Every Apple silicon Mac has a Neural Engine, and most of the time it sits idle. I wanted it to do my embeddings. It took a lot longer than I expected, and most of what I learned isn't in the docs, so here it is.

The context is a local code search tool I build in Rust. It embeds every chunk of code and docs it indexes, which on my biggest test setup is about 879k chunks. The machine is a 2020 M1 MacBook Air, 8 GB, fanless. So embedding speed is the whole game for the first index.

The end result on that laptop, sustained.

Lane Chunks per second
CPU only, ONNX Runtime int8 30.7
CPU and GPU about 63
Neural Engine, native Core ML 236.5

The model

I use IBM's granite-embedding-small-english-r2. 384 dimensions, Apache-2.0, about 52 MB quantised to int8. On the CPU it runs through ONNX Runtime with the ort crate, and that part was easy.

Attempt one, the Core ML execution provider

ONNX Runtime ships a Core ML execution provider, so the obvious move is to flip it on and let it place ops on the Neural Engine. I did, and the throughput barely moved.

The reason is that the Core ML provider places almost nothing of this model on the Neural Engine. The ANE is picky about tensor shapes, data types and op types. When a graph doesn't fit what it likes, Core ML quietly runs those parts on the CPU or GPU, and you pay for the hand-offs between them. Nothing errors. It just isn't fast.

If you try this yourself, don't trust "it ran with the Core ML provider" as proof that it ran on the Neural Engine. Check where the ops actually landed, Xcode's Core ML performance report shows the compute unit for each op.

Attempt two, rewrite the model for the Neural Engine

What worked was a native Core ML model built in the layout the Neural Engine wants. Roughly, that meant three things.

  • Tensors in (B, C, 1, S) shape, batch, channels, a dummy height of 1, then sequence. That's the shape Apple's own ANE transformer work uses.
  • Linear layers as 1x1 convolutions, since the ANE is built around convolution.
  • Attention split per head instead of one big batched matmul.

The ANE also wants fixed shapes. So the model is converted at batch 4 with six sequence-length buckets, 128, 192, 256, 320, 384 and 512 tokens, and a batch runs at a bucket long enough for it. After all that, 3,375 of the model's 3,380 ops run on the Neural Engine.

Two fp16 problems came up along the way, because the ANE computes in fp16.

  • LayerNorm overflowed. Squaring activations to get the variance can blow past fp16's max of 65,504. The fix was to square a sixty-fourth of the input instead and scale back after, so the sum of squares stays in range.
  • Fully padded rows produced NaNs. If every key in an attention row is masked, softmax divides zero by zero. A padding row now keeps one key unmasked.

After both fixes every vector matches fp32 at a cosine of 0.9999 or better.

Calling Core ML from Rust

There's no need for a Swift shim. The objc2-core-ml crate exposes the Core ML API directly. Loading a model for one bucket looks like this.

let config = unsafe { MLModelConfiguration::new() }; unsafe { config.setComputeUnits(match kind { Kind::NeuralEngine => MLComputeUnits::CPUAndNeuralEngine, Kind::Gpu => MLComputeUnits::CPUAndGPU, }); if let Some(function) = &group.function { let name = function.replace("{S}", &bucket.to_string()); config.setFunctionName(Some(&NSString::from_str(&name))); } } let url = NSURL::fileURLWithPath(&NSString::from_str(&path.to_string_lossy())); unsafe { MLModel::modelWithContentsOfURL_configuration_error(&url, &config) }

Note CPUAndNeuralEngine. You can't ask for the Neural Engine alone, Core ML always keeps a fallback. That's another reason the model has to be built so nothing needs the fallback.

Run it in a separate process

The embedding work doesn't run inside the main daemon. Each accelerator lane is a child process, semlith __embed-worker ane, that reads length-framed batches on stdin and writes vectors to stdout.

The reason is crashes. Accelerator drivers can fail in ways you can't catch, and when one does I want to lose a worker, not the daemon that's serving search to every agent. A worker that exits, errors or misses its per-batch deadline fails only its own lane. Its batch goes back in the queue and the run finishes on whatever lanes are left, the CPU if nothing else.

Diagram of the daemon sending batches to a separate __embed-worker ane process over stdin and stdout, which runs Core ML on the Neural Engine. Notes below cover the known-answer check, crash isolation, the stable worker copy and the compile lock.

Each accelerator lane is its own process, so a driver crash costs a worker, not the search daemon.

A new worker also has to prove itself before it gets real work. It embeds 32 known chunks, and its vectors have to match the fp32 reference at cosine 0.999 or better. A lane that gives subtly wrong vectors is worse than a slow lane, because nothing downstream would notice.

The compile cache trap

The first time Core ML loads a model for the Neural Engine, macOS compiles it. On the M1 that's 30 to 43 seconds per bucket, one bucket at a time, so about three minutes for all six. A worker can start embedding after about 39 seconds from cold, and about 2 seconds once cached.

Here's what I didn't expect. macOS caches the compiled model under the name of the program that loaded it. Ship a new build and the binary changes, so every user recompiles all six buckets on upgrade.

The fix was to run the worker from a stable copy of the binary, accel/coreml-worker-v<N>, and only bump N when the model or the worker protocol actually changes. Upgrades now reuse the cache.

A few more things I hit around that cache.

  • Killing a program doesn't cancel its compile. The system compiler service works through a dead program's requests anyway, and once a backlog of them kept a lane compiling for over two hours. So compiles now take a file lock, coreml-compile.lock, which stops a second copy from queueing the same compile and stops the daemon retiring a worker mid-compile.
  • A killed worker leaves a half-written compile of about 94 MB behind, so the daemon cleans those up.
  • Load order matters. A session is ready once its longest bucket loads, then it loads the other five on its own thread at user-request priority. That took 27 seconds a bucket against 44 at a spawned thread's default priority.
  • macOS will clear the cache on its own when disk gets low. I saw it happen with 13 GB free. Plan for a recompile at any time.
  • It wrecks benchmarks. If you A/B two builds by alternating them run by run, every run is a fresh compile and your timings measure the compiler. Run each build in a block with a warm-up first.

Don't run the CPU beside it

My first instinct was to run every lane at once, Neural Engine plus CPU plus GPU, and add up the throughput. It went the other way.

Two int8 CPU threads running beside the Neural Engine cut it from 236 to 79 chunks a second. They take the cores that tokenise the text and feed the ANE, and starving it costs far more than the CPU lane adds. A GPU lane beside it added about 30% in bursts but only 4% sustained. So the CPU lane is off by default next to the Neural Engine now.

What it adds up to

The 70-repo test setup, 66,333 files and 879,439 chunks, took 70.7 minutes for its first index, a mean of 207 chunks a second, mostly on battery. Memory stays around 600 MB however much you index, with the daemon and workers peaking at 686 MB during a Neural Engine run. After the first index only changed chunks are re-embedded, so saving a file is searchable again in about a second.

None of it needs a GPU server or an API key, and the chip was already in the laptop.

The code is in https://github.com/semlith/semlith if you want to look, the Core ML loading is in src/coreml.rs and the worker lanes in src/accel.rs. The measurements and how to rerun them are on semlith.com/benchmarks.

Has anyone got the Core ML execution provider in ONNX Runtime placing a transformer on the ANE properly? I'd like to know what I missed.