Skip to content

Commit

Permalink
Update
Browse files Browse the repository at this point in the history
  • Loading branch information
Nicholas Sarkauskas committed Jul 9, 2024
1 parent ead14c3 commit 0a72a73
Show file tree
Hide file tree
Showing 7 changed files with 191 additions and 57 deletions.
3 changes: 3 additions & 0 deletions src/coll_score/ucc_coll_score_map.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,9 @@ static ucc_status_t ucc_coll_score_map_lookup(ucc_score_map_t *map,
"host" range list */
mt = UCC_MEMORY_TYPE_HOST;
}
if (!ucc_coll_args_is_mem_symmetric(&bargs->args, map->team_rank)) {
return UCC_ERR_INVALID_PARAM;
}
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
Expand Down
17 changes: 13 additions & 4 deletions src/components/base/ucc_base_iface.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,11 +159,20 @@ enum {
UCC_BASE_CARGS_MAX_FRAG_COUNT = UCC_BIT(0)
};

typedef struct ucc_asymmetric_save_info {
union {
ucc_coll_buffer_info_t info; /*!< Buffer info for the collective */
ucc_coll_buffer_info_v_t info_v; /*!< Buffer info for the collective */
} old_asymmetric_buffer;
ucc_mc_buffer_header_t *scratch;
} ucc_asymmetric_save_info_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_asymmetric_save_info_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,
Expand Down
24 changes: 20 additions & 4 deletions src/core/ucc_coll.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,8 @@ UCC_CORE_PROFILE_FUNC(ucc_status_t, ucc_collective_init,
ucc_ee_type_t coll_ee_type;
size_t coll_size;

printf("nick in ucc_collective_init: dst buffer=%p\n", coll_args->dst.info.buffer);

if (ucc_unlikely(team->state != UCC_TEAM_ACTIVE)) {
ucc_error("team %p is used before team create is completed", team);
return UCC_ERR_INVALID_PARAM;
Expand Down Expand Up @@ -230,19 +232,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)) {
status = ucc_coll_args_update_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;
}

coll_mem_type = ucc_coll_args_mem_type(&op_args.args, team->rank);

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);
switch(coll_mem_type) {
case UCC_MEMORY_TYPE_CUDA:
case UCC_MEMORY_TYPE_CUDA_MANAGED:
Expand All @@ -251,7 +263,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:
Expand Down Expand Up @@ -299,6 +311,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;
}

Expand Down
20 changes: 20 additions & 0 deletions src/schedule/ucc_schedule.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -185,6 +186,25 @@ 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_asymmetric_save_info_t *save = &task->bargs.asymmetric_save_info;
if (save->scratch != NULL) {
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));
}
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 */
Expand Down
127 changes: 122 additions & 5 deletions src/utils/ucc_coll_utils.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -94,6 +94,126 @@ ucc_coll_args_is_mem_symmetric(const ucc_coll_args_t *args,
return 0;
}


