From 3f1fa0796079e8eeabf4de3b2cd238b6776d03b7 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Tue, 22 Sep 2026 11:10:36 +0000 Subject: [PATCH 1/8] UCP/CORE: Add shared topology assignment policy Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 18 ++ src/ucp/core/ucp_context.h | 2 + src/ucp/core/ucp_gpu_nic_assignment.c | 65 +++++- src/ucp/core/ucp_gpu_nic_assignment.h | 18 ++ src/ucp/core/ucp_types.h | 14 +- test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 191 ++++++++++++++---- 6 files changed, 258 insertions(+), 50 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index c0a62ceb920..0a431fd4f1b 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -108,6 +108,14 @@ static const char *ucp_fence_modes[] = { [UCP_FENCE_MODE_LAST] = NULL }; +static const char *ucp_gpu_net_device_modes[] = { + [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED] = "not_shared", + [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP] = "not_shared_flip", + [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN] = "not_shared_round_robin", + [UCP_GPU_NET_DEVICE_MODE_SHARED] = "shared", + [UCP_GPU_NET_DEVICE_MODE_LAST] = NULL +}; + static const char *ucp_rndv_modes[] = { [UCP_RNDV_MODE_AUTO] = "auto", [UCP_RNDV_MODE_GET_ZCOPY] = "get_zcopy", @@ -615,6 +623,16 @@ static ucs_config_field_t ucp_context_config_table[] = { ucs_offsetof(ucp_context_config_t, connect_all_to_all), UCS_CONFIG_TYPE_BOOL}, + {"GPU_NET_DEVICE_MODE", "not_shared", + "Select how network devices are assigned to GPUs.\n" + " - not_shared : use the default FLIP assignment.\n" + " - not_shared_flip : use the FLIP assignment.\n" + " - not_shared_round_robin : use the round-robin assignment.\n" + " - shared : assign every network device in a topology\n" + " group to every GPU in that group.", + ucs_offsetof(ucp_context_config_t, gpu_net_device_mode), + UCS_CONFIG_TYPE_ENUM(ucp_gpu_net_device_modes)}, + {"SINGLE_NET_DEVICE", "n", "Restrict each protocol's lanes to one network device.\n" "The device is picked among the closest eligible network devices using \n" diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 2ffbbb67739..edae667a466 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -255,6 +255,8 @@ typedef struct ucp_context_config { /** Extend endpoint lanes connections of each local device to all remote * devices */ int connect_all_to_all; + /** GPU network device assignment mode */ + ucp_gpu_net_device_mode_t gpu_net_device_mode; /** Restrict lanes to one network device per protocol */ int proto_use_single_net_device; /** Max HCAs for GPU memory registration: auto=closest, N=limit, inf=all */ diff --git a/src/ucp/core/ucp_gpu_nic_assignment.c b/src/ucp/core/ucp_gpu_nic_assignment.c index 728f4cf843f..72c4ed360b1 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.c +++ b/src/ucp/core/ucp_gpu_nic_assignment.c @@ -33,6 +33,22 @@ typedef struct { } ucp_nics_string_buffers_t; +ucp_gpu_nic_assignment_policy_t +ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode) +{ + switch (mode) { + case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED: + case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP: + return UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; + case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN: + return UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN; + case UCP_GPU_NET_DEVICE_MODE_SHARED: + return UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED; + default: + ucs_fatal("invalid gpu network device mode %d", (int)mode); + } +} + static size_t ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_policy_t policy, size_t num_gpus, size_t nic_idx) @@ -327,16 +343,39 @@ ucp_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t *assignment, } static void -ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, - const ucs_topo_group_t *group, - ucp_gpu_nic_assignment_policy_t policy) +ucp_gpu_nic_assignment_add_group_shared(ucp_gpu_nic_assignment_t *assignment, + const ucs_topo_group_t *group) +{ + ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap; + const ucs_topo_group_element_t *gpu, *nic; + + ucs_array_for_each(gpu, &group->gpus) { + ucs_assert((gpu->num_sys_devs > 0) && + (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN)); + + nic_sys_dev_bitmap = ucs_const_cast( + ucp_gpu_nic_sys_dev_bitmap_t*, + ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0])); + ucs_assert(nic_sys_dev_bitmap != NULL); + + /* Assign all NICs in the group to the GPU. */ + ucs_array_for_each(nic, &group->nics) { + ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic); + } + } +} + +static void ucp_gpu_nic_assignment_add_group_not_shared( + ucp_gpu_nic_assignment_t *assignment, const ucs_topo_group_t *group, + ucp_gpu_nic_assignment_policy_t policy) { size_t num_gpus = ucs_array_length(&group->gpus); ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap; - const ucs_topo_group_element_t *nic; - const ucs_topo_group_element_t *gpu; - size_t nic_idx; - size_t gpu_idx; + const ucs_topo_group_element_t *gpu, *nic; + size_t nic_idx, gpu_idx; + + ucs_assertv(policy != UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, + "invalid policy: %d", (int)policy); ucs_array_for_each_index(nic, nic_idx, &group->nics) { gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, nic_idx); @@ -354,6 +393,18 @@ ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, } } +static void +ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, + const ucs_topo_group_t *group, + ucp_gpu_nic_assignment_policy_t policy) +{ + if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { + ucp_gpu_nic_assignment_add_group_shared(assignment, group); + } else { + ucp_gpu_nic_assignment_add_group_not_shared(assignment, group, policy); + } +} + ucs_status_t ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups, ucp_gpu_nic_assignment_policy_t policy, diff --git a/src/ucp/core/ucp_gpu_nic_assignment.h b/src/ucp/core/ucp_gpu_nic_assignment.h index 3965ae74bbf..4246d303901 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.h +++ b/src/ucp/core/ucp_gpu_nic_assignment.h @@ -7,6 +7,8 @@ #ifndef UCP_GPU_NIC_ASSIGNMENT_H_ #define UCP_GPU_NIC_ASSIGNMENT_H_ +#include "ucp_types.h" + #include #include @@ -52,10 +54,26 @@ typedef enum { */ UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + /** + * Assign every NIC, with all of its ports, to every GPU in the group. + */ + UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST } ucp_gpu_nic_assignment_policy_t; +/** + * Get the assignment policy of a GPU network device mode. + * + * @param [in] mode GPU network device mode. + * + * @return Assignment policy corresponding to @a mode. + */ +ucp_gpu_nic_assignment_policy_t +ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode); + + /** * Build GPU-to-NIC assignments from topology groups. * diff --git a/src/ucp/core/ucp_types.h b/src/ucp/core/ucp_types.h index baf544c592d..00c122c52db 100644 --- a/src/ucp/core/ucp_types.h +++ b/src/ucp/core/ucp_types.h @@ -1,5 +1,5 @@ /** - * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2017. ALL RIGHTS RESERVED. + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED. * * See file LICENSE for terms. */ @@ -233,6 +233,18 @@ typedef enum { } ucp_fence_mode_t; +/** + * GPU network device assignment mode. + */ +typedef enum { + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED, + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP, + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN, + UCP_GPU_NET_DEVICE_MODE_SHARED, + UCP_GPU_NET_DEVICE_MODE_LAST +} ucp_gpu_net_device_mode_t; + + /** * Communication scheme in RNDV protocol. */ diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index f22e76d2ea9..d78bdca0eff 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -219,9 +219,10 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void check_assignment(const topology_shape_t &config, - ucp_gpu_nic_assignment_policy_t policy, - const std::vector &expected_owners) + void + check_assignment_not_shared_inner(const topology_shape_t &config, + ucp_gpu_nic_assignment_policy_t policy, + const std::vector &expected_owners) { ASSERT_NE(config.num_groups, 0); ASSERT_NE(config.num_nic_ports, 0); @@ -255,10 +256,11 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void check_clique_assignment(size_t num_groups, size_t num_gpus_per_group, - size_t num_nics_per_group, - ucp_gpu_nic_assignment_policy_t policy, - const std::vector &expected_owners) + void check_assignment_not_shared(size_t num_groups, + size_t num_gpus_per_group, + size_t num_nics_per_group, + ucp_gpu_nic_assignment_policy_t policy, + const std::vector &expected_owners) { topology_shape_t config; @@ -274,7 +276,75 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ++num_nic_ports) { config.num_gpu_devices = num_gpu_devices; config.num_nic_ports = num_nic_ports; - check_assignment(config, policy, expected_owners); + check_assignment_not_shared_inner(config, policy, + expected_owners); + } + } + } + + void check_assignment_shared_inner(const topology_shape_t &config) + { + ASSERT_NE(config.num_groups, 0); + ASSERT_NE(config.num_nic_ports, 0); + ASSERT_NE(config.num_gpu_devices, 0); + ASSERT_LE(config.num_gpu_devices, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); + ASSERT_LE(config.num_nic_ports, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); + ASSERT_LE(config.num_sys_devs(), UCS_SYS_DEVICE_ID_COUNT - 1); + + build_groups(config); + build_assignment(UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED); + + if (config.num_gpus() == 0) { + EXPECT_EQ(nullptr, m_assignment.nic_sys_dev_bitmaps); + EXPECT_EQ(0ul, m_assignment.num_bitmaps); + return; + } + + check_gpu_device_aliases(config); + + for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { + const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( + &m_assignment, gpu_sys_dev(config, gpu_idx, 0)); + const size_t gpu_group_idx = gpu_idx / config.num_gpus_per_group; + + ASSERT_NE(nullptr, nic_sys_dev_bitmap); + + for (size_t nic_idx = 0; nic_idx < config.num_nics(); ++nic_idx) { + const size_t nic_group_idx = nic_idx / + config.num_nics_per_group; + + for (size_t port_idx = 0; port_idx < config.num_nic_ports; + ++port_idx) { + EXPECT_EQ(gpu_group_idx == nic_group_idx, + ucp_gpu_nic_bitmap_get( + nic_sys_dev_bitmap, + nic_port_sys_dev(config, nic_idx, + port_idx))); + } + } + + EXPECT_EQ(config.num_nics_per_group * config.num_nic_ports, + static_cast( + UCS_STATIC_BITMAP_POPCOUNT(*nic_sys_dev_bitmap))); + } + } + + void check_assignment_shared(size_t num_groups, size_t num_gpus_per_group, + size_t num_nics_per_group) + { + topology_shape_t config; + + config.num_groups = num_groups; + config.num_gpus_per_group = num_gpus_per_group; + config.num_nics_per_group = num_nics_per_group; + + for (size_t num_gpu_devices = 1; num_gpu_devices <= 2; + ++num_gpu_devices) { + for (size_t num_nic_ports = 1; num_nic_ports <= 2; + ++num_nic_ports) { + config.num_gpu_devices = num_gpu_devices; + config.num_nic_ports = num_nic_ports; + check_assignment_shared_inner(config); } } } @@ -284,49 +354,86 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ucp_gpu_nic_assignment_t m_assignment; }; -UCS_TEST_F(test_ucp_gpu_nic_assignment, no_nics) { - check_clique_assignment(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, mode_to_policy) { + EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + ucp_gpu_net_device_mode_get_policy( + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED)); + EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + ucp_gpu_net_device_mode_get_policy( + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP)); + EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + ucp_gpu_net_device_mode_get_policy( + UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN)); + EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, + ucp_gpu_net_device_mode_get_policy( + UCP_GPU_NET_DEVICE_MODE_SHARED)); +} + +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_no_nics) { + check_assignment_not_shared(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, no_gpus) { - check_clique_assignment(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_no_gpus) { + check_assignment_not_shared(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, clique_flip_divisible) { +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_flip_divisible) { /* Number of NICs is divisible by the number of GPUs. */ - check_clique_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 1, 0, /**/ - 2, 3, 3, 2, /**/ - 4, 5, 5, 4}); - check_clique_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 2, 2, 1, 0, /**/ - 3, 4, 5, 5, 4, 3}); + check_assignment_not_shared(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 1, 0, /**/ + 2, 3, 3, 2, /**/ + 4, 5, 5, 4}); + check_assignment_not_shared(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 2, 2, 1, 0, /**/ + 3, 4, 5, 5, 4, 3}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, clique_flip_not_divisible) { - check_clique_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 1, 0, 0, /**/ - 2, 3, 3, 2, 2}); - check_clique_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, /**/ - 3, 4}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_flip_not_divisible) { + check_assignment_not_shared(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 1, 0, 0, /**/ + 2, 3, 3, 2, 2}); + check_assignment_not_shared(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, /**/ + 3, 4}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, clique_round_robin_divisible) { - check_clique_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 0, 1, /**/ - 2, 3, 2, 3, /**/ - 4, 5, 4, 5}); - check_clique_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 2, 0, 1, 2, /**/ - 3, 4, 5, 3, 4, 5}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_round_robin_divisible) { + check_assignment_not_shared(3, 2, 4, + UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 0, 1, /**/ + 2, 3, 2, 3, /**/ + 4, 5, 4, 5}); + check_assignment_not_shared(2, 3, 6, + UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 2, 0, 1, 2, /**/ + 3, 4, 5, 3, 4, 5}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, clique_round_robin_not_divisible) { - check_clique_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 0, 1, 0, /**/ - 2, 3, 2, 3, 2}); - check_clique_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, /**/ - 3, 4}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_round_robin_not_divisible) { + check_assignment_not_shared(2, 2, 5, + UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 0, 1, 0, /**/ + 2, 3, 2, 3, 2}); + check_assignment_not_shared(2, 3, 2, + UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, /**/ + 3, 4}); +} + +UCS_TEST_F(test_ucp_gpu_nic_assignment, shared_no_nics) { + check_assignment_shared(2, 3, 0); +} + +UCS_TEST_F(test_ucp_gpu_nic_assignment, shared_no_gpus) { + check_assignment_shared(2, 0, 3); +} + +UCS_TEST_F(test_ucp_gpu_nic_assignment, shared) { + /* 2 GPUs, 4 NICs. */ + check_assignment_shared(3, 2, 4); + + /* More GPUs than NICs. */ + check_assignment_shared(2, 3, 2); } From 8514fbe79b367ff0c4d3643570a362100fd21bad Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Tue, 22 Sep 2026 12:45:21 +0000 Subject: [PATCH 2/8] TEST/UCP: Add check_assignment template Signed-off-by: Guy Ealey Morag --- test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 58 +++++++++---------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index d78bdca0eff..4ea8c250d0c 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -256,32 +256,6 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void check_assignment_not_shared(size_t num_groups, - size_t num_gpus_per_group, - size_t num_nics_per_group, - ucp_gpu_nic_assignment_policy_t policy, - const std::vector &expected_owners) - { - topology_shape_t config; - - config.num_groups = num_groups; - config.num_gpus_per_group = num_gpus_per_group; - config.num_nics_per_group = num_nics_per_group; - - /* Each GPU may appear as two devices if MPS MLOPart is enabled. */ - for (size_t num_gpu_devices = 1; num_gpu_devices <= 2; - ++num_gpu_devices) { - /* Each NIC may appear as two ports if dual-port mode is enabled. */ - for (size_t num_nic_ports = 1; num_nic_ports <= 2; - ++num_nic_ports) { - config.num_gpu_devices = num_gpu_devices; - config.num_nic_ports = num_nic_ports; - check_assignment_not_shared_inner(config, policy, - expected_owners); - } - } - } - void check_assignment_shared_inner(const topology_shape_t &config) { ASSERT_NE(config.num_groups, 0); @@ -329,8 +303,9 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void check_assignment_shared(size_t num_groups, size_t num_gpus_per_group, - size_t num_nics_per_group) + template + void check_assignment(size_t num_groups, size_t num_gpus_per_group, + size_t num_nics_per_group, Check check) { topology_shape_t config; @@ -338,17 +313,42 @@ class test_ucp_gpu_nic_assignment : public ucs::test { config.num_gpus_per_group = num_gpus_per_group; config.num_nics_per_group = num_nics_per_group; + /* Each GPU may appear as two devices if MPS MLOPart is enabled. */ for (size_t num_gpu_devices = 1; num_gpu_devices <= 2; ++num_gpu_devices) { + /* Each NIC may appear as two ports if dual-port mode is enabled. */ for (size_t num_nic_ports = 1; num_nic_ports <= 2; ++num_nic_ports) { config.num_gpu_devices = num_gpu_devices; config.num_nic_ports = num_nic_ports; - check_assignment_shared_inner(config); + check(config); } } } + void check_assignment_not_shared(size_t num_groups, + size_t num_gpus_per_group, + size_t num_nics_per_group, + ucp_gpu_nic_assignment_policy_t policy, + const std::vector &expected_owners) + { + check_assignment(num_groups, num_gpus_per_group, num_nics_per_group, + [this, policy, + &expected_owners](const topology_shape_t &config) { + check_assignment_not_shared_inner(config, policy, + expected_owners); + }); + } + + void check_assignment_shared(size_t num_groups, size_t num_gpus_per_group, + size_t num_nics_per_group) + { + check_assignment(num_groups, num_gpus_per_group, num_nics_per_group, + [this](const topology_shape_t &config) { + check_assignment_shared_inner(config); + }); + } + private: ucs_topo_groups_t m_groups; ucp_gpu_nic_assignment_t m_assignment; From e66bc85359a82313aa14acbb43901e1c9c09dec8 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Tue, 22 Sep 2026 14:01:30 +0000 Subject: [PATCH 3/8] UCP/CORE: Add GPU_NIC_ASSIGNMENT_MODE with 'auto' and 'off' Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 33 +++++---- src/ucp/core/ucp_context.h | 4 +- src/ucp/core/ucp_gpu_nic_assignment.c | 33 ++++++--- src/ucp/core/ucp_gpu_nic_assignment.h | 15 ++-- src/ucp/core/ucp_types.h | 15 ++-- test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 73 +++++++++++++++---- 6 files changed, 121 insertions(+), 52 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 0a431fd4f1b..264d41b69b5 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -108,12 +108,13 @@ static const char *ucp_fence_modes[] = { [UCP_FENCE_MODE_LAST] = NULL }; -static const char *ucp_gpu_net_device_modes[] = { - [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED] = "not_shared", - [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP] = "not_shared_flip", - [UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN] = "not_shared_round_robin", - [UCP_GPU_NET_DEVICE_MODE_SHARED] = "shared", - [UCP_GPU_NET_DEVICE_MODE_LAST] = NULL +static const char *ucp_gpu_nic_assignment_modes[] = { + [UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO] = "auto", + [UCP_GPU_NIC_ASSIGNMENT_MODE_OFF] = "off", + [UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP] = "flip", + [UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN] = "round_robin", + [UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED] = "shared", + [UCP_GPU_NIC_ASSIGNMENT_MODE_LAST] = NULL }; static const char *ucp_rndv_modes[] = { @@ -170,6 +171,7 @@ const char *ucp_extra_op_attr_flags_names[] = { static UCS_CONFIG_DEFINE_ARRAY(memunit_sizes, sizeof(size_t), UCS_CONFIG_TYPE_MEMUNITS); +/* clang-format off */ static ucs_config_field_t ucp_context_config_table[] = { {"SELECT_DISTANCE_MD", "cuda_cpy", "MD whose distance is queried when evaluating transport selection score", @@ -623,15 +625,15 @@ static ucs_config_field_t ucp_context_config_table[] = { ucs_offsetof(ucp_context_config_t, connect_all_to_all), UCS_CONFIG_TYPE_BOOL}, - {"GPU_NET_DEVICE_MODE", "not_shared", - "Select how network devices are assigned to GPUs.\n" - " - not_shared : use the default FLIP assignment.\n" - " - not_shared_flip : use the FLIP assignment.\n" - " - not_shared_round_robin : use the round-robin assignment.\n" - " - shared : assign every network device in a topology\n" - " group to every GPU in that group.", - ucs_offsetof(ucp_context_config_t, gpu_net_device_mode), - UCS_CONFIG_TYPE_ENUM(ucp_gpu_net_device_modes)}, + {"GPU_NIC_ASSIGNMENT_MODE", "auto", + "Select how NICs are assigned to GPUs.\n" + " - auto : enable automatically on supported hardware, otherwise disable.\n" + " - off : disable GPU-to-NIC assignment, fallback to free lane selection.\n" + " - flip : require the not shared flip assignment policy.\n" + " - round_robin : require the not shared round-robin assignment policy.\n" + " - shared : require the shared assignment policy.", + ucs_offsetof(ucp_context_config_t, gpu_nic_assignment_mode), + UCS_CONFIG_TYPE_ENUM(ucp_gpu_nic_assignment_modes)}, {"SINGLE_NET_DEVICE", "n", "Restrict each protocol's lanes to one network device.\n" @@ -668,6 +670,7 @@ static ucs_config_field_t ucp_context_config_table[] = { {NULL} }; +/* clang-format on */ static ucs_config_field_t ucp_config_table[] = { {"NET_DEVICES", UCP_RSC_CONFIG_ALL, diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index edae667a466..e7132d90ccf 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -255,8 +255,8 @@ typedef struct ucp_context_config { /** Extend endpoint lanes connections of each local device to all remote * devices */ int connect_all_to_all; - /** GPU network device assignment mode */ - ucp_gpu_net_device_mode_t gpu_net_device_mode; + /** GPU-to-NIC assignment mode */ + ucp_gpu_nic_assignment_mode_t gpu_nic_assignment_mode; /** Restrict lanes to one network device per protocol */ int proto_use_single_net_device; /** Max HCAs for GPU memory registration: auto=closest, N=limit, inf=all */ diff --git a/src/ucp/core/ucp_gpu_nic_assignment.c b/src/ucp/core/ucp_gpu_nic_assignment.c index 72c4ed360b1..c0717981f95 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.c +++ b/src/ucp/core/ucp_gpu_nic_assignment.c @@ -33,19 +33,32 @@ typedef struct { } ucp_nics_string_buffers_t; -ucp_gpu_nic_assignment_policy_t -ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode) +int ucp_gpu_nic_assignment_policy_resolve( + ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model, + ucp_gpu_nic_assignment_policy_t *policy_p) { + ucs_assert(policy_p != NULL); + switch (mode) { - case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED: - case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP: - return UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; - case UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN: - return UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN; - case UCP_GPU_NET_DEVICE_MODE_SHARED: - return UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED; + case UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO: + if (cpu_model != UCS_CPU_MODEL_NVIDIA_VERA) { + return 0; + } + *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; + return 1; + case UCP_GPU_NIC_ASSIGNMENT_MODE_OFF: + return 0; + case UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP: + *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; + return 1; + case UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN: + *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN; + return 1; + case UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED: + *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED; + return 1; default: - ucs_fatal("invalid gpu network device mode %d", (int)mode); + ucs_fatal("invalid gpu-nic assignment mode %d", (int)mode); } } diff --git a/src/ucp/core/ucp_gpu_nic_assignment.h b/src/ucp/core/ucp_gpu_nic_assignment.h index 4246d303901..08bdf223c57 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.h +++ b/src/ucp/core/ucp_gpu_nic_assignment.h @@ -9,6 +9,7 @@ #include "ucp_types.h" +#include #include #include @@ -64,14 +65,18 @@ typedef enum { /** - * Get the assignment policy of a GPU network device mode. + * Resolve the GPU-to-NIC assignment policy. * - * @param [in] mode GPU network device mode. + * @param [in] mode GPU-to-NIC assignment mode. + * @param [in] cpu_model CPU model used to resolve automatic mode. + * @param [out] policy_p Resolved assignment policy, + * populated only when assignment is enabled. * - * @return Assignment policy corresponding to @a mode. + * @return Nonzero if assignment is enabled, or zero if it is disabled. */ -ucp_gpu_nic_assignment_policy_t -ucp_gpu_net_device_mode_get_policy(ucp_gpu_net_device_mode_t mode); +int ucp_gpu_nic_assignment_policy_resolve( + ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model, + ucp_gpu_nic_assignment_policy_t *policy_p); /** diff --git a/src/ucp/core/ucp_types.h b/src/ucp/core/ucp_types.h index 00c122c52db..d145a969241 100644 --- a/src/ucp/core/ucp_types.h +++ b/src/ucp/core/ucp_types.h @@ -234,15 +234,16 @@ typedef enum { /** - * GPU network device assignment mode. + * GPU-to-NIC assignment mode. */ typedef enum { - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED, - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP, - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN, - UCP_GPU_NET_DEVICE_MODE_SHARED, - UCP_GPU_NET_DEVICE_MODE_LAST -} ucp_gpu_net_device_mode_t; + UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, + UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, + UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, + UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, + UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, + UCP_GPU_NIC_ASSIGNMENT_MODE_LAST +} ucp_gpu_nic_assignment_mode_t; /** diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index 4ea8c250d0c..b75f59257b0 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -276,6 +276,23 @@ class test_ucp_gpu_nic_assignment : public ucs::test { check_gpu_device_aliases(config); + ASSERT_NE(config.num_gpus_per_group, 0); + + if (config.num_nics() == 0) { + for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { + const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( + &m_assignment, gpu_sys_dev(config, gpu_idx, 0)); + + ASSERT_NE(nullptr, nic_sys_dev_bitmap); + EXPECT_EQ(0ul, static_cast(UCS_STATIC_BITMAP_POPCOUNT( + *nic_sys_dev_bitmap))); + } + + return; + } + + ASSERT_NE(config.num_nics_per_group, 0); + for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( &m_assignment, gpu_sys_dev(config, gpu_idx, 0)); @@ -354,19 +371,49 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ucp_gpu_nic_assignment_t m_assignment; }; -UCS_TEST_F(test_ucp_gpu_nic_assignment, mode_to_policy) { - EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - ucp_gpu_net_device_mode_get_policy( - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED)); - EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - ucp_gpu_net_device_mode_get_policy( - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_FLIP)); - EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - ucp_gpu_net_device_mode_get_policy( - UCP_GPU_NET_DEVICE_MODE_NOT_SHARED_ROUND_ROBIN)); - EXPECT_EQ(UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, - ucp_gpu_net_device_mode_get_policy( - UCP_GPU_NET_DEVICE_MODE_SHARED)); +UCS_TEST_F(test_ucp_gpu_nic_assignment, mode_resolution) { + struct test_param_t { + ucp_gpu_nic_assignment_mode_t mode; + ucs_cpu_model_t cpu_model; + int expected_enabled; + ucp_gpu_nic_assignment_policy_t expected_policy; + }; + + const test_param_t test_params[] = { + /* Auto mode with VERA CPU. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_NVIDIA_VERA, 1, + UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP}, + /* Auto mode with non-VERA CPU. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_UNKNOWN, 0, + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, + {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_NVIDIA_GRACE, 0, + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, + /* Off mode. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, UCS_CPU_MODEL_NVIDIA_VERA, 0, + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, + {UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, UCS_CPU_MODEL_UNKNOWN, 0, + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, + /* Flip mode. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, UCS_CPU_MODEL_UNKNOWN, 1, + UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP}, + /* Round-robin mode. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, UCS_CPU_MODEL_UNKNOWN, 1, + UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN}, + /* Shared mode. */ + {UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, UCS_CPU_MODEL_UNKNOWN, 1, + UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED}, + }; + + for (const auto ¶ms : test_params) { + ucp_gpu_nic_assignment_policy_t policy = + UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST; + + EXPECT_EQ(params.expected_enabled, + ucp_gpu_nic_assignment_policy_resolve(params.mode, + params.cpu_model, + &policy)); + EXPECT_EQ(params.expected_policy, policy); + } } UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_no_nics) { From 8d616a0cc672b5aa4c00d85d85642f0846acf842 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Tue, 22 Sep 2026 16:59:59 +0000 Subject: [PATCH 4/8] TEST/UCP: Fix coverity error Signed-off-by: Guy Ealey Morag --- test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index b75f59257b0..227ca42552e 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -276,7 +276,10 @@ class test_ucp_gpu_nic_assignment : public ucs::test { check_gpu_device_aliases(config); - ASSERT_NE(config.num_gpus_per_group, 0); + if (config.num_gpus_per_group == 0) { + ADD_FAILURE() << "num_gpus_per_group must be non-zero"; + return; + } if (config.num_nics() == 0) { for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { @@ -291,7 +294,10 @@ class test_ucp_gpu_nic_assignment : public ucs::test { return; } - ASSERT_NE(config.num_nics_per_group, 0); + if (config.num_nics_per_group == 0) { + ADD_FAILURE() << "num_nics_per_group must be non-zero"; + return; + } for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( From abc8ac4b000941f7d3b1e7564a3bf147f7670e72 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Wed, 23 Sep 2026 07:17:08 +0000 Subject: [PATCH 5/8] UCP/CORE: Use ucp_gpu_nic_assignment_policy_resolve Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 20 +++++++++++--------- src/ucp/core/ucp_gpu_nic_assignment.c | 1 + 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 0eab914e586..618a8020aef 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -2776,19 +2776,21 @@ ucp_version_check(unsigned api_major_version, unsigned api_minor_version) } static ucs_status_t -ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t **assignment_p) +ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode, + ucp_gpu_nic_assignment_t **assignment_p) { + ucp_gpu_nic_assignment_policy_t policy; ucp_gpu_nic_assignment_t *assignment; ucs_topo_groups_t groups; ucs_status_t status; *assignment_p = NULL; - /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */ - if (ucs_arch_get_cpu_model() != UCS_CPU_MODEL_NVIDIA_VERA) { - ucs_debug("gpu-nic assignment is not supported on %s architecture, " + if (!ucp_gpu_nic_assignment_policy_resolve(mode, ucs_arch_get_cpu_model(), + &policy)) { + ucs_debug("gpu-nic assignment is disabled (mode %s, cpu model %s), " "skipping", - ucs_cpu_model_name()); + ucp_gpu_nic_assignment_modes[mode], ucs_cpu_model_name()); return UCS_OK; } @@ -2809,9 +2811,7 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t **assignment_p) goto out_release_groups; } - status = ucp_gpu_nic_assignment_build(&groups, - UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - assignment); + status = ucp_gpu_nic_assignment_build(&groups, policy, assignment); if (status != UCS_OK) { ucs_free(assignment); goto out_release_groups; @@ -2878,7 +2878,9 @@ ucs_status_t ucp_init_version(unsigned api_major_version, unsigned api_minor_ver goto err_thread_lock_finalize; } - status = ucp_context_gpu_nic_assignment_init(&context->gpu_nic_assignment); + status = ucp_context_gpu_nic_assignment_init( + context->config.ext.gpu_nic_assignment_mode, + &context->gpu_nic_assignment); if (status != UCS_OK) { goto err_free_res; } diff --git a/src/ucp/core/ucp_gpu_nic_assignment.c b/src/ucp/core/ucp_gpu_nic_assignment.c index c0717981f95..bb0ef4b749d 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.c +++ b/src/ucp/core/ucp_gpu_nic_assignment.c @@ -41,6 +41,7 @@ int ucp_gpu_nic_assignment_policy_resolve( switch (mode) { case UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO: + /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */ if (cpu_model != UCS_CPU_MODEL_NVIDIA_VERA) { return 0; } From 9413b2aa1c4b3183b1c5ce0e9d278e075ae45697 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Wed, 23 Sep 2026 10:31:22 +0000 Subject: [PATCH 6/8] UCP/CORE: PR comments fixes Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 22 +- src/ucp/core/ucp_gpu_nic_assignment.c | 71 ++--- src/ucp/core/ucp_gpu_nic_assignment.h | 10 +- src/ucp/core/ucp_types.h | 5 + test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 299 ++++++------------ 5 files changed, 145 insertions(+), 262 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 618a8020aef..551c7555388 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -628,12 +628,22 @@ static ucs_config_field_t ucp_context_config_table[] = { UCS_CONFIG_TYPE_BOOL}, {"GPU_NIC_ASSIGNMENT_MODE", "auto", - "Select how NICs are assigned to GPUs.\n" - " - auto : enable automatically on supported hardware, otherwise disable.\n" - " - off : disable GPU-to-NIC assignment, fallback to free lane selection.\n" - " - flip : require the not shared flip assignment policy.\n" - " - round_robin : require the not shared round-robin assignment policy.\n" - " - shared : require the shared assignment policy.", + "Assign NICs to GPUs within each topology group, and restrict the lanes\n" + "for a GPU's memory to the NICs assigned to that GPU.\n" + "All ports of a NIC are assigned together.\n" + "The 'flip', 'round_robin' and 'shared' modes apply on any hardware.\n" + "With 'flip' and 'round_robin', a group with fewer NICs than GPUs leaves\n" + "some GPUs without any NICs.\n" + " - auto : use 'flip' on hardware with a known GPU-NIC topology,\n" + " otherwise 'off'.\n" + " - off : do not assign; select lanes from all NICs.\n" + " - flip : assign each NIC to a single GPU, walking the N GPUs of\n" + " the group forward then backward:\n" + " 0, 1, .., N-1, N-1, .., 1, 0, 0, 1, ..\n" + " - round_robin : assign each NIC to a single GPU, walking the N GPUs of\n" + " the group in ascending order:\n" + " 0, 1, .., N-1, 0, 1, .., N-1, 0, ..\n" + " - shared : assign all NICs of a group to every GPU of that group.", ucs_offsetof(ucp_context_config_t, gpu_nic_assignment_mode), UCS_CONFIG_TYPE_ENUM(ucp_gpu_nic_assignment_modes)}, diff --git a/src/ucp/core/ucp_gpu_nic_assignment.c b/src/ucp/core/ucp_gpu_nic_assignment.c index bb0ef4b749d..f14c363e337 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.c +++ b/src/ucp/core/ucp_gpu_nic_assignment.c @@ -357,65 +357,44 @@ ucp_gpu_nic_assignment_init(ucp_gpu_nic_assignment_t *assignment, } static void -ucp_gpu_nic_assignment_add_group_shared(ucp_gpu_nic_assignment_t *assignment, - const ucs_topo_group_t *group) +ucp_gpu_nic_assignment_add_nic_to_gpu(ucp_gpu_nic_assignment_t *assignment, + const ucs_topo_group_element_t *gpu, + const ucs_topo_group_element_t *nic) { ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap; - const ucs_topo_group_element_t *gpu, *nic; - ucs_array_for_each(gpu, &group->gpus) { - ucs_assert((gpu->num_sys_devs > 0) && - (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN)); + ucs_assert((gpu->num_sys_devs > 0) && + (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN)); - nic_sys_dev_bitmap = ucs_const_cast( - ucp_gpu_nic_sys_dev_bitmap_t*, - ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0])); - ucs_assert(nic_sys_dev_bitmap != NULL); + nic_sys_dev_bitmap = ucs_const_cast( + ucp_gpu_nic_sys_dev_bitmap_t*, + ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0])); + ucs_assert(nic_sys_dev_bitmap != NULL); - /* Assign all NICs in the group to the GPU. */ - ucs_array_for_each(nic, &group->nics) { - ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic); - } - } + ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic); } -static void ucp_gpu_nic_assignment_add_group_not_shared( - ucp_gpu_nic_assignment_t *assignment, const ucs_topo_group_t *group, - ucp_gpu_nic_assignment_policy_t policy) +static void +ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, + const ucs_topo_group_t *group, + ucp_gpu_nic_assignment_policy_t policy) { size_t num_gpus = ucs_array_length(&group->gpus); - ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap; const ucs_topo_group_element_t *gpu, *nic; size_t nic_idx, gpu_idx; - ucs_assertv(policy != UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, - "invalid policy: %d", (int)policy); - ucs_array_for_each_index(nic, nic_idx, &group->nics) { - gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, nic_idx); - gpu = &ucs_array_elem(&group->gpus, gpu_idx); - - ucs_assert((gpu->num_sys_devs > 0) && - (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN)); - - nic_sys_dev_bitmap = ucs_const_cast( - ucp_gpu_nic_sys_dev_bitmap_t*, - ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0])); - ucs_assert(nic_sys_dev_bitmap != NULL); - - ucp_gpu_nic_bitmap_add_nic(nic_sys_dev_bitmap, nic); - } -} - -static void -ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, - const ucs_topo_group_t *group, - ucp_gpu_nic_assignment_policy_t policy) -{ - if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { - ucp_gpu_nic_assignment_add_group_shared(assignment, group); - } else { - ucp_gpu_nic_assignment_add_group_not_shared(assignment, group, policy); + if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { + /* Assign the NIC to every GPU in the group. */ + ucs_array_for_each(gpu, &group->gpus) { + ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic); + } + } else { + gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, + nic_idx); + gpu = &ucs_array_elem(&group->gpus, gpu_idx); + ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic); + } } } diff --git a/src/ucp/core/ucp_gpu_nic_assignment.h b/src/ucp/core/ucp_gpu_nic_assignment.h index 1cd8f9c161e..a681504056b 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.h +++ b/src/ucp/core/ucp_gpu_nic_assignment.h @@ -41,17 +41,17 @@ struct ucp_gpu_nic_assignment { typedef enum { /** * Assign each NIC, with all of its ports, going forward and then backward - * across the GPUs. + * across the N GPUs of the group. * The group's NICs are assigned to the following GPU indices in order: - * 0, 1, ..., num_gpus-1, num_gpus-1, ..., 1, 0, 0, 1, ... + * 0, 1, ..., N-1, N-1, ..., 1, 0, 0, 1, ... */ UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, /** - * Assign each NIC, with all of its ports, to GPUs repeatedly in ascending - * order. + * Assign each NIC, with all of its ports, to the N GPUs of the group + * repeatedly in ascending order. * The group's NICs are assigned to the following GPU indices in order: - * 0, 1, ..., num_gpus-1, 0, 1, ..., num_gpus-1, 0, 1, ... + * 0, 1, ..., N-1, 0, 1, ..., N-1, 0, 1, ... */ UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, diff --git a/src/ucp/core/ucp_types.h b/src/ucp/core/ucp_types.h index 5ccabaec60f..0b21a8bc115 100644 --- a/src/ucp/core/ucp_types.h +++ b/src/ucp/core/ucp_types.h @@ -238,10 +238,15 @@ typedef enum { * GPU-to-NIC assignment mode. */ typedef enum { + /* Enable on hardware with a known GPU-NIC topology, otherwise disable */ UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, + /* Do not assign NICs to GPUs, select lanes from all NICs */ UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, + /* Assign each NIC to a single GPU, walking GPUs forward then backward */ UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, + /* Assign each NIC to a single GPU, walking the GPUs in ascending order */ UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, + /* Assign every NIC of a group to every GPU of that group */ UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, UCP_GPU_NIC_ASSIGNMENT_MODE_LAST } ucp_gpu_nic_assignment_mode_t; diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index 227ca42552e..75051256661 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -161,6 +161,13 @@ class test_ucp_gpu_nic_assignment : public ucs::test { void check_nic_owners(const topology_shape_t &config, const std::vector &expected_owners) { + /* Validate the expected owners vector. */ + ASSERT_EQ(config.num_nics(), expected_owners.size()); + for (size_t owner_idx = 0; owner_idx < expected_owners.size(); + ++owner_idx) { + ASSERT_LT(expected_owners[owner_idx], config.num_gpus()); + } + /* Expected number of assigned NIC ports per physical GPU. */ std::vector expected_port_counts(config.num_gpus(), 0); @@ -219,116 +226,82 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void - check_assignment_not_shared_inner(const topology_shape_t &config, - ucp_gpu_nic_assignment_policy_t policy, - const std::vector &expected_owners) + void check_nic_shared(const topology_shape_t &config) { - ASSERT_NE(config.num_groups, 0); - ASSERT_NE(config.num_nic_ports, 0); - ASSERT_NE(config.num_gpu_devices, 0); + /* Verify that each GPU owns all ports of every NIC in its group. */ + for (size_t group_idx = 0; group_idx < config.num_groups; ++group_idx) { + const size_t first_nic_idx = group_idx * config.num_nics_per_group; + const size_t end_nic_idx = first_nic_idx + + config.num_nics_per_group; - ASSERT_LE(config.num_gpu_devices, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); - ASSERT_LE(config.num_nic_ports, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); + for (size_t local_idx = 0; local_idx < config.num_gpus_per_group; + ++local_idx) { + const size_t gpu_idx = (group_idx * config.num_gpus_per_group) + + local_idx; + const ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap = + ucp_gpu_nic_assignment_lookup(&m_assignment, + gpu_sys_dev(config, + gpu_idx, 0)); + ASSERT_NE(nullptr, nic_sys_dev_bitmap); - /* Leave (UCS_SYS_DEVICE_ID_UNKNOWN - 1) unmapped for the ungrouped lookup. */ - ASSERT_LE(config.num_sys_devs(), UCS_SYS_DEVICE_ID_COUNT - 1); + for (size_t nic_idx = 0; nic_idx < config.num_nics(); + ++nic_idx) { + const bool same_group = (nic_idx >= first_nic_idx) && + (nic_idx < end_nic_idx); + + for (size_t port_idx = 0; port_idx < config.num_nic_ports; + ++port_idx) { + EXPECT_EQ(same_group, + ucp_gpu_nic_bitmap_get( + nic_sys_dev_bitmap, + nic_port_sys_dev(config, nic_idx, + port_idx))); + } + } - if (config.num_gpus() == 0) { - ASSERT_TRUE(expected_owners.empty()); - } else { - ASSERT_EQ(config.num_nics(), expected_owners.size()); - for (size_t owner_idx = 0; owner_idx < expected_owners.size(); - ++owner_idx) { - ASSERT_LT(expected_owners[owner_idx], config.num_gpus()); + EXPECT_EQ(config.num_nics_per_group * config.num_nic_ports, + static_cast(UCS_STATIC_BITMAP_POPCOUNT( + *nic_sys_dev_bitmap))); } } - - build_groups(config); - build_assignment(policy); - - if (config.num_gpus() == 0) { - EXPECT_EQ(nullptr, m_assignment.nic_sys_dev_bitmaps); - EXPECT_EQ(0ul, m_assignment.num_bitmaps); - } else { - check_gpu_device_aliases(config); - check_nic_owners(config, expected_owners); - } } - void check_assignment_shared_inner(const topology_shape_t &config) + void check_assignment(const topology_shape_t &config, + ucp_gpu_nic_assignment_policy_t policy, + const std::vector &expected_owners) { ASSERT_NE(config.num_groups, 0); ASSERT_NE(config.num_nic_ports, 0); ASSERT_NE(config.num_gpu_devices, 0); + ASSERT_LE(config.num_gpu_devices, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); ASSERT_LE(config.num_nic_ports, UCS_TOPO_MAX_SYS_DEVS_PER_ELEMENT); + + /* Leave (UCS_SYS_DEVICE_ID_UNKNOWN - 1) unmapped for the ungrouped lookup. */ ASSERT_LE(config.num_sys_devs(), UCS_SYS_DEVICE_ID_COUNT - 1); build_groups(config); - build_assignment(UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED); + build_assignment(policy); if (config.num_gpus() == 0) { + ASSERT_TRUE(expected_owners.empty()); EXPECT_EQ(nullptr, m_assignment.nic_sys_dev_bitmaps); EXPECT_EQ(0ul, m_assignment.num_bitmaps); - return; - } - - check_gpu_device_aliases(config); - - if (config.num_gpus_per_group == 0) { - ADD_FAILURE() << "num_gpus_per_group must be non-zero"; - return; - } - - if (config.num_nics() == 0) { - for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { - const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( - &m_assignment, gpu_sys_dev(config, gpu_idx, 0)); - - ASSERT_NE(nullptr, nic_sys_dev_bitmap); - EXPECT_EQ(0ul, static_cast(UCS_STATIC_BITMAP_POPCOUNT( - *nic_sys_dev_bitmap))); + } else { + check_gpu_device_aliases(config); + if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { + ASSERT_TRUE(expected_owners.empty()); + check_nic_shared(config); + } else { + check_nic_owners(config, expected_owners); } - - return; - } - - if (config.num_nics_per_group == 0) { - ADD_FAILURE() << "num_nics_per_group must be non-zero"; - return; - } - - for (size_t gpu_idx = 0; gpu_idx < config.num_gpus(); ++gpu_idx) { - const auto *nic_sys_dev_bitmap = ucp_gpu_nic_assignment_lookup( - &m_assignment, gpu_sys_dev(config, gpu_idx, 0)); - const size_t gpu_group_idx = gpu_idx / config.num_gpus_per_group; - - ASSERT_NE(nullptr, nic_sys_dev_bitmap); - - for (size_t nic_idx = 0; nic_idx < config.num_nics(); ++nic_idx) { - const size_t nic_group_idx = nic_idx / - config.num_nics_per_group; - - for (size_t port_idx = 0; port_idx < config.num_nic_ports; - ++port_idx) { - EXPECT_EQ(gpu_group_idx == nic_group_idx, - ucp_gpu_nic_bitmap_get( - nic_sys_dev_bitmap, - nic_port_sys_dev(config, nic_idx, - port_idx))); - } - } - - EXPECT_EQ(config.num_nics_per_group * config.num_nic_ports, - static_cast( - UCS_STATIC_BITMAP_POPCOUNT(*nic_sys_dev_bitmap))); } } - template void check_assignment(size_t num_groups, size_t num_gpus_per_group, - size_t num_nics_per_group, Check check) + size_t num_nics_per_group, + ucp_gpu_nic_assignment_policy_t policy, + const std::vector &expected_owners) { topology_shape_t config; @@ -344,149 +317,65 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ++num_nic_ports) { config.num_gpu_devices = num_gpu_devices; config.num_nic_ports = num_nic_ports; - check(config); + check_assignment(config, policy, expected_owners); } } } - void check_assignment_not_shared(size_t num_groups, - size_t num_gpus_per_group, - size_t num_nics_per_group, - ucp_gpu_nic_assignment_policy_t policy, - const std::vector &expected_owners) - { - check_assignment(num_groups, num_gpus_per_group, num_nics_per_group, - [this, policy, - &expected_owners](const topology_shape_t &config) { - check_assignment_not_shared_inner(config, policy, - expected_owners); - }); - } - - void check_assignment_shared(size_t num_groups, size_t num_gpus_per_group, - size_t num_nics_per_group) - { - check_assignment(num_groups, num_gpus_per_group, num_nics_per_group, - [this](const topology_shape_t &config) { - check_assignment_shared_inner(config); - }); - } - private: ucs_topo_groups_t m_groups; ucp_gpu_nic_assignment_t m_assignment; }; -UCS_TEST_F(test_ucp_gpu_nic_assignment, mode_resolution) { - struct test_param_t { - ucp_gpu_nic_assignment_mode_t mode; - ucs_cpu_model_t cpu_model; - int expected_enabled; - ucp_gpu_nic_assignment_policy_t expected_policy; - }; - - const test_param_t test_params[] = { - /* Auto mode with VERA CPU. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_NVIDIA_VERA, 1, - UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP}, - /* Auto mode with non-VERA CPU. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_UNKNOWN, 0, - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, - {UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, UCS_CPU_MODEL_NVIDIA_GRACE, 0, - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, - /* Off mode. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, UCS_CPU_MODEL_NVIDIA_VERA, 0, - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, - {UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, UCS_CPU_MODEL_UNKNOWN, 0, - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST}, - /* Flip mode. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, UCS_CPU_MODEL_UNKNOWN, 1, - UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP}, - /* Round-robin mode. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, UCS_CPU_MODEL_UNKNOWN, 1, - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN}, - /* Shared mode. */ - {UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, UCS_CPU_MODEL_UNKNOWN, 1, - UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED}, - }; - - for (const auto ¶ms : test_params) { - ucp_gpu_nic_assignment_policy_t policy = - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST; - - EXPECT_EQ(params.expected_enabled, - ucp_gpu_nic_assignment_policy_resolve(params.mode, - params.cpu_model, - &policy)); - EXPECT_EQ(params.expected_policy, policy); - } -} - -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_no_nics) { - check_assignment_not_shared(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, no_nics) { + check_assignment(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_no_gpus) { - check_assignment_not_shared(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, no_gpus) { + check_assignment(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_flip_divisible) { +UCS_TEST_F(test_ucp_gpu_nic_assignment, flip_divisible) { /* Number of NICs is divisible by the number of GPUs. */ - check_assignment_not_shared(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 1, 0, /**/ - 2, 3, 3, 2, /**/ - 4, 5, 5, 4}); - check_assignment_not_shared(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 2, 2, 1, 0, /**/ - 3, 4, 5, 5, 4, 3}); + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 1, 0, /**/ + 2, 3, 3, 2, /**/ + 4, 5, 5, 4}); + check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 2, 2, 1, 0, /**/ + 3, 4, 5, 5, 4, 3}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_flip_not_divisible) { - check_assignment_not_shared(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, 1, 0, 0, /**/ - 2, 3, 3, 2, 2}); - check_assignment_not_shared(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - {0, 1, /**/ - 3, 4}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, flip_not_divisible) { + check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, 1, 0, 0, /**/ + 2, 3, 3, 2, 2}); + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + {0, 1, /**/ + 3, 4}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_round_robin_divisible) { - check_assignment_not_shared(3, 2, 4, - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 0, 1, /**/ - 2, 3, 2, 3, /**/ - 4, 5, 4, 5}); - check_assignment_not_shared(2, 3, 6, - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 2, 0, 1, 2, /**/ - 3, 4, 5, 3, 4, 5}); +UCS_TEST_F(test_ucp_gpu_nic_assignment, round_robin_divisible) { + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 0, 1, /**/ + 2, 3, 2, 3, /**/ + 4, 5, 4, 5}); + check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 2, 0, 1, 2, /**/ + 3, 4, 5, 3, 4, 5}); } -UCS_TEST_F(test_ucp_gpu_nic_assignment, not_shared_round_robin_not_divisible) { - check_assignment_not_shared(2, 2, 5, - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, 0, 1, 0, /**/ - 2, 3, 2, 3, 2}); - check_assignment_not_shared(2, 3, 2, - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - {0, 1, /**/ - 3, 4}); -} - -UCS_TEST_F(test_ucp_gpu_nic_assignment, shared_no_nics) { - check_assignment_shared(2, 3, 0); -} - -UCS_TEST_F(test_ucp_gpu_nic_assignment, shared_no_gpus) { - check_assignment_shared(2, 0, 3); +UCS_TEST_F(test_ucp_gpu_nic_assignment, round_robin_not_divisible) { + check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, 0, 1, 0, /**/ + 2, 3, 2, 3, 2}); + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + {0, 1, /**/ + 3, 4}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, shared) { - /* 2 GPUs, 4 NICs. */ - check_assignment_shared(3, 2, 4); - + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, {}); /* More GPUs than NICs. */ - check_assignment_shared(2, 3, 2); + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, {}); } From fe3025b2b1d65da9f5bd88b51f503c46202853e4 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Wed, 23 Sep 2026 10:39:22 +0000 Subject: [PATCH 7/8] UCP/CORE: Merge assignment policy and mode enums Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 22 +++++--- src/ucp/core/ucp_gpu_nic_assignment.c | 54 +++++-------------- src/ucp/core/ucp_gpu_nic_assignment.h | 46 +--------------- src/ucp/core/ucp_types.h | 25 +++++++-- test/gtest/ucp/test_ucp_gpu_nic_assignment.cc | 38 ++++++------- 5 files changed, 69 insertions(+), 116 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 551c7555388..a5950894e2f 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -2789,21 +2789,29 @@ static ucs_status_t ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode, ucp_gpu_nic_assignment_t **assignment_p) { - ucp_gpu_nic_assignment_policy_t policy; ucp_gpu_nic_assignment_t *assignment; ucs_topo_groups_t groups; ucs_status_t status; *assignment_p = NULL; - if (!ucp_gpu_nic_assignment_policy_resolve(mode, ucs_arch_get_cpu_model(), - &policy)) { - ucs_debug("gpu-nic assignment is disabled (mode %s, cpu model %s), " - "skipping", - ucp_gpu_nic_assignment_modes[mode], ucs_cpu_model_name()); + if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO) { + /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */ + if (ucs_arch_get_cpu_model() != UCS_CPU_MODEL_NVIDIA_VERA) { + ucs_debug("gpu-nic assignment is not supported on %s " + "architecture, skipping", + ucs_cpu_model_name()); + return UCS_OK; + } + + mode = UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP; + } else if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_OFF) { + ucs_debug("gpu-nic assignment is disabled by configuration"); return UCS_OK; } + ucs_debug("gpu-nic assignment mode %s", ucp_gpu_nic_assignment_modes[mode]); + status = ucs_topo_build_groups(&groups); if (status != UCS_OK) { return status; @@ -2821,7 +2829,7 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode, goto out_release_groups; } - status = ucp_gpu_nic_assignment_build(&groups, policy, assignment); + status = ucp_gpu_nic_assignment_build(&groups, mode, assignment); if (status != UCS_OK) { ucs_free(assignment); goto out_release_groups; diff --git a/src/ucp/core/ucp_gpu_nic_assignment.c b/src/ucp/core/ucp_gpu_nic_assignment.c index f14c363e337..682fc9f770e 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.c +++ b/src/ucp/core/ucp_gpu_nic_assignment.c @@ -33,38 +33,8 @@ typedef struct { } ucp_nics_string_buffers_t; -int ucp_gpu_nic_assignment_policy_resolve( - ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model, - ucp_gpu_nic_assignment_policy_t *policy_p) -{ - ucs_assert(policy_p != NULL); - - switch (mode) { - case UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO: - /* TODO: Improve Vera Rubin detection by checking NICs/GPUs models. */ - if (cpu_model != UCS_CPU_MODEL_NVIDIA_VERA) { - return 0; - } - *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; - return 1; - case UCP_GPU_NIC_ASSIGNMENT_MODE_OFF: - return 0; - case UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP: - *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP; - return 1; - case UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN: - *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN; - return 1; - case UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED: - *policy_p = UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED; - return 1; - default: - ucs_fatal("invalid gpu-nic assignment mode %d", (int)mode); - } -} - static size_t -ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_policy_t policy, +ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_mode_t mode, size_t num_gpus, size_t nic_idx) { size_t gpu_idx, leg; @@ -72,12 +42,12 @@ ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_policy_t policy, ucs_assert(num_gpus > 0); gpu_idx = nic_idx % num_gpus; - if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN) { + if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN) { return gpu_idx; } - ucs_assertv(policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - "invalid policy: %d", (int)policy); + ucs_assertv(mode == UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, "invalid mode: %d", + (int)mode); leg = nic_idx / num_gpus; if ((leg % 2) != 0) { gpu_idx = num_gpus - 1 - gpu_idx; @@ -377,20 +347,20 @@ ucp_gpu_nic_assignment_add_nic_to_gpu(ucp_gpu_nic_assignment_t *assignment, static void ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, const ucs_topo_group_t *group, - ucp_gpu_nic_assignment_policy_t policy) + ucp_gpu_nic_assignment_mode_t mode) { size_t num_gpus = ucs_array_length(&group->gpus); const ucs_topo_group_element_t *gpu, *nic; size_t nic_idx, gpu_idx; ucs_array_for_each_index(nic, nic_idx, &group->nics) { - if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { + if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED) { /* Assign the NIC to every GPU in the group. */ ucs_array_for_each(gpu, &group->gpus) { ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic); } } else { - gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(policy, num_gpus, + gpu_idx = ucp_gpu_nic_assignment_get_gpu_idx(mode, num_gpus, nic_idx); gpu = &ucs_array_elem(&group->gpus, gpu_idx); ucp_gpu_nic_assignment_add_nic_to_gpu(assignment, gpu, nic); @@ -400,7 +370,7 @@ ucp_gpu_nic_assignment_add_group(ucp_gpu_nic_assignment_t *assignment, ucs_status_t ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups, - ucp_gpu_nic_assignment_policy_t policy, + ucp_gpu_nic_assignment_mode_t mode, ucp_gpu_nic_assignment_t *assignment_p) { ucp_gpu_nic_assignment_t assignment; @@ -408,8 +378,10 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups, ucs_status_t status; ucs_assert(groups != NULL); - ucs_assertv(policy < UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST, - "invalid gpu-nic assignment policy %d", (int)policy); + ucs_assertv((mode == UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP) || + (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN) || + (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED), + "invalid gpu-nic assignment mode %d", (int)mode); status = ucp_gpu_nic_assignment_init(&assignment, groups); if (status != UCS_OK) { @@ -421,7 +393,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups, continue; } - ucp_gpu_nic_assignment_add_group(&assignment, group, policy); + ucp_gpu_nic_assignment_add_group(&assignment, group, mode); } ucp_gpu_nic_assignment_log(&assignment, groups); diff --git a/src/ucp/core/ucp_gpu_nic_assignment.h b/src/ucp/core/ucp_gpu_nic_assignment.h index a681504056b..ff39dbe476a 100644 --- a/src/ucp/core/ucp_gpu_nic_assignment.h +++ b/src/ucp/core/ucp_gpu_nic_assignment.h @@ -9,7 +9,6 @@ #include "ucp_types.h" -#include #include #include @@ -38,47 +37,6 @@ struct ucp_gpu_nic_assignment { }; -typedef enum { - /** - * Assign each NIC, with all of its ports, going forward and then backward - * across the N GPUs of the group. - * The group's NICs are assigned to the following GPU indices in order: - * 0, 1, ..., N-1, N-1, ..., 1, 0, 0, 1, ... - */ - UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, - - /** - * Assign each NIC, with all of its ports, to the N GPUs of the group - * repeatedly in ascending order. - * The group's NICs are assigned to the following GPU indices in order: - * 0, 1, ..., N-1, 0, 1, ..., N-1, 0, 1, ... - */ - UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, - - /** - * Assign every NIC, with all of its ports, to every GPU in the group. - */ - UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, - - UCP_GPU_NIC_ASSIGNMENT_POLICY_LAST -} ucp_gpu_nic_assignment_policy_t; - - -/** - * Resolve the GPU-to-NIC assignment policy. - * - * @param [in] mode GPU-to-NIC assignment mode. - * @param [in] cpu_model CPU model used to resolve automatic mode. - * @param [out] policy_p Resolved assignment policy, - * populated only when assignment is enabled. - * - * @return Nonzero if assignment is enabled, or zero if it is disabled. - */ -int ucp_gpu_nic_assignment_policy_resolve( - ucp_gpu_nic_assignment_mode_t mode, ucs_cpu_model_t cpu_model, - ucp_gpu_nic_assignment_policy_t *policy_p); - - /** * Build GPU-to-NIC assignments from topology groups. * @@ -86,7 +44,7 @@ int ucp_gpu_nic_assignment_policy_resolve( * topology group. * * @param [in] groups Topology groups to build assignments from. - * @param [in] policy Assignment policy. + * @param [in] mode Assignment mode: flip, round-robin or shared. * @param [out] assignment_p Completed assignment. Updated only on success. * * @return UCS_OK on success, or an error status if assignment construction @@ -94,7 +52,7 @@ int ucp_gpu_nic_assignment_policy_resolve( */ ucs_status_t ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups, - ucp_gpu_nic_assignment_policy_t policy, + ucp_gpu_nic_assignment_mode_t mode, ucp_gpu_nic_assignment_t *assignment_p); diff --git a/src/ucp/core/ucp_types.h b/src/ucp/core/ucp_types.h index 0b21a8bc115..161ddb4601b 100644 --- a/src/ucp/core/ucp_types.h +++ b/src/ucp/core/ucp_types.h @@ -238,16 +238,31 @@ typedef enum { * GPU-to-NIC assignment mode. */ typedef enum { - /* Enable on hardware with a known GPU-NIC topology, otherwise disable */ + /** Use flip on hardware with a known GPU-NIC topology, otherwise off */ UCP_GPU_NIC_ASSIGNMENT_MODE_AUTO, - /* Do not assign NICs to GPUs, select lanes from all NICs */ + + /** Do not assign NICs to GPUs, select lanes from all NICs */ UCP_GPU_NIC_ASSIGNMENT_MODE_OFF, - /* Assign each NIC to a single GPU, walking GPUs forward then backward */ + + /** + * Assign each NIC, with all of its ports, going forward and then backward + * across the N GPUs of the group. + * The group's NICs are assigned to the following GPU indices in order: + * 0, 1, ..., N-1, N-1, ..., 1, 0, 0, 1, ... + */ UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, - /* Assign each NIC to a single GPU, walking the GPUs in ascending order */ + + /** + * Assign each NIC, with all of its ports, to the N GPUs of the group + * repeatedly in ascending order. + * The group's NICs are assigned to the following GPU indices in order: + * 0, 1, ..., N-1, 0, 1, ..., N-1, 0, 1, ... + */ UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, - /* Assign every NIC of a group to every GPU of that group */ + + /** Assign every NIC, with all of its ports, to every GPU in the group */ UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, + UCP_GPU_NIC_ASSIGNMENT_MODE_LAST } ucp_gpu_nic_assignment_mode_t; diff --git a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc index 75051256661..712bcccb39d 100644 --- a/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc +++ b/test/gtest/ucp/test_ucp_gpu_nic_assignment.cc @@ -118,12 +118,12 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } } - void build_assignment(ucp_gpu_nic_assignment_policy_t policy) + void build_assignment(ucp_gpu_nic_assignment_mode_t mode) { ucs_status_t status; ucp_gpu_nic_assignment_release(&m_assignment); - status = ucp_gpu_nic_assignment_build(&m_groups, policy, &m_assignment); + status = ucp_gpu_nic_assignment_build(&m_groups, mode, &m_assignment); ASSERT_UCS_OK(status); /* Unknown device lookup */ @@ -267,7 +267,7 @@ class test_ucp_gpu_nic_assignment : public ucs::test { } void check_assignment(const topology_shape_t &config, - ucp_gpu_nic_assignment_policy_t policy, + ucp_gpu_nic_assignment_mode_t mode, const std::vector &expected_owners) { ASSERT_NE(config.num_groups, 0); @@ -281,7 +281,7 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ASSERT_LE(config.num_sys_devs(), UCS_SYS_DEVICE_ID_COUNT - 1); build_groups(config); - build_assignment(policy); + build_assignment(mode); if (config.num_gpus() == 0) { ASSERT_TRUE(expected_owners.empty()); @@ -289,7 +289,7 @@ class test_ucp_gpu_nic_assignment : public ucs::test { EXPECT_EQ(0ul, m_assignment.num_bitmaps); } else { check_gpu_device_aliases(config); - if (policy == UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED) { + if (mode == UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED) { ASSERT_TRUE(expected_owners.empty()); check_nic_shared(config); } else { @@ -300,7 +300,7 @@ class test_ucp_gpu_nic_assignment : public ucs::test { void check_assignment(size_t num_groups, size_t num_gpus_per_group, size_t num_nics_per_group, - ucp_gpu_nic_assignment_policy_t policy, + ucp_gpu_nic_assignment_mode_t mode, const std::vector &expected_owners) { topology_shape_t config; @@ -317,7 +317,7 @@ class test_ucp_gpu_nic_assignment : public ucs::test { ++num_nic_ports) { config.num_gpu_devices = num_gpu_devices; config.num_nic_ports = num_nic_ports; - check_assignment(config, policy, expected_owners); + check_assignment(config, mode, expected_owners); } } } @@ -328,54 +328,54 @@ class test_ucp_gpu_nic_assignment : public ucs::test { }; UCS_TEST_F(test_ucp_gpu_nic_assignment, no_nics) { - check_assignment(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); + check_assignment(2, 3, 0, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, no_gpus) { - check_assignment(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, {}); + check_assignment(2, 0, 3, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, flip_divisible) { /* Number of NICs is divisible by the number of GPUs. */ - check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {0, 1, 1, 0, /**/ 2, 3, 3, 2, /**/ 4, 5, 5, 4}); - check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {0, 1, 2, 2, 1, 0, /**/ 3, 4, 5, 5, 4, 3}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, flip_not_divisible) { - check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {0, 1, 1, 0, 0, /**/ 2, 3, 3, 2, 2}); - check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_FLIP, + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP, {0, 1, /**/ 3, 4}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, round_robin_divisible) { - check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, {0, 1, 0, 1, /**/ 2, 3, 2, 3, /**/ 4, 5, 4, 5}); - check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + check_assignment(2, 3, 6, UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, {0, 1, 2, 0, 1, 2, /**/ 3, 4, 5, 3, 4, 5}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, round_robin_not_divisible) { - check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + check_assignment(2, 2, 5, UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, {0, 1, 0, 1, 0, /**/ 2, 3, 2, 3, 2}); - check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_ROUND_ROBIN, + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_MODE_ROUND_ROBIN, {0, 1, /**/ 3, 4}); } UCS_TEST_F(test_ucp_gpu_nic_assignment, shared) { - check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, {}); + check_assignment(3, 2, 4, UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, {}); /* More GPUs than NICs. */ - check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_POLICY_SHARED, {}); + check_assignment(2, 3, 2, UCP_GPU_NIC_ASSIGNMENT_MODE_SHARED, {}); } From 173d40f6462b2595ac7a45c21ae1b4f6306f7cb7 Mon Sep 17 00:00:00 2001 From: Guy Ealey Morag Date: Wed, 23 Sep 2026 10:55:46 +0000 Subject: [PATCH 8/8] UCP/CORE: Revert clang format on/off Signed-off-by: Guy Ealey Morag --- src/ucp/core/ucp_context.c | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index a5950894e2f..81bd0220d3a 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -173,7 +173,6 @@ const char *ucp_extra_op_attr_flags_names[] = { static UCS_CONFIG_DEFINE_ARRAY(memunit_sizes, sizeof(size_t), UCS_CONFIG_TYPE_MEMUNITS); -/* clang-format off */ static ucs_config_field_t ucp_context_config_table[] = { {"SELECT_DISTANCE_MD", "cuda_cpy", "MD whose distance is queried when evaluating transport selection score", @@ -682,7 +681,6 @@ static ucs_config_field_t ucp_context_config_table[] = { {NULL} }; -/* clang-format on */ static ucs_config_field_t ucp_config_table[] = { {"NET_DEVICES", UCP_RSC_CONFIG_ALL,