Skip to content

[ROCm] Migrate SMI device queries from rocm_smi to amd_smi - #1144

Merged
mminutoli merged 1 commit into
rocm-jaxlib-0.11.1-alphafrom
users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-0.11.1-alpha
Aug 28, 2026
Merged

mminutoli merged 1 commit into
rocm-jaxlib-0.11.1-alphafrom
users/gulsumgudukbay/migrate-rocm-smi-to-amd-smi-0.11.1-alpha

Conversation

@gulsumgudukbay

Copy link
Copy Markdown

Ports the SMI migration (#1113) to rocm-jaxlib-0.11.1-alpha, ahead of
TheRock#6852, which removes
librocm_smi64 from the ROCm SDK.

This is the ref that ROCm/jax rocm-jaxlib-v0.11.1 actually consumes. That
branch builds with bzlmod (common --enable_bzlmod --noenable_workspace), so
its XLA comes from the archive_override in MODULE.bazel, which pins this
branch — not from third_party/xla/revision.bzl. The earlier port to
rocm-jaxlib-v0.11.1 (#1140) therefore never reached the JAX wheel build, and
TheRock CI still failed with:

SolibSymlink .../librocm_smi64.so failed: missing input file
@@xla++rocm_configure_ext+local_config_rocm//rocm:rocm_dist/lib/librocm_smi64.so

Changes

Device queries move behind smi_util.h, with one backend compiled in per
build. The :smi alias and backend selection follow the existing
rocm_version_number() >= 71300 gate, so ROCm >= 7.13 (including ROCm 10.x,
which encodes as 100100) resolves to amd_smi and never pulls
:rocm_smi, whose interface_library hard-references librocm_smi64.so.

Relative to the raw cherry-pick, the only conflict resolution was in
third_party/gpus/rocm/BUILD.tpl: this branch already carried amdsmi_libs as
a plain cc_library, so it is converted to a full rocm_lib_import with an
interface_library, keeping the existing :system_libs dep.

Verification

Built against a full ROCm 7.14 SDK with librocm_smi64.so* deleted, simulating
TheRock#6852:

  • smi_util, rocm_pcie_bandwidth, rocm_xgmi_topology, rocm_executor all build.
  • jax-rocm-pjrt wheel builds from ROCm/jax rocm-jaxlib-v0.11.1.
  • The resulting xla_rocm_plugin.so lists libamd_smi.so.26 as NEEDED, with
    no librocm_smi64 entry.

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).

(cherry picked from commit 45f7885)
@mminutoli
mminutoli merged commit 65f2f02 into rocm-jaxlib-0.11.1-alpha Aug 28, 2026
6 of 14 checks passed
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.

2 participants