From 45f7885d3642a2f2f65c8b32abbde986db1f5683 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Thu, 27 Aug 2026 02:31:49 +0000 Subject: [PATCH] [ROCm] Migrate SMI device queries from rocm_smi to amd_smi Cherry-pick of ROCm/xla#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). --- third_party/gpus/rocm/BUILD.tpl | 11 +- xla/stream_executor/rocm/BUILD | 46 ++++- xla/stream_executor/rocm/rocm_executor.cc | 23 ++- .../rocm/rocm_pcie_bandwidth.cc | 51 ++--- .../rocm/rocm_pcie_bandwidth.h | 7 +- xla/stream_executor/rocm/rocm_smi_util.h | 51 ----- .../rocm/rocm_xgmi_topology.cc | 76 ++++---- xla/stream_executor/rocm/rocm_xgmi_topology.h | 7 +- .../rocm/rocm_xgmi_topology_test.cc | 15 +- .../rocm/{rocm_smi_util.cc => smi_util.cc} | 86 +++------ xla/stream_executor/rocm/smi_util.h | 100 ++++++++++ xla/stream_executor/rocm/smi_util_amd_smi.cc | 176 ++++++++++++++++++ xla/stream_executor/rocm/smi_util_rocm_smi.cc | 150 +++++++++++++++ 13 files changed, 574 insertions(+), 225 deletions(-) delete mode 100644 xla/stream_executor/rocm/rocm_smi_util.h rename xla/stream_executor/rocm/{rocm_smi_util.cc => smi_util.cc} (53%) create mode 100644 xla/stream_executor/rocm/smi_util.h create mode 100644 xla/stream_executor/rocm/smi_util_amd_smi.cc create mode 100644 xla/stream_executor/rocm/smi_util_rocm_smi.cc diff --git a/third_party/gpus/rocm/BUILD.tpl b/third_party/gpus/rocm/BUILD.tpl index a723d728c32613..c017334792f44c 100644 --- a/third_party/gpus/rocm/BUILD.tpl +++ b/third_party/gpus/rocm/BUILD.tpl @@ -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( diff --git a/xla/stream_executor/rocm/BUILD b/xla/stream_executor/rocm/BUILD index 800e9543bc0a64..62dbb8d5668159 100644 --- a/xla/stream_executor/rocm/BUILD +++ b/xla/stream_executor/rocm/BUILD @@ -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") @@ -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", ], ) @@ -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", ], ) @@ -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"], diff --git a/xla/stream_executor/rocm/rocm_executor.cc b/xla/stream_executor/rocm/rocm_executor.cc index ea0726b6bd40a5..c223e16cbc663d 100644 --- a/xla/stream_executor/rocm/rocm_executor.cc +++ b/xla/stream_executor/rocm/rocm_executor.cc @@ -1168,26 +1168,31 @@ RocmExecutor::CreateDeviceDescription(int device_ordinal) { } { - std::optional pcie_bw = gpu::GetRocmPcieBandwidth(pci_bus_id); - if (pcie_bw.has_value()) { + absl::StatusOr 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 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 << ")"; } } diff --git a/xla/stream_executor/rocm/rocm_pcie_bandwidth.cc b/xla/stream_executor/rocm/rocm_pcie_bandwidth.cc index 11c9ed02125395..03a2afd14e0ea7 100644 --- a/xla/stream_executor/rocm/rocm_pcie_bandwidth.cc +++ b/xla/stream_executor/rocm/rocm_pcie_bandwidth.cc @@ -13,12 +13,14 @@ limitations under the License. #include "xla/stream_executor/rocm/rocm_pcie_bandwidth.h" #include -#include +#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 { @@ -49,43 +51,22 @@ constexpr int64_t ComputePcieBandwidthFromSpeedAndWidth( } // namespace -std::optional GetRocmPcieBandwidth(absl::string_view pci_bus_id) { - absl::MutexLock lock(rocm_smi_mutex); +absl::StatusOr GetRocmPcieBandwidth(absl::string_view pci_bus_id) { + absl::MutexLock lock(smi_mutex); - if (!InitRocmSmi()) return std::nullopt; + ABSL_RETURN_IF_ERROR(InitSmi()); - std::optional 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 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(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 = diff --git a/xla/stream_executor/rocm/rocm_pcie_bandwidth.h b/xla/stream_executor/rocm/rocm_pcie_bandwidth.h index 7b2d6fbbb91c56..c72de149b093fd 100644 --- a/xla/stream_executor/rocm/rocm_pcie_bandwidth.h +++ b/xla/stream_executor/rocm/rocm_pcie_bandwidth.h @@ -14,15 +14,16 @@ limitations under the License. #define XLA_STREAM_EXECUTOR_ROCM_ROCM_PCIE_BANDWIDTH_H_ #include -#include +#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 GetRocmPcieBandwidth(absl::string_view pci_bus_id); +// Returns bandwidth in bytes/second. Does not log; the caller decides what a +// failure means. +absl::StatusOr GetRocmPcieBandwidth(absl::string_view pci_bus_id); } // namespace stream_executor::gpu diff --git a/xla/stream_executor/rocm/rocm_smi_util.h b/xla/stream_executor/rocm/rocm_smi_util.h deleted file mode 100644 index e9bbd94479fe8c..00000000000000 --- a/xla/stream_executor/rocm/rocm_smi_util.h +++ /dev/null @@ -1,51 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_STREAM_EXECUTOR_ROCM_ROCM_SMI_UTIL_H_ -#define XLA_STREAM_EXECUTOR_ROCM_ROCM_SMI_UTIL_H_ - -#include -#include - -#include "absl/base/attributes.h" -#include "absl/strings/string_view.h" -#include "absl/synchronization/mutex.h" - -namespace stream_executor::gpu { - -struct BdfComponents { - uint64_t domain; - uint64_t bus; - uint64_t device; - uint64_t function; -}; - -// Process-global lock serializing all rocm_smi (rsmi_*) access from XLA. -// rocm_smi only guards state with a per-device mutex, but some of it is global -// (e.g. the shared gpu_metrics object), so concurrent queries on different -// devices race. Hold this across each full rsmi call sequence. -ABSL_CONST_INIT extern absl::Mutex rocm_smi_mutex; - -// Returns true if rocm_smi was successfully initialized. -bool InitRocmSmi(); - -// Parses a PCI bus ID string (e.g., "0000:41:00.0") into its BDF components. -// Returns std::nullopt on parse failure. -std::optional ParseBdf(absl::string_view pci_bus_id); - -// Finds the rocm_smi device index that matches the given PCI bus ID. -// Returns std::nullopt if not found. -std::optional FindDeviceIndex(const BdfComponents& target_bdf); - -} // namespace stream_executor::gpu - -#endif // XLA_STREAM_EXECUTOR_ROCM_ROCM_SMI_UTIL_H_ diff --git a/xla/stream_executor/rocm/rocm_xgmi_topology.cc b/xla/stream_executor/rocm/rocm_xgmi_topology.cc index d113b62f6ebece..0f921df31f22db 100644 --- a/xla/stream_executor/rocm/rocm_xgmi_topology.cc +++ b/xla/stream_executor/rocm/rocm_xgmi_topology.cc @@ -13,72 +13,64 @@ limitations under the License. #include "xla/stream_executor/rocm/rocm_xgmi_topology.h" #include -#include +#include +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.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 { -XgmiTopologyInfo GetRocmXgmiTopology(absl::string_view pci_bus_id) { - XgmiTopologyInfo info; - - absl::MutexLock lock(rocm_smi_mutex); +absl::StatusOr GetRocmXgmiTopology( + absl::string_view pci_bus_id) { + absl::MutexLock lock(smi_mutex); - if (!InitRocmSmi()) return info; + ABSL_RETURN_IF_ERROR(InitSmi()); - std::optional bdf = ParseBdf(pci_bus_id); - if (!bdf.has_value()) { - LOG(WARNING) << "Failed to parse PCI bus ID for xGMI query: " << pci_bus_id; - return info; - } + ABSL_ASSIGN_OR_RETURN(BdfComponents bdf, ParseBdf(pci_bus_id)); + ABSL_ASSIGN_OR_RETURN(SmiDeviceHandle device, FindDevice(bdf)); - std::optional dev_idx = FindDeviceIndex(*bdf); - if (!dev_idx.has_value()) { - LOG(WARNING) << "rocm_smi: could not find device for PCI bus ID " - << pci_bus_id << " (xGMI query)"; - return info; - } + XgmiTopologyInfo info; - // Query xGMI hive ID. - uint64_t hive_id = 0; - rsmi_status_t status = rsmi_dev_xgmi_hive_id_get(*dev_idx, &hive_id); - if (status == RSMI_STATUS_SUCCESS) { - info.hive_id = hive_id; + // A device outside an xGMI hive is a normal configuration, and SMI reports + // it the same way it reports a failed query, so this is not fatal. + absl::StatusOr hive_id = QueryHiveId(device); + if (hive_id.ok()) { + info.hive_id = *hive_id; } else { - VLOG(1) << "rsmi_dev_xgmi_hive_id_get failed for " << pci_bus_id + VLOG(2) << "xGMI hive ID query failed for " << pci_bus_id << ": " + << hive_id.status().message() << "; device may not be in an xGMI hive."; } - // Count xGMI links by querying link type to every other device. - uint32_t num_devices = 0; - status = rsmi_num_monitor_devices(&num_devices); - if (status != RSMI_STATUS_SUCCESS || num_devices <= 1) { - return info; - } + // Count peers reachable over xGMI by querying the link type to every other + // device. This counts peer GPUs, not physical links. + ABSL_ASSIGN_OR_RETURN(std::vector devices, + EnumerateDevices()); + if (devices.size() <= 1) return info; int xgmi_links = 0; - for (uint32_t i = 0; i < num_devices; ++i) { - if (i == *dev_idx) continue; - - uint64_t hops = 0; - RSMI_IO_LINK_TYPE link_type = RSMI_IOLINK_TYPE_UNDEFINED; - status = rsmi_topo_get_link_type(*dev_idx, i, &hops, &link_type); - if (status != RSMI_STATUS_SUCCESS) continue; - - if (link_type == RSMI_IOLINK_TYPE_XGMI) { - ++xgmi_links; + for (SmiDeviceHandle peer : devices) { + if (peer == device) continue; + absl::StatusOr is_peer = IsXgmiPeer(device, peer); + if (!is_peer.ok()) { + VLOG(2) << "xGMI link type query failed for " << pci_bus_id << ": " + << is_peer.status().message(); + continue; } + if (*is_peer) ++xgmi_links; } info.active_links = xgmi_links; VLOG(1) << "xGMI topology for " << pci_bus_id << ": " << xgmi_links << " active xGMI links" - << " (hive_id=" << hive_id << ", num_devices=" << num_devices << ")"; + << " (hive_id=" << info.hive_id << ", num_devices=" << devices.size() + << ")"; return info; } diff --git a/xla/stream_executor/rocm/rocm_xgmi_topology.h b/xla/stream_executor/rocm/rocm_xgmi_topology.h index 55b5ba1b4e856c..065baeb6da6eab 100644 --- a/xla/stream_executor/rocm/rocm_xgmi_topology.h +++ b/xla/stream_executor/rocm/rocm_xgmi_topology.h @@ -15,6 +15,7 @@ limitations under the License. #include +#include "absl/status/statusor.h" #include "absl/strings/string_view.h" namespace stream_executor::gpu { @@ -26,8 +27,10 @@ struct XgmiTopologyInfo { // Queries xGMI topology information for a ROCm device. // pci_bus_id is the PCI bus ID string from HIP (e.g. "0000:41:00.0"). -// Returns topology info. On failure, returns default (0 links, no hive). -XgmiTopologyInfo GetRocmXgmiTopology(absl::string_view pci_bus_id); +// A device with no xGMI links is not a failure; it reports zero active links. +// Does not log; the caller decides what a failure means. +absl::StatusOr GetRocmXgmiTopology( + absl::string_view pci_bus_id); } // namespace stream_executor::gpu diff --git a/xla/stream_executor/rocm/rocm_xgmi_topology_test.cc b/xla/stream_executor/rocm/rocm_xgmi_topology_test.cc index ea41460a4759b9..fed4ce625ad0dc 100644 --- a/xla/stream_executor/rocm/rocm_xgmi_topology_test.cc +++ b/xla/stream_executor/rocm/rocm_xgmi_topology_test.cc @@ -23,17 +23,14 @@ TEST(XgmiTopologyInfoTest, DefaultValues) { EXPECT_EQ(info.hive_id, 0); } -TEST(GetRocmXgmiTopologyTest, InvalidBdfReturnsDefault) { - // An invalid PCI bus ID should return default topology (0 links). - XgmiTopologyInfo info = GetRocmXgmiTopology("invalid"); - EXPECT_EQ(info.active_links, 0); - EXPECT_EQ(info.hive_id, 0); +TEST(GetRocmXgmiTopologyTest, InvalidBdfFails) { + // An unparsable PCI bus ID cannot name a device, so the query fails rather + // than reporting an empty topology. + EXPECT_FALSE(GetRocmXgmiTopology("invalid").ok()); } -TEST(GetRocmXgmiTopologyTest, EmptyBdfReturnsDefault) { - XgmiTopologyInfo info = GetRocmXgmiTopology(""); - EXPECT_EQ(info.active_links, 0); - EXPECT_EQ(info.hive_id, 0); +TEST(GetRocmXgmiTopologyTest, EmptyBdfFails) { + EXPECT_FALSE(GetRocmXgmiTopology("").ok()); } } // namespace diff --git a/xla/stream_executor/rocm/rocm_smi_util.cc b/xla/stream_executor/rocm/smi_util.cc similarity index 53% rename from xla/stream_executor/rocm/rocm_smi_util.cc rename to xla/stream_executor/rocm/smi_util.cc index 838be674c064a3..a2afb68e9c3ae7 100644 --- a/xla/stream_executor/rocm/rocm_smi_util.cc +++ b/xla/stream_executor/rocm/smi_util.cc @@ -10,50 +10,42 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/stream_executor/rocm/rocm_smi_util.h" +// Backend independent part of the SMI utilities. The SMI calls themselves live +// in smi_util_amd_smi.cc and smi_util_rocm_smi.cc. + +#include "xla/stream_executor/rocm/smi_util.h" #include -#include -#include #include "absl/base/attributes.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/numbers.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/tsl/platform/logging.h" namespace stream_executor::gpu { -ABSL_CONST_INIT absl::Mutex rocm_smi_mutex(absl::kConstInit); - -bool InitRocmSmi() { - static bool initialized = []() { - rsmi_status_t status = rsmi_init(0); - if (status != RSMI_STATUS_SUCCESS) { - const char* err_str = nullptr; - rsmi_status_string(status, &err_str); - LOG(WARNING) << "rsmi_init failed: " - << (err_str ? err_str : "unknown error"); - return false; - } - return true; - }(); - return initialized; -} +ABSL_CONST_INIT absl::Mutex smi_mutex(absl::kConstInit); // Parses a PCI bus/device/function ID string into its numeric components. // Accepts two formats: // DDDD:BB:DD.F - domain:bus:device.function (e.g. "0000:41:00.0") // BB:DD.F - bus:device.function, domain defaults to 0 // All fields are hex -std::optional ParseBdf(absl::string_view pci_bus_id) { +absl::StatusOr ParseBdf(absl::string_view pci_bus_id) { + const auto invalid = [&] { + return absl::InvalidArgumentError( + absl::StrCat("cannot parse PCI bus ID: ", pci_bus_id)); + }; + BdfComponents bdf = {}; // Determine which format we have by counting colons. // Two colons -> DDDD:BB:DD.F, one colon -> BB:DD.F. size_t first_colon = pci_bus_id.find(':'); - if (first_colon == absl::string_view::npos) return std::nullopt; + if (first_colon == absl::string_view::npos) return invalid(); size_t second_colon = pci_bus_id.find(':', first_colon + 1); size_t dot; @@ -61,67 +53,37 @@ std::optional ParseBdf(absl::string_view pci_bus_id) { if (second_colon != absl::string_view::npos) { // DDDD:BB:DD.F format dot = pci_bus_id.find('.', second_colon + 1); - if (dot == absl::string_view::npos) return std::nullopt; + if (dot == absl::string_view::npos) return invalid(); if (!absl::SimpleHexAtoi(pci_bus_id.substr(0, first_colon), &bdf.domain)) - return std::nullopt; + return invalid(); if (!absl::SimpleHexAtoi( pci_bus_id.substr(first_colon + 1, second_colon - first_colon - 1), &bdf.bus)) - return std::nullopt; + return invalid(); if (!absl::SimpleHexAtoi( pci_bus_id.substr(second_colon + 1, dot - second_colon - 1), &bdf.device)) - return std::nullopt; + return invalid(); if (!absl::SimpleHexAtoi(pci_bus_id.substr(dot + 1), &bdf.function)) - return std::nullopt; + return invalid(); } else { // BB:DD.F format (domain = 0) dot = pci_bus_id.find('.', first_colon + 1); - if (dot == absl::string_view::npos) return std::nullopt; + if (dot == absl::string_view::npos) return invalid(); bdf.domain = 0; if (!absl::SimpleHexAtoi(pci_bus_id.substr(0, first_colon), &bdf.bus)) - return std::nullopt; + return invalid(); if (!absl::SimpleHexAtoi( pci_bus_id.substr(first_colon + 1, dot - first_colon - 1), &bdf.device)) - return std::nullopt; + return invalid(); if (!absl::SimpleHexAtoi(pci_bus_id.substr(dot + 1), &bdf.function)) - return std::nullopt; + return invalid(); } return bdf; } -std::optional FindDeviceIndex(const BdfComponents& target_bdf) { - uint32_t num_devices = 0; - rsmi_status_t status = rsmi_num_monitor_devices(&num_devices); - if (status != RSMI_STATUS_SUCCESS || num_devices == 0) { - return std::nullopt; - } - - for (uint32_t i = 0; i < num_devices; ++i) { - uint64_t bdfid = 0; - status = rsmi_dev_pci_id_get(i, &bdfid); - if (status != RSMI_STATUS_SUCCESS) continue; - - // Unpack rocm_smi's 64-bit BDF format into individual fields. - // See - // rocm-systems/projects/rocm-smi-lib/src/rocm_smi.cc:rsmi_dev_pci_id_get - // for details on the packing. - uint32_t domain = (bdfid >> 32) & 0xFFFFFFFF; - uint8_t bus = (bdfid >> 8) & 0xFF; - uint8_t device = (bdfid >> 3) & 0x1F; - uint8_t function = bdfid & 0x7; - - if (domain == target_bdf.domain && bus == target_bdf.bus && - device == target_bdf.device && function == target_bdf.function) { - return i; - } - } - - return std::nullopt; -} - } // namespace stream_executor::gpu diff --git a/xla/stream_executor/rocm/smi_util.h b/xla/stream_executor/rocm/smi_util.h new file mode 100644 index 00000000000000..aeca31ca24d58e --- /dev/null +++ b/xla/stream_executor/rocm/smi_util.h @@ -0,0 +1,100 @@ +/* Copyright 2026 The OpenXLA Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_STREAM_EXECUTOR_ROCM_SMI_UTIL_H_ +#define XLA_STREAM_EXECUTOR_ROCM_SMI_UTIL_H_ + +#include +#include + +#include "absl/base/attributes.h" +#include "absl/base/thread_annotations.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" + +// Backend independent view of the SMI queries XLA needs. Exactly one backend +// implements them, smi_util_amd_smi.cc or smi_util_rocm_smi.cc, and +// the BUILD file picks the source and the library to link against. Nothing +// here exposes an SMI type, so callers never need to know which one they got. +// The rationale for the ROCm 7.13 boundary lives next to that choice, in +// xla/stream_executor/rocm/BUILD. + +namespace stream_executor::gpu { + +struct BdfComponents { + uint64_t domain; + uint64_t bus; + uint64_t device; + uint64_t function; +}; + +// Opaque identifier of an SMI visible GPU. Only the backend interprets the +// value: rocm_smi identifies a GPU by a dense monitor device index, amd_smi by +// an opaque processor handle, and both fit in a uintptr_t. +struct SmiDeviceHandle { + uintptr_t value = 0; + + friend bool operator==(SmiDeviceHandle lhs, SmiDeviceHandle rhs) { + return lhs.value == rhs.value; + } + friend bool operator!=(SmiDeviceHandle lhs, SmiDeviceHandle rhs) { + return !(lhs == rhs); + } +}; + +// Current PCIe link state of a device. +struct PcieLinkStatus { + uint32_t speed_mt_per_sec; + uint16_t width; +}; + +// Process-global lock serializing all SMI access from XLA. rocm_smi only +// guards state with a per-device mutex, but some of it is global (e.g. the +// shared gpu_metrics object), so concurrent queries on different devices race. +// amd_smi embeds rocm_smi and inherits the same problem. Hold this across each +// full SMI call sequence. +ABSL_CONST_INIT extern absl::Mutex smi_mutex; + +// Initializes the SMI library at most once per process. Every later call +// reports the result of that one attempt. +absl::Status InitSmi() ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +// Parses a PCI bus ID string (e.g., "0000:41:00.0") into its BDF components. +// Touches no SMI state, so it needs no lock. +absl::StatusOr ParseBdf(absl::string_view pci_bus_id); + +// Returns handles for every SMI visible GPU, in SMI enumeration order. +absl::StatusOr> EnumerateDevices() + ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +// Finds the SMI device that matches the given PCI bus ID. +absl::StatusOr FindDevice(const BdfComponents& target_bdf) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +// Returns the current PCIe link speed and width of the device. +absl::StatusOr QueryPcieLinkStatus(SmiDeviceHandle device) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +// Returns the xGMI hive ID of the device. Fails if it is not in a hive, which +// SMI does not distinguish from a failed query. +absl::StatusOr QueryHiveId(SmiDeviceHandle device) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +// Returns true if src reaches dst over an xGMI link. +absl::StatusOr IsXgmiPeer(SmiDeviceHandle src, SmiDeviceHandle dst) + ABSL_EXCLUSIVE_LOCKS_REQUIRED(smi_mutex); + +} // namespace stream_executor::gpu + +#endif // XLA_STREAM_EXECUTOR_ROCM_SMI_UTIL_H_ diff --git a/xla/stream_executor/rocm/smi_util_amd_smi.cc b/xla/stream_executor/rocm/smi_util_amd_smi.cc new file mode 100644 index 00000000000000..001c3c6821acd0 --- /dev/null +++ b/xla/stream_executor/rocm/smi_util_amd_smi.cc @@ -0,0 +1,176 @@ +/* Copyright 2026 The OpenXLA Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// amd_smi backend for the SMI queries declared in smi_util.h. Compiled in +// from ROCm 7.13 on; smi_util_rocm_smi.cc takes its place below that. + +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "rocm/include/amd_smi/amdsmi.h" +#include "xla/stream_executor/rocm/smi_util.h" +#include "xla/tsl/platform/logging.h" + +namespace stream_executor::gpu { +namespace { + +amdsmi_processor_handle ToProcessorHandle(SmiDeviceHandle device) { + return reinterpret_cast(device.value); +} + +SmiDeviceHandle ToDeviceHandle(amdsmi_processor_handle processor) { + return SmiDeviceHandle{reinterpret_cast(processor)}; +} + +absl::Status SmiError(absl::string_view api, amdsmi_status_t status) { + const char* err_str = nullptr; + amdsmi_status_code_to_string(status, &err_str); + return absl::InternalError( + absl::StrCat(api, " failed: ", err_str ? err_str : "unknown error")); +} + +amdsmi_status_t InitLibrary() { + amdsmi_status_t status = amdsmi_init(AMDSMI_INIT_AMD_GPUS); + if (status == AMDSMI_STATUS_SUCCESS) { + VLOG(1) << "SMI device queries go through amd_smi."; + } + return status; +} + +} // namespace + +absl::Status InitSmi() { + // Caching the raw status, not an absl::Status, keeps this a plain static. + static const amdsmi_status_t status = InitLibrary(); + if (status != AMDSMI_STATUS_SUCCESS) return SmiError("amdsmi_init", status); + return absl::OkStatus(); +} + +absl::StatusOr> EnumerateDevices() { + uint32_t num_sockets = 0; + if (amdsmi_status_t status = amdsmi_get_socket_handles(&num_sockets, nullptr); + status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_get_socket_handles", status); + } + + if (num_sockets == 0) return std::vector(); + + std::vector sockets(num_sockets); + if (amdsmi_status_t status = + amdsmi_get_socket_handles(&num_sockets, sockets.data()); + status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_get_socket_handles", status); + } + + // A socket that cannot be read is skipped rather than failing the whole + // enumeration, so the other GPUs still get described. + std::vector devices; + for (amdsmi_socket_handle socket : sockets) { + uint32_t num_processors = 0; + if (amdsmi_status_t status = + amdsmi_get_processor_handles(socket, &num_processors, nullptr); + status != AMDSMI_STATUS_SUCCESS) { + VLOG(2) << "Skipping socket: " + << SmiError("amdsmi_get_processor_handles", status).message(); + continue; + } + if (num_processors == 0) continue; + + std::vector processors(num_processors); + if (amdsmi_status_t status = amdsmi_get_processor_handles( + socket, &num_processors, processors.data()); + status != AMDSMI_STATUS_SUCCESS) { + VLOG(2) << "Skipping socket: " + << SmiError("amdsmi_get_processor_handles", status).message(); + continue; + } + + // amdsmi_init(AMDSMI_INIT_AMD_GPUS) already restricts enumeration to + // sockets holding AMD GPUs, but a socket can still expose processors that + // are not GPUs, so filter by type. + for (amdsmi_processor_handle processor : processors) { + processor_type_t type = AMDSMI_PROCESSOR_TYPE_UNKNOWN; + if (amdsmi_get_processor_type(processor, &type) == + AMDSMI_STATUS_SUCCESS && + type == AMDSMI_PROCESSOR_TYPE_AMD_GPU) { + devices.push_back(ToDeviceHandle(processor)); + } + } + } + + return devices; +} + +absl::StatusOr FindDevice(const BdfComponents& target_bdf) { + 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; + + amdsmi_processor_handle handle = nullptr; + amdsmi_status_t status = amdsmi_get_processor_handle_from_bdf(bdf, &handle); + if (status == AMDSMI_STATUS_NOT_FOUND || handle == nullptr) { + return absl::NotFoundError( + absl::StrFormat("amd_smi exposes no device with BDF %04x:%02x:%02x.%x", + target_bdf.domain, target_bdf.bus, target_bdf.device, + target_bdf.function)); + } + if (status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_get_processor_handle_from_bdf", status); + } + + return ToDeviceHandle(handle); +} + +absl::StatusOr QueryPcieLinkStatus(SmiDeviceHandle device) { + amdsmi_pcie_info_t pcie_info = {}; + if (amdsmi_status_t status = + amdsmi_get_pcie_info(ToProcessorHandle(device), &pcie_info); + status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_get_pcie_info", status); + } + + // amdsmi.h documents pcie_metric.pcie_speed as "current PCIe speed in MT/s", + // so no scaling here, unlike rocm_smi's 0.1 GT/s field. + return PcieLinkStatus{pcie_info.pcie_metric.pcie_speed, + pcie_info.pcie_metric.pcie_width}; +} + +absl::StatusOr QueryHiveId(SmiDeviceHandle device) { + amdsmi_xgmi_info_t xgmi_info = {}; + if (amdsmi_status_t status = + amdsmi_get_xgmi_info(ToProcessorHandle(device), &xgmi_info); + status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_get_xgmi_info", status); + } + return xgmi_info.xgmi_hive_id; +} + +absl::StatusOr IsXgmiPeer(SmiDeviceHandle src, SmiDeviceHandle dst) { + // The API rejects a null hops pointer; only the link type is used. + uint64_t hops = 0; + amdsmi_link_type_t link_type = AMDSMI_LINK_TYPE_UNKNOWN; + if (amdsmi_status_t status = amdsmi_topo_get_link_type( + ToProcessorHandle(src), ToProcessorHandle(dst), &hops, &link_type); + status != AMDSMI_STATUS_SUCCESS) { + return SmiError("amdsmi_topo_get_link_type", status); + } + return link_type == AMDSMI_LINK_TYPE_XGMI; +} + +} // namespace stream_executor::gpu diff --git a/xla/stream_executor/rocm/smi_util_rocm_smi.cc b/xla/stream_executor/rocm/smi_util_rocm_smi.cc new file mode 100644 index 00000000000000..01fcff2ddde07e --- /dev/null +++ b/xla/stream_executor/rocm/smi_util_rocm_smi.cc @@ -0,0 +1,150 @@ +/* Copyright 2026 The OpenXLA Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// rocm_smi backend for the SMI queries declared in smi_util.h. Compiled +// in below ROCm 7.13; smi_util_amd_smi.cc takes its place from 7.13 on. + +#include +#include + +#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/str_format.h" +#include "absl/strings/string_view.h" +#include "rocm/include/rocm_smi/rocm_smi.h" +#include "xla/stream_executor/rocm/smi_util.h" +#include "xla/tsl/platform/logging.h" + +namespace stream_executor::gpu { +namespace { + +uint32_t ToDeviceIndex(SmiDeviceHandle device) { + return static_cast(device.value); +} + +SmiDeviceHandle ToDeviceHandle(uint32_t device_index) { + return SmiDeviceHandle{device_index}; +} + +absl::Status SmiError(absl::string_view api, rsmi_status_t status) { + const char* err_str = nullptr; + rsmi_status_string(status, &err_str); + return absl::InternalError( + absl::StrCat(api, " failed: ", err_str ? err_str : "unknown error")); +} + +rsmi_status_t InitLibrary() { + rsmi_status_t status = rsmi_init(0); + if (status == RSMI_STATUS_SUCCESS) { + VLOG(1) << "SMI device queries go through rocm_smi."; + } + return status; +} + +} // namespace + +absl::Status InitSmi() { + // Caching the raw status, not an absl::Status, keeps this a plain static. + static const rsmi_status_t status = InitLibrary(); + if (status != RSMI_STATUS_SUCCESS) return SmiError("rsmi_init", status); + return absl::OkStatus(); +} + +absl::StatusOr> EnumerateDevices() { + uint32_t num_devices = 0; + if (rsmi_status_t status = rsmi_num_monitor_devices(&num_devices); + status != RSMI_STATUS_SUCCESS) { + return SmiError("rsmi_num_monitor_devices", status); + } + + std::vector devices; + devices.reserve(num_devices); + for (uint32_t i = 0; i < num_devices; ++i) { + devices.push_back(ToDeviceHandle(i)); + } + return devices; +} + +absl::StatusOr FindDevice(const BdfComponents& target_bdf) { + ABSL_ASSIGN_OR_RETURN(std::vector devices, + EnumerateDevices()); + + for (SmiDeviceHandle device : devices) { + uint64_t bdfid = 0; + if (rsmi_status_t status = + rsmi_dev_pci_id_get(ToDeviceIndex(device), &bdfid); + status != RSMI_STATUS_SUCCESS) { + VLOG(2) << "Skipping device " << ToDeviceIndex(device) << ": " + << SmiError("rsmi_dev_pci_id_get", status).message(); + continue; + } + + // Unpack rocm_smi's 64-bit BDF format into individual fields. + // See + // rocm-systems/projects/rocm-smi-lib/src/rocm_smi.cc:rsmi_dev_pci_id_get + // for details on the packing. + uint32_t domain = (bdfid >> 32) & 0xFFFFFFFF; + uint8_t bus = (bdfid >> 8) & 0xFF; + uint8_t device_number = (bdfid >> 3) & 0x1F; + uint8_t function = bdfid & 0x7; + + if (domain == target_bdf.domain && bus == target_bdf.bus && + device_number == target_bdf.device && function == target_bdf.function) { + return device; + } + } + + return absl::NotFoundError( + absl::StrFormat("rocm_smi exposes no device with BDF %04x:%02x:%02x.%x", + target_bdf.domain, target_bdf.bus, target_bdf.device, + target_bdf.function)); +} + +absl::StatusOr QueryPcieLinkStatus(SmiDeviceHandle device) { + rsmi_gpu_metrics_t gpu_metrics = {}; + if (rsmi_status_t status = + rsmi_dev_gpu_metrics_info_get(ToDeviceIndex(device), &gpu_metrics); + status != RSMI_STATUS_SUCCESS) { + return SmiError("rsmi_dev_gpu_metrics_info_get", status); + } + + // rocm_smi reports pcie_link_speed in units of 0.1 GT/s, so scale to MT/s. + return PcieLinkStatus{ + static_cast(gpu_metrics.pcie_link_speed) * 100, + gpu_metrics.pcie_link_width}; +} + +absl::StatusOr QueryHiveId(SmiDeviceHandle device) { + uint64_t hive_id = 0; + if (rsmi_status_t status = + rsmi_dev_xgmi_hive_id_get(ToDeviceIndex(device), &hive_id); + status != RSMI_STATUS_SUCCESS) { + return SmiError("rsmi_dev_xgmi_hive_id_get", status); + } + return hive_id; +} + +absl::StatusOr IsXgmiPeer(SmiDeviceHandle src, SmiDeviceHandle dst) { + // The API rejects a null hops pointer; only the link type is used. + uint64_t hops = 0; + RSMI_IO_LINK_TYPE link_type = RSMI_IOLINK_TYPE_UNDEFINED; + if (rsmi_status_t status = rsmi_topo_get_link_type( + ToDeviceIndex(src), ToDeviceIndex(dst), &hops, &link_type); + status != RSMI_STATUS_SUCCESS) { + return SmiError("rsmi_topo_get_link_type", status); + } + return link_type == RSMI_IOLINK_TYPE_XGMI; +} + +} // namespace stream_executor::gpu