Fix GPU device index resolution and init/shutdown lifecycle races

cmake/hwsConfig.cmake.in

@@ -112,6 +112,15 @@ if (HWS_HAS_CPU_SUPPORT)
     endif ()
 endif ()
 
+# check whether MPI support is enabled
+string_contains("${HWS_COMPILE_DEFINITIONS}" "HWS_MPI_SUPPORT_ENABLED" HWS_HAS_MPI_SUPPORT)
+if (HWS_HAS_MPI_SUPPORT)
+    find_dependency(MPI COMPONENTS CXX)
+    if (NOT hws_FIND_QUIETLY)
+        message(STATUS "Enabled MPI support via hws.")
+    endif ()
+endif ()
+
 # check whether NVIDIA GPUs are supported
 string_contains("${HWS_COMPILE_DEFINITIONS}" "HWS_FOR_NVIDIA_GPUS_ENABLED" HWS_HAS_GPU_NVIDIA_SUPPORT)
 if (HWS_HAS_GPU_NVIDIA_SUPPORT)

include/hws/gpu_amd/hardware_sampler.hpp

@@ -18,11 +18,11 @@
 
 #include "fmt/ostream.h"  // fmt::formatter, fmt::ostream_formatter
 
-#include <atomic>   // std::atomic
 #include <chrono>   // std::chrono::milliseconds, std::chrono_literals namespace
 #include <cstddef>  // std::size_t
 #include <cstdint>  // std::uint32_t
 #include <iosfwd>   // std::ostream forward declaration
