TL/UCP: Add stubs for sliding window allreduce

nsarka · Jan 29, 2024 · 06fc8d3 · 06fc8d3
1 parent c9f1286
commit 06fc8d3
Show file tree

Hide file tree

Showing 15 changed files with 593 additions and 57 deletions.
diff --git a/src/components/tl/ucp/Makefile.am b/src/components/tl/ucp/Makefile.am
@@ -36,10 +36,12 @@ alltoallv =                        \
 	alltoallv/alltoallv_onesided.c
 
 allreduce =                           \
-	allreduce/allreduce.h             \
-	allreduce/allreduce.c             \
-	allreduce/allreduce_knomial.c     \
-	allreduce/allreduce_sra_knomial.c \
+	allreduce/allreduce.h                      \
+	allreduce/allreduce.c                      \
+	allreduce/allreduce_knomial.c              \
+	allreduce/allreduce_sra_knomial.c          \
+	allreduce/allreduce_sliding_window.c       \
+	allreduce/allreduce_sliding_window_setup.c \
 	allreduce/allreduce_dbt.c
 
 barrier =                     \

diff --git a/src/components/tl/ucp/allreduce/allreduce.c b/src/components/tl/ucp/allreduce/allreduce.c
@@ -25,6 +25,10 @@ ucc_base_coll_alg_info_t
              .name = "dbt",
              .desc = "alreduce over double binary tree where a leaf in one tree "
                      "will be intermediate in other (optimized for BW)"},
+        [UCC_TL_UCP_ALLREDUCE_ALG_SLIDING_WINDOW] =
+            {.id   = UCC_TL_UCP_ALLREDUCE_ALG_SLIDING_WINDOW,
+             .name = "sliding_window",
+             .desc = "sliding window allreduce (optimized for running on DPU)"},
         [UCC_TL_UCP_ALLREDUCE_ALG_LAST] = {
             .id = 0, .name = NULL, .desc = NULL}};
 
@@ -51,3 +55,11 @@ ucc_status_t ucc_tl_ucp_allreduce_knomial_init(ucc_base_coll_args_t *coll_args,
 out:
     return status;
 }
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_init(ucc_base_coll_args_t *coll_args,
+                                         ucc_base_team_t *     team,
+                                         ucc_coll_task_t **    task_h)
+{
+    return UCC_OK;
+}
diff --git a/src/components/tl/ucp/allreduce/allreduce.h b/src/components/tl/ucp/allreduce/allreduce.h
@@ -11,6 +11,7 @@
 enum {
     UCC_TL_UCP_ALLREDUCE_ALG_KNOMIAL,
     UCC_TL_UCP_ALLREDUCE_ALG_SRA_KNOMIAL,
+    UCC_TL_UCP_ALLREDUCE_ALG_SLIDING_WINDOW,
     UCC_TL_UCP_ALLREDUCE_ALG_DBT,
     UCC_TL_UCP_ALLREDUCE_ALG_LAST
 };
@@ -36,16 +37,110 @@ ucc_status_t ucc_tl_ucp_allreduce_init(ucc_tl_ucp_task_t *task);
 #define ALLREDUCE_TASK_CHECK(_args, _team)                                     \
     CHECK_SAME_MEMTYPE((_args), (_team));
 
