Docker imagewith ROCM support

[jalberto]

There are docker images ready for cuda, it would be great if there was also official containers with AMD ROCM support.

[jonatanklosko]

We may consider these once we have precompiled XLA for ROCM, but the truth is that figuring out the versions to precompile with and testing that things work are too much in terms of maintainability (both in terms of time and hardware), given the low demand :<

[josevalim]

@jalberto exactly. This is not a Livebook or elixir-nx/xla issue, it is an openxla/xla issue. All elixir-nx/xla does is to precompile assets and we could precompile ROCM if someone provides a reliable recipe for precompiling openxla/xla with ROCM. The answer to most of your questions is achieved by building elixir-nx/xla locally (which in turns just builds openxla/xla). gfx11 support would also have to be added to openxla/xla (it may already have been added).

[jalberto]

I understand, still unable to build it locally or using docker, even just the base code (not trying to add upgraded rocm or even gfx1100).
So I wonder if cuda/cpu versions has the same reproducibility issues

I have a lot of experience with containers and still I am unable to reproduce the builds, and I am glad to add my 2c, just point me in the right direction and I will dig there :)

[jonatanklosko]

So I wonder if cuda/cpu versions has the same reproducibility issues

Usually the hard part is getting through the build, with the precompiled binaries that we have for cuda/cpu it generally works smoothly. We use the build scripts for precompilation, so in that sense they are more reliable.

The ROCM build is much more off the beaten track, we try to help whenever someone attempts to build it. I added the Docker-based build, because otherwise getting the right environment for built was not straightforward; I am still confused why it stopped working within Docker.

[jonatanklosko]

For the record, Jax (the Python counterpart to Nx + EXLA) integrates with openxla/xla and they also say that ROCm support is experimental and they don't provide precompiled binaries either (which they do for cpu, cuda and tpu). That's what I mean when I say that ROCm support is second-class.

That said, the main argument is still that maintaining precompiled binaries and validating that ROCm works under different configurations is just too much maintenance overhead at the moment (time and hardware wise).

[jalberto]

I understand, thanks for your time!

I will continue trying and will post any progress

[jalberto]

@jonatanklosko sorry for bothering you, can you tell me if the current binary provided already implements the fixes and gfx11 support mentioned here?

elixir-nx/xla#63

If it doesn't, could you point me where to find it?

Thanks

[jonatanklosko]

@jalberto the binary at https://static.jonatanklosko.com/builds/0.6.0/xla_extension-x86_64-linux-gnu-rocm.tar.gz is built on elixir-nx/xla v0.6.0, so it doesn't have the gfx11 support. The discussion at elixir-nx/xla#63 outlines a few necessary changes, but it's definitely hacky. Again, the best bet is for us to update to latest openxla/xla (the revision we use is from 6 months ago); but there is other work that this would affect, that's why I'm not certain about when we ship the update.