Skip to content

[ROCm] Migrate SMI device queries from rocm_smi to amd_smi (v0.11.1) #1140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
i-chaochen merged 1 commit into from
Aug 27, 2026
Merged

[ROCm] Migrate SMI device queries from rocm_smi to amd_smi (v0.11.1) #1140

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions third_party/gpus/rocm/BUILD.tpl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,9 +326,16 @@ rocm_lib_import(
],
)

cc_library(
name = "amdsmi_libs",
# Both SMI libraries are exposed here; consumers pick one. They must not be
# mixed in a single process, see the :smi alias in
# xla/stream_executor/rocm/BUILD.
rocm_lib_import(
name = "amdsmi",
data = glob(["%{rocm_root}/lib/libamd_smi.so*"]),
interface_library = "%{rocm_root}/lib/libamd_smi.so",
deps = [
":system_libs",
],
)

rocm_lib_import(
Expand Down
46 changes: 36 additions & 10 deletions xla/stream_executor/rocm/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
load(
"@local_config_rocm//rocm:build_defs.bzl",
"rocm_library",
"rocm_version_number",
)
load("//xla:xla.default.bzl", "xla_cc_test")
load("//xla/service/gpu:build_defs.bzl", "get_cub_sort_kernel_types")
Expand Down Expand Up @@ -371,21 +372,43 @@ cc_library(
deps = ["//xla/stream_executor:platform_id"],
)

# The SMI library XLA queries devices with. RCCL links librocm_smi64 below
# ROCm 7.13 and libamd_smi from 7.13 on, and XLA follows the same boundary so
# that only one SMI library is ever present in the process: libamd_smi embeds a
# copy of rocm_smi and the two export hundreds of identical amd::smi:: and
# rsmi_ symbols, so loading both lets the dynamic loader interpose ABI
# incompatible definitions.
# See RCCL change: https://github.com/ROCm/rocm-systems/commit/244047310e
_USE_AMD_SMI = rocm_version_number() >= 71300

alias(
name = "smi",
actual = "@local_config_rocm//rocm:amdsmi" if _USE_AMD_SMI else "@local_config_rocm//rocm:rocm_smi",
)

# Backend implementing the SMI queries declared in smi_util.h. Exactly one
# is compiled in, matching the library :smi resolves to.
_SMI_BACKEND_SRCS = ["smi_util_amd_smi.cc"] if _USE_AMD_SMI else ["smi_util_rocm_smi.cc"]

cc_library(
name = "rocm_smi_util",
srcs = ["rocm_smi_util.cc"],
hdrs = ["rocm_smi_util.h"],
name = "smi_util",
srcs = ["smi_util.cc"] + _SMI_BACKEND_SRCS,
hdrs = ["smi_util.h"],
tags = [
"gpu",
"rocm-only",
],
deps = [
":smi",
"//xla/tsl/platform:logging",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@local_config_rocm//rocm:rocm_headers",
"@local_config_rocm//rocm:rocm_smi",
],
)

Expand All @@ -398,12 +421,13 @@ cc_library(
"rocm-only",
],
deps = [
":rocm_smi_util",
":smi_util",
"//xla/tsl/platform:logging",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@local_config_rocm//rocm:rocm_headers",
"@local_config_rocm//rocm:rocm_smi",
],
)

Expand All @@ -416,15 +440,17 @@ cc_library(
"rocm-only",
],
deps = [
":rocm_smi_util",
":smi_util",
"//xla/tsl/platform:logging",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_macros",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@local_config_rocm//rocm:rocm_headers",
"@local_config_rocm//rocm:rocm_smi",
],
)


