[WIP] [ROCm] Migrate SMI device queries from rocm_smi to amd_smi by nurmukhametov · Pull Request #1113 · ROCm/xla

nurmukhametov · 2026-08-18T10:15:24Z

Motivation

To get review.

Technical Details

RCCL dropped librocm_smi64 in ROCm 7.13 and links libamd_smi instead. The two libraries cannot coexist in one process. XLA now switches SMI libraries at the same boundary RCCL does, so only one is ever present. Selection is compile and link time only, via a version gated :smi target in BUILD.tpl and #if (TF_ROCM_VERSION >= 71300) in the three existing sources. The ROCm floor stays at 7.1; bumping it to 7.13 later reduces to deleting the #else branches.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

i-chaochen · 2026-08-18T10:30:05Z

+  uint16_t width;
+};
+
+#if (TF_ROCM_VERSION >= 71300)


this macro guarad on rocm version will be challenged by upstream to change as the runtime check..

I don't think so. Why should they challenge that?

they against these new added #ifdef e.g., openxla#47447 (review) and openxla#46895 (comment)

they against these new added #ifdef e.g., openxla#47447 (review) and openxla#46895 (comment)

IMO, these are completely different cases.

Split the implementation into separate files and choose the impl via bazel.

I'm not sure how this will play with you work to use amd_smi to get memory bandwidth. Not sure what we have here is simpler than just moving to amd_smi and having dlopen workaround for pre 7.14.

Split the implementation into separate files and choose the impl via bazel.

Done.

I'm not sure how this will play with you work to use amd_smi to get memory bandwidth. Not sure what we have here is simpler than just moving to amd_smi and having dlopen workaround for pre 7.14.

I will implement it for amd_smi only. Below 7.13, it will just keep today's behavior, the per-gfx constant.

claude · 2026-08-18T20:15:07Z

+  // amd_smi reports pcie_speed directly in MT/s, unlike rocm_smi below.
+  return PcieLinkStatus{pcie_info.pcie_metric.pcie_speed,
+                        pcie_info.pcie_metric.pcie_width};


nit / fragility: The comment says amd_smi reports speed "directly in MT/s." Could you add a reference to the specific amd_smi header or doc that confirms this? The rocm_smi path below documents its 0.1 GT/s unit, which is great — having the same clarity here would protect against silent regressions if the semantics ever change across amd_smi versions.

A sanity check (e.g., LOG_IF(WARNING, speed > 128000)) would also catch unexpected unit changes early, but that's optional.

Put the reference to amd_smi header.

Resolved - the amdsmi.h reference in rocm_smi_util_amd_smi.cc:132 documents the unit clearly.

Resolved — amd_smi header reference added. Thanks.

claude · 2026-08-18T20:15:11Z

+  amdsmi_bdf_t bdf = {};
+  bdf.bdf.domain_number = target_bdf.domain;
+  bdf.bdf.bus_number = target_bdf.bus;
+  bdf.bdf.device_number = target_bdf.device;
+  bdf.bdf.function_number = target_bdf.function;


nit: BdfComponents uses uint64_t for all fields (domain, bus, device, function), but amdsmi_bdf_t's internal fields are narrower (typically uint32_t/uint16_t/bitfields). The assignments here involve implicit narrowing conversions. While PCI BDF values are always small enough in practice, explicit static_casts would make the assumption visible and prevent compiler warnings under stricter settings.

The field types match at uint64_t, and static_cast neither documents nor silences anything

Resolved - understood, BdfComponents fields are uint64_t and the amdsmi_bdf_t accessors return matching types, so no narrowing occurs.

Resolved — agreed, the field types already match at uint64_t so no narrowing occurs and a static_cast would be noise.

claude · 2026-08-18T20:15:13Z

