CORE: Implement weak asymmetric mem with gtests

src/coll_score/ucc_coll_score_map.c

@@ -87,15 +87,12 @@ static ucc_status_t ucc_coll_score_map_lookup(ucc_score_map_t *map,
     ucc_list_link_t *list;
     ucc_msg_range_t *r;
 
-    if (mt == UCC_MEMORY_TYPE_ASYMMETRIC) {
-        /* TODO */
-        ucc_debug("asymmetric memory type is not supported");
-        return UCC_ERR_NOT_SUPPORTED;
-    } else if (mt == UCC_MEMORY_TYPE_NOT_APPLY) {
+    if (mt == UCC_MEMORY_TYPE_NOT_APPLY) {
         /* Temporary solution: for Barrier, Fanin, Fanout - use
            "host" range list */
         mt = UCC_MEMORY_TYPE_HOST;
     }
+    ucc_assert(ucc_coll_args_is_mem_symmetric(&bargs->args, map->team_rank));
     if (msgsize == UCC_MSG_SIZE_INVALID || msgsize == UCC_MSG_SIZE_ASYMMETRIC) {
         /* These algorithms require global communication to get the same msgsize estimation.
            Can't use msg ranges. Use msize 0 (assuming the range list should only contain 1

src/components/base/ucc_base_iface.h

@@ -159,11 +159,20 @@ enum {
     UCC_BASE_CARGS_MAX_FRAG_COUNT = UCC_BIT(0)
 };
 
+typedef struct ucc_buffer_info_asymmetric_memtype {
+    union {
+        ucc_coll_buffer_info_t      info;
+        ucc_coll_buffer_info_v_t    info_v;
+    } old_asymmetric_buffer;
+    ucc_mc_buffer_header_t *scratch;
+} ucc_buffer_info_asymmetric_memtype_t;
+
 typedef struct ucc_base_coll_args {
-    uint64_t         mask;
-    ucc_coll_args_t  args;
-    ucc_team_t      *team;
-    size_t           max_frag_count;
+    uint64_t                             mask;
+    ucc_coll_args_t                      args;
+    ucc_team_t                          *team;
+    size_t                               max_frag_count;
+    ucc_buffer_info_asymmetric_memtype_t asymmetric_save_info;
 } ucc_base_coll_args_t;
 
 typedef ucc_status_t (*ucc_base_coll_init_fn_t)(ucc_base_coll_args_t *coll_args,

src/core/ucc_coll.c

@@ -230,19 +230,29 @@ UCC_CORE_PROFILE_FUNC(ucc_status_t, ucc_collective_init,
     UCC_COPY_PARAM_BY_FIELD(&op_args.args, coll_args, UCC_COLL_ARGS_FIELD_FLAGS,
                             flags);
 
+    if (!ucc_coll_args_is_mem_symmetric(&op_args.args, team->rank) &&
+        ucc_coll_args_is_rooted(op_args.args.coll_type)) {
+        status = ucc_coll_args_init_asymmetric_buffer(&op_args.args, team,
+                    &op_args.asymmetric_save_info);
+        if (ucc_unlikely(status != UCC_OK)) {
+            ucc_error("handling asymmetric memory failed");
+            return status;
+        }
+    }
+
     status = ucc_coll_init(team->score_map, &op_args, &task);
     if (UCC_ERR_NOT_SUPPORTED == status) {
         ucc_debug("failed to init collective: not supported");
-        return status;
+        goto free_scratch;
     } else if (ucc_unlikely(status < 0)) {
         ucc_error("failed to init collective: %s", ucc_status_string(status));
-        return status;
+        goto free_scratch;
     }
 
     task->flags |= UCC_COLL_TASK_FLAG_TOP_LEVEL;
     if (task->flags & UCC_COLL_TASK_FLAG_EXECUTOR) {
         task->flags |= UCC_COLL_TASK_FLAG_EXECUTOR_STOP;
-        coll_mem_type = ucc_coll_args_mem_type(coll_args, team->rank);
+        coll_mem_type = ucc_coll_args_mem_type(&op_args.args, team->rank);
         switch(coll_mem_type) {
         case UCC_MEMORY_TYPE_CUDA:
         case UCC_MEMORY_TYPE_CUDA_MANAGED:
@@ -251,7 +261,7 @@ UCC_CORE_PROFILE_FUNC(ucc_status_t, ucc_collective_init,
         case UCC_MEMORY_TYPE_ROCM:
             coll_ee_type = UCC_EE_ROCM_STREAM;
             break;
-       case UCC_MEMORY_TYPE_HOST:
+        case UCC_MEMORY_TYPE_HOST:
             coll_ee_type = UCC_EE_CPU_THREAD;
             break;
         default:
@@ -299,6 +309,10 @@ UCC_CORE_PROFILE_FUNC(ucc_status_t, ucc_collective_init,
 
 coll_finalize:
     task->finalize(task);
+free_scratch:
+    if (op_args.asymmetric_save_info.scratch != NULL) {
+        ucc_mc_free(op_args.asymmetric_save_info.scratch);
+    }
     return status;
 }
 

src/schedule/ucc_schedule.h

@@ -14,6 +14,7 @@
 #include "utils/ucc_coll_utils.h"
 #include "components/base/ucc_base_iface.h"
 #include "components/ec/ucc_ec.h"
+#include "components/mc/ucc_mc.h"
 
 #define MAX_LISTENERS 4
 
@@ -165,6 +166,36 @@ ucc_status_t ucc_dependency_handler(ucc_coll_task_t *parent,
 ucc_status_t ucc_triggered_post(ucc_ee_h ee, ucc_ev_t *ev,
                                 ucc_coll_task_t *task);
 
+static inline
+ucc_status_t ucc_copy_asymmetric_buffer_out(ucc_coll_task_t *task)
+{
+    ucc_status_t                          status    = UCC_OK;
+    ucc_coll_args_t                      *coll_args = &task->bargs.args;
+    ucc_buffer_info_asymmetric_memtype_t *save      = &task->bargs.asymmetric_save_info;
+    ucc_rank_t                            size      = task->team->params.size;
+
+    if(task->bargs.args.coll_type == UCC_COLL_TYPE_GATHERV) {
+        status = ucc_mc_memcpy(save->old_asymmetric_buffer.info_v.buffer,
+                        save->scratch->addr,
+                        ucc_coll_args_get_total_count(coll_args,
+                            coll_args->dst.info_v.counts, size),
+                        save->old_asymmetric_buffer.info_v.mem_type,
+                        save->scratch->mt);
+    } else {
+        status = ucc_mc_memcpy(save->old_asymmetric_buffer.info.buffer,
+                        save->scratch->addr,
+                        ucc_dt_size(save->old_asymmetric_buffer.info.datatype) *
+                            save->old_asymmetric_buffer.info.count,
+                        save->old_asymmetric_buffer.info.mem_type,
+                        save->scratch->mt);
+    }
+    if (ucc_unlikely(status != UCC_OK)) {
+        ucc_error("error copying back to old asymmetric buffer: %s",
+                    ucc_status_string(status));
+    }
+    return status;
+}
+
 static inline ucc_status_t ucc_task_complete(ucc_coll_task_t *task)
 {
     ucc_status_t        status    = task->status;
@@ -185,6 +216,20 @@ static inline ucc_status_t ucc_task_complete(ucc_coll_task_t *task)
        with schedules are not released during a callback (if set). */
 
     if (ucc_likely(status == UCC_OK)) {
+        ucc_buffer_info_asymmetric_memtype_t *save = &task->bargs.asymmetric_save_info;
+        if (save->scratch != NULL) {
+            status = ucc_copy_asymmetric_buffer_out(task);
+            if (status != UCC_OK) {
+                ucc_error("failure copying out asymmetric buffer: %s",
+                          ucc_status_string(status));
+            }
+            status = ucc_mc_free(save->scratch);
+            if (ucc_unlikely(status != UCC_OK)) {
+                ucc_error("error freeing scratch asymmetric buffer: %s",
+                            ucc_status_string(status));
+            }
+            save->scratch = NULL;
+        }
         status = ucc_event_manager_notify(task, UCC_EVENT_COMPLETED);
     } else {
         /* error in task status */

src/utils/ucc_coll_utils.c

@@ -52,7 +52,7 @@ ucc_memory_type_t ucc_mem_type_from_str(const char *str)
     return UCC_MEMORY_TYPE_LAST;
 }
 
-static inline int
+int
 ucc_coll_args_is_mem_symmetric(const ucc_coll_args_t *args,
                                ucc_rank_t rank)
 {
@@ -94,6 +94,68 @@ ucc_coll_args_is_mem_symmetric(const ucc_coll_args_t *args,
     return 0;
 }
 
+
+/* If this is the root and the src/dst buffers are asymmetric, the dst needs
+   to have a new allocation to make the mem types match. On task completion,
+   copy the result back into the old dst */
+ucc_status_t
+ucc_coll_args_init_asymmetric_buffer(ucc_coll_args_t *args,
+                                     ucc_team_h team,
+                                     ucc_buffer_info_asymmetric_memtype_t *save_info)
+{
+    ucc_status_t status = UCC_OK;
+
+    if (UCC_IS_INPLACE(*args)) {
+        return UCC_ERR_INVALID_PARAM;
+    }
+    switch (args->coll_type) {
+    case UCC_COLL_TYPE_REDUCE:
+    case UCC_COLL_TYPE_GATHER:
+    case UCC_COLL_TYPE_SCATTER:
+    case UCC_COLL_TYPE_SCATTERV:
+    {
+        ucc_memory_type_t mem_type = args->src.info.mem_type;
+        if (args->coll_type == UCC_COLL_TYPE_SCATTERV) {
+            mem_type = args->src.info_v.mem_type;
+        }
+        memcpy(&save_info->old_asymmetric_buffer.info,
+               &args->dst.info, sizeof(ucc_coll_buffer_info_t));
+        status = ucc_mc_alloc(&save_info->scratch,
+                              ucc_dt_size(args->dst.info.datatype) *
+                                args->dst.info.count,
+                                mem_type);
+        if (ucc_unlikely(UCC_OK != status)) {
+            ucc_error("failed to allocate replacement "
+                      "memory for asymmetric buffer");
+            return status;
+        }
+        args->dst.info.buffer = save_info->scratch->addr;
+        args->dst.info.mem_type = mem_type;
+        return UCC_OK;
+    }
+    case UCC_COLL_TYPE_GATHERV:
+    {
+        memcpy(&save_info->old_asymmetric_buffer.info_v,
+               &args->dst.info_v, sizeof(ucc_coll_buffer_info_v_t));
+        status = ucc_mc_alloc(&save_info->scratch,
+                              ucc_coll_args_get_total_count(args,
+                                args->dst.info_v.counts, team->size),
+                              args->src.info.mem_type);
+        if (ucc_unlikely(UCC_OK != status)) {
+            ucc_error("failed to allocate replacement "
+                      "memory for asymmetric buffer");
+            return status;
+        }
+        args->dst.info_v.buffer = save_info->scratch->addr;
+        args->dst.info_v.mem_type = args->src.info.mem_type;
+        return UCC_OK;
+    }
+    default:
+        break;
+    }
+    return UCC_ERR_INVALID_PARAM;
+}
+
 int ucc_coll_args_is_predefined_dt(const ucc_coll_args_t *args, ucc_rank_t rank)
 {
     switch (args->coll_type) {
@@ -163,9 +225,6 @@ ucc_memory_type_t ucc_coll_args_mem_type(const ucc_coll_args_t *args,
 {
     ucc_rank_t root = args->root;
 
-    if (!ucc_coll_args_is_mem_symmetric(args, rank)) {
-        return UCC_MEMORY_TYPE_ASYMMETRIC;
-    }
     switch (args->coll_type) {
     case UCC_COLL_TYPE_BARRIER:
     case UCC_COLL_TYPE_FANIN:
@@ -180,7 +239,6 @@ ucc_memory_type_t ucc_coll_args_mem_type(const ucc_coll_args_t *args,
         return args->dst.info.mem_type;
     case UCC_COLL_TYPE_ALLGATHERV:
     case UCC_COLL_TYPE_REDUCE_SCATTERV:
-        return args->dst.info_v.mem_type;
     case UCC_COLL_TYPE_ALLTOALLV:
         return args->dst.info_v.mem_type;
     case UCC_COLL_TYPE_REDUCE:
@@ -323,7 +381,7 @@ ucc_ep_map_t ucc_ep_map_from_array_64(uint64_t **array, ucc_rank_t size,
                                          need_free, 1);
 }
 
-static inline int ucc_coll_args_is_rooted(ucc_coll_type_t ct)
+int ucc_coll_args_is_rooted(ucc_coll_type_t ct)
 {
     if (ct == UCC_COLL_TYPE_REDUCE || ct == UCC_COLL_TYPE_BCAST ||
         ct == UCC_COLL_TYPE_GATHER || ct == UCC_COLL_TYPE_SCATTER ||

src/utils/ucc_coll_utils.h

@@ -335,4 +335,16 @@ static inline size_t ucc_buffer_block_offset_aligned(size_t total_count,
                            operations. */
 int ucc_coll_args_is_predefined_dt(const ucc_coll_args_t *args, ucc_rank_t rank);
 
+int ucc_coll_args_is_mem_symmetric(const ucc_coll_args_t *args, ucc_rank_t rank);
+
+int ucc_coll_args_is_rooted(ucc_coll_type_t ct);
+
+typedef struct ucc_buffer_info_asymmetric_memtype ucc_buffer_info_asymmetric_memtype_t;
+typedef struct ucc_mc_buffer_header ucc_mc_buffer_header_t;
+
+ucc_status_t
+ucc_coll_args_init_asymmetric_buffer(ucc_coll_args_t *args,
+                                       ucc_team_h team,
+                                       ucc_buffer_info_asymmetric_memtype_t *save_info);
+
 #endif

test/gtest/Makefile.am

@@ -106,7 +106,8 @@ gtest_SOURCES =                           \
 	coll_score/test_score.cc              \
 	coll_score/test_score_str.cc          \
 	coll_score/test_score_update.cc       \
-	active_set/test_active_set.cc
+	active_set/test_active_set.cc         \
+	asym_mem/test_asymmetric_memory.cc
 
 if TL_MLX5_ENABLED
     gtest_SOURCES  += tl/mlx5/test_tl_mlx5.cc \