Skip to content

Commit

Permalink
Gatherv working
Browse files Browse the repository at this point in the history
  • Loading branch information
@nsarka
nsarka committed Jul 12, 2024
1 parent 0c34341 commit 6289d70
Show file tree
Hide file tree
Showing 3 changed files with 267 additions and 133 deletions.
65 changes: 51 additions & 14 deletions src/schedule/ucc_schedule.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,50 @@ 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_asymmetric_save_info_t *save)
{
ucc_status_t status = UCC_OK;
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));
}
return status;
}

static inline ucc_status_t ucc_copy_asymmetric_buffer_v_out(ucc_coll_task_t *task)
{
ucc_status_t status = UCC_OK;
ucc_coll_args_t *coll_args = &task->bargs.args;
ucc_asymmetric_save_info_t *save = &task->bargs.asymmetric_save_info;
ucc_rank_t size = task->team->params.size;
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);
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));
}
return status;
}

static inline ucc_status_t ucc_task_complete(ucc_coll_task_t *task)
{
ucc_status_t status = task->status;
Expand All @@ -188,20 +232,13 @@ static inline ucc_status_t ucc_task_complete(ucc_coll_task_t *task)
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));
status = (task->bargs.args.coll_type == UCC_COLL_TYPE_GATHERV ||
task->bargs.args.coll_type == UCC_COLL_TYPE_SCATTERV) ?
ucc_copy_asymmetric_buffer_v_out(task) :
ucc_copy_asymmetric_buffer_out(save);
if (status != UCC_OK) {
ucc_error("failure copying out asymmetric buffer: %s",
ucc_status_string(status));
}
save->scratch = NULL;
}
Expand Down
93 changes: 16 additions & 77 deletions src/utils/ucc_coll_utils.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,85 +102,12 @@ 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:
Expand All @@ -200,10 +127,22 @@ ucc_coll_args_update_asymmetric_buffer(ucc_coll_args_t *args,
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_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;
}
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;
Expand Down
Loading

0 comments on commit 6289d70

Please sign in to comment.