TL/UCP: Add Sliding Window allreduce impl

openucx · Apr 18, 2024 · 139ee8d · 139ee8d
1 parent 8a465a5
commit 139ee8d
Show file tree

Hide file tree

Showing 5 changed files with 762 additions and 64 deletions.
diff --git a/src/components/tl/ucp/allreduce/allreduce.c b/src/components/tl/ucp/allreduce/allreduce.c
@@ -57,11 +57,38 @@ ucc_status_t ucc_tl_ucp_allreduce_knomial_init(ucc_base_coll_args_t *coll_args,
 }
 
 ucc_status_t
-ucc_tl_ucp_allreduce_sliding_window_init(ucc_base_coll_args_t __attribute__((unused)) *coll_args, //NOLINT
-                                         ucc_base_team_t __attribute__((unused)) *team, //NOLINT
-                                         ucc_coll_task_t __attribute__((unused)) **task_h) //NOLINT
+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_coll_task_t *coll_task = NULL;
-    ucc_tl_ucp_allreduce_sliding_window_progress(coll_task);
-    return UCC_OK;
+    ucc_status_t             status  = UCC_OK;
+    ucc_tl_ucp_team_t       *tl_team = ucc_derived_of(team, ucc_tl_ucp_team_t);
+    ucc_tl_ucp_task_t       *task;
+    ucc_ee_executor_params_t params;
+
+    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;
+
+    ucc_tl_ucp_allreduce_sliding_window_task_init(coll_args, team, task);
+
+    params.mask    = UCC_EE_EXECUTOR_PARAM_FIELD_TYPE;
+    params.ee_type = UCC_EE_CPU_THREAD;
+    status =
+        ucc_ee_executor_init(&params, &task->allreduce_sliding_window.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
@@ -44,18 +44,18 @@ ucc_status_t ucc_tl_ucp_allreduce_knomial_init(ucc_base_coll_args_t *coll_args,
 
 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_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_task_init(ucc_base_coll_args_t *coll_args,
-                                              ucc_base_team_t *     team,
-                                              ucc_tl_ucp_task_t *   task);
+                                              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_tl_ucp_task_t *sw_task, int inplace);
 
 ucc_status_t
 ucc_tl_ucp_allreduce_sliding_window_free_gwbi(ucc_coll_task_t *coll_task);