+#include <mutex>    // std::mutex
 
 namespace hws {
 
@@ -132,8 +132,12 @@ class gpu_amd_hardware_sampler : public hardware_sampler {
      */
     void sampling_loop() final;
 
-    /// The ID of the device to sample.
+    /// The ROCm SMI device index to sample, resolved from hip_device_id_ via a PCI bus ID match (ROCm SMI's own
+    /// device enumeration isn't affected by HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES, unlike HIP's).
     std::uint32_t device_id_{};
+    /// The HIP-relative device index this hardware sampler was constructed with; only used for the one HIP call
+    /// (hipGetDeviceProperties) that needs a HIP-space rather than a ROCm-SMI-space index.
+    std::uint32_t hip_device_id_{};
 
     /// The general AMD GPU samples.
     rocm_smi_general_samples general_samples_{};
@@ -146,10 +150,14 @@ class gpu_amd_hardware_sampler : public hardware_sampler {
     /// The temperature related AMD GPU samples.
     rocm_smi_temperature_samples temperature_samples_{};
 
-    /// The total number of currently active AMD GPU hardware samplers.
-    inline static std::atomic<int> instances_{ 0 };
-    /// True if the ROCm SMI environment has been successfully initialized (only done by a single hardware sampler).
-    inline static std::atomic<bool> init_finished_{ false };
+    /// Guards `instances_` and every `rsmi_init()`/`rsmi_shut_down()` call, so that the "first instance
+    /// initializes, last instance shuts down" decision and the actual init/shutdown call happen as one atomic
+    /// step - a busy-wait on a plain flag can't do this: a failing first `rsmi_init()` never sets it, permanently
+    /// stranding every waiter, and nothing prevents a shutdown from racing a concurrent init.
+    inline static std::mutex lifecycle_mutex_{};
+    /// The total number of currently active AMD GPU hardware samplers; only ever read/written while holding
+    /// `lifecycle_mutex_`.
+    inline static int instances_{ 0 };
 };
 
 /**

include/hws/gpu_amd/utility.hpp

@@ -17,11 +17,10 @@
 
 #include <stdexcept>  // std::runtime_error
 #include <string>     // std::string
+#include <vector>     // std::vector
 
 #if defined(HWS_MPI_SUPPORT_ENABLED)
     #include "hws/visible_gpu_device.hpp"  // hws::detail::visible_gpu_device
-
-    #include <vector>  // std::vector
 #endif
 
 namespace hws::detail {
@@ -74,6 +73,12 @@ namespace hws::detail {
  */
 [[nodiscard]] std::string performance_level_to_string(rsmi_dev_perf_level_t perf_level);
 
+/**
+ * @brief Return the PCI bus ID (e.g. `"0000:c1:00.0"`) of the AMD GPU device with the given HIP @p local_index.
+ * @param[in] local_index the local HIP device index
+ * @return the PCI bus ID (`[[nodiscard]]`)
+ */
+[[nodiscard]] std::string amd_device_pci_bus_id(int local_index);
 
 #if defined(HWS_MPI_SUPPORT_ENABLED)
 

include/hws/gpu_intel/hardware_sampler.hpp

@@ -19,10 +19,10 @@
 
 #include "fmt/format.h"  // fmt::formatter, fmt::ostream_formatter
 
-#include <atomic>   // std::atomic
 #include <chrono>   // std::chrono::milliseconds, std::chrono_literals namespace
 #include <cstddef>  // std::size_t
 #include <iosfwd>   // std::ostream forward declaration
+#include <mutex>    // std::mutex
 #include <string>   // std::string
 
 namespace hws {
@@ -146,10 +146,14 @@ class gpu_intel_hardware_sampler : public hardware_sampler {
     /// The temperature related Intel GPU samples.
     level_zero_temperature_samples temperature_samples_{};
 
-    /// The total number of currently active Intel GPU hardware samplers.
-    inline static std::atomic<int> instances_{ 0 };
-    /// True if the Level Zero environment has been successfully initialized (only done by a single hardware sampler).
-    inline static std::atomic<bool> init_finished_{ false };
+    /// Guards `initialized_` and the `zeInit()` call, so that the "has anyone already initialized?" check and the
+    /// actual `zeInit()` call happen as one atomic step - a busy-wait on a plain flag can't do this: a failing
+    /// first `zeInit()` never sets it, permanently stranding every waiter. There is no matching shut down call
+    /// (Level Zero has none), so unlike the AMD/NVIDIA backends this only ever needs to run `zeInit()` once
+    /// successfully, never again for the lifetime of the process.
+    inline static std::mutex lifecycle_mutex_{};
+    /// True once `zeInit()` has completed successfully; only ever read/written while holding `lifecycle_mutex_`.
+    inline static bool initialized_{ false };
 };
 
 /**

include/hws/gpu_nvidia/hardware_sampler.hpp

@@ -19,10 +19,10 @@
 
 #include "fmt/format.h"  // fmt::formatter, fmt::ostream_formatter
 
-#include <atomic>   // std::atomic
 #include <chrono>   // std::chrono::milliseconds, std::chrono_literals namespace
 #include <cstddef>  // std::size_t
 #include <iosfwd>   // std::ostream forward declaration
+#include <mutex>    // std::mutex
 #include <string>   // std::string
 
 namespace hws {
@@ -147,10 +147,14 @@ class gpu_nvidia_hardware_sampler : public hardware_sampler {
     /// The temperature related NVIDIA GPU samples.
     nvml_temperature_samples temperature_samples_{};
 
-    /// The total number of currently active NVIDIA GPU hardware samplers.
-    inline static std::atomic<int> instances_{ 0 };
-    /// True if the NVML environment has been successfully initialized (only done by a single hardware sampler).
-    inline static std::atomic<bool> init_finished_{ false };
+    /// Guards `instances_` and every `nvmlInit()`/`nvmlShutdown()` call, so that the "first instance initializes,
+    /// last instance shuts down" decision and the actual init/shutdown call happen as one atomic step - a
+    /// busy-wait on a plain flag can't do this: a failing first `nvmlInit()` never sets it, permanently stranding
+    /// every waiter, and nothing prevents a shutdown from racing a concurrent init.
+    inline static std::mutex lifecycle_mutex_{};
+    /// The total number of currently active NVIDIA GPU hardware samplers; only ever read/written while holding
+    /// `lifecycle_mutex_`.
+    inline static int instances_{ 0 };
 };
 
 /**

include/hws/gpu_nvidia/utility.hpp

@@ -69,6 +69,13 @@ namespace hws::detail {
 
 #endif
 
+/**
+ * @brief Return the PCI bus ID (e.g. `"0000:c1:00.0"`) of the NVIDIA GPU device with the given CUDA @p local_index.
+ * @param[in] local_index the local CUDA device index
+ * @return the PCI bus ID (`[[nodiscard]]`)
+ */
+[[nodiscard]] std::string nvidia_device_pci_bus_id(int local_index);
+
 #if defined(HWS_MPI_SUPPORT_ENABLED)
 
 /**

include/hws/utility.hpp

@@ -19,6 +19,7 @@
 #include <chrono>         // std::chrono::duration
 #include <cmath>          // std::trunc
 #include <cstddef>        // std::size_t
+#include <cstdint>        // std::uint32_t
 #include <optional>       // std::optional
 #include <stdexcept>      // std::runtime_error
 #include <string>         // std::string, std::stof, std::stod, std::stold
@@ -255,6 +256,22 @@ template <typename T>
  */
 [[nodiscard]] std::string indent_lines(const std::string &text, std::string_view prefix);
 
+/**
+ * @brief Format a PCI domain/bus/device triplet as the canonical Linux sysfs PCI bus ID string
+ *        `"<domain>:<bus>:<device>.0"` (e.g. `"0000:c1:00.0"`), lowercase hex, 4/2/2 digits.
+ * @details The function is fixed to `0`: on multi-die/multi-partition accelerators (e.g. AMD MI300-series
+ *          "partitions") the PCI function field is repurposed by the vendor's management library for
+ *          partition/die identification, but the function seen by the OS/sysfs for the *device* itself is
+ *          always `0` - so vendor-provided domain/bus/device values should be combined with a hardcoded `0`
+ *          function here rather than a vendor-reported function value, to stay comparable with sysfs PCI bus IDs
+ *          (see e.g. `hws::detail::enumerate_all_amd_gpu_pci_bus_ids()`/`enumerate_all_nvidia_gpu_pci_bus_ids()`).
+ * @param[in] domain the PCI domain
+ * @param[in] bus the PCI bus number
+ * @param[in] device the PCI device (slot) number
+ * @return the formatted PCI bus ID string (`[[nodiscard]]`)
+ */
+[[nodiscard]] std::string format_pci_bus_id(std::uint32_t domain, std::uint32_t bus, std::uint32_t device);
+
 /*****************************************************************************************************/
 /**                                      other free functions                                       **/
 /*****************************************************************************************************/

src/hws/gpu_amd/hardware_sampler.cpp

@@ -35,6 +35,62 @@
 
 namespace hws {
 
+namespace {
+
+/**
+ * @brief Convert a ROCm SMI BDFID (as returned by `rsmi_dev_pci_id_get()`) to a sysfs-style PCI bus ID string.
+ * @details BDFID = (DOMAIN << 32) | (PARTITION << 28) | (BUS << 8) | (DEVICE << 3) | FUNCTION (see ROCm SMI's
+ *          `rsmi_dev_pci_id_get` documentation). On MI-series partitioned devices the function bits are
+ *          repurposed for the partition ID instead of a real PCI function - but the OS/sysfs-visible PCI address
+ *          for the device itself always has function 0, so the function is intentionally not extracted here; see
+ *          `hws::detail::format_pci_bus_id()`.
+ */
+[[nodiscard]] std::string bdfid_to_pci_bus_id(const std::uint64_t bdfid) {
+    const auto domain = static_cast<std::uint32_t>((bdfid >> 32) & 0xffffffffull);
+    const auto bus = static_cast<std::uint32_t>((bdfid >> 8) & 0xffull);
+    const auto device = static_cast<std::uint32_t>((bdfid >> 3) & 0x1full);
+    return detail::format_pci_bus_id(domain, bus, device);
+}
+
+/**
+ * @brief Resolve the ROCm SMI device index that corresponds to the physical device HIP considers index
+ *        @p hip_device_id, by matching PCI bus IDs.
+ * @details Necessary because ROCm SMI enumerates every physical AMD GPU on the node unconditionally, while HIP's
+ *          enumeration is filtered/reordered by `HIP_VISIBLE_DEVICES`/`ROCR_VISIBLE_DEVICES` - the same index
+ *          number in both APIs can refer to different physical devices. Requires `rsmi_init()` to have already
+ *          been called.
+ * @throws std::runtime_error if ROCm SMI's device count can't be queried, if none of its devices' PCI bus IDs
+ *         match @p hip_device_id's, or if more than one does - `bdfid_to_pci_bus_id()` deliberately drops the
+ *         BDFID's partition bits (see its docs), so on a partitioned MI-series accelerator several ROCm SMI
+ *         entries can share one normalized bus ID; silently returning the first match there would be exactly the
+ *         HIP-index-used-as-RSMI-index bug this function exists to avoid, just triggered by partition mode
+ *         instead of a visibility mask.
+ */
+[[nodiscard]] std::uint32_t resolve_rsmi_device_id(const std::uint32_t hip_device_id) {
+    std::uint32_t rsmi_count{};
+    if (rsmi_num_monitor_devices(&rsmi_count) != RSMI_STATUS_SUCCESS) {
+        throw std::runtime_error{ "gpu_amd_hardware_sampler: couldn't query the number of ROCm SMI devices while resolving the physical device for HIP index " + std::to_string(hip_device_id) + "!" };
+    }
+
+    const std::string hip_bus_id = detail::amd_device_pci_bus_id(static_cast<int>(hip_device_id));
+    std::optional<std::uint32_t> resolved{};
+    for (std::uint32_t rsmi_idx = 0; rsmi_idx < rsmi_count; ++rsmi_idx) {
+        std::uint64_t bdfid{};
+        if (rsmi_dev_pci_id_get(rsmi_idx, &bdfid) == RSMI_STATUS_SUCCESS && bdfid_to_pci_bus_id(bdfid) == hip_bus_id) {
+            if (resolved.has_value()) {
+                throw std::runtime_error{ "gpu_amd_hardware_sampler: found more than one ROCm SMI device with PCI bus ID " + hip_bus_id + " (HIP index " + std::to_string(hip_device_id) + ") - likely a partitioned accelerator, which isn't supported yet!" };
+            }
+            resolved = rsmi_idx;
+        }
+    }
+    if (!resolved.has_value()) {
+        throw std::runtime_error{ "gpu_amd_hardware_sampler: couldn't find a ROCm SMI device with PCI bus ID " + hip_bus_id + " (HIP index " + std::to_string(hip_device_id) + ")!" };
+    }
+    return resolved.value();
+}
+
+}  // namespace
+
 gpu_amd_hardware_sampler::gpu_amd_hardware_sampler(const sample_category category) :
     gpu_amd_hardware_sampler{ 0, HWS_SAMPLING_INTERVAL, category } { }
 
@@ -46,15 +102,29 @@ gpu_amd_hardware_sampler::gpu_amd_hardware_sampler(const std::chrono::millisecon
 
 gpu_amd_hardware_sampler::gpu_amd_hardware_sampler(const std::size_t device_id, const std::chrono::milliseconds sampling_interval, const sample_category category) :
     hardware_sampler{ sampling_interval, category },
-    device_id_{ static_cast<std::uint32_t>(device_id) } {
-    // make sure that rsmi_init is only called once for all instances
-    if (instances_++ == 0) {
-        HWS_ROCM_SMI_ERROR_CHECK(rsmi_init(std::uint64_t{ 0 }))
-        // notify that initialization has been finished
-        init_finished_ = true;
-    } else {
-        // wait until init has been finished!
-        while (!init_finished_) { }
+    hip_device_id_{ static_cast<std::uint32_t>(device_id) } {
+    // make sure that rsmi_init is only called once for all instances; holding lifecycle_mutex_ for the whole
+    // "am I first?" decision plus the rsmi_init() call itself serializes it against every other constructor and
+    // destructor, so a failing rsmi_init() can never strand a waiter the way a busy-wait on a flag could - the
+    // next constructor to acquire the mutex simply sees instances_ still 0 and retries rsmi_init() itself
+    {
+        const std::lock_guard<std::mutex> lock{ lifecycle_mutex_ };
+        if (instances_ == 0) {
+            HWS_ROCM_SMI_ERROR_CHECK(rsmi_init(std::uint64_t{ 0 }))
+        }
+        ++instances_;
+    }
+
+    // resolve device_id_ only after rsmi_init() has definitely run (by this instance or a previous one, guaranteed
+    // since we're now a counted instance); if resolution throws, roll the count back under the same mutex
+    try {
+        device_id_ = resolve_rsmi_device_id(hip_device_id_);
+    } catch (...) {
+        const std::lock_guard<std::mutex> lock{ lifecycle_mutex_ };
+        if (--instances_ == 0) {
+            rsmi_shut_down();
+        }
+        throw;
     }
 }
 
@@ -65,12 +135,11 @@ gpu_amd_hardware_sampler::~gpu_amd_hardware_sampler() {
             this->stop_sampling();
         }
 
-        // the last instance must shut down the ROCm SMI runtime
-        // make sure that rsmi_shut_down is only called once
+        // the last instance must shut down the ROCm SMI runtime; guarded by the same mutex as the constructor so
+        // this can't race a concurrent constructor's "am I first?" check
+        const std::lock_guard<std::mutex> lock{ lifecycle_mutex_ };
         if (--instances_ == 0) {
             HWS_ROCM_SMI_ERROR_CHECK(rsmi_shut_down())
-            // reset init_finished flag
-            init_finished_ = false;
         }
     } catch (const std::exception &e) {
         std::cerr << e.what() << std::endl;
@@ -94,7 +163,7 @@ void gpu_amd_hardware_sampler::sampling_loop() {
         general_samples_.byte_order_ = "Little Endian";
 
         hipDeviceProp_t prop{};
-        if (hipGetDeviceProperties(&prop, static_cast<int>(device_id_)) == hipSuccess) {
+        if (hipGetDeviceProperties(&prop, static_cast<int>(hip_device_id_)) == hipSuccess) {
             const std::string architecture{ prop.gcnArchName };
             general_samples_.architecture_ = architecture.substr(0, architecture.find_first_of('\0'));
         }

src/hws/gpu_amd/utility.cpp

@@ -9,13 +9,14 @@
 
 #include "rocm_smi/rocm_smi.h"  // ROCm SMI runtime functions
 
+#include "hip/hip_runtime_api.h"  // hipGetDeviceCount, hipDeviceGetPCIBusId
+
 #include <string>  // std::string
-#include <vector>  // std::vector
 
 #if defined(HWS_MPI_SUPPORT_ENABLED) && defined(HWS_FOR_AMD_GPUS_ENABLED)
     #include "hws/visible_gpu_device.hpp"  // hws::detail::visible_gpu_device, hws::detail::device_backend_kind
 
-    #include "hip/hip_runtime_api.h"  // hipGetDeviceCount, hipDeviceGetPCIBusId
+    #include <vector>  // std::vector
 #endif
 
 namespace hws::detail {
@@ -46,6 +47,12 @@ std::string performance_level_to_string(const rsmi_dev_perf_level_t perf_level)
     }
 }
 
+std::string amd_device_pci_bus_id(const int local_index) {
+    char bus_id[64] = {};
+    HWS_HIP_ERROR_CHECK(hipDeviceGetPCIBusId(bus_id, sizeof(bus_id), local_index));
+    return std::string{ bus_id };
+}
+
 #if defined(HWS_MPI_SUPPORT_ENABLED) && defined(HWS_FOR_AMD_GPUS_ENABLED)
 
 namespace {
@@ -58,9 +65,7 @@ namespace {
  * @return the physical ID of the AMD GPU device
  */
 [[nodiscard]] std::string amd_physical_id(const int local_index) {
-    char bus_id[64] = {};
-    HWS_HIP_ERROR_CHECK(hipDeviceGetPCIBusId(bus_id, sizeof(bus_id), local_index));
-    return std::string{ "amd:" } + bus_id;
+    return std::string{ "amd:" } + amd_device_pci_bus_id(local_index);
 }
 
 }  // namespace

src/hws/gpu_intel/hardware_sampler.cpp

@@ -44,17 +44,19 @@ gpu_intel_hardware_sampler::gpu_intel_hardware_sampler(const std::chrono::millis
 
 gpu_intel_hardware_sampler::gpu_intel_hardware_sampler(const std::size_t device_id, const std::chrono::milliseconds sampling_interval, const sample_category category) :
     hardware_sampler{ sampling_interval, category } {
-    // make sure that zeInit is only called once for all instances
-    if (instances_++ == 0) {
-        HWS_LEVEL_ZERO_ERROR_CHECK(zeInit(ZE_INIT_FLAG_GPU_ONLY))
-        // notify that initialization has been finished
-        init_finished_ = true;
-    } else {
-        // wait until init has been finished!
-        while (!init_finished_) { }
+    // make sure that zeInit is only called once for all instances; holding lifecycle_mutex_ for the whole
+    // "already initialized?" check plus the zeInit() call itself serializes it against every other constructor,
+    // so a failing zeInit() can never strand a waiter the way a busy-wait on a flag could - the next constructor
+    // to acquire the mutex simply sees initialized_ still false and retries zeInit() itself
+    {
+        const std::lock_guard<std::mutex> lock{ lifecycle_mutex_ };
+        if (!initialized_) {
+            HWS_LEVEL_ZERO_ERROR_CHECK(zeInit(ZE_INIT_FLAG_GPU_ONLY))
+            initialized_ = true;
+        }
     }
 
-    // initialize samples -> can't be done beforehand since the device handle can only be initialized after a call to nvmlInit
+    // initialize samples -> can't be done beforehand since the device handle can only be initialized after a call to zeInit
     device_ = detail::level_zero_device_handle{ device_id };
 }