amx VS neural-engine

I am talking about the matrix/vector coprocessor (AMX). You can find some reverse-engineered documentation here: https://github.com/corsix/amx

On M3 a singe matrix block can achieve ~ 1TFLOP on DGEMM, I assume it will be closer to 4TFLOPS for SGEMM. The Max variants have two such blocks. Didn't do precise benchmarking myself, but switching Python/R matrix libraries to use Apple's BLAS result in 5-6x perf improvement on matrix heavy code for me.

It's really cool. I hope it becomes more common for training/inference/numerics capable accelerators to be included in consumer hardware.

Apple's AMX is really under-documented, while the instructions were reverse engineered, Virtually no benchmarks are available comparing current chip generations, models and variants.

https://github.com/corsix/amx

Visible in the unofficial documentation for AMX instructions too - M2 only bf16 functionality - https://github.com/corsix/amx/blob/main/matfp.md

Confusingly there are 2 mechanisms to do matrix operations on the new apple hardware - AMX (https://github.com/corsix/amx) - and the ANE (apple neural engine) - which is enabled by CoreML. This code does not run on the neural engine but the author has a branch for his whisper.cpp project which uses it here: https://github.com/ggerganov/whisper.cpp/pull/566 - so it may not be long before we see it applied here as well. All of this is to say that it actually could get significantly faster if some of this work was able to be handed to the ANE with CoreML.

really? seems pretty well documented here: https://github.com/corsix/amx

Article is almost two years old, and has a huge correction at the bottom. It's just a proprietary ISA extension, there's even a repo documenting what's been reverse engineered.

You are correct, in that those are the four

My understanding is that the AMX is more tightly wound with the CPU, ultimately being accessible via an instruction set (https://github.com/corsix/amx), and it is useful if you need to do matrix multiplications interleaved with other CPU tasks. A common example would be a VIO loop or something where you want that data in the CPU caches.

The GPU and Neural Engine are not that – they take some time to set up and initialize. They also can parallelize tasks to a much higher degree. The GPU is more generalizable, because you can write compute shaders to do anything in parallel, but it uses a lot of resources. I'll have to check out the PR to see how exactly the MPS shaders match up with the task at hand, because you could also consider writing Metal compute shaders by hand.

I know the least about the ANE, but it has specific hardware for running ML models, and you have to process the weights ahead of time to make sure they are in the right format. It can run ML models very efficiently and is the most battery friendly.

If someone were to build a Linux system with proprietary ISA extensions, how would they do it given Linux is open source? Are there any examples of this being done? Would it be possible at all?

I got inspiration from this (https://github.com/corsix/amx) and I wondered if someone has done it before on a Linux-based system. I understand a userspace library could be created to access those instructions from userspace, but how would then they be implemented in the kernel? Through a proprietary kernel module built using a custom compiler? Or is that not needed at all and the library could just run on the processor taking advantage of the proprietary extensions?

yep. they have a neural engine that is separate from the CPU and GPU that does really fast matmuls https://github.com/hollance/neural-engine. it's basically completely undocumented.

Looks like the ANE ("Apple Neural Engine") cores are powerful but not as flexible/programmable as the GPU cores. There is no sign that LLM inference is possible with them or ever will be unless Apple either opens up the closed ANE software framework for extensibility or they extend the ANE framework to support modern LLMs themselves. I would not hold my breath.

If Apple would wake up to what's happening with llama.cpp etc then I don't see such a big role for paying for remote access to big models via API

Currently a Macbook has a Neural Engine that is sitting idle 99% of the time and only suitable for running limited models (poorly documented, opaque rules about what ops can be accelerated, a black box compiler [1] and an apparent 3GB model size limit [2])

OTOH you can buy a Macbook with 64GB 'unified' memory and a Neural Engine today

If you squint a bit and look into the near future it's not so hard to imagine a future Mx chip with a more capable Neural Engine and yet more RAM, and able to run the largest GPT3 class models locally. (Ideally with better developer tools so other compilers can target the NE)

And then imagine it does that while leaving the CPU+GPU mostly free to run apps/games ... the whole experience of using a computer could change radically in that case.

I find it hard not to think this is coming within 5 years (although equally, I can imagine this is not on Apple's roadmap at all currently)

[1] https://github.com/hollance/neural-engine

My question too. This semi-answer on the page seems to contradict itself (source: https://github.com/hollance/neural-engine/blob/master/docs/p... ):

"> Can I program the ANE directly?

Unfortunately not. You can only use the Neural Engine through Core ML at the moment.

There currently is no public framework for programming the ANE. There are several private, undocumented frameworks but obviously we cannot use them as Apple rejects apps that use private frameworks.

(Perhaps in the future Apple will provide a public version of AppleNeuralEngine.framework.)"

The last part links to this bunch of headers:

https://github.com/nst/iOS-Runtime-Headers/tree/master/Priva...

So might it be more accurate to say you can program it directly, but won't end up with something that can be distributed on the app store?

Read between their buzzwords. Apple's Neural Engine does nothing for training. It's purely for inference and it still requires the developers to go through their API. If a model uses a type layer Apple doesn't support, it's back to the CPU/GPU.