Play nicely with Elixir's releases.

[hickscorp]

I've noticed that setting XLA_TARGET=xxx mix release compute will result in the libexla_extension.so being built with this name always, regardless of xxx.
This is a bit problematic, because it implies having as many releases defined as we want to target different architectures.
For example, our mix.exs file looks like this:

  defp releases(),
    do:
      [
        web: [
          version: "0.0.1",
          applications: [...]
        ],
        worker: [
          version: "0.0.1",
          applications: [...]
        ]
      ]
      |> add_compute_releases([:cuda120, :rocm])

  defp add_compute_releases(acc, xla_targets),
    do:
      xla_targets
      |> Enum.reduce(acc, fn xla_target, acc ->
        acc ++
          [
            {:"compute_#{xla_target}",
             [
               version: "0.0.1",
               applications: [...]
             ]}
          ]
      end)

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 two compute_xxx releases seems to be libexla_extension.so...

A few ideas here:

• It would be really nice being able to indicate more than one XLA target architecture. Then upon starting the release, we would specify which one to use, and a different libexla_extension_xxx.so could be picked (via symlink?).
• Not having to toy too much with ENV variables (XLA_TARGET) could be really nice. For example, being able to indicate the architectures directly in mix.exs, maybe even in the releases section...
• Why not pluralizing XLA_RELEASE into XLA_RELEASES, telling XLA to build as many as requested, so that a release could embed many libexla_extension_xxx.so and do the switcheroo at boot?

Let me know if I misunderstood anything!

[jonatanklosko]

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.

[polvalente]

@jonatanklosko do you still see this plugin path moving forward?

[jonatanklosko]

Yes, I think eventually we will want to migrate towards plugins. Looks like Jax has been only using plugins now for a while.

Ideally we would encapsulate all platform-specific logic in the plugins (e.g. cuda custom calls), so that we only have a single variant of EXLA. Then we could precompile EXLA itself per architecture + NIF version, and plugins would be separate.

We actually considered this in the past, but since XLA kept both the StreamExecutor GPU API and the plugin API, it was unclear what was going to happen long term (not sure if it is more clear now, but at least Jax seems to have moved clearly in the plugins direction). One challenge with moving to the plugin API is that it doesn't support infeed/outfeed, so for Nx hooks (print) we need a different mechanism, and it would be per platform. Now that we have runtime_call, we are closer to that, but we would need to implement runtime_call at least for CUDA, and possibly have non-blocking variants. There is a bit more details in elixir-nx/nx#1541 (comment).

So it's not necessarily a priority, but as some point it would be worth revisiting the XLA integration, and look more specifically how Jax does it now.