/* If the src/dst buffers are asymmetric, one of them needs to have a new
allocation */
ucc_status_t
ucc_coll_args_update_asymmetric_buffer(ucc_coll_args_t *args,
ucc_team_h team,
ucc_asymmetric_save_info_t *save_info)
{
//ucc_rank_t root = args->root;
//ucc_rank_t rank = team->rank;
ucc_status_t status = UCC_OK;

if (UCC_IS_INPLACE(*args)) {
return UCC_ERR_INVALID_PARAM;
}
switch (args->coll_type) {
case UCC_COLL_TYPE_BARRIER:
case UCC_COLL_TYPE_BCAST:
case UCC_COLL_TYPE_FANIN:
case UCC_COLL_TYPE_FANOUT:
return UCC_ERR_INVALID_PARAM;
case UCC_COLL_TYPE_ALLTOALL:
case UCC_COLL_TYPE_ALLREDUCE:
case UCC_COLL_TYPE_ALLGATHER:
case UCC_COLL_TYPE_REDUCE_SCATTER:
/*{
if (args->dst.info.mem_type == mem_type) {
*old_buffer = args->src.info.buffer;
status = ucc_mc_alloc(args->src.info.buffer, ucc_dt_size(args->src.info.datatype) * args->src.info.count, mem_type);
if (ucc_unlikely(UCC_OK != status)) {
ucc_error("failed to allocate replacement memory for asymmetric buffer");
return status;
}
} else {
*old_buffer = args->dst.info.buffer;
status = ucc_mc_alloc(args->dst.info.buffer, 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;
}
}
}*/
case UCC_COLL_TYPE_ALLGATHERV:
case UCC_COLL_TYPE_REDUCE_SCATTERV:
/*{
if (args->dst.info_v.mem_type == mem_type) {
*old_buffer = args->src.info.buffer;
status = ucc_mc_alloc(args->src.info.buffer, ucc_dt_size(args->src.info.datatype) * args->src.info.count, mem_type);
if (ucc_unlikely(UCC_OK != status)) {
ucc_error("failed to allocate replacement memory for asymmetric buffer");
return status;
}
} else {
*old_buffer = args->dst.info_v.buffer;
status = ucc_mc_alloc(args->dst.info_v.buffer, ucc_dt_size(args->dst.info_v.datatype) * args->dst.info_v.counts[rank], mem_type);
if (ucc_unlikely(UCC_OK != status)) {
ucc_error("failed to allocate replacement memory for asymmetric buffer");
return status;
}
}
}*/
case UCC_COLL_TYPE_ALLTOALLV:
/*{
ucc_count_t sum_count = 0;
ucc_rank_t i;
if (args->dst.info_v.mem_type == mem_type) {
*old_buffer = args->src.info_v.buffer;
for(i = 0; i < team->size; i++) {
sum_count += args->src.info_v.counts[i];
}
status = ucc_mc_alloc(args->src.info_v.buffer, ucc_dt_size(args->src.info_v.datatype) * sum_count, mem_type);
if (ucc_unlikely(UCC_OK != status)) {
ucc_error("failed to allocate replacement memory for asymmetric buffer");
return status;
}
} else {
*old_buffer = args->dst.info_v.buffer;
for(i = 0; i < team->size; i++) {
sum_count += args->dst.info_v.counts[i];
}
status = ucc_mc_alloc(args->dst.info_v.buffer, ucc_dt_size(args->dst.info_v.datatype) * sum_count, mem_type);
if (ucc_unlikely(UCC_OK != status)) {
ucc_error("failed to allocate replacement memory for asymmetric buffer");
return status;
}
}
}*/
case UCC_COLL_TYPE_REDUCE:
case UCC_COLL_TYPE_GATHER:
case UCC_COLL_TYPE_SCATTER:
{
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,
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.buffer = save_info->scratch->addr;
args->dst.info.mem_type = args->src.info.mem_type;
return UCC_OK;
}
case UCC_COLL_TYPE_GATHERV: /*
return (root != rank ? NULL : (
args->dst.info_v.mem_type == mem_type ? args->src.info.buffer : args->dst.info_v.buffer;
)) */
case UCC_COLL_TYPE_SCATTERV: /*
return (root != rank ? NULL : (
args->dst.info.mem_type == mem_type ? args->src.info_v.buffer : args->dst.info.buffer;
))*/
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) {
Expand Down Expand Up @@ -163,18 +283,15 @@ 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:
case UCC_COLL_TYPE_FANOUT:
return UCC_MEMORY_TYPE_NOT_APPLY;
case UCC_COLL_TYPE_BCAST:
return args->src.info.mem_type;
case UCC_COLL_TYPE_ALLTOALL:
case UCC_COLL_TYPE_ALLREDUCE:
case UCC_COLL_TYPE_ALLTOALL:
case UCC_COLL_TYPE_ALLGATHER:
case UCC_COLL_TYPE_REDUCE_SCATTER:
return args->dst.info.mem_type;
Expand Down
10 changes: 10 additions & 0 deletions src/utils/ucc_coll_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -335,4 +335,14 @@ 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);

typedef struct ucc_asymmetric_save_info ucc_asymmetric_save_info_t;
typedef struct ucc_mc_buffer_header ucc_mc_buffer_header_t;

ucc_status_t
ucc_coll_args_update_asymmetric_buffer(ucc_coll_args_t *args,
ucc_team_h team,
ucc_asymmetric_save_info_t *save_info);

#endif
Loading

0 comments on commit 0a72a73

Please sign in to comment.