Bad interaction between JAX and Reactant on aarch64-linux

[mofeing]

We are currently running into this error on aarch64-linux CI for the Python-JAX test:

2024-12-30 21:52:33.649409: I external/xla/xla/service/llvm_ir/llvm_command_line_options.cc:50] XLA (re)initializing LLVM with options fingerprint: 12488328760584806018
--
  | 'neoverse-n1' is not a recognized processor for this target (ignoring processor)
  | '+crc' is not a recognized feature for this target (ignoring feature)
  | '+fp-armv8' is not a recognized feature for this target (ignoring feature)
  | '+lse' is not a recognized feature for this target (ignoring feature)
  | '+neon' is not a recognized feature for this target (ignoring feature)
  | '+sha2' is not a recognized feature for this target (ignoring feature)
...
LLVM ERROR: 64-bit code requested on a subtarget that doesn't support it!

[740] signal (6.-6): Aborted
--
  | in expression starting at /cache/build/default-aws-aarch64-ci-0-4/julialang/reactant-dot-jl/test/integration/python.jl:7

[wsmoses]

Incidentally this seems like jax itself doesn’t support aarch64 is the issue (and the same llvm initialization we use gets broken).

If we don’t initialize jax I presume this succeeds?

[wsmoses]

Also on a debug build I’m curious what the backtrace is

[mofeing]

mmm but there are aarch64-linux wheels for jaxlib https://pypi.org/project/jaxlib/#files

Also on a debug build I’m curious what the backtrace is

too bad i don't have access to a aarch64-linux machine right now 🥲 i should get access to Fugaku at some moment, but maybe @giordano has sth?

[giordano]

What should I try to do (but probably I won't have much time before a couple of days)

[mofeing]

yeah, don't worry. there is no rush.

basically...

1. set the preference for Reactant_jll such that we use the variant mode+dbg. i think LocalPreferences.toml should look similar to this

[Reactant_jll]
mode = "dbg"

2. try running this test https://github.com/EnzymeAD/Reactant.jl/blob/main/test/integration/python.jl

[giordano]

julia> using Reactant_jll

julia> Reactant_jll.is_available()
false

julia> Reactant_jll.host_platform
Linux aarch64 {cxxstring_abi=cxx11, gpu=cuda, julia_version=1.11.2, libc=glibc, libgfortran_version=5.0.0, libstdcxx_version=3.4.30, mode=dbg}

Reactant isn't available for me. Is that because this system has got an Nvidia GPU and we don't have a CUDA build for for aarch64?

[wsmoses]

Yeah we don't make debug cuda builds because there's so much debug info that it breaks things

[giordano]

I had the same also without debug mode. OK, next time I'll need to try a machine without a GPU

[wsmoses]

You should also be able to just force the jll preferences to not auto detect the cuda version

[mofeing]

i think you just need to put

cuda = "none"

in the LocalPreferences.toml

[giordano]

i think you just need to put

cuda = "none"

in the LocalPreferences.toml

I presume you meant

gpu = "none"

but I used

julia> using Preferences

julia> set_preferences!(
           "LocalPreferences.toml",
           "Reactant_jll",
           "mode" => "dbg",
           "gpu" => "none",
       )

and after restart I still get

julia> Reactant_jll.is_available()
false

julia> Reactant_jll.host_platform
Linux aarch64 {cxxstring_abi=cxx11, gpu=cuda, julia_version=1.11.2, libc=glibc, libgfortran_version=5.0.0, libstdcxx_version=3.4.30, mode=dbg}

I believe this is because platform augmentation is double flawed:

• you force gpu discovery even if the user set gpu="none": https://github.com/JuliaPackaging/Yggdrasil/blob/614abd45579152a7f40a1566c294f5abe2f9f61c/R/Reactant/build_tarballs.jl#L341-L351
• you do this also on platforms where GPU support isn't available at all

[wsmoses]

Oh yeah we should limit that to just x86 Linux (in fairness at the time that was written only x86 and macOS existed and we clearly didn’t update it)

[wsmoses]

This is what we get for trying to be helpful and auto enable the cuda jll when possible xD

[giordano]

If I set the env var REACTANT_GPU=none then the augmentation is correct:

julia> Reactant_jll.host_platform
Linux aarch64 {cxxstring_abi=cxx11, gpu=none, julia_version=1.11.2, libc=glibc, libgfortran_version=5.0.0, libstdcxx_version=3.4.30, mode=dbg}

but it still doesn't make the artifact available, which now puzzles me:

julia> Reactant_jll.is_available()
false

🤔