+
+// Returns true if src reaches dst over an xGMI link.
+bool IsXgmiPeer(SmiDeviceHandle src, SmiDeviceHandle dst) {
+  uint64_t hops = 0;


nit: hops is written by both amdsmi_topo_get_link_type and rsmi_topo_get_link_type but is never read afterward. Moving it inside each #if/#else block (or at minimum annotating with [[maybe_unused]]) would be cleaner.

It is not unused, it is used. Made a comment about it.

Resolved - the comment at line 148 makes clear hops is a required out-parameter. Thanks for clarifying.

Resolved — hops is a required out-parameter for both amdsmi_topo_get_link_type and rsmi_topo_get_link_type. Clarifying comment added.

claude · 2026-08-18T20:15:18Z

+    data = glob([
+        "%{rocm_root}/lib/libamd_smi.so*" if _use_amd_smi else "%{rocm_root}/lib/librocm_smi64.so*",


observation: glob() will silently produce an empty list if the expected library is missing from the ROCm installation (e.g., ROCm 7.13 installed but libamd_smi.so not present). The error would only surface at runtime. This is inherent to glob-based patterns and not new to this PR — just flagging in case a more explicit fail-fast check (e.g., a Starlark fail() when the glob returns empty) is desired.

The glob is indeed silent when empty (XLA sets --incompatible_disallow_empty_glob=false in .bazelrc:5), but it is not what gates this. rocm_lib_import passes interface_library to cc_import(shared_library = ...) as a bare label, never through a glob, so a missing library is a build-time error naming the exact path, not a runtime surprise.

Resolved - correct, interface_library is a bare label so a missing .so is a build-time error, not a silent runtime issue.

Resolved — good point about interface_library in cc_import making a missing .so a build-time error regardless of the glob behavior. The safety net is at the right layer.

claude · 2026-08-18T20:15:27Z

Review Summary

Clean migration from rocm_smi to amd_smi for ROCm >= 7.13. The version gating, Bazel plumbing, and SmiDeviceHandle abstraction are well structured, and the rationale for the 7.13 boundary (matching RCCL's switch, avoiding symbol conflicts) is clearly documented.

One item worth verifying before merge: the peer == *device pointer comparison in rocm_xgmi_topology.cc relies on amdsmi_get_processor_handle_from_bdf and amdsmi_get_processor_handles returning the same opaque pointer for the same physical device. If they don't, the self-skip won't fire.

The remaining inline comments are minor (implicit narrowing, unused variable, documentation asks).

draganmladjenovic · 2026-08-19T14:40:37Z

-    static_library = "%{rocm_root}/lib/libhsakmt.a",
-    visibility = ["//visibility:public"],
-)
+# RCCL links librocm_smi64 below ROCm 7.13 and libamd_smi from 7.13 on. XLA


leave both libs and use alias at point of usage in xla/stream_executor/rocm/BUILD. The point of adding both is to not care which one is actuality used.

claude · 2026-08-21T10:22:52Z

Re-review Summary

All 5 previous findings have been addressed by the developer:

Pointer identity for device comparison — verified as the library's own equality definition
PCIe speed units (MT/s) — amd_smi header reference added to the comment
BdfComponents narrowing — field types already match at uint64_t, no narrowing occurs
hops out-parameter — clarifying comment added; it is a required parameter, not unused
glob() silent empty list — interface_library in cc_import provides the build-time safety net

No new issues found. The refactoring cleanly splits the SMI backends into separate compilation units with correct build-time selection via rocm_version_number() >= 71300.

draganmladjenovic · 2026-08-21T15:32:14Z

-// Returns std::nullopt if not found.
-std::optional<uint32_t> FindDeviceIndex(const BdfComponents& target_bdf);
+// Returns handles for every SMI visible GPU, in SMI enumeration order. Empty
+// on failure. Callers must hold rocm_smi_mutex.


Good comment on mutex. Maybe use TF_EXCLUSIVE_LOCKS_REQUIRED here. Not that we check of it. But it is somewhat selfdocumenting.

draganmladjenovic · 2026-08-21T15:39:46Z

+
+// Finds the SMI device that matches the given PCI bus ID.
+// Returns std::nullopt if not found. Callers must hold rocm_smi_mutex.
+std::optional<SmiDeviceHandle> FindDevice(const BdfComponents& target_bdf);


Does it make sense for this to be StatusOr? I guess that would avoid most of the need for kSmiLibraryName and caller can chose if to just vlog the message or treat it as something serious. Also we can ensure we only log the message once per device. I won't hold of PR on this. This can be a follow up. This std::optional somehow passed the initial review on upstream.

Yes, it makes sense. I have removed kSmiLibraryName completely as well.

draganmladjenovic · 2026-08-25T16:37:31Z

-                 << pci_bus_id;
-    return std::nullopt;
-  }
+  absl::StatusOr<SmiDeviceHandle> device = FindDevice(*bdf);


*ASIGN_OR_RETURN macro?

Actually this reminders me @nurmukhametov shall we rm this part https://github.com/ROCm/xla/blob/rocm-dev-infra/.claude/skills/review-xla-pr/SKILL.md#11-xla-project-specific-c-conventions due to openxla@eb278d3 ?

These stale skills might hurt claude code review quality? Instead of giving the specific instruction, maybe we just let Claude to follow other existing error handle APIs in the codebase.

In addition, we're still using opus-4.6, wondering can we use opus-4.8?

*ASIGN_OR_RETURN macro?

Done.

Actually this reminders me @nurmukhametov shall we rm this part https://github.com/ROCm/xla/blob/rocm-dev-infra/.claude/skills/review-xla-pr/SKILL.md#11-xla-project-specific-c-conventions due to openxla@eb278d3 ?

These stale skills might hurt claude code review quality? Instead of giving the specific instruction, maybe we just let Claude to follow other existing error handle APIs in the codebase.

Yeah, could be. The stale things need to be removed for sure.

In addition, we're still using opus-4.6, wondering can we use opus-4.8?

We need to use the latest model.

draganmladjenovic · 2026-08-25T16:39:07Z


-  if (!InitRocmSmi()) return std::nullopt;
+  if (!InitSmi()) return absl::UnavailableError("SMI is not available");


Maybe make InitSmi return status.

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

) Cherry-pick of ROCm#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)

nurmukhametov · 2026-09-28T16:23:33Z

It was merged upstream in openxla@ca92135

nurmukhametov force-pushed the anurmukh/migrate-rocm-smi-to-amd-smi branch from 1cba3a7 to 1d0de6d Compare August 18, 2026 10:21

nurmukhametov requested a review from draganmladjenovic August 18, 2026 10:22

nurmukhametov added the claude-review Request a Claude AI code review for this PR label Aug 18, 2026

nurmukhametov requested a review from i-chaochen August 18, 2026 10:22

github-actions Bot removed the claude-review Request a Claude AI code review for this PR label Aug 18, 2026

i-chaochen reviewed Aug 18, 2026

View reviewed changes

nurmukhametov force-pushed the anurmukh/migrate-rocm-smi-to-amd-smi branch from 1d0de6d to 8b44a37 Compare August 18, 2026 11:37

claude Bot reviewed Aug 18, 2026

View reviewed changes

Comment thread xla/stream_executor/rocm/rocm_xgmi_topology.cc Resolved

claude Bot reviewed Aug 18, 2026

View reviewed changes

nurmukhametov force-pushed the anurmukh/migrate-rocm-smi-to-amd-smi branch from 8b44a37 to 0ce954d Compare August 19, 2026 11:11

draganmladjenovic reviewed Aug 19, 2026

View reviewed changes

nurmukhametov requested a review from draganmladjenovic August 21, 2026 10:09

nurmukhametov added the claude-review Request a Claude AI code review for this PR label Aug 21, 2026

github-actions Bot removed the claude-review Request a Claude AI code review for this PR label Aug 21, 2026

draganmladjenovic reviewed Aug 21, 2026

View reviewed changes

draganmladjenovic approved these changes Aug 21, 2026

View reviewed changes

nurmukhametov requested a review from draganmladjenovic August 24, 2026 15:05

nurmukhametov force-pushed the anurmukh/migrate-rocm-smi-to-amd-smi branch from 927e3cc to 40dd19d Compare August 25, 2026 14:30

draganmladjenovic reviewed Aug 25, 2026

View reviewed changes

gulsumgudukbay mentioned this pull request Aug 28, 2026

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

Merged

nurmukhametov closed this Sep 28, 2026

		data = glob([
		"%{rocm_root}/lib/libamd_smi.so" if _use_amd_smi else "%{rocm_root}/lib/librocm_smi64.so",


		if (!InitRocmSmi()) return std::nullopt;
		if (!InitSmi()) return absl::UnavailableError("SMI is not available");

Conversation

nurmukhametov commented Aug 18, 2026

Motivation

Technical Details

Submission Checklist

Uh oh!

Uh oh!

i-chaochen Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

i-chaochen Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

claude Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot commented Aug 18, 2026

Review Summary

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot commented Aug 21, 2026

i-chaochen Aug 18, 2026 •

edited

Loading

i-chaochen Aug 18, 2026 •

edited

Loading

i-chaochen Aug 25, 2026 •

edited

Loading