xla_cc_test(
name = "rocm_xgmi_topology_test",
srcs = ["rocm_xgmi_topology_test.cc"],
Expand Down
23 changes: 14 additions & 9 deletions xla/stream_executor/rocm/rocm_executor.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1168,26 +1168,31 @@ RocmExecutor::CreateDeviceDescription(int device_ordinal) {
}

{
std::optional<int64_t> pcie_bw = gpu::GetRocmPcieBandwidth(pci_bus_id);
if (pcie_bw.has_value()) {
absl::StatusOr<int64_t> pcie_bw = gpu::GetRocmPcieBandwidth(pci_bus_id);
if (pcie_bw.ok()) {
desc.set_pcie_bandwidth(*pcie_bw);
} else {
LOG(WARNING) << "Could not determine PCIe bandwidth for device "
<< device_ordinal
<< " via rocm_smi. Assuming PCIe Gen4 x16.";
<< device_ordinal << " via SMI ("
<< pcie_bw.status().message()
<< "). Assuming PCIe Gen4 x16.";
desc.set_pcie_bandwidth(32LL * 1024 * 1024 * 1024);
}
}

{
gpu::XgmiTopologyInfo xgmi = gpu::GetRocmXgmiTopology(pci_bus_id);
if (xgmi.active_links > 0) {
absl::StatusOr<gpu::XgmiTopologyInfo> xgmi =
gpu::GetRocmXgmiTopology(pci_bus_id);
if (!xgmi.ok()) {
LOG(WARNING) << "Could not determine xGMI topology for device "
<< device_ordinal << " via SMI: " << xgmi.status().message();
} else if (xgmi->active_links > 0) {
DeviceInterconnectInfo info;
info.active_links = xgmi.active_links;
info.active_links = xgmi->active_links;
desc.set_device_interconnect_info(info);
VLOG(1) << "Device " << device_ordinal << ": detected "
<< xgmi.active_links << " active xGMI links"
<< " (hive_id=" << xgmi.hive_id << ")";
<< xgmi->active_links << " active xGMI links"
<< " (hive_id=" << xgmi->hive_id << ")";
}
}

Expand Down
51 changes: 16 additions & 35 deletions xla/stream_executor/rocm/rocm_pcie_bandwidth.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,14 @@ limitations under the License.
#include "xla/stream_executor/rocm/rocm_pcie_bandwidth.h"

#include <cstdint>
#include <optional>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "rocm/include/rocm_smi/rocm_smi.h"
#include "xla/stream_executor/rocm/rocm_smi_util.h"
#include "xla/stream_executor/rocm/smi_util.h"
#include "xla/tsl/platform/logging.h"

namespace stream_executor::gpu {
Expand Down Expand Up @@ -49,43 +51,22 @@ constexpr int64_t ComputePcieBandwidthFromSpeedAndWidth(

} // namespace

std::optional<int64_t> GetRocmPcieBandwidth(absl::string_view pci_bus_id) {
absl::MutexLock lock(rocm_smi_mutex);
absl::StatusOr<int64_t> GetRocmPcieBandwidth(absl::string_view pci_bus_id) {
absl::MutexLock lock(smi_mutex);

if (!InitRocmSmi()) return std::nullopt;
ABSL_RETURN_IF_ERROR(InitSmi());

std::optional<BdfComponents> bdf = ParseBdf(pci_bus_id);
if (!bdf.has_value()) {
LOG(WARNING) << "Failed to parse PCI bus ID: " << pci_bus_id;
return std::nullopt;
}

std::optional<uint32_t> dev_idx = FindDeviceIndex(*bdf);
if (!dev_idx.has_value()) {
LOG(WARNING) << "rocm_smi: could not find device for PCI bus ID "
<< pci_bus_id;
return std::nullopt;
}

rsmi_gpu_metrics_t gpu_metrics = {};
rsmi_status_t status = rsmi_dev_gpu_metrics_info_get(*dev_idx, &gpu_metrics);
if (status != RSMI_STATUS_SUCCESS) {
const char* err_str = nullptr;
rsmi_status_string(status, &err_str);
LOG(WARNING) << "rsmi_dev_gpu_metrics_info_get failed for " << pci_bus_id
<< ": " << (err_str ? err_str : "unknown error");
return std::nullopt;
}
ABSL_ASSIGN_OR_RETURN(BdfComponents bdf, ParseBdf(pci_bus_id));
ABSL_ASSIGN_OR_RETURN(SmiDeviceHandle device, FindDevice(bdf));
ABSL_ASSIGN_OR_RETURN(PcieLinkStatus link, QueryPcieLinkStatus(device));

uint32_t speed_mt_per_sec =
static_cast<uint32_t>(gpu_metrics.pcie_link_speed) * 100;
uint16_t width = gpu_metrics.pcie_link_width;
uint32_t speed_mt_per_sec = link.speed_mt_per_sec;
uint16_t width = link.width;

if (speed_mt_per_sec == 0 || width == 0) {
LOG(WARNING) << "rocm_smi gpu_metrics reported zero PCIe speed ("
<< speed_mt_per_sec << " MT/s) or width (" << width
<< " lanes) for " << pci_bus_id;
return std::nullopt;
return absl::InternalError(
absl::StrCat("SMI reported zero PCIe speed (", speed_mt_per_sec,
" MT/s) or width (", width, " lanes)"));
}

int64_t bandwidth =
Expand Down
7 changes: 4 additions & 3 deletions xla/stream_executor/rocm/rocm_pcie_bandwidth.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,16 @@ limitations under the License.
#define XLA_STREAM_EXECUTOR_ROCM_ROCM_PCIE_BANDWIDTH_H_

#include <cstdint>
#include <optional>

#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"

namespace stream_executor::gpu {

// pci_bus_id is the PCI bus ID string from HIP (e.g. "0000:41:00.0").
// Returns bandwidth in bytes/second, or std::nullopt if the query fails.
std::optional<int64_t> GetRocmPcieBandwidth(absl::string_view pci_bus_id);
// Returns bandwidth in bytes/second. Does not log; the caller decides what a
// failure means.
absl::StatusOr<int64_t> GetRocmPcieBandwidth(absl::string_view pci_bus_id);

} // namespace stream_executor::gpu

Expand Down
51 changes: 0 additions & 51 deletions xla/stream_executor/rocm/rocm_smi_util.h
View file
Open in desktop

This file was deleted.

Loading
Loading