CORE: handle nonrooted asym colls

openucx · Jul 22, 2024 · 5e24e86 · 5e24e86
1 parent 62d54e9
commit 5e24e86
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/src/core/ucc_coll.c b/src/core/ucc_coll.c
@@ -230,7 +230,8 @@ 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)) {
+    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)) {

diff --git a/src/utils/ucc_coll_utils.c b/src/utils/ucc_coll_utils.c
@@ -381,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 ||

diff --git a/src/utils/ucc_coll_utils.h b/src/utils/ucc_coll_utils.h
@@ -337,6 +337,8 @@ 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;