How to compile the jaxlib with debug symbols?

[juuso-oskari]

The problem is that building the jaxlib as:

python build/build.py --enable_cuda --bazel_options=--override_repository=xla=/xla --bazel_options=--jobs=4

works, but when adding (for debugging symbols):

python build/build.py --enable_cuda --bazel_options=--override_repository=xla=/xla --bazel_options=--jobs=4 --bazel_options=--compilation_mode=dbg

it crashes with:

ERROR: /root/.cache/bazel/_bazel_root/bee4ad1fd43279be7a03b33426e824d5/external/xla/xla/python/BUILD:1257:21: Linking external/xla/xla/python/xla_extension.so failed: (Exit 1): crosstool_wrapper_driver_is_not_gcc failed: error executing command (from target @xla//xla/python:xla_extension.so) 

[bmaxdk]

Normally when you add debug symbols --compilation_mode=dbg, the build can get pretty heavy. If you have enough memory and cpu/gpu power, can you try this?

python build/build.py --enable_cuda --bazel_options=--override_repository=xla=/xla --bazel_options=--jobs=8 --bazel_options=--compilation_mode=dbg

[juuso-oskari]

I managed to build by pulling the latest xla and jax repos. But I have another problem:

The reason why I wanted to build the jaxlib from source is that I want to debug the xla backend that gets called when we run a jax script.

I try to debug:

gdb python
(gdb) break InitializePjitFunction # add breakpoint at one of the cpp functions that gets called
(gdb) run simple_jax_script.py

...

Thread 1 "python" hit Breakpoint 1, jax::(anonymous namespace)::InitializePjitFunction (fn_obj=0x7fffd6f3c310, function_name="fft", fun=std::optional<nanobind::callable> = {...}, cache_miss=..., static_argnums=std::vector of length 2, capacity 2 = {...}, static_argnames=std::vector of length 2, capacity 2 = {...}, donate_argnums=std::vector of length 0, capacity 0, pytree_registry=std::shared_ptr<xla::PyTreeRegistry> (use count 1, weak count 1) = {...}, shard_arg_fallback=..., cache=std::shared_ptr<jax::(anonymous namespace)::PjitFunctionCache> (use count 1, weak count 0) = {...}) at external/xla/xla/python/pjit.cc:1044
1044    external/xla/xla/python/pjit.cc: No such file or directory.
(gdb) 

So it cannot find the xla source files (which would be located in /xla in my case, instead it tries to search that external (?)). How could I build the jaxlib so that it would link to the original c++ source files still?

[jaro-sevcik]

Have you tried set substitute-path external /xla in gdb?