Skip to content

[ROCm] Migrate SMI device queries from rocm_smi to amd_smi (v0.11.0) - #1139

Merged
i-chaochen merged 1 commit into
rocm-jaxlib-v0.11.0from
users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-v0.11.0
Aug 27, 2026
Merged

i-chaochen merged 1 commit into
rocm-jaxlib-v0.11.0from
users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-v0.11.0

Conversation

@gulsumgudukbay

@gulsumgudukbay gulsumgudukbay commented Aug 26, 2026 •

Copy link
Copy Markdown

Summary

Minimal port of ROCm/xla#1113 (02298de) onto rocm-jaxlib-v0.11.0. This is not a full cherry-pick: only the SMI migration lands, scoped to what this release branch needs to build once librocm_smi64 is absent from the SDK (ROCm/TheRock#6852).

What changes (13 files):

  • third_party/gpus/rocm/BUILD.tpl: replace amdsmi_libs with rocm_lib_import(name = "amdsmi", ...); leave rocm_smi and the rest of the release-branch template untouched.
  • xla/stream_executor/rocm/BUILD: add :smi alias (amd_smi from ROCm 7.13+, rocm_smi below), :smi_util, and wire pcie/xgmi deps; no unrelated release-branch BUILD churn.
  • C++ SMI layer from [WIP] [ROCm] Migrate SMI device queries from rocm_smi to amd_smi #1113: smi_util.{h,cc}, smi_util_amd_smi.cc, smi_util_rocm_smi.cc; update pcie/xgmi callers and tests; delete rocm_smi_util.*.
  • rocm_executor.cc: only the pcie/xgmi caller hunks from [WIP] [ROCm] Migrate SMI device queries from rocm_smi to amd_smi #1113 (~23 lines), not main's wider TF→ABSL status-macro refactor.

Release-branch adaptations vs main #1113:

  • Bazel deps use //xla/tsl/platform:status_macros (release branches lack @com_google_absl//absl/status:status_macros).
  • Status macros in the three adapted .cc files use RETURN_IF_ERROR / ASSIGN_OR_RETURN from that header.

Test plan

Based on: #1113 (02298de)

@gulsumgudukbay
gulsumgudukbay force-pushed the users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-v0.11.0 branch 6 times, most recently from b43f7be to 96c0969 Compare August 27, 2026 02:19
Cherry-pick of #1113 (02298de) with SMI-only Bazel changes
on this release branch. Unblocks JAX wheel builds when librocm_smi64
is absent from the ROCm SDK (ROCm/TheRock#6852).
@gulsumgudukbay
gulsumgudukbay force-pushed the users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-v0.11.0 branch from 96c0969 to 2225114 Compare August 27, 2026 02:31

@nurmukhametov nurmukhametov left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@i-chaochen

Copy link
Copy Markdown
Collaborator

do we need to merge these PRs?

@gulsumgudukbay

Copy link
Copy Markdown
Author

do we need to merge these PRs?

yes please

@i-chaochen
i-chaochen merged commit 2cab052 into rocm-jaxlib-v0.11.0 Aug 27, 2026
5 of 11 checks passed
@gulsumgudukbay
gulsumgudukbay deleted the users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-v0.11.0 branch August 27, 2026 16:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants