TL/UCP: Add all-reduce ring alogrithm

Signed-off-by: Armen Ratner <[email protected]>

Author: armratner

src/components/tl/ucp/allreduce/allreduce.c

@@ -32,7 +32,7 @@ ucc_base_coll_alg_info_t
         [UCC_TL_UCP_ALLREDUCE_ALG_RING] =
             {.id   = UCC_TL_UCP_ALLREDUCE_ALG_RING,
              .name = "ring",
-             .desc = "ring-based allreduce (optimized for large messages and simple topologies)"},
+             .desc = "ring-based allreduce (optimized for BW and simple topologies)"},
         [UCC_TL_UCP_ALLREDUCE_ALG_LAST] = {
             .id = 0, .name = NULL, .desc = NULL}};
 

src/components/tl/ucp/allreduce/allreduce.h

@@ -82,8 +82,6 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task);
 
 ucc_status_t ucc_tl_ucp_allreduce_ring_start(ucc_coll_task_t *coll_task);
 
-ucc_status_t ucc_tl_ucp_allreduce_ring_init_common(ucc_tl_ucp_task_t *task);
-
 ucc_status_t ucc_tl_ucp_allreduce_ring_init(ucc_base_coll_args_t *coll_args,
                                             ucc_base_team_t *team,
                                             ucc_coll_task_t **task_h);

src/components/tl/ucp/allreduce/allreduce_ring.c

@@ -20,13 +20,15 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task)
     size_t             count     = TASK_ARGS(task).dst.info.count;
     ucc_datatype_t     dt        = TASK_ARGS(task).dst.info.datatype;
     size_t             data_size = count * ucc_dt_size(dt);
+    int                num_chunks = tsize; // Number of chunks equals number of ranks
     size_t             chunk_size, offset, remaining;
     ucc_rank_t         sendto, recvfrom;
     void              *recv_buf, *send_buf, *reduce_buf;
     ucc_status_t       status;
+    int                step, chunk;
 
-    int num_chunks = tsize; // Use the number of ranks as the number of chunks (this is dynamic)
-    chunk_size = (data_size + num_chunks - 1) / num_chunks; // Ensure chunks fit into data evenly
+    // Divide data into chunks, rounding up to ensure we cover all data
+    chunk_size = ucc_div_round_up(data_size, num_chunks);
 
     if (UCC_IS_INPLACE(TASK_ARGS(task))) {
         sbuf = rbuf;
@@ -39,16 +41,22 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task)
     sendto   = ucc_ep_map_eval(task->subset.map, (trank + 1) % tsize);
     recvfrom = ucc_ep_map_eval(task->subset.map, (trank - 1 + tsize) % tsize);
 
-    while (task->tagged.send_posted < tsize - 1) {
-        int step = task->tagged.send_posted;
-
-        for (int chunk = 0; chunk < num_chunks; chunk++) {
+    /* 
+     * In the ring algorithm, each process sends/receives tsize-1 times
+     * This is because after tsize-1 steps, each piece of data has traversed 
+     * the entire ring and completed its reduction
+     */
+    while (task->allreduce_ring.step < tsize - 1) {
+        step = task->allreduce_ring.step;
+        
+        /* Resume from the last processed chunk */
+        for (chunk = task->allreduce_ring.chunk; chunk < num_chunks; chunk++) {
             offset = chunk * chunk_size;
             remaining = (chunk == num_chunks - 1) ? data_size - offset : chunk_size;
 
-            send_buf  = (step == 0) ? sbuf + offset : rbuf + offset;
-            recv_buf  = task->allreduce_ring.scratch + offset;
-            reduce_buf = rbuf + offset;
+            send_buf  = (step == 0) ? PTR_OFFSET(sbuf, offset) : PTR_OFFSET(rbuf, offset);
+            recv_buf  = PTR_OFFSET(task->allreduce_ring.scratch, offset);
+            reduce_buf = PTR_OFFSET(rbuf, offset);
 
             UCPCHECK_GOTO(
                 ucc_tl_ucp_send_nb(send_buf, remaining, mem_type, sendto, team, task),
@@ -57,7 +65,11 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task)
                 ucc_tl_ucp_recv_nb(recv_buf, remaining, mem_type, recvfrom, team, task),
                 task, out);
 
+            /* Save current chunk position before testing progress */
+            task->allreduce_ring.chunk = chunk;
+            
             if (UCC_INPROGRESS == ucc_tl_ucp_test(task)) {
+                /* Return and resume from this chunk next time */
                 return;
             }
 
@@ -73,7 +85,9 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task)
             }
         }
 
-        task->tagged.send_posted++;
+        task->allreduce_ring.step++;
+        /* Reset chunk counter for the next step */
+        task->allreduce_ring.chunk = 0;
     }
 
     ucc_assert(UCC_TL_UCP_TASK_P2P_COMPLETE(task));
