TL/UCP: add exact mixed-radix allgather schedules

src/coll_patterns/recursive_knomial.h

@@ -7,8 +7,9 @@
 #ifndef RECURSIVE_KNOMIAL_H_
 #define RECURSIVE_KNOMIAL_H_
 
+#include "utils/ucc_datastruct.h"
+
 #define UCC_KN_PEER_NULL ((ucc_rank_t)-1)
-typedef uint16_t ucc_kn_radix_t;
 
 enum {
     KN_NODE_BASE,  /* Participates in the main loop of the recursive KN algorithm */
@@ -29,7 +30,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 */
@@ -54,6 +55,7 @@ typedef struct ucc_knomial_pattern {
     ucc_rank_t     block_size;
     ptrdiff_t      block_offset;
     int            is64;
+    const ucc_kn_radix_t *radices;
 } ucc_knomial_pattern_t;
 
 /**
@@ -101,6 +103,7 @@ ucc_knomial_pattern_init_impl(ucc_rank_t size, ucc_rank_t rank,
     p->rank          = rank;
     p->backward      = backward;
     p->iteration     = 0;
+    p->radices       = NULL;
     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) ?
@@ -229,6 +232,9 @@ ucc_knomial_pattern_next_iteration(ucc_knomial_pattern_t *p)
 {
     p->iteration++;
     p->radix_pow *= p->radix;
+    if (p->radices && !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)

src/coll_patterns/sra_knomial.h

@@ -261,15 +261,37 @@ static inline void ucc_kn_g_pattern_next_iter(ucc_knomial_pattern_t *p)
 }
 
 static inline void
-ucc_kn_ag_pattern_init(ucc_rank_t size, ucc_rank_t rank, ucc_kn_radix_t radix,
+ucc_kn_ag_pattern_init(ucc_rank_t size, ucc_rank_t rank,
+                       const ucc_kn_radix_t *radices, uint8_t nradices,
                        size_t count, ucc_knomial_pattern_t *p)
 {
-    ucc_knomial_pattern_init(size, rank, radix, p);
+    ucc_assert(nradices > 0 && nradices <= UCC_KN_MAX_RADIX_PHASES);
+    if (nradices == 1) {
+        ucc_knomial_pattern_init(size, rank, radices[0], p);
+        p->type         = KN_PATTERN_ALLGATHER;
+        p->count        = count;
+        p->block_size   = p->radix_pow * p->radix;
+        p->block_offset = ucc_knomial_pattern_loop_rank(p, rank) /
+                          p->block_size * p->block_size;
+        return;
+    }
+
+    p->radix        = radices[0];
     p->type         = KN_PATTERN_ALLGATHER;
+    p->iteration    = 0;
+    p->n_iters      = nradices;
+    p->node_type    = KN_NODE_BASE;
+    p->backward     = 0;
+    p->radix_pow    = 1;
+    p->full_pow_size = size;
+    p->size         = size;
+    p->rank         = rank;
+    p->n_extra      = 0;
     p->count        = count;
-    p->block_size   = p->radix_pow * radix;
-    p->block_offset = ucc_knomial_pattern_loop_rank(p, rank) / p->block_size *
-                      p->block_size;
+    p->block_size   = p->radix;
+    p->block_offset = ucc_knomial_pattern_loop_rank(p, rank) /
+                      p->block_size * p->block_size;
+    p->radices      = radices;
 }
 
 static inline void

src/components/tl/ucp/allgather/allgather_knomial.c

@@ -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"
@@ -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) {
@@ -243,8 +245,10 @@ 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);
+        ucc_kn_ag_pattern_init(size, rank,
+                               p->radices ? p->radices : &radix,
+                               p->radices ? p->n_iters : 1,
+                               args->dst.info.count, p);
         offset = ucc_buffer_block_offset(args->dst.info.count, size, rank) *
                  ucc_dt_size(args->dst.info.datatype);
         rbuf   = args->dst.info.buffer;
@@ -306,9 +310,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_schedule_t *schedule)
 {
     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);
@@ -324,6 +329,13 @@ ucc_status_t ucc_tl_ucp_allgather_knomial_init_r(
         task->subset.map    = sbgp->map;
     }
     task->allgather_kn.p.radix = radix;
+    if (schedule && schedule->n_radices > 1) {
+        task->allgather_kn.p.radices = schedule->radices;
+        task->allgather_kn.p.n_iters = schedule->n_radices;
+    } else {
+        task->allgather_kn.p.radices = NULL;
+        task->allgather_kn.p.n_iters = 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;
@@ -357,19 +369,64 @@ 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);
+}
+
+static int ucc_tl_ucp_allgather_knomial_schedule_matches(
+    const ucc_kn_radix_schedule_t *schedule, ucc_rank_t team_size)
+{
+    ucc_rank_t product = 1;
+    uint8_t    i;
+
+    if (schedule->n_radices <= 1) {
+        return 1;
+    }
+    for (i = 0; i < schedule->n_radices; i++) {
+        if (product > team_size / schedule->radices[i]) {
+            return 0;
+        }
+        product *= schedule->radices[i];
+    }
+    return product == team_size;
+}
+
 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)
 {
     ucc_tl_ucp_team_t *tl_team = ucc_derived_of(team, ucc_tl_ucp_team_t);
-    ucc_mrange_uint_t *p       = &tl_team->cfg.allgather_kn_radix;
     ucc_rank_t         tsize   = UCC_TL_TEAM_SIZE(tl_team);
     ucc_memory_type_t  mtype   = GET_MT(&coll_args->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;
-
-    radix = ucc_tl_ucp_get_knomial_radix(tl_team, count, dtype, mtype, p, 0);
-
-    return ucc_tl_ucp_allgather_knomial_init_r(coll_args, team, task_h, radix);
+    size_t             msgsize = count * ucc_dt_size(dtype);
+    const ucc_kn_radix_schedule_t *schedule;
+    ucc_kn_radix_t                 radix;
+
+    schedule = ucc_mrange_kn_radix_get(&tl_team->cfg.allgather_kn_radix,
+                                       msgsize, mtype);
+    if (schedule->n_radices == 0) {
+        radix = mtype == UCC_MEMORY_TYPE_HOST ? tl_team->opt_radix_host :
+                                                tl_team->opt_radix;
+    } else {
+        radix = schedule->radices[0];
+    }
+    radix = ucc_min(radix, tsize);
+    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 (!ucc_tl_ucp_allgather_knomial_schedule_matches(schedule, tsize)) {
+        tl_error(UCC_TL_TEAM_LIB(tl_team),
+                 "ALLGATHER_KN_RADIX schedule product must equal team size %u",
+                 tsize);
+        return UCC_ERR_INVALID_PARAM;
+    }
+    return ucc_tl_ucp_allgather_knomial_init_common(
+        coll_args, team, task_h, radix, schedule);
 }

src/components/tl/ucp/tl_ucp.c

@@ -168,9 +168,11 @@ ucc_config_field_t ucc_tl_ucp_lib_config_table[] = {
      ucc_offsetof(ucc_tl_ucp_lib_config_t, reduce_scatter_kn_radix),
      UCC_CONFIG_TYPE_UINT},
 
-    {"ALLGATHER_KN_RADIX", "auto", "Radix of the knomial allgather algorithm",
+    {"ALLGATHER_KN_RADIX", "auto",
+     "Radix or ordered exact radix schedule of the knomial allgather "
+     "algorithm, for example 4 or 8x6",
      ucc_offsetof(ucc_tl_ucp_lib_config_t, allgather_kn_radix),
-     UCC_CONFIG_TYPE_UINT_RANGED},
+     UCC_CONFIG_TYPE_KN_RADIX},
 
     {"BCAST_KN_RADIX", "4", "Radix of the recursive-knomial bcast algorithm",
      ucc_offsetof(ucc_tl_ucp_lib_config_t, bcast_kn_radix),

src/components/tl/ucp/tl_ucp.h

@@ -62,7 +62,7 @@ typedef struct ucc_tl_ucp_lib_config {
     ucc_mrange_uint_t                  allreduce_kn_radix;
     ucc_mrange_uint_t                  allreduce_sra_kn_radix;
     uint32_t                           reduce_scatter_kn_radix;
-    ucc_mrange_uint_t                  allgather_kn_radix;
+    ucc_mrange_kn_radix_t              allgather_kn_radix;
     uint32_t                           bcast_kn_radix;
     ucc_mrange_uint_t                  bcast_sag_kn_radix;
     uint32_t                           reduce_kn_radix;

src/utils/ucc_datastruct.c

@@ -45,3 +45,32 @@ void ucc_mrange_uint_destroy(ucc_mrange_uint_t *param)
         ucc_free(r);
     }
 }
+
+ucc_status_t ucc_mrange_kn_radix_copy(ucc_mrange_kn_radix_t       *dst,
+                                      const ucc_mrange_kn_radix_t *src)
+{
+    ucc_mrange_kn_radix_entry_t *r, *r_dup;
+
+    dst->default_value = src->default_value;
+    ucc_list_head_init(&dst->ranges);
+    ucc_list_for_each(r, &src->ranges, list_elem) {
+        r_dup = ucc_malloc(sizeof(*r_dup), "kn radix range dup");
+        if (ucc_unlikely(!r_dup)) {
+            ucc_mrange_kn_radix_destroy(dst);
+            return UCC_ERR_NO_MEMORY;
+        }
+        *r_dup = *r;
+        ucc_list_add_tail(&dst->ranges, &r_dup->list_elem);
+    }
+    return UCC_OK;
+}
+
+void ucc_mrange_kn_radix_destroy(ucc_mrange_kn_radix_t *param)
+{
+    ucc_mrange_kn_radix_entry_t *r, *r_tmp;
+
+    ucc_list_for_each_safe(r, r_tmp, &param->ranges, list_elem) {
+        ucc_list_del(&r->list_elem);
+        ucc_free(r);
+    }
+}

src/utils/ucc_datastruct.h

@@ -7,13 +7,22 @@
 #define UCC_DATASTRUCT_H_
 #include <ucc/api/ucc.h>
 #include "ucc_list.h"
+#include <limits.h>
 #include <stdint.h>
 
 #define UCC_LIST_HEAD UCS_LIST_HEAD
 typedef uint32_t ucc_rank_t;
 #define UCC_RANK_INVALID UINT32_MAX
 #define UCC_RANK_MAX UCC_RANK_INVALID - 1
 
+typedef uint16_t ucc_kn_radix_t;
+#define UCC_KN_MAX_RADIX_PHASES (sizeof(ucc_rank_t) * CHAR_BIT - 1)
+
+typedef struct ucc_kn_radix_schedule {
+    ucc_kn_radix_t radices[UCC_KN_MAX_RADIX_PHASES];
+    uint8_t        n_radices;
+} ucc_kn_radix_schedule_t;
+
 typedef struct ucc_subset {
     ucc_ep_map_t map;
     ucc_rank_t   myrank;
@@ -37,11 +46,29 @@ typedef struct ucc_mrange_uint {
     unsigned        default_value;
 } ucc_mrange_uint_t;
 
+typedef struct ucc_mrange_kn_radix_entry {
+    ucc_list_link_t         list_elem;
+    size_t                  start;
+    size_t                  end;
+    uint32_t                mtypes;
+    ucc_kn_radix_schedule_t value;
+} ucc_mrange_kn_radix_entry_t;
+
+typedef struct ucc_mrange_kn_radix {
+    ucc_list_link_t         ranges;
+    ucc_kn_radix_schedule_t default_value;
+} ucc_mrange_kn_radix_t;
+
 ucc_status_t ucc_mrange_uint_copy(ucc_mrange_uint_t       *dst,
                                   const ucc_mrange_uint_t *src);
 
 void ucc_mrange_uint_destroy(ucc_mrange_uint_t *param);
 
+ucc_status_t ucc_mrange_kn_radix_copy(ucc_mrange_kn_radix_t       *dst,
+                                      const ucc_mrange_kn_radix_t *src);
+
+void ucc_mrange_kn_radix_destroy(ucc_mrange_kn_radix_t *param);
+
 static inline unsigned ucc_mrange_uint_get(
     ucc_mrange_uint_t *param, size_t range_value, ucc_memory_type_t mem_type)
 {
@@ -56,4 +83,19 @@ static inline unsigned ucc_mrange_uint_get(
     return param->default_value;
 }
 
+static inline const ucc_kn_radix_schedule_t *ucc_mrange_kn_radix_get(
+    const ucc_mrange_kn_radix_t *param, size_t range_value,
+    ucc_memory_type_t mem_type)
+{
+    ucc_mrange_kn_radix_entry_t *r;
+
+    ucc_list_for_each(r, &param->ranges, list_elem) {
+        if (r->start <= range_value && range_value <= r->end &&
+            (UCC_BIT(mem_type) & r->mtypes)) {
+            return &r->value;
+        }
+    }
+    return &param->default_value;
+}
+
 #endif