Cherry-pick PR #983 ROCm fixes to rocm-jaxlib-v0.10.2#985
Closed
magaonka-amd wants to merge 4 commits into
Closed
Conversation
Imported from GitHub PR openxla#40385 📝 Summary of Changes Move from linking to rocm libraries trough DsoLoader to direct linking. Make sure we only link to specified rocm library, not to any other library said library may use. This change leaves enough wrapper (namespace wrap) headers needed for TF. 🎯 Justification Makes our build targets not tainted by transitive dependencies of rocm libraries. DsoLoader use has little benefit for rocm since the libraries get loaded ahead of time by jax plugin anyways. 🚀 Kind of Contribution ♻️ Cleanup Copybara import of the project: -- 38470a3 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: [ROCm] Streamline bazel targets for rocm libraries -- 7f2ec34 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: [ROCm] Remove rocm_solver_context -- 1e34f26 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: [ROCm] Remove hipblaslt_wrapper.h -- 2a35a53 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: [ROCm] Remove hipsparse_wrapper.h -- 37338f2 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove rocsolver_wrapper.h -- 545039f by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove rocblas_wrapper.h -- c7e959e by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove rocm_driver_wrapper.h -- e41363d by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove wrap namespace usage from rocm_fft -- 94dc13e by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove wrap namespace usage from rocm_dnn -- 0885e40 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove TryDlopenROCmLibraries -- 983cc3f by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove :all_files -- aa5b6c3 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Partialy restore *wrapper.h files needed by ROCm TF -- d2e89c2 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Use more filegroups -- 3abd237 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Remove DsoLoader ROCm support -- 474acd2 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Unblock X86 build CI -- 57e1100 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Fix RBE build -- 8a55bfb by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: Move to rocm_lib_import macro Merging this change closes openxla#40385 COPYBARA_INTEGRATE_REVIEW=openxla#40385 from ROCm:draganm/rocm_libs 8a55bfb PiperOrigin-RevId: 904252712 (cherry picked from commit 4f993de) (cherry picked from commit 16f1d8d)
Imported from GitHub PR openxla#41591 📝 Summary of Changes Default for hermetic paths in non-hermetic build too 🎯 Justification Makes CI pick up rocm libs when doing local build + remote testing 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- 2ab9ead by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>: [ROCm] Unblock CI after PR openxla#40385 Merging this change closes openxla#41591 COPYBARA_INTEGRATE_REVIEW=openxla#41591 from ROCm:draganm/rocm_libs_followup 2ab9ead PiperOrigin-RevId: 909965137 (cherry picked from commit d9922e7) (cherry picked from commit 83380a5)
(cherry picked from commit 8ef3bb4)
magaonka-amd
added a commit
to ROCm/jax
that referenced
this pull request
Jun 23, 2026
Fetch the pinned XLA archive from github.com/ROCm/xla instead of openxla/xla so the release build pulls the ROCm release branch (rocm-jaxlib-v0.10.2). XLA_COMMIT/XLA_SHA256 are unchanged for now (5a9e73cb, content-identical across both repos); the SHA will be bumped to the rocm-jaxlib-v0.10.2 HEAD once ROCm/xla#985 merges.
i-chaochen
reviewed
Jun 23, 2026
i-chaochen
left a comment
i-chaochen
left a comment
Collaborator
There was a problem hiding this comment.
Thanks! Please also cherry-pick this openxla#44389 , openxla#41116 and openxla#44428 to v0.10.1 release and waiting for the CI result!
Collaborator
|
0.10.2 CI needs this commit 13dc711 @magaonka-amd please cherry-pick it as well. Thanks! |
This was referenced Jun 24, 2026
Author
|
@i-chaochen cherry-picked all applicable changes |
Collaborator
|
:) You don't need to open so many PRs on our side, just put these in your this PR for v0.10.2 is easier to track From this PR https://github.com/ROCm/xla/pull/988/changes could you please rm |
Author
|
Superseded by #993, which combines all four ROCm 0.10.2 cherry-pick PRs into a single PR against rocm-jaxlib-v0.10.2. |
Author
|
I tested locally it works , but in CI it looks like test path injected from dev infra not working in release branch and build fails. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Brings the ROCm release deltas from #983 (
Cherrypick pr932 to v0.10.1) onto the newrocm-jaxlib-v0.10.2branch, which was cut from the JAX 0.10.2 pinned XLA commit5a9e73cbd92530cac2ac36f4736a774b2412afe2.Commits (cherry-picked with
-x)PR #40385: [ROCm] Streamline bazel targets for rocm librariesPR #41591: [ROCm] Unblock CI after PR #40385[ROCm] Add librocm_kpack and libclang-cpp to runfiles globsFix formatting error in rocm_blas.ccAll four applied cleanly onto the pinned base (no empty/dropped commits).
Files changed
third_party/gpus/rocm/build_defs.bzl.tplthird_party/gpus/rocm/BUILD.tplxla/stream_executor/rocm/rocm_blas.ccxla/tsl/platform/default/dso_loader.ccTest Plan
release/0.10.2release-validation workflow.