Skip to content

TL/UCP: add exact mixed-radix allgather schedules #1327

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
jeffnvidia wants to merge 8 commits into
base: master
Choose a base branch
from
Open

TL/UCP: add exact mixed-radix allgather schedules #1327

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ec51644
TL/UCP: add exact mixed-radix allgather
jeffnvidia Jul 19, 2026
94d3bd9
TEST: cover mixed-radix allgather schedules
jeffnvidia Jul 19, 2026
d17161b
TL/UCP: address mixed-radix review
jeffnvidia Aug 20, 2026
155b53b
TL/UCP: remove dead radix assignment
jeffnvidia Aug 20, 2026
8e47d8d
TL/UCP: address follow-up mixed-radix review
jeffnvidia Aug 25, 2026
1efec89
TL/UCP: refine exact radix sequence handling
jeffnvidia Aug 31, 2026
1f41ba6
TL/UCP: unify radix sequence state
jeffnvidia Sep 1, 2026
3bd32dd
CODESTYLE: clean up mixed-radix formatting
jeffnvidia Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion src/coll_patterns/recursive_knomial.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@
#ifndef RECURSIVE_KNOMIAL_H_
#define RECURSIVE_KNOMIAL_H_

#include <limits.h>

#define UCC_KN_PEER_NULL ((ucc_rank_t)-1)
/* An exact schedule of radix >= 2 cannot have more phases than this. */
#define UCC_KN_MAX_RADIX_PHASES (sizeof(ucc_rank_t) * CHAR_BIT - 1)
typedef uint16_t ucc_kn_radix_t;

