Open hickscorp opened 5 months ago
As far as I understand, the OpenXLA goal is eventually have each platform abstracted away into a self-contained plugin, that can be compiled separately and loaded dynamically. Once that's stable we may have a single base binary and several plugin binaries that can be downloaded, which should make it easy to support multiple as well. This is an ongoing work though, by a brief look, JAX did it for CUDA, but they reverted it from suggested setup instructions because there were issues, so it still seems pretty experimental.
We could work out a solution with the current binaries, but honestly I think it's worth waiting for the plugins. FWIW it's a rare use case to build for both CUDA and ROCm (note that we don't even have precompiled archives for ROCm), and using CPU and CUDA is possible with the single CUDA binary.
I've noticed that setting
XLA_TARGET=xxx mix release compute
will result in thelibexla_extension.so
being built with this name always, regardless ofxxx
. This is a bit problematic, because it implies having as many releases defined as we want to target different architectures. For example, ourmix.exs
file looks like this:This results in a list of releases available such as
web, worker, compute_cuda120, compute_rocm
. This is a bit cumbersome, given that the only thing that changes between the twocompute_xxx
releases seems to belibexla_extension.so
...A few ideas here:
libexla_extension_xxx.so
could be picked (via symlink?).XLA_TARGET
) could be really nice. For example, being able to indicate the architectures directly inmix.exs
, maybe even in thereleases
section...XLA_RELEASE
intoXLA_RELEASES
, telling XLA to build as many as requested, so that a release could embed manylibexla_extension_xxx.so
and do the switcheroo at boot?Let me know if I misunderstood anything!