TL/CUDA: Linear Broadcast for GPU

src/components/tl/cuda/Makefile.am

@@ -27,6 +27,11 @@ alltoallv = \
  alltoallv/alltoallv.c \
  alltoallv/alltoallv_ce.c
 
+bcast = \
+ bcast/bcast.h \
+ bcast/bcast.c \
+ bcast/bcast_linear.c
+
 reduce_scatter = \
  reduce_scatter/reduce_scatter.h \
  reduce_scatter/reduce_scatter.c \
@@ -54,6 +59,7 @@ sources = \
  $(allgatherv) \
  $(alltoall) \
  $(alltoallv) \
+ $(bcast) \
  $(reduce_scatter) \
  $(reduce_scatterv)
 

src/components/tl/cuda/allgather/allgather.c

@@ -44,7 +44,7 @@ ucc_status_t ucc_tl_cuda_allgather_init(ucc_base_coll_args_t *coll_args,
 {
  ucc_tl_cuda_team_t *team = ucc_derived_of(tl_team, ucc_tl_cuda_team_t);
 
- if (ucc_tl_cuda_team_topo_is_fully_conntected(team->topo)) {
+ if (ucc_tl_cuda_team_topo_is_fully_connected(team->topo)) {
  return ucc_tl_cuda_allgather_linear_init(coll_args, tl_team, task_p);
  } else {
  return ucc_tl_cuda_allgather_ring_init(coll_args, tl_team, task_p);

src/components/tl/cuda/allgather/allgather_linear.c

@@ -15,7 +15,7 @@ ucc_status_t ucc_tl_cuda_allgather_linear_init(ucc_base_coll_args_t *coll_args,
  ucc_tl_cuda_task_t *task;
  ucc_status_t status;
 
- if (ucc_unlikely(!ucc_tl_cuda_team_topo_is_fully_conntected(team->topo) ||
+ if (ucc_unlikely(!ucc_tl_cuda_team_topo_is_fully_connected(team->topo) ||
  UCC_TL_TEAM_SIZE(team) - 1 > UCC_EE_EXECUTOR_MULTI_OP_NUM_BUFS)) {
  return UCC_ERR_NOT_SUPPORTED;
  }

src/components/tl/cuda/allgatherv/allgatherv.c

@@ -47,7 +47,7 @@ ucc_status_t ucc_tl_cuda_allgatherv_init(ucc_base_coll_args_t *coll_args,
 {
  ucc_tl_cuda_team_t *team = ucc_derived_of(tl_team, ucc_tl_cuda_team_t);
 
- if (ucc_tl_cuda_team_topo_is_fully_conntected(team->topo)) {
+ if (ucc_tl_cuda_team_topo_is_fully_connected(team->topo)) {
  return ucc_tl_cuda_allgatherv_linear_init(coll_args, tl_team, task_p);
  } else {
  return ucc_tl_cuda_allgatherv_ring_init(coll_args, tl_team, task_p);

src/components/tl/cuda/allgatherv/allgatherv_linear.c

@@ -55,22 +55,6 @@ enum
  * other ranks to finish */
 };
 
-static inline int get_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- return sync->seq_num[step_id];
-}
-
-static inline void set_rank_step(ucc_tl_cuda_task_t *task, ucc_rank_t rank,
- int step, int step_id)
-{
- ucc_tl_cuda_sync_t *sync = TASK_SYNC(task, rank);
-
- sync->seq_num[step_id] = step;
-}
-
 ucc_status_t ucc_tl_cuda_allgatherv_linear_finalize(ucc_coll_task_t *coll_task)
 {
  ucc_tl_cuda_task_t *task = ucc_derived_of(coll_task, ucc_tl_cuda_task_t);
@@ -432,7 +416,7 @@ ucc_status_t ucc_tl_cuda_allgatherv_linear_init(ucc_base_coll_args_t *coll_args,
  ucc_tl_cuda_task_t *task;
  ucc_status_t status;
 
- if (ucc_unlikely(!ucc_tl_cuda_team_topo_is_fully_conntected(team->topo) ||
+ if (ucc_unlikely(!ucc_tl_cuda_team_topo_is_fully_connected(team->topo) ||
  UCC_TL_TEAM_SIZE(team) - 1 > UCC_EE_EXECUTOR_MULTI_OP_NUM_BUFS)) {
  return UCC_ERR_NOT_SUPPORTED;
  }

src/components/tl/cuda/bcast/bcast.c

@@ -0,0 +1,28 @@
+/**
+ * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See file LICENSE for terms.
+ */
+
+#include "bcast.h"
+#include "components/mc/ucc_mc.h"
+
+ucc_base_coll_alg_info_t
+ ucc_tl_cuda_bcast_algs[UCC_TL_CUDA_BCAST_ALG_LAST + 1] = {
+ [UCC_TL_CUDA_BCAST_ALG_LINEAR] = {.id = UCC_TL_CUDA_BCAST_ALG_LINEAR,
+ .name = "linear",
+ .desc = "linear bcast algorithm"},
+ [UCC_TL_CUDA_BCAST_ALG_LAST] = {.id = 0, .name = NULL, .desc = NULL}};
+
+ucc_status_t ucc_tl_cuda_bcast_init(ucc_base_coll_args_t *coll_args,
+ ucc_base_team_t *tl_team,
+ ucc_coll_task_t **task_p)
+{
+ ucc_tl_cuda_team_t *team = ucc_derived_of(tl_team, ucc_tl_cuda_team_t);
+
+ if (ucc_tl_cuda_team_topo_is_fully_connected(team->topo)) {
+ return ucc_tl_cuda_bcast_linear_init(coll_args, tl_team, task_p);
+ } else {
+ return UCC_ERR_NOT_SUPPORTED;
+ }
+}

src/components/tl/cuda/bcast/bcast.h

@@ -0,0 +1,43 @@
+/**
+ * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See file LICENSE for terms.
+ */
+
+#ifndef BCAST_H_
+#define BCAST_H_
+
+#include "tl_cuda.h"
+#include "tl_cuda_coll.h"
+
+enum
+{
+ UCC_TL_CUDA_BCAST_ALG_LINEAR,
+ UCC_TL_CUDA_BCAST_ALG_LAST
+};
+
+extern ucc_base_coll_alg_info_t
+ ucc_tl_cuda_bcast_algs[UCC_TL_CUDA_BCAST_ALG_LAST + 1];
+
+#define UCC_TL_CUDA_BCAST_DEFAULT_ALG_SELECT_STR "bcast:cuda:@0"
+
+ucc_status_t ucc_tl_cuda_bcast_init(ucc_base_coll_args_t *coll_args,
+ ucc_base_team_t *tl_team,
+ ucc_coll_task_t **task_p);
+
+ucc_status_t ucc_tl_cuda_bcast_linear_init(ucc_base_coll_args_t *coll_args,
+ ucc_base_team_t *tl_team,
+ ucc_coll_task_t **task_p);
+
+static inline int ucc_tl_cuda_bcast_alg_from_str(const char *str)
+{
+ int i;
+ for (i = 0; i < UCC_TL_CUDA_BCAST_ALG_LAST; i++) {
+ if (0 == strcasecmp(str, ucc_tl_cuda_bcast_algs[i].name)) {
+ break;
+ }
+ }
+ return i;
+}
+
+#endif