@@ -84,48 +98,50 @@ void ucc_tl_ucp_allreduce_ring_progress(ucc_coll_task_t *coll_task)
 
 ucc_status_t ucc_tl_ucp_allreduce_ring_start(ucc_coll_task_t *coll_task)
 {
-    ucc_tl_ucp_task_t *task      = ucc_derived_of(coll_task, ucc_tl_ucp_task_t);
-    ucc_tl_ucp_team_t *team      = TASK_TEAM(task);
-    size_t             count     = TASK_ARGS(task).dst.info.count;
-    ucc_datatype_t     dt        = TASK_ARGS(task).dst.info.datatype;
-    size_t             data_size = count * ucc_dt_size(dt);
-    ucc_status_t       status;
+    ucc_tl_ucp_task_t *task = ucc_derived_of(coll_task, ucc_tl_ucp_task_t);
+    ucc_tl_ucp_team_t *team = TASK_TEAM(task);
 
     UCC_TL_UCP_PROFILE_REQUEST_EVENT(coll_task, "ucp_allreduce_ring_start", 0);
     ucc_tl_ucp_task_reset(task, UCC_INPROGRESS);
 
-    /* Allocate scratch space for the receive buffer */
-    status = ucc_mc_alloc(&task->allreduce_ring.scratch_mc_header,
-                          data_size, TASK_ARGS(task).dst.info.mem_type);
-    task->allreduce_ring.scratch = task->allreduce_ring.scratch_mc_header->addr;
-    if (ucc_unlikely(status != UCC_OK)) {
-        tl_error(UCC_TASK_LIB(task), "failed to allocate scratch buffer");
-        return status;
-    }
-
     return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super);
 }
 
 ucc_status_t ucc_tl_ucp_allreduce_ring_init_common(ucc_tl_ucp_task_t *task)
 {
     ucc_tl_ucp_team_t *team = TASK_TEAM(task);
     ucc_sbgp_t        *sbgp;
+    size_t             count     = TASK_ARGS(task).dst.info.count;
+    ucc_datatype_t     dt        = TASK_ARGS(task).dst.info.datatype;
+    size_t             data_size = count * ucc_dt_size(dt);
+    ucc_status_t       status;
 
     if (!ucc_coll_args_is_predefined_dt(&TASK_ARGS(task), UCC_RANK_INVALID)) {
         tl_error(UCC_TASK_LIB(task), "user defined datatype is not supported");
         return UCC_ERR_NOT_SUPPORTED;
     }
 
-    if (!(task->flags & UCC_TL_UCP_TASK_FLAG_SUBSET)) {
-        if (team->cfg.use_reordering) {
-            sbgp = ucc_topo_get_sbgp(team->topo, UCC_SBGP_FULL_HOST_ORDERED);
-            task->subset.myrank = sbgp->group_rank;
-            task->subset.map    = sbgp->map;
-        }
+    if (!(task->flags & UCC_TL_UCP_TASK_FLAG_SUBSET) && team->cfg.use_reordering) {
+        sbgp = ucc_topo_get_sbgp(team->topo, UCC_SBGP_FULL_HOST_ORDERED);
+        task->subset.myrank = sbgp->group_rank;
+        task->subset.map    = sbgp->map;
+    }
+
+    /* Allocate scratch space for the receive buffer */
+    status = ucc_mc_alloc(&task->allreduce_ring.scratch_mc_header,
+                          data_size, TASK_ARGS(task).dst.info.mem_type);
+    if (ucc_unlikely(status != UCC_OK)) {
+        tl_error(UCC_TASK_LIB(task), "failed to allocate scratch buffer");
+        return status;
     }
+    task->allreduce_ring.scratch = task->allreduce_ring.scratch_mc_header->addr;
 
+    task->allreduce_ring.step = 0;  /* Initialize step counter */
+    task->allreduce_ring.chunk = 0;  /* Initialize chunk counter */
+    
     task->super.post     = ucc_tl_ucp_allreduce_ring_start;
     task->super.progress = ucc_tl_ucp_allreduce_ring_progress;
+    task->super.finalize = ucc_tl_ucp_allreduce_ring_finalize;
 
     return UCC_OK;
 }

src/components/tl/ucp/tl_ucp_coll.h

@@ -268,13 +268,15 @@ typedef struct ucc_tl_ucp_task {
             ucc_rank_t              iteration;
             int                     phase;
         } alltoall_bruck;
-        char                        plugin_data[UCC_TL_UCP_TASK_PLUGIN_MAX_DATA];
         struct {
             void                   *scratch;
             ucc_mc_buffer_header_t *scratch_mc_header;
             ucc_ee_executor_task_t *etask;
             ucc_ee_executor_t      *executor;
+            int                     step;  /* Track algorithm steps separately */
+            int                     chunk; /* Track current chunk being processed */
         } allreduce_ring;
+        char                        plugin_data[UCC_TL_UCP_TASK_PLUGIN_MAX_DATA];
     };
 } ucc_tl_ucp_task_t;