Tales of the M1 GPU is a very nicely written discussion of GPUs in general and the Apple M1 GPU in particular. It would be great if someone could suggest a companion article to discuss how Axon and Nx make use of GPU capabilities.
-r
3 Likes
Note I’m not an expert so the following may not be completely accurate although I think the broad brushstrokes aren’t too far off.
'Nx` can operate with different backends and some of those backends, like tensorflow utilise GPUs.
To make that work efficiently, GPU code needs to be loaded onto the GPU and the data needs to be efficiently accessed on it too.
The way that happens is with the defn (define numerical) functions which very cleverly compile your Elixir code (a subset of the full language) into backend code. And if that backend executes on the GPU then you get the appropriate performance outcome. Much of Nx (and therefore libraries which depend on it like Axon) is written as defn.
2 Likes
Note no expert as well
elixir-nx can use different backends… the torchx backend has beta support for the m1 gpu - and in a few weeks/months it will be super easy to get going - right now the install with gpu support is a bit cumbersome afaik…
1 Like
