TL/UCP: Convert sliding window to schedule-based

openucx · Jun 6, 2024 · 3f13f77 · 3f13f77
1 parent 9e4bcc4
commit 3f13f77
Show file tree

Hide file tree

Showing 9 changed files with 508 additions and 416 deletions.
diff --git a/src/components/tl/ucp/Makefile.am b/src/components/tl/ucp/Makefile.am
@@ -113,6 +113,8 @@ sources =                 \
 	tl_ucp_ep.c           \
 	tl_ucp_coll.c         \
 	tl_ucp_service_coll.c \
+	tl_ucp_dpu_offload.h  \
+	tl_ucp_dpu_offload.c  \
 	$(allgather)          \
 	$(allgatherv)         \
 	$(alltoall)           \

diff --git a/src/components/tl/ucp/allreduce/allreduce.c b/src/components/tl/ucp/allreduce/allreduce.c
@@ -7,6 +7,8 @@
 #include "tl_ucp.h"
 #include "allreduce.h"
 #include "utils/ucc_coll_utils.h"
+#include "tl_ucp_dpu_offload.h"
+#include "../allgather/allgather.h"
 
 ucc_base_coll_alg_info_t
     ucc_tl_ucp_allreduce_algs[UCC_TL_UCP_ALLREDUCE_ALG_LAST + 1] = {
@@ -55,40 +57,3 @@ 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)
-{
-    ucc_tl_ucp_task_t       *task;
-    ucc_status_t             status  = UCC_OK;
-    ucc_tl_ucp_team_t       *tl_team = ucc_derived_of(team, ucc_tl_ucp_team_t);
-
-    ALLREDUCE_TASK_CHECK(coll_args->args, tl_team);
-
-    task = ucc_tl_ucp_init_task(coll_args, team);
-    if (ucc_unlikely(!task)) {
-        ucc_error("couldnt allocate task");
-        return UCC_ERR_NO_MEMORY;
-    }
-    *task_h              = &task->super;
-    task->super.post     = ucc_tl_ucp_allreduce_sliding_window_start;
-    task->super.progress = ucc_tl_ucp_allreduce_sliding_window_progress;
-    task->super.finalize = ucc_tl_ucp_allreduce_sliding_window_finalize;
-
-    status = ucc_tl_ucp_allreduce_sliding_window_task_init(coll_args, team, task);
-    if (status != UCC_OK) {
-        ucc_tl_ucp_put_task(task);
-        ucc_error("failed to init task: %s", ucc_status_string(status));
-    }
-
-    task->super.flags |= UCC_COLL_TASK_FLAG_EXECUTOR;
-
-    if (UCC_OK != status) {
-        ucc_error("failed to init executor: %s", ucc_status_string(status));
-    }
-
-out:
-    return status;
-}
diff --git a/src/components/tl/ucp/allreduce/allreduce.h b/src/components/tl/ucp/allreduce/allreduce.h
@@ -70,11 +70,14 @@ 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);
+void ucc_tl_ucp_allreduce_sliding_window_rdma_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_sliding_window_rdma_task_post(ucc_coll_task_t *coll_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,