jax
af2dcef8 - [ROCm] Rollup of changes for rocm7 support

Commit

160 days ago

[ROCm] Rollup of changes for rocm7 support - Add rocm7 plugin namespace to jaxlib ext's This adds support to load kernels from the ROCM 7.x version of the plugin which will be released imminently, but still load the "rocm60" version, if the rocm7 version is not found. - Add missing dependency on amdhip64 runtime _hybrid, _prng, and _triton were missing links to the HIP runtime library which is has symbols that are required for these shared objects to load properly. rocm:hip target from XLA unfortunately has a dependency on amd_comgr which is incorrect, but will have to be fixed in XLA sources. The proposed fix for this is at https://github.com/openxla/xla/pull/27498 - Change rocm jax plugin to use new hip target This is a minimally linked target for HIP runtime (libamdhip64.so), such that we can avoid having ld overlink to things we do not need in JAX-rocm. Namely, comgr and rocm-smi.

References

rocm7-plugin-updates

#29281 - [ROCm] ROCm7 Plugin Updates

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#579 - Create rocm-test-requirements.txt

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency

#585 - update a test for checking zero ROCm GPU event