enum {
Expand All @@ -29,7 +33,7 @@ enum {

typedef struct ucc_knomial_pattern {

ucc_kn_radix_t radix; /* knomial tree radix */
ucc_kn_radix_t radix; /* radix for current iteration */
uint8_t type; /* pattern type */
uint8_t iteration; /* current iteration */
uint8_t n_iters; /* number of iterations in knomial algorithm */
Expand All @@ -54,6 +58,8 @@ typedef struct ucc_knomial_pattern {
ucc_rank_t block_size;
ptrdiff_t block_offset;
int is64;
uint8_t is_mixed;
ucc_kn_radix_t radices[UCC_KN_MAX_RADIX_PHASES];
Comment thread
jeffnvidia marked this conversation as resolved.
Outdated
} ucc_knomial_pattern_t;

/**
Expand Down Expand Up @@ -101,6 +107,7 @@ ucc_knomial_pattern_init_impl(ucc_rank_t size, ucc_rank_t rank,
p->rank = rank;
p->backward = backward;
p->iteration = 0;
p->is_mixed = 0;
n_full_subtrees = ucc_kn_pattern_n_full(p);
p->n_extra = has_extra ? size - n_full_subtrees * p->full_pow_size : 0;
p->n_iters = (p->n_extra && n_full_subtrees == 1) ?
Expand Down Expand Up @@ -229,6 +236,9 @@ ucc_knomial_pattern_next_iteration(ucc_knomial_pattern_t *p)
{
p->iteration++;
p->radix_pow *= p->radix;
if (p->is_mixed && !ucc_knomial_pattern_loop_done(p)) {
p->radix = p->radices[p->iteration];
}
}

static inline void ucc_knomial_pattern_prev_iteration(ucc_knomial_pattern_t *p)
Expand Down
35 changes: 35 additions & 0 deletions src/coll_patterns/sra_knomial.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
static inline
ucc_rank_t ucc_kn_compute_step_radix(ucc_knomial_pattern_t *p)
{
if (p->is_mixed) {
return p->radix;
}
int n_full = ucc_kn_pattern_n_full(p);

return p->radix_pow * p->radix >= p->size ? (n_full > 1 ? n_full : p->radix)
Expand Down Expand Up @@ -272,6 +275,38 @@ ucc_kn_ag_pattern_init(ucc_rank_t size, ucc_rank_t rank, ucc_kn_radix_t radix,
p->block_size;
}

static inline void ucc_kn_ag_pattern_reset(ucc_knomial_pattern_t *p)
{
p->iteration = 0;
p->radix = p->radices[0];
p->radix_pow = 1;
p->block_size = p->radix;
p->block_offset = ucc_knomial_pattern_loop_rank(p, p->rank) /
p->block_size * p->block_size;
}

static inline void ucc_kn_ag_pattern_init_mixed(
ucc_rank_t size, ucc_rank_t rank, size_t count,
const ucc_kn_radix_t *radices, uint8_t nradices, ucc_knomial_pattern_t *p)
{
uint8_t i;

ucc_assert(nradices > 0 && nradices <= UCC_KN_MAX_RADIX_PHASES);
p->type = KN_PATTERN_ALLGATHER;
p->n_iters = nradices;
p->node_type = KN_NODE_BASE;
p->backward = 0;
p->size = size;
p->rank = rank;
p->n_extra = 0;
p->count = count;
p->is_mixed = 1;
for (i = 0; i < nradices; i++) {
p->radices[i] = radices[i];
}
ucc_kn_ag_pattern_reset(p);
}

static inline void
ucc_kn_agx_pattern_init(ucc_rank_t size, ucc_rank_t rank, ucc_kn_radix_t radix,
size_t count, ucc_knomial_pattern_t *p)
Expand Down
1 change: 1 addition & 0 deletions src/components/tl/ucp/Makefile.am
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ endif

allgather = \
allgather/allgather.h \
allgather/allgather_knomial_schedule.c \
allgather/allgather.c \
allgather/allgather_ring.c \
allgather/allgather_neighbor.c \
Expand Down
4 changes: 4 additions & 0 deletions src/components/tl/ucp/allgather/allgather.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,4 +100,8 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_init(ucc_base_coll_args_t *coll_args,
ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
ucc_base_coll_args_t *coll_args, ucc_base_team_t *team,
ucc_coll_task_t **task_h, ucc_kn_radix_t radix);

ucc_status_t ucc_tl_ucp_allgather_knomial_parse_radices(
const char *value, ucc_rank_t team_size, ucc_kn_radix_t *radices,
uint8_t *nradices);
#endif
61 changes: 55 additions & 6 deletions src/components/tl/ucp/allgather/allgather_knomial.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "tl_ucp_sendrecv.h"
#include "tl_ucp_copy.h"
#include "core/ucc_progress_queue.h"
#include "allgather.h"
#include "coll_patterns/sra_knomial.h"
#include "tl_ucp_task.h"
#include "ucc/api/ucc.h"
Expand Down Expand Up @@ -198,6 +199,7 @@ void ucc_tl_ucp_allgather_knomial_progress(ucc_coll_task_t *coll_task)
return;
}
ucc_kn_ag_pattern_next_iter(p);
radix = p->radix;
}

if (KN_NODE_PROXY == node_type) {
Expand Down Expand Up @@ -243,8 +245,11 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_start(ucc_coll_task_t *coll_task)
task->allgather_kn.copy_task = NULL;
task->allgather_kn.phase = UCC_KN_PHASE_INIT;
if (ct == UCC_COLL_TYPE_ALLGATHER) {
ucc_kn_ag_pattern_init(size, rank, radix, args->dst.info.count,
&task->allgather_kn.p);
if (p->is_mixed) {
ucc_kn_ag_pattern_reset(p);
} else {
ucc_kn_ag_pattern_init(size, rank, radix, args->dst.info.count, p);
}
Comment thread
jeffnvidia marked this conversation as resolved.
Outdated
offset = ucc_buffer_block_offset(args->dst.info.count, size, rank) *
ucc_dt_size(args->dst.info.datatype);
rbuf = args->dst.info.buffer;
Expand Down Expand Up @@ -306,9 +311,10 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_start(ucc_coll_task_t *coll_task)
return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super);
}

ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
static ucc_status_t ucc_tl_ucp_allgather_knomial_init_common(
ucc_base_coll_args_t *coll_args, ucc_base_team_t *team,
ucc_coll_task_t **task_h, ucc_kn_radix_t radix)
ucc_coll_task_t **task_h, ucc_kn_radix_t radix,
const ucc_kn_radix_t *radices, uint8_t nradices)
{
ucc_tl_ucp_team_t *tl_team = ucc_derived_of(team, ucc_tl_ucp_team_t);
ucc_tl_ucp_context_t *ctx = UCC_TL_UCP_TEAM_CTX(tl_team);
Expand All @@ -323,7 +329,20 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
task->subset.myrank = sbgp->group_rank;
task->subset.map = sbgp->map;
}
task->allgather_kn.p.radix = radix;
if (nradices > 0) {
ucc_rank_t tsize = UCC_TL_TEAM_SIZE(tl_team);

ucc_kn_ag_pattern_init_mixed(
tsize,
task->subset.myrank,
GET_TOTAL_COUNT(&coll_args->args, tsize),
radices,
nradices,
&task->allgather_kn.p);
} else {
task->allgather_kn.p.radix = radix;
task->allgather_kn.p.is_mixed = 0;
}
if (!UCC_IS_INPLACE(coll_args->args)) {
if (ctx->cfg.local_copy_type == UCC_TL_UCP_LOCAL_COPY_TYPE_EC) {
task->super.flags |= UCC_COLL_TASK_FLAG_EXECUTOR;
Expand Down Expand Up @@ -357,6 +376,14 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
return UCC_OK;
}

ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
ucc_base_coll_args_t *coll_args, ucc_base_team_t *team,
ucc_coll_task_t **task_h, ucc_kn_radix_t radix)
{
return ucc_tl_ucp_allgather_knomial_init_common(
coll_args, team, task_h, radix, NULL, 0);
}

ucc_status_t ucc_tl_ucp_allgather_knomial_init(ucc_base_coll_args_t *coll_args,
ucc_base_team_t *team,
ucc_coll_task_t **task_h)
Expand All @@ -368,8 +395,30 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_init(ucc_base_coll_args_t *coll_args,
size_t count = GET_TOTAL_COUNT(&coll_args->args, tsize);
ucc_datatype_t dtype = GET_DT(&coll_args->args);
ucc_kn_radix_t radix;
ucc_kn_radix_t radices[UCC_KN_MAX_RADIX_PHASES];
uint8_t nradices = 0;
ucc_status_t status;

radix = ucc_tl_ucp_get_knomial_radix(tl_team, count, dtype, mtype, p, 0);
if (coll_args->args.coll_type != UCC_COLL_TYPE_ALLGATHER) {
return ucc_tl_ucp_allgather_knomial_init_r(
coll_args, team, task_h, radix);
}

if (tl_team->cfg.allgather_kn_mixed_radices != NULL &&
Comment thread
jeffnvidia marked this conversation as resolved.
Outdated
tl_team->cfg.allgather_kn_mixed_radices[0] != '\0') {
status = ucc_tl_ucp_allgather_knomial_parse_radices(
tl_team->cfg.allgather_kn_mixed_radices, tsize, radices, &nradices);
if (status != UCC_OK) {
tl_error(
UCC_TL_TEAM_LIB(tl_team),
"invalid ALLGATHER_KN_MIXED_RADICES '%s' for team size %u",
tl_team->cfg.allgather_kn_mixed_radices,
tsize);
return status;
}
}

return ucc_tl_ucp_allgather_knomial_init_r(coll_args, team, task_h, radix);
return ucc_tl_ucp_allgather_knomial_init_common(
coll_args, team, task_h, radix, radices, nradices);
}
50 changes: 50 additions & 0 deletions src/components/tl/ucp/allgather/allgather_knomial_schedule.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
/**
* Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* See file LICENSE for terms.
*/

#include "allgather.h"

#include <stdlib.h>

ucc_status_t ucc_tl_ucp_allgather_knomial_parse_radices(
Comment thread
jeffnvidia marked this conversation as resolved.
Outdated
const char *value, ucc_rank_t team_size, ucc_kn_radix_t *radices,
uint8_t *nradices)
{
const char *p = value;
char *end;
unsigned long parsed;
ucc_rank_t product = 1;
uint8_t n = 0;

*nradices = 0;
if (value == NULL || value[0] == '\0') {
return UCC_ERR_NOT_FOUND;
}

while (*p != '\0') {
if (n == UCC_KN_MAX_RADIX_PHASES) {
return UCC_ERR_INVALID_PARAM;
}
parsed = strtoul(p, &end, 10);
if (end == p || parsed < 2 || parsed > UINT16_MAX ||
product > UCC_RANK_MAX / parsed) {
return UCC_ERR_INVALID_PARAM;
}
radices[n++] = (ucc_kn_radix_t)parsed;
product *= (ucc_rank_t)parsed;
if (*end == '\0') {
break;
}
if (*end != ',' || end[1] == '\0') {
return UCC_ERR_INVALID_PARAM;
}
p = end + 1;
}

if (product != team_size) {
return UCC_ERR_INVALID_PARAM;
}
*nradices = n;
return UCC_OK;
}
7 changes: 7 additions & 0 deletions src/components/tl/ucp/tl_ucp.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,13 @@ ucc_config_field_t ucc_tl_ucp_lib_config_table[] = {
ucc_offsetof(ucc_tl_ucp_lib_config_t, allgather_kn_radix),
UCC_CONFIG_TYPE_UINT_RANGED},

{"ALLGATHER_KN_MIXED_RADICES", "",
"Optional exact mixed-radix schedule for knomial allgather, for example "

@Sergei-Lebedev Sergei-Lebedev Aug 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i would maybe change semantic here, instead of exact decompostion consider it as allowed factors in decomposition so it generalises to any team size in a run.

@jeffnvidia jeffnvidia Aug 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

don't we want the option for the user to provide the exact order ? (between 8x6 and 6x8 for example) I'm thinking that maybe the user would for example want to start with the first radix being the number of ranks in one node.

In the second PR, I provide an auto tuner that selects the radices and the order anyways but if the user chooses to select, I was thinking we ggive him "full control"

@Sergei-Lebedev Sergei-Lebedev Aug 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree, but at the same time if we consider this config as exact decomposition then it is less useful. If context creates multiple teams of different size then most of them will skip this value. Maybe extend with extra "=", for instance
RADIX="2x3x4" tells UCC to decompose using provided values
RADIX="=2x3x4" use exact decomposition and works for team size 24 only

@jeffnvidia jeffnvidia Aug 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I kept the x-separated value as an exact ordered sequence. We already have auto for team-size-independent selection and the explicit sequence for users who need exact ordering. To handle contexts containing teams of different sizes, a sequence-product mismatch now falls back to auto instead of failing. This addresses the multi-team concern without introducing a third allowed-factors mode or additional = syntax.

"4,4,6. Every radix must be at least 2 and their product must equal the "
"team size. Invalid schedules are rejected",
ucc_offsetof(ucc_tl_ucp_lib_config_t, allgather_kn_mixed_radices),
UCC_CONFIG_TYPE_STRING},

Comment thread
jeffnvidia marked this conversation as resolved.
Outdated
{"BCAST_KN_RADIX", "4", "Radix of the recursive-knomial bcast algorithm",
ucc_offsetof(ucc_tl_ucp_lib_config_t, bcast_kn_radix),
UCC_CONFIG_TYPE_UINT},
Expand Down
1 change: 1 addition & 0 deletions src/components/tl/ucp/tl_ucp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ typedef struct ucc_tl_ucp_lib_config {
ucc_mrange_uint_t allreduce_sra_kn_radix;
uint32_t reduce_scatter_kn_radix;
ucc_mrange_uint_t allgather_kn_radix;
char *allgather_kn_mixed_radices;
uint32_t bcast_kn_radix;
ucc_mrange_uint_t bcast_sag_kn_radix;
uint32_t reduce_kn_radix;
Expand Down
6 changes: 6 additions & 0 deletions test/gtest/Makefile.am
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,12 @@ gtest_LDFLAGS += $(UCX_LDFLAGS)
gtest_LDADD += $(UCX_LIBS) $(UCX_LIBADD)
endif

if TL_UCP_ENABLED

@Sergei-Lebedev Sergei-Lebedev Sep 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why test knomial_schedule depends on TL UCP?

@jeffnvidia jeffnvidia Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It no longer depends on TL/UCP. The test was renamed to test_knomial_seq.cc, uses only the generic K-nomial pattern code, and is now included unconditionally in the gtest sources.

gtest_SOURCES += \
coll/test_knomial_schedule.cc \
$(top_srcdir)/src/components/tl/ucp/allgather/allgather_knomial_schedule.c
endif

noinst_HEADERS = \
common/gtest.h \
common/test.h \
Expand Down
Loading
Loading