UCS/SYS/TOPO: Allow multiple sibling devices per memory device

src/ucs/sys/topo/base/topo.c

@@ -510,18 +510,18 @@ ucs_topo_find_device_by_bus_id_value(const ucs_sys_bus_id_t *bus_id,
         *sys_dev_p = ucs_topo_global_ctx.num_devices++;
         device     = &ucs_topo_global_ctx.devices[*sys_dev_p];
 
-        device->bus_id          = *bus_id;
-        device->name            = name;
-        device->name_priority   = 0;
-        device->numa_node       = numa_node;
-        device->pci_id          = pci_id;
-        device->user_value      = user_value;
-        device->device_class    = UCS_TOPO_DEVICE_CLASS_UNKNOWN;
-        device->flags           = 0;
-        device->class_ordinal   = UCS_SYS_DEVICE_ORDINAL_INVALID;
-        device->sibling_role    = UCS_TOPO_SIBLING_ROLE_NONE;
-        device->sibling_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
-        device->sys_dev_aux     = UCS_SYS_DEVICE_ID_UNKNOWN;
+        device->bus_id        = *bus_id;
+        device->name          = name;
+        device->name_priority = 0;
+        device->numa_node     = numa_node;
+        device->pci_id        = pci_id;
+        device->user_value    = user_value;
+        device->device_class  = UCS_TOPO_DEVICE_CLASS_UNKNOWN;
+        device->flags         = 0;
+        device->class_ordinal = UCS_SYS_DEVICE_ORDINAL_INVALID;
+        device->sibling_role  = UCS_TOPO_SIBLING_ROLE_NONE;
+        device->sys_dev_aux   = UCS_SYS_DEVICE_ID_UNKNOWN;
+        UCS_STATIC_BITMAP_RESET_ALL(&device->siblings);
 
         if (user_value == UCS_SYS_DEVICE_USER_VALUE_EMPTY) {
             ucs_debug(
@@ -803,22 +803,23 @@ ucs_topo_is_reachable(ucs_sys_device_t sys_dev, ucs_sys_device_t sys_dev_mem)
     result =
             (sys_dev >= ucs_topo_global_ctx.num_devices) ||
             (sys_dev_mem >= ucs_topo_global_ctx.num_devices) ||
-            /*
-             * Memory device was never matched with a sibling, it does not
-             * mandate an auxiliary path.
-             */
-            (ucs_topo_global_ctx.devices[sys_dev_mem].sibling_sys_dev ==
-             UCS_SYS_DEVICE_ID_UNKNOWN) ||
             /* The device itself never uses auxiliary path */
             (ucs_topo_global_ctx.devices[sys_dev].sibling_role !=
              UCS_TOPO_SIBLING_ROLE_DEV) ||
             /*
-             * MPS MLOParts create multiple MEM aliases for one DEV. Each MEM
-             * alias records its matched DEV; the DEV stores only one
-             * representative MEM sibling.
+             * MPS MLOParts create multiple MEM aliases for one DEV, and
+             * multiplane NICs create multiple DEVs for one MEM. Each MEM
+             * records all of its matched DEVs.
+             */
+            UCS_STATIC_BITMAP_GET(
+                    ucs_topo_global_ctx.devices[sys_dev_mem].siblings,
+                    sys_dev) ||
+            /*
+             * Memory device was never matched with a sibling, it does not
+             * mandate an auxiliary path.
              */
-            (ucs_topo_global_ctx.devices[sys_dev_mem].sibling_sys_dev ==
-             sys_dev);
+            UCS_STATIC_BITMAP_IS_ZERO(
+                    ucs_topo_global_ctx.devices[sys_dev_mem].siblings);
     ucs_spin_unlock(&ucs_topo_global_ctx.lock);
 
     return result;
@@ -836,11 +837,12 @@ int ucs_topo_is_sibling(ucs_sys_device_t sys_dev, ucs_sys_device_t sys_dev_mem)
                  (sys_dev_mem != UCS_SYS_DEVICE_ID_UNKNOWN) &&
                  (sys_dev_mem < ucs_topo_global_ctx.num_devices) &&
                  /*
-                  * MLOPart allows multiple MEM devices to point at one DEV, so
-                  * the MEM-side sibling is the canonical relationship.
+                  * The MEM side holds every matched DEV, including all planes
+                  * of a multiplane NIC.
                   */
-                 (ucs_topo_global_ctx.devices[sys_dev_mem].sibling_sys_dev ==
-                  sys_dev);
+                 UCS_STATIC_BITMAP_GET(
+                         ucs_topo_global_ctx.devices[sys_dev_mem].siblings,
+                         sys_dev);
 
     if (is_sibling) {
         role_dev     = ucs_topo_global_ctx.devices[sys_dev].sibling_role;
@@ -1372,28 +1374,40 @@ ucs_topo_is_pci_bridge(ucs_sys_device_t sys_dev1, ucs_sys_device_t sys_dev2)
     return result;
 }
 
-static int ucs_topo_sys_device_sibling_match(ucs_sys_device_t sys_dev,
-                                             ucs_sys_device_t sys_dev_mem)
+static void ucs_topo_sys_device_sibling_match(ucs_sys_device_t sys_dev,
+                                              ucs_sys_device_t sys_dev_mem)
 {
     ucs_sys_device_t sys_dev_aux;
 
     if ((ucs_topo_global_ctx.devices[sys_dev_mem].sibling_role !=
          UCS_TOPO_SIBLING_ROLE_MEM) ||
         (ucs_topo_global_ctx.devices[sys_dev].sibling_role !=
          UCS_TOPO_SIBLING_ROLE_DEV)) {
-        return 0;
+        return;
     }
 
     sys_dev_aux = ucs_topo_global_ctx.devices[sys_dev].sys_dev_aux;
     if (ucs_topo_is_pci_bridge(sys_dev_mem, sys_dev_aux)) {
-        ucs_topo_global_ctx.devices[sys_dev_mem].sibling_sys_dev = sys_dev;
-        ucs_topo_global_ctx.devices[sys_dev].sibling_sys_dev     = sys_dev_mem;
+        UCS_STATIC_BITMAP_SET(
+                &ucs_topo_global_ctx.devices[sys_dev_mem].siblings, sys_dev);
+        UCS_STATIC_BITMAP_SET(&ucs_topo_global_ctx.devices[sys_dev].siblings,
+                              sys_dev_mem);
         ucs_trace("sys_dev=%u with sys_dev_aux=%u matched with sys_dev_mem=%u",
                   sys_dev, sys_dev_aux, sys_dev_mem);
-        return 1;
+    }
+}
+
+static void ucs_topo_sys_device_clear_siblings(ucs_sys_device_t sys_dev)
+{
+    ucs_topo_sys_device_info_t *device = &ucs_topo_global_ctx.devices[sys_dev];
+    unsigned peer;
+
+    UCS_STATIC_BITMAP_FOR_EACH_BIT(peer, &device->siblings) {
+        UCS_STATIC_BITMAP_RESET(&ucs_topo_global_ctx.devices[peer].siblings,
+                                sys_dev);
     }
 
-    return 0;
+    UCS_STATIC_BITMAP_RESET_ALL(&device->siblings);
 }
 
 /* Memory device supports auxiliary path if a sibling is found */
@@ -1410,13 +1424,19 @@ ucs_status_t ucs_topo_sys_device_enable_aux_path(ucs_sys_device_t sys_dev)
         goto out;
     }
 
+    if (ucs_topo_global_ctx.devices[sys_dev].sibling_role ==
+        UCS_TOPO_SIBLING_ROLE_MEM) {
+        status = UCS_OK;
+        goto out;
+    }
+
+    ucs_topo_sys_device_clear_siblings(sys_dev);
     ucs_topo_global_ctx.devices[sys_dev].sibling_role =
             UCS_TOPO_SIBLING_ROLE_MEM;
 
+    /* All devices sharing the auxiliary device match the memory device */
     for (dev = 0; dev < ucs_topo_global_ctx.num_devices; dev++) {
-        if (ucs_topo_sys_device_sibling_match(dev, sys_dev)) {
-            break;
-        }
+        ucs_topo_sys_device_sibling_match(dev, sys_dev);
     }
 
     status = UCS_OK;
@@ -1440,6 +1460,14 @@ ucs_status_t ucs_topo_sys_device_set_sys_dev_aux(ucs_sys_device_t sys_dev,
         goto out;
     }
 
+    if ((ucs_topo_global_ctx.devices[sys_dev].sibling_role ==
+         UCS_TOPO_SIBLING_ROLE_DEV) &&
+        (ucs_topo_global_ctx.devices[sys_dev].sys_dev_aux == sys_dev_aux)) {
+        status = UCS_OK;
+        goto out;
+    }
+
+    ucs_topo_sys_device_clear_siblings(sys_dev);
     ucs_topo_global_ctx.devices[sys_dev].sys_dev_aux = sys_dev_aux;
     ucs_topo_global_ctx.devices[sys_dev].sibling_role =
             UCS_TOPO_SIBLING_ROLE_DEV;
@@ -1462,8 +1490,8 @@ int ucs_topo_device_has_sibling(ucs_sys_device_t sys_dev)
 
     ucs_spin_lock(&ucs_topo_global_ctx.lock);
     result = (sys_dev < ucs_topo_global_ctx.num_devices) &&
-             (ucs_topo_global_ctx.devices[sys_dev].sibling_sys_dev !=
-              UCS_SYS_DEVICE_ID_UNKNOWN);
+             !UCS_STATIC_BITMAP_IS_ZERO(
+                     ucs_topo_global_ctx.devices[sys_dev].siblings);
     ucs_spin_unlock(&ucs_topo_global_ctx.lock);
 
     return result;

src/ucs/sys/topo/base/topo_groups.h

@@ -50,13 +50,6 @@ UCS_ARRAY_DECLARE_TYPE(ucs_topo_group_array_t, size_t, ucs_topo_group_t);
 typedef ucs_topo_group_array_t ucs_topo_groups_t;
 
 
-/**
- * @ingroup UCS_RESOURCE
- * Bitmap of system devices.
- */
-typedef ucs_static_bitmap_s(UCS_SYS_DEVICE_ID_COUNT) ucs_sys_device_bitmap_t;
-
-
 /**
  * Build system topology groups.
  *

src/ucs/sys/topo/base/topo_int.h

@@ -9,6 +9,8 @@
 
 #include "topo.h"
 
+#include <ucs/datastruct/static_bitmap.h>
+
 BEGIN_C_DECLS
 
 /* Possible role of a current device wrt its sibling */
@@ -24,6 +26,13 @@ typedef enum {
 } ucs_topo_sibling_role_t;
 
 
+/**
+ * @ingroup UCS_RESOURCE
+ * Bitmap of system devices.
+ */
+typedef ucs_static_bitmap_s(UCS_SYS_DEVICE_ID_COUNT) ucs_sys_device_bitmap_t;
+
+
 typedef struct {
     ucs_sys_bus_id_t        bus_id;
     char                    *name;
@@ -43,8 +52,10 @@ typedef struct {
     ucs_sys_device_t        sys_dev_aux;
 
     ucs_topo_sibling_role_t sibling_role; /* Role of the current device */
-    /* MEM role: matched DEV. DEV role: one representative matched MEM. */
-    ucs_sys_device_t        sibling_sys_dev;
+    /* Matched siblings: all DEVs for the MEM role, all MEMs for the DEV role.
+     * Several DEVs can share one auxiliary device (e.g. planes of a multiplane
+     * NIC), and several MEMs can sit behind it (e.g. MPS MLOPart aliases). */
+    ucs_sys_device_bitmap_t siblings;
 } ucs_topo_sys_device_info_t;
 
 

test/gtest/ucs/test_topo.cc

@@ -1001,3 +1001,163 @@ UCS_TEST_F(test_topo, sibling_multiple_memory_devices) {
     expect_sibling(gpu_dev1);
     expect_sibling(gpu_dev2);
 }
+
+// Multiple network devices sharing a single auxiliary device, with multiple
+// memory devices behind it.
+class test_topo_multi_dev_sibling : public test_topo {
+protected:
+    static constexpr int num_planes = 3;
+    static constexpr int num_mems   = 2;
+
+    // Test both registration orders: memory devices first (matched by
+    // set_sys_dev_aux), or network devices first (matched by enable_aux_path)
+    void run_test(bool mem_first)
+    {
+        read_pcie_devices();
+        if ((m_hcas.size() < (num_planes + 1)) || m_gpus.empty() ||
+            m_dmas.empty()) {
+            UCS_TEST_SKIP_R("Not enough HCA, GPU and DMA PCIe device");
+        }
+
+        std::string sibling_gpu, sibling_dma;
+        get_siblings(m_hcas[0], sibling_gpu, sibling_dma);
+        if (sibling_dma.empty()) {
+            UCS_TEST_SKIP_R("No sibling DMA device found");
+        }
+
+        // Memory devices: the same GPU BDF with distinct user values
+        std::vector<ucs_sys_device_t> mems;
+        auto gpu_dev = register_device("gpu0", sibling_gpu);
+        ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, gpu_dev);
+        mems.push_back(gpu_dev);
+
+        ucs_sys_bus_id_t gpu_bus_id;
+        ASSERT_UCS_OK(ucs_topo_get_device_bus_id(gpu_dev, &gpu_bus_id));
+        for (int i = 1; i < num_mems; ++i) {
+            ucs_sys_device_t mem_alias;
+            ASSERT_UCS_OK(ucs_topo_find_device_by_bus_id_and_user_value(
+                    &gpu_bus_id, i, &mem_alias));
+            ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, mem_alias);
+            mems.push_back(mem_alias);
+        }
+
+        auto dma_dev = register_device("dma", sibling_dma);
+        ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, dma_dev);
+
+        if (mem_first) {
+            for (auto mem : mems) {
+                ASSERT_UCS_OK(ucs_topo_sys_device_enable_aux_path(mem));
+            }
+        }
+
+        // All planes report the same auxiliary device
+        std::vector<ucs_sys_device_t> planes;
+        for (int i = 0; i < num_planes; ++i) {
+            auto plane = register_device("plane" + std::to_string(i),
+                                         m_hcas[i]);
+            ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, plane);
+            ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(plane, dma_dev));
+            planes.push_back(plane);
+        }
+
+        // A device whose auxiliary device is not behind the GPU's bridge
+        auto other = register_device("other", m_hcas[num_planes]);
+        ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, other);
+        ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(other, other));
+
+        if (!mem_first) {
+            for (auto mem : mems) {
+                ASSERT_UCS_OK(ucs_topo_sys_device_enable_aux_path(mem));
+            }
+        }
+
+        // Every memory device is matched with every plane, and vice versa
+        for (auto mem : mems) {
+            EXPECT_TRUE(ucs_topo_device_has_sibling(mem));
+            for (auto plane : planes) {
+                EXPECT_TRUE(ucs_topo_is_sibling(plane, mem));
+                EXPECT_TRUE(ucs_topo_is_reachable(plane, mem));
+            }
+        }
+
+        for (auto plane : planes) {
+            EXPECT_TRUE(ucs_topo_device_has_sibling(plane));
+        }
+
+        // Skip if the other HCA happens to be behind the GPU's bridge as well
+        if (ucs_topo_is_sibling(other, mems[0])) {
+            UCS_TEST_SKIP_R("Other HCA is behind the GPU bridge");
+        }
+
+        EXPECT_FALSE(ucs_topo_device_has_sibling(other));
+        for (auto mem : mems) {
+            EXPECT_FALSE(ucs_topo_is_sibling(other, mem));
+            EXPECT_FALSE(ucs_topo_is_reachable(other, mem));
+        }
+    }
+};
+
+UCS_TEST_F(test_topo_multi_dev_sibling, mem_first)
+{
+    run_test(true);
+}
+
+UCS_TEST_F(test_topo_multi_dev_sibling, devs_first)
+{
+    run_test(false);
+}
+
+// A pairing must not outlive the role or the auxiliary device it was made for
+UCS_TEST_F(test_topo_multi_dev_sibling, stale_siblings)
+{
+    read_pcie_devices();
+    if ((m_hcas.size() < 2) || m_gpus.empty() || m_dmas.empty()) {
+        UCS_TEST_SKIP_R("Not enough HCA, GPU and DMA PCIe device");
+    }
+
+    std::string sibling_gpu, sibling_dma;
+    get_siblings(m_hcas[0], sibling_gpu, sibling_dma);
+    if (sibling_dma.empty()) {
+        UCS_TEST_SKIP_R("No sibling DMA device found");
+    }
+
+    auto gpu_dev = register_device("gpu0", sibling_gpu);
+    ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, gpu_dev);
+    auto dma_dev = register_device("dma", sibling_dma);
+    ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, dma_dev);
+    auto plane = register_device("plane", m_hcas[0]);
+    ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, plane);
+
+    // An auxiliary device which is not behind the GPU's bridge. This depends
+    // on the PCIe topology of the machine.
+    auto other = register_device("other", m_hcas[1]);
+    ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, other);
+    ASSERT_UCS_OK(ucs_topo_sys_device_enable_aux_path(gpu_dev));
+    ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(other, other));
+    if (ucs_topo_is_sibling(other, gpu_dev)) {
+        UCS_TEST_SKIP_R("Other HCA is behind the GPU bridge");
+    }
+
+    ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(plane, dma_dev));
+    ASSERT_TRUE(ucs_topo_is_sibling(plane, gpu_dev));
+
+    // Setting the same auxiliary device again keeps the pairing
+    ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(plane, dma_dev));
+    EXPECT_TRUE(ucs_topo_is_sibling(plane, gpu_dev));
+
+    // Role flip DEV -> MEM: the memory device forgets the former DEV
+    ASSERT_UCS_OK(ucs_topo_sys_device_enable_aux_path(plane));
+    EXPECT_FALSE(ucs_topo_is_sibling(plane, gpu_dev));
+    EXPECT_FALSE(ucs_topo_device_has_sibling(gpu_dev));
+
+    // Role flip MEM -> DEV pairs it again
+    ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(plane, dma_dev));
+    EXPECT_TRUE(ucs_topo_is_sibling(plane, gpu_dev));
+    EXPECT_TRUE(ucs_topo_device_has_sibling(gpu_dev));
+
+    // Moving the device to another auxiliary device drops the old pairing
+    ASSERT_UCS_OK(ucs_topo_sys_device_set_sys_dev_aux(plane, other));
+    EXPECT_FALSE(ucs_topo_is_sibling(plane, gpu_dev));
+    EXPECT_FALSE(ucs_topo_device_has_sibling(plane));
+    EXPECT_FALSE(ucs_topo_device_has_sibling(gpu_dev));
+}