Skip to content

Commit

Permalink
CORE: asymmetric mem feedback
Browse files Browse the repository at this point in the history
  • Loading branch information
@nsarka
nsarka committed Jul 22, 2024
1 parent 9d8c131 commit c88eea3
Show file tree
Hide file tree
Showing 7 changed files with 30 additions and 40 deletions.
6 changes: 1 addition & 5 deletions src/coll_score/ucc_coll_score_map.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,7 @@ 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)) {
ucc_debug("Mem was asymmetric even though asymmetric memory should "
"be handled prior to finding coll score");
return UCC_ERR_INVALID_PARAM;
}
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
Expand Down
14 changes: 7 additions & 7 deletions src/components/base/ucc_base_iface.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,20 +159,20 @@ enum {
UCC_BASE_CARGS_MAX_FRAG_COUNT = UCC_BIT(0)
};

typedef struct ucc_asymmetric_save_info {
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_asymmetric_save_info_t;
} 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;
ucc_asymmetric_save_info_t asymmetric_save_info;
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,
Expand Down
2 changes: 1 addition & 1 deletion src/core/ucc_coll.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -231,7 +231,7 @@ UCC_CORE_PROFILE_FUNC(ucc_status_t, ucc_collective_init,
flags);

if (!ucc_coll_args_is_mem_symmetric(&op_args.args, team->rank)) {
status = ucc_coll_args_update_asymmetric_buffer(&op_args.args, team,
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");
Expand Down
27 changes: 11 additions & 16 deletions src/schedule/ucc_schedule.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ 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 ucc_copy_asymmetric_buffer_out(ucc_buffer_info_asymmetric_memtype_t *save)
{
ucc_status_t status = UCC_OK;
status = ucc_mc_memcpy(save->old_asymmetric_buffer.info.buffer,
Expand All @@ -180,21 +180,16 @@ ucc_status_t ucc_copy_asymmetric_buffer_out(ucc_asymmetric_save_info_t *save)
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;
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;
status = ucc_mc_memcpy(save->old_asymmetric_buffer.info_v.buffer,
save->scratch->addr,
ucc_coll_args_get_total_count(coll_args,
Expand All @@ -205,11 +200,6 @@ ucc_status_t ucc_copy_asymmetric_buffer_v_out(ucc_coll_task_t *task)
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;
}

Expand All @@ -233,7 +223,7 @@ 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;
ucc_buffer_info_asymmetric_memtype_t *save = &task->bargs.asymmetric_save_info;
if (save->scratch != NULL) {
status = (task->bargs.args.coll_type == UCC_COLL_TYPE_GATHERV) ?
ucc_copy_asymmetric_buffer_v_out(task) :
Expand All @@ -242,6 +232,11 @@ static inline ucc_status_t ucc_task_complete(ucc_coll_task_t *task)
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);
Expand Down
6 changes: 3 additions & 3 deletions src/utils/ucc_coll_utils.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,9 +99,9 @@ ucc_coll_args_is_mem_symmetric(const ucc_coll_args_t *args,
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_update_asymmetric_buffer(ucc_coll_args_t *args,
ucc_team_h team,
ucc_asymmetric_save_info_t *save_info)
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;

Expand Down
6 changes: 3 additions & 3 deletions src/utils/ucc_coll_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -337,12 +337,12 @@ 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_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_update_asymmetric_buffer(ucc_coll_args_t *args,
ucc_coll_args_init_asymmetric_buffer(ucc_coll_args_t *args,
ucc_team_h team,
ucc_asymmetric_save_info_t *save_info);
ucc_buffer_info_asymmetric_memtype_t *save_info);

#endif
9 changes: 4 additions & 5 deletions test/gtest/asym_mem/test_asymmetric_memory.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,12 @@ class test_asymmetric_memory : public ucc::test,
ucc_coll_args_t *coll = (ucc_coll_args_t*)
calloc(1, sizeof(ucc_coll_args_t));

ctxs[i]->args = coll;

coll->coll_type = coll_type;
ctxs[i]->args = coll;
coll->coll_type = coll_type;
coll->src.info.mem_type = src_mem_type;
coll->src.info.count = (ucc_count_t)msglen * src_modifier;
coll->src.info.count = (ucc_count_t)msglen * src_modifier;
coll->src.info.datatype = UCC_DT_INT8;
coll->root = root;
coll->root = root;

UCC_CHECK(ucc_mc_alloc(&ctxs[i]->src_mc_header,
msglen * src_modifier, src_mem_type));
Expand Down

0 comments on commit c88eea3

Please sign in to comment.