XLA_TARGET for Apple Metal Silicon?

[shaqq]

Apologies for the n00b question

If I want BumbleBee/Axon/EXLA to use GPUs on Apple M1 Silicon via Metal, what should I set the XLA_TARGET to?

[josevalim]

XLA has no official support afaik. Apple release some plugin for old versions but afaik that is already outdated.

[shaqq]

I see - I think #51 is the reference issue, and maybe jax-ml/jax#8074.

I can send up a PR to the README here to clarify this (if you'd like, @josevalim). I'm sure others are looking to use Bumblebee on their Macs and want to know if it works on Apple Metal.

[joelberkeley]

Any progress with this? I'm trying to build the PJRT plugin myself but don't know where to start. It's far more opaque than building for aarch64 CPUs, which isn't particularly easy either

[polvalente]

We've released elixir-nx/emlx for Metal support!

There's still some work needed for compiler acceleration, but I was already able to see some improvements on workflows that don't make use of the EXLA compiler, just the backend.

[josevalim]

The plugins are out of scope for this repository but the good news is that you should be able to just import existing precompiled ones. Here is a PR where we add support for Metal plugins in Nx: elixir-nx/nx#1504 - but it wasn't merged yet because several operations are missing on Apple's plugins and unfortunately the plugin is closed source.

[joelberkeley]

OK so the jax metal plugin packaged in the python wheel isn't python-specific (as many jax things are, I think)? And it can be used as a normal PjRt plugin?

[jonatanklosko]

@joelberkeley correct, it's just a PjRt plugin!

[joelberkeley]

Do you know if it's an XLA plugin? The website says

The Metal plugin uses the OpenXLA compiler

but OpenXLA is an ecosystem not a compiler. It's not important but I am curious.

[jonatanklosko]

I'm not sure what you mean by XLA plugin :D There are many parts, so I may be missing the right wording, but the Jax Metal provides a PjRt plugin API implementation, in the sense that you load a dynamic library via pjrt::LoadPjrtPlugin.

[joelberkeley]

I'm not sure what you mean by XLA plugin

A PjRt plugin that uses the XLA compiler.

in the sense that you load a dynamic library via pjrt::LoadPjrtPlugin

ah, so this plugin doesn't have a C API, for use with the PjRt C API? I'm assuming the presence of an extern "C" function GetPjrtApi like in the CPU plugin

[joelberkeley]

ah, ok so that function calls GetPjrtApi anyway, so we're on the same page (assuming an extern "C")