+#define ALLREDUCE_PACKED_KEY_MAX_LEN 1024
+
+typedef struct ucc_tl_ucp_allreduce_sw_global_work_buf_info {
+    void *packed_src_memh;
+    void *packed_dst_memh;
+} ucc_tl_ucp_allreduce_sw_global_work_buf_info;
+
+typedef enum ucc_tl_ucp_allreduce_sw_buf_state
+{
+    FREE,
+    RECVING,
+    REDUCING,
+    REDUCED,
+    SENDING,
+    IDLE,
+} ucc_tl_ucp_allreduce_sw_buf_state;
+
+typedef struct ucc_tl_ucp_allreduce_sw_buf {
+    void *                            buf;
+    ucc_tl_ucp_allreduce_sw_buf_state state;
+    ucs_status_ptr_t                  ucp_req;
+    size_t                            count;
+    size_t                            bytes;
+} ucc_tl_ucp_allreduce_sw_buf;
+
+typedef struct ucc_tl_ucp_allreduce_sw_pipeline {
+    ucc_tl_ucp_allreduce_sw_buf  accbuf;
+    ucc_tl_ucp_allreduce_sw_buf *getbuf;
+    ucs_status_ptr_t *           put_requests;
+    size_t                       buffer_size;
+    size_t                       num_buffers;
+    size_t                       avail_buffs;
+    size_t                       my_count;
+    size_t                       my_offset;
+    size_t                       count_issued;
+    size_t                       count_received;
+    size_t                       count_reduced;
+    size_t                       count_serviced;
+    size_t                       get_idx;
+    size_t                       red_idx;
+    ucc_rank_t                   src_rank;
+    ucc_rank_t                   dst_rank;
+    int                          done_get;
+    int                          done_red;
+    int                          done_put;
+    int                          posted_put;
+} ucc_tl_ucp_allreduce_sw_pipeline;
+
+struct ucc_tl_ucp_allreduce_sw_export_buf {
+    ucp_context_h ucp_context;
+    ucp_mem_h     memh;
+    void *        packed_memh;
+    size_t        packed_memh_len;
+    void *        packed_key;
+    size_t        packed_key_len;
+    uint64_t      memh_id;
+};
+
+typedef struct ucc_tl_ucp_allreduce_sw_host_allgather {
+    void *src_buf;
+    void *dst_buf;
+    char  packed_src_key[ALLREDUCE_PACKED_KEY_MAX_LEN];
+    char  packed_dst_key[ALLREDUCE_PACKED_KEY_MAX_LEN];
+} ucc_tl_ucp_allreduce_sw_host_allgather;
+
 ucc_status_t ucc_tl_ucp_allreduce_knomial_init(ucc_base_coll_args_t *coll_args,
                                                ucc_base_team_t *team,
                                                ucc_coll_task_t **task_h);
 
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_init(ucc_base_coll_args_t *coll_args,
+                                         ucc_base_team_t *     team,
+                                         ucc_coll_task_t **    task_h);
+
 ucc_status_t ucc_tl_ucp_allreduce_knomial_init_common(ucc_tl_ucp_task_t *task);
 
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_alloc_pipe(ucc_base_coll_args_t *coll_args,
+                                               ucc_base_team_t *     team,
+                                               ucc_tl_ucp_task_t *   task);
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_task_init(ucc_base_coll_args_t *coll_args,
+                                              ucc_base_team_t *     team,
+                                              ucc_tl_ucp_task_t *   task);
+
+ucc_status_t ucc_tl_ucp_allreduce_sliding_window_allgather_info_finalize(
+    ucc_service_coll_req_t *scoll_req, ucc_tl_ucp_task_t *sw_task);
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_free_gwbi(ucc_coll_task_t *coll_task);
+
 ucc_status_t ucc_tl_ucp_allreduce_knomial_start(ucc_coll_task_t *task);
 
 void ucc_tl_ucp_allreduce_knomial_progress(ucc_coll_task_t *task);
 
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_start(ucc_coll_task_t *coll_task);
+
+void ucc_tl_ucp_allreduce_sliding_window_progress(ucc_coll_task_t *task);
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_finalize(ucc_coll_task_t *task);
+
 ucc_status_t ucc_tl_ucp_allreduce_knomial_finalize(ucc_coll_task_t *task);
 
 ucc_status_t ucc_tl_ucp_allreduce_sra_knomial_init(ucc_base_coll_args_t *coll_args,

diff --git a/src/components/tl/ucp/allreduce/allreduce_sliding_window.c b/src/components/tl/ucp/allreduce/allreduce_sliding_window.c
@@ -0,0 +1,71 @@
+/**
+ * Copyright(c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See file LICENSE for terms.
+ */
+
+#include "allreduce.h"
+#include "../allgather/allgather.h"
+#include "../barrier/barrier.h"
+#include "utils/ucc_dt_reduce.h"
+#include "tl_ucp_ep.h"
+
+static inline void
+ucc_tl_ucp_allreduce_sliding_window_reset_buf(ucc_tl_ucp_allreduce_sw_buf *buf)
+{
+}
+
+static inline void ucc_tl_ucp_allreduce_sliding_window_reset_pipeline(
+    ucc_tl_ucp_allreduce_sw_pipeline *pipe, ucc_rank_t rank,
+    size_t put_window_size)
+{
+}
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_start(ucc_coll_task_t *coll_task)
+{
+    return UCC_OK;
+}
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_finalize(ucc_coll_task_t *coll_task)
+{
+    return UCC_OK;
+}
+
+static inline void ucc_tl_ucp_allreduce_sliding_window_reduction(
+    ucc_coll_task_t *coll_task, ucc_tl_ucp_allreduce_sw_buf *accbuf,
+    ucc_tl_ucp_allreduce_sw_buf *getbuf)
+{
+}
+
+static inline void
+ucc_tl_ucp_allreduce_sliding_window_test_reduction(ucc_tl_ucp_task_t *task)
+{
+}
+
+static inline ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_req_test(ucs_status_ptr_t   request,
+                                             ucc_tl_ucp_task_t *task)
+{
+    return UCC_OK;
+}
+
+static inline void ucc_tl_ucp_allreduce_sliding_window_allgather_info_test(
+    ucc_coll_task_t *coll_task)
+{
+}
+
+static inline void ucc_tl_ucp_allreduce_sliding_window_allgather_free_rkeys(
+    ucc_coll_task_t *coll_task)
+{
+}
+
+static inline void
+ucc_tl_ucp_allreduce_sliding_window_barrier(ucc_coll_task_t *coll_task)
+{
+}
+
+void ucc_tl_ucp_allreduce_sliding_window_progress(ucc_coll_task_t *coll_task)
+{
+}
diff --git a/src/components/tl/ucp/allreduce/allreduce_sliding_window_setup.c b/src/components/tl/ucp/allreduce/allreduce_sliding_window_setup.c
@@ -0,0 +1,38 @@
+/**
+ * Copyright(c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See file LICENSE for terms.
+ */
+
+#include "allreduce.h"
+#include "../allgather/allgather.h"
+#include "utils/ucc_dt_reduce.h"
+#include "tl_ucp_ep.h"
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_task_init(ucc_base_coll_args_t *coll_args,
+                                              ucc_base_team_t *     team,
+                                              ucc_tl_ucp_task_t *   task)
+{
+    return UCC_OK;
+}
+
+ucc_status_t ucc_tl_ucp_allreduce_sliding_window_allgather_info_finalize(
+    ucc_service_coll_req_t *scoll_req, ucc_tl_ucp_task_t *sw_task)
+{
+    return UCC_OK;
+}
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_alloc_pipe(ucc_base_coll_args_t *coll_args,
+                                               ucc_base_team_t *     team,
+                                               ucc_tl_ucp_task_t *   task)
+{
+    return UCC_OK;
+}
+
+ucc_status_t
+ucc_tl_ucp_allreduce_sliding_window_free_gwbi(ucc_coll_task_t *coll_task)
+{
+    return UCC_OK;
+}
diff --git a/src/components/tl/ucp/barrier/barrier.c b/src/components/tl/ucp/barrier/barrier.c
@@ -7,9 +7,6 @@
 #include "tl_ucp.h"
 #include "barrier.h"
 
-ucc_status_t ucc_tl_ucp_barrier_knomial_start(ucc_coll_task_t *task);
-void ucc_tl_ucp_barrier_knomial_progress(ucc_coll_task_t *task);
-
 ucc_base_coll_alg_info_t
     ucc_tl_ucp_barrier_algs[UCC_TL_UCP_BARRIER_ALG_LAST + 1] = {
         [UCC_TL_UCP_BARRIER_ALG_KNOMIAL] =

diff --git a/src/components/tl/ucp/barrier/barrier.h b/src/components/tl/ucp/barrier/barrier.h
@@ -18,4 +18,7 @@ extern ucc_base_coll_alg_info_t
 
 ucc_status_t ucc_tl_ucp_barrier_init(ucc_tl_ucp_task_t *task);
 
+ucc_status_t ucc_tl_ucp_barrier_knomial_start(ucc_coll_task_t *task);
+void         ucc_tl_ucp_barrier_knomial_progress(ucc_coll_task_t *task);
+
 #endif
diff --git a/src/components/tl/ucp/tl_ucp.c b/src/components/tl/ucp/tl_ucp.c
@@ -104,6 +104,23 @@ ucc_config_field_t ucc_tl_ucp_lib_config_table[] = {
      ucc_offsetof(ucc_tl_ucp_lib_config_t, allreduce_kn_radix),
      UCC_CONFIG_TYPE_UINT_RANGED},
 
+    {"ALLREDUCE_SLIDING_WIN_BUF_SIZE", "65536",
+     "Buffer size of the sliding window allreduce algorithm",
+     ucc_offsetof(ucc_tl_ucp_lib_config_t, allreduce_sliding_window_buf_size),
+     UCC_CONFIG_TYPE_MEMUNITS},
+
+    {"ALLREDUCE_SLIDING_WIN_PUT_WINDOW_SIZE", "0",
+     "Max concurrent puts in SW Allreduce. <= 0 means set to team size",
+     ucc_offsetof(ucc_tl_ucp_lib_config_t,
+                  allreduce_sliding_window_put_window_size),
+     UCC_CONFIG_TYPE_UINT},
+
+    {"ALLREDUCE_SLIDING_WIN_NUM_GET_BUFS", "0",
+     "Number of get buffers for sliding window AR. <= 0 means set to team size",
+     ucc_offsetof(ucc_tl_ucp_lib_config_t,
+                  allreduce_sliding_window_num_get_bufs),
+     UCC_CONFIG_TYPE_UINT},
+
     {"ALLREDUCE_SRA_KN_RADIX", "auto",
      "Radix of the scatter-reduce-allgather (SRA) knomial allreduce algorithm",
      ucc_offsetof(ucc_tl_ucp_lib_config_t, allreduce_sra_kn_radix),

diff --git a/src/components/tl/ucp/tl_ucp.h b/src/components/tl/ucp/tl_ucp.h
@@ -48,6 +48,9 @@ typedef struct ucc_tl_ucp_lib_config {
     uint32_t                 fanin_kn_radix;
     uint32_t                 fanout_kn_radix;
     uint32_t                 barrier_kn_radix;
+    size_t                   allreduce_sliding_window_buf_size;
+    uint32_t                 allreduce_sliding_window_put_window_size;
+    uint32_t                 allreduce_sliding_window_num_get_bufs;
     ucc_mrange_uint_t        allreduce_kn_radix;
     ucc_mrange_uint_t        allreduce_sra_kn_radix;
     uint32_t                 reduce_scatter_kn_radix;

diff --git a/src/components/tl/ucp/tl_ucp_coll.c b/src/components/tl/ucp/tl_ucp_coll.c
@@ -278,6 +278,9 @@ ucc_status_t ucc_tl_ucp_alg_id_to_init(int alg_id, const char *alg_id_str,
         case UCC_TL_UCP_ALLREDUCE_ALG_DBT:
             *init = ucc_tl_ucp_allreduce_dbt_init;
             break;
+        case UCC_TL_UCP_ALLREDUCE_ALG_SLIDING_WINDOW:
+            *init = ucc_tl_ucp_allreduce_sliding_window_init;
+            break;
         default:
             status = UCC_ERR_INVALID_PARAM;
             break;

diff --git a/src/components/tl/ucp/tl_ucp_coll.h b/src/components/tl/ucp/tl_ucp_coll.h
@@ -88,6 +88,11 @@ enum ucc_tl_ucp_task_flags {
     UCC_TL_UCP_TASK_FLAG_SUBSET = UCC_BIT(0),
 };
 
+typedef struct ucc_tl_ucp_allreduce_sw_pipeline
+    ucc_tl_ucp_allreduce_sw_pipeline;
+typedef struct ucc_tl_ucp_allreduce_sw_host_allgather
+    ucc_tl_ucp_allreduce_sw_host_allgather;
+
 typedef struct ucc_tl_ucp_task {
     ucc_coll_task_t super;
     uint32_t        flags;
@@ -121,6 +126,27 @@ typedef struct ucc_tl_ucp_task {
             ucc_ee_executor_task_t *etask;
             ucc_ee_executor_t      *executor;
         } allreduce_kn;
+        struct {
+            int                                        reduce_in_progress;
+            ucp_rkey_h *                               src_rkeys; //unpacked
+            ucp_rkey_h *                               dst_rkeys; //unpacked
+            ucp_ep_h *                                 eps;
+            void **                                    sbufs;
+            void **                                    rbufs;
+            ucc_coll_task_t *                          allreduce_task_h;
+            ucc_tl_ucp_allreduce_sw_pipeline *         pipe;
+            ucc_ee_executor_task_t *                   etask;
+            ucc_ee_executor_t *                        executor;
+            int                                        put_window_size;
+            int                                        num_get_bufs;
+            ucs_status_ptr_t *                         put_requests;
+            ucc_service_coll_req_t *                   allgather_scoll_req;
+            ucc_tl_ucp_allreduce_sw_host_allgather *   allgather_data;
+            ucc_coll_task_t *                          barrier_task;
+            struct ucc_tl_ucp_allreduce_sw_export_buf *src_ebuf;
+            struct ucc_tl_ucp_allreduce_sw_export_buf *dst_ebuf;
+            int                                        inplace;
+        } allreduce_sliding_window;
         struct {
             int                     phase;
             ucc_knomial_pattern_t   p;