XLA_TARGET for Apple Metal Silicon?

elixir-nx / xla

Pre-compiled XLA extension

Apache License 2.0

83 stars 21 forks source link

XLA_TARGET for Apple Metal Silicon? #74

Closed shaqq closed 4 months ago

shaqq commented 4 months ago

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 commented 4 months ago

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

shaqq commented 4 months ago

I see - I think https://github.com/elixir-nx/xla/pull/51 is the reference issue, and maybe https://github.com/google/jax/issues/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.