openucx · Sergei-Lebedev · Feb 25, 2025 · samnordmann · Feb 26, 2025 · samnordmann
diff --git a/config/m4/sharp.m4 b/config/m4/sharp.m4
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # See file LICENSE for terms.
 #
 
@@ -45,6 +45,7 @@ AS_IF([test "x$with_sharp" != "xno"],
                 AC_CHECK_DECLS([SHARP_COLL_HIDE_ERRORS], [], [], [[#include <sharp/api/sharp_coll.h>]])
                 AC_CHECK_DECLS([SHARP_COLL_DISABLE_LAZY_GROUP_RESOURCE_ALLOC], [], [], [[#include <sharp/api/sharp_coll.h>]])
                 AC_CHECK_DECLS([sharp_coll_do_reduce_scatter], [], [], [[#include <sharp/api/sharp_coll.h>]])
+                AC_CHECK_DECLS([sharp_coll_do_allgather], [], [], [[#include <sharp/api/sharp_coll.h>]])
             ],
             [
                 AS_IF([test "x$with_sharp" != "xguess"],

diff --git a/src/components/tl/sharp/tl_sharp.h b/src/components/tl/sharp/tl_sharp.h
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -104,6 +104,10 @@ typedef struct ucc_tl_sharp_task {
     ucc_coll_task_t             super;
     void                       *req_handle;
     union {
+        struct {
+            ucc_tl_sharp_reg_t *s_mem_h;
+            ucc_tl_sharp_reg_t *r_mem_h;
+        } allgather;
         struct {
             ucc_tl_sharp_reg_t *s_mem_h;
             ucc_tl_sharp_reg_t *r_mem_h;
@@ -138,12 +142,11 @@ ucc_status_t sharp_status_to_ucc_status(int status);
 #define UCC_TL_BASIC_SHARP_SUPPORTED_COLLS                                     \
     (UCC_COLL_TYPE_ALLREDUCE | UCC_COLL_TYPE_BARRIER | UCC_COLL_TYPE_BCAST)
 
-#if HAVE_DECL_SHARP_COLL_DO_REDUCE_SCATTER
-#define UCC_TL_SHARP_SUPPORTED_COLLS                                           \
-    (UCC_TL_BASIC_SHARP_SUPPORTED_COLLS | UCC_COLL_TYPE_REDUCE_SCATTER)
-#else
-#define UCC_TL_SHARP_SUPPORTED_COLLS (UCC_TL_BASIC_SHARP_SUPPORTED_COLLS)
-#endif
+
+#define UCC_TL_SHARP_SUPPORTED_COLLS \
+    (UCC_TL_BASIC_SHARP_SUPPORTED_COLLS | \
+    (HAVE_DECL_SHARP_COLL_DO_REDUCE_SCATTER ? UCC_COLL_TYPE_REDUCE_SCATTER : 0) | \
+    (HAVE_DECL_SHARP_COLL_DO_ALLGATHER ? UCC_COLL_TYPE_ALLGATHER : 0))
 
 UCC_CLASS_DECLARE(ucc_tl_sharp_team_t, ucc_base_context_t *,
                   const ucc_base_team_params_t *);

diff --git a/src/components/tl/sharp/tl_sharp_coll.c b/src/components/tl/sharp/tl_sharp_coll.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -160,18 +160,39 @@ void ucc_tl_sharp_collective_progress(ucc_coll_task_t *coll_task)
     if (task->req_handle != NULL) {
         completed = sharp_coll_req_test(task->req_handle);
         if (completed) {
-            if (TASK_ARGS(task).coll_type == UCC_COLL_TYPE_ALLREDUCE) {
+            switch(TASK_ARGS(task).coll_type){
+            case UCC_COLL_TYPE_ALLREDUCE:
                 if (!UCC_IS_INPLACE(TASK_ARGS(task))) {
                     ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
                                                 task->allreduce.s_mem_h);
                 }
                 ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
                                             task->allreduce.r_mem_h);
-            }
-            if (TASK_ARGS(task).coll_type == UCC_COLL_TYPE_BCAST) {
+                break;
+            case UCC_COLL_TYPE_BCAST:
                 ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
                                             task->bcast.mem_h);
+                break;
+            case UCC_COLL_TYPE_ALLGATHER:
+                if (!UCC_IS_INPLACE(TASK_ARGS(task))) {
+                    ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
+                                                task->allgather.s_mem_h);
+                }
+                ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
+                                            task->allgather.r_mem_h);
+                break;
+            case UCC_COLL_TYPE_REDUCE_SCATTER:
+                if (!UCC_IS_INPLACE(TASK_ARGS(task))) {
+                    ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
+                                                task->reduce_scatter.s_mem_h);
+                }
+                ucc_tl_sharp_mem_deregister(TASK_TEAM(task),
+                                            task->reduce_scatter.r_mem_h);
+                break;
+            default:
+                break;
             }
+
             sharp_coll_req_free(task->req_handle);
             coll_task->status = UCC_OK;
             UCC_TL_SHARP_PROFILE_REQUEST_EVENT(coll_task,
@@ -450,3 +471,88 @@ ucc_status_t ucc_tl_sharp_barrier_init(ucc_tl_sharp_task_t *task)
     task->super.progress = ucc_tl_sharp_collective_progress;
     return UCC_OK;
 };
+
+#if HAVE_DECL_SHARP_COLL_DO_ALLGATHER
+ucc_status_t ucc_tl_sharp_allgather_start(ucc_coll_task_t *coll_task)
+{
+    ucc_tl_sharp_task_t *task  = ucc_derived_of(coll_task, ucc_tl_sharp_task_t);
+    ucc_tl_sharp_team_t *team  = TASK_TEAM(task);
+    ucc_coll_args_t     *args  = &TASK_ARGS(task);
+    size_t               count = args->dst.info.count;
+    ucc_datatype_t       dt    = args->dst.info.datatype;
+    struct sharp_coll_gather_spec gather_spec;
+    size_t                        src_data_size, dst_data_size;
+    int                           ret;
+
+    UCC_TL_SHARP_PROFILE_REQUEST_EVENT(coll_task, "sharp_allgather_start", 0);
+
+    src_data_size = ucc_dt_size(dt) * count / UCC_TL_TEAM_SIZE(team);
+    dst_data_size = ucc_dt_size(dt) * count;
+
+    if (!UCC_IS_INPLACE(*args)) {
+        ucc_tl_sharp_mem_register(TASK_CTX(task), team, args->src.info.buffer,
+                                  src_data_size, &task->allgather.s_mem_h);
+    }
+    ucc_tl_sharp_mem_register(TASK_CTX(task), team, args->dst.info.buffer,
+                              dst_data_size, &task->allgather.r_mem_h);
+
+    if (!UCC_IS_INPLACE(*args)) {
+        gather_spec.sbuf_desc.buffer.ptr = args->src.info.buffer;
+        gather_spec.sbuf_desc.buffer.mem_handle = task->allgather.s_mem_h->mr;
+        gather_spec.sbuf_desc.mem_type =
+            ucc_to_sharp_memtype[args->src.info.mem_type];
+    } else {
+        gather_spec.sbuf_desc.buffer.ptr = PTR_OFFSET(args->dst.info.buffer,
+                                                      UCC_TL_TEAM_RANK(team) *
+                                                      src_data_size);
+        gather_spec.sbuf_desc.buffer.mem_handle = task->allgather.r_mem_h->mr;
+        gather_spec.sbuf_desc.mem_type =
+            ucc_to_sharp_memtype[args->dst.info.mem_type];
+    }
+    gather_spec.sbuf_desc.buffer.length = src_data_size;
+    gather_spec.sbuf_desc.type = SHARP_DATA_BUFFER;
+
+    gather_spec.rbuf_desc.buffer.ptr        = args->dst.info.buffer;
+    gather_spec.rbuf_desc.buffer.length     = dst_data_size;
+    gather_spec.rbuf_desc.buffer.mem_handle = task->allgather.r_mem_h->mr;
+    gather_spec.rbuf_desc.type              = SHARP_DATA_BUFFER;
+    gather_spec.rbuf_desc.mem_type =
+        ucc_to_sharp_memtype[args->dst.info.mem_type];
+    gather_spec.offset    = 0;
+    gather_spec.dtype = SHARP_DTYPE_INT8;
+
+    ret = sharp_coll_do_allgather_nb(team->sharp_comm, &gather_spec,
+                                     &task->req_handle);
+    if (ret != SHARP_COLL_SUCCESS) {
+        tl_error(UCC_TASK_LIB(task),
+                 "sharp_coll_do_allgather_nb failed:%s",
+                 sharp_coll_strerror(ret));
+        coll_task->status = ucc_tl_sharp_status_to_ucc(ret);
+        return ucc_task_complete(coll_task);
+    }
+    coll_task->status = UCC_INPROGRESS;
+
+    return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super);
+}
+
+ucc_status_t ucc_tl_sharp_allgather_init(ucc_tl_sharp_task_t *task)
+{
+    ucc_coll_args_t *args = &TASK_ARGS(task);
+
+    if (!ucc_coll_args_is_predefined_dt(args, UCC_RANK_INVALID)) {
+        return UCC_ERR_NOT_SUPPORTED;
+    }
+
+    if ((!UCC_IS_INPLACE(*args) &&
+        ucc_to_sharp_memtype[args->src.info.mem_type] == SHARP_MEM_TYPE_LAST) ||
+        ucc_to_sharp_memtype[args->dst.info.mem_type] == SHARP_MEM_TYPE_LAST ||
+        ucc_to_sharp_dtype[UCC_DT_PREDEFINED_ID(args->dst.info.datatype)] ==
+            SHARP_DTYPE_NULL) {
+        return UCC_ERR_NOT_SUPPORTED;
+    }
+
+    task->super.post     = ucc_tl_sharp_allgather_start;
+    task->super.progress = ucc_tl_sharp_collective_progress;
+    return UCC_OK;
+};
+#endif
diff --git a/src/components/tl/sharp/tl_sharp_coll.h b/src/components/tl/sharp/tl_sharp_coll.h
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -23,4 +23,9 @@ ucc_status_t ucc_tl_sharp_bcast_init(ucc_tl_sharp_task_t *task);
 #if HAVE_DECL_SHARP_COLL_DO_REDUCE_SCATTER
 ucc_status_t ucc_tl_sharp_reduce_scatter_init(ucc_tl_sharp_task_t *task);
 #endif
+
+#if HAVE_DECL_SHARP_COLL_DO_ALLGATHER
+ucc_status_t ucc_tl_sharp_allgather_init(ucc_tl_sharp_task_t *task);
+#endif
+
 #endif
diff --git a/src/components/tl/sharp/tl_sharp_team.c b/src/components/tl/sharp/tl_sharp_team.c
@@ -1,5 +1,5 @@
 /**
- * Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  *
  * See file LICENSE for terms.
  */
@@ -254,6 +254,11 @@ ucc_status_t ucc_tl_sharp_coll_init(ucc_base_coll_args_t *coll_args,
     case UCC_COLL_TYPE_REDUCE_SCATTER:
         status = ucc_tl_sharp_reduce_scatter_init(task);
         break;
+#endif
+#if HAVE_DECL_SHARP_COLL_DO_ALLGATHER
+    case UCC_COLL_TYPE_ALLGATHER:
+        status = ucc_tl_sharp_allgather_init(task);
+        break;
 #endif
     default:
         tl_debug(UCC_TASK_LIB(task),