Skip to content

Commit 56a3874

Browse files
wfaderhold21wfaderhold21
committed
CORE: various memh fixes
1 parent e7c538e commit 56a3874

1 file changed

Lines changed: 73 additions & 20 deletions

File tree

‎src/core/ucc_context.c‎

Lines changed: 73 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@
1717
#include "utils/ucc_debug.h"
1818
#include "ucc_progress_queue.h"
1919

20+
_Static_assert(UCC_MEM_MAP_TL_NAME_LEN == sizeof(size_t),
21+
"UCC_MEM_MAP_TL_NAME_LEN must equal sizeof(size_t) for pack layout");
22+
2023
static uint32_t ucc_context_seq_num = 0;
2124
static ucc_config_field_t ucc_context_config_table[] = {
2225
{"ESTIMATED_NUM_EPS", "0",
@@ -1226,6 +1229,14 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context,
12261229
ucc_error("cannot import NULL memory handle");
12271230
return UCC_ERR_INVALID_PARAM;
12281231
}
1232+
if (!*memh) {
1233+
ucc_error("cannot import NULL *memh");
1234+
return UCC_ERR_INVALID_PARAM;
1235+
}
1236+
if (!memh_size) {
1237+
ucc_error("memh_size cannot be NULL");
1238+
return UCC_ERR_INVALID_PARAM;
1239+
}
12291240
if (!params) {
12301241
ucc_error("params cannot be NULL");
12311242
return UCC_ERR_INVALID_PARAM;
@@ -1244,7 +1255,8 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context,
12441255
/* memh should have been used in exchanges or from a remote process,
12451256
addresses, etc. likely garbage. fix it */
12461257
local_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc(
1247-
ctx->n_tl_ctx, sizeof(ucc_mem_map_tl_t), "tl memh");
1258+
ucc_max(ctx->n_tl_ctx, (int)local_memh->num_tls),
1259+
sizeof(ucc_mem_map_tl_t), "tl memh");
12481260
if (!local_memh->tl_h) {
12491261
ucc_error("failed to allocate tl memh for import");
12501262
return UCC_ERR_NO_MEMORY;
@@ -1263,6 +1275,24 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context,
12631275
local_memh, &local_memh->tl_h[j]);
12641276
if (status < UCC_ERR_NOT_IMPLEMENTED) {
12651277
ucc_error("failed to import mem map memh %d", status);
1278+
/* unmap slots that were successfully mapped using
1279+
name-based matching (data is indexed by exporter slot) */
1280+
for (int m = 0; m < (int)local_memh->num_tls; m++) {
1281+
if (!local_memh->tl_h[m].tl_data) {
1282+
continue;
1283+
}
1284+
for (int k = 0; k < ctx->n_tl_ctx; k++) {
1285+
if (strcmp(local_memh->tl_h[m].tl_name,
1286+
tls->names[k]) == 0) {
1287+
ucc_tl_lib_t *rl = ucc_derived_of(
1288+
ctx->tl_ctx[k]->super.lib, ucc_tl_lib_t);
1289+
rl->iface->context.mem_unmap(
1290+
(const ucc_base_context_t *)ctx->tl_ctx[k],
1291+
mode, &local_memh->tl_h[m]);
1292+
break;
1293+
}
1294+
}
1295+
}
12661296
ucc_free(local_memh->tl_h);
12671297
return status;
12681298
}
@@ -1354,7 +1384,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context,
13541384
}
13551385
}
13561386
packed_buffers =
1357-
(void **)ucc_calloc(ctx->n_tl_ctx, sizeof(void *), "packed buffers");
1387+
(void **)ucc_calloc(ucc_max(ctx->n_tl_ctx, (int)local_memh->num_tls),
1388+
sizeof(void *), "packed buffers");
13581389
if (!packed_buffers) {
13591390
if (mode == UCC_MEM_MAP_MODE_EXPORT) {
13601391
ucc_free(local_memh->tl_h);
@@ -1403,6 +1434,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context,
14031434
ucc_free(local_memh->tl_h);
14041435
ucc_free(local_memh);
14051436
}
1437+
*memh = NULL;
1438+
*memh_size = 0;
14061439
return UCC_OK;
14071440
}
14081441

@@ -1448,6 +1481,10 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context,
14481481
exported_memh->tl_h = local_memh->tl_h;
14491482

14501483
for (i = 0; i < local_memh->num_tls; i++) {
1484+
if (!local_memh->tl_h[i].packed_size) {
1485+
ucc_free(packed_buffers[i]);
1486+
continue;
1487+
}
14511488
/* these are already packed in order */
14521489
strncpy(PTR_OFFSET(exported_memh->pack_buffer, offset),
14531490
local_memh->tl_h[i].tl_name, UCC_MEM_MAP_TL_NAME_LEN);
@@ -1482,10 +1519,10 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context,
14821519
i = ctx->n_tl_ctx;
14831520
failed_mem_map:
14841521
for (int j = 0; j < i; j++) {
1485-
tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t);
1486-
tl_lib->iface->context.mem_unmap((const ucc_base_context_t *)ctx,
1487-
UCC_MEM_MAP_MODE_EXPORT,
1488-
&local_memh->tl_h[j]);
1522+
tl_lib = ucc_derived_of(ctx->tl_ctx[j]->super.lib, ucc_tl_lib_t);
1523+
tl_lib->iface->context.mem_unmap(
1524+
(const ucc_base_context_t *)ctx->tl_ctx[j],
1525+
mode, &local_memh->tl_h[j]);
14891526
}
14901527
if (mode == UCC_MEM_MAP_MODE_EXPORT) {
14911528
ucc_free(local_memh->tl_h);
@@ -1507,13 +1544,28 @@ ucc_status_t ucc_mem_map(ucc_context_h context, ucc_mem_map_mode_t mode,
15071544
if (mode == UCC_MEM_MAP_MODE_IMPORT || mode == UCC_MEM_MAP_MODE_IMPORT_OFFLOAD) {
15081545
return ucc_mem_map_import(context, mode, params, memh_size, memh);
15091546
}
1510-
if (!params) {
1511-
ucc_error("params cannot be NULL");
1547+
if (!memh_size) {
1548+
ucc_error("memh_size cannot be NULL");
15121549
return UCC_ERR_INVALID_PARAM;
15131550
}
1514-
if (params->n_segments > 1) {
1515-
ucc_error("UCC only supports one mapping per call");
1516-
return UCC_ERR_INVALID_PARAM;
1551+
if (mode == UCC_MEM_MAP_MODE_EXPORT_OFFLOAD) {
1552+
if (!memh || !*memh) {
1553+
ucc_error("EXPORT_OFFLOAD requires a non-NULL *memh");
1554+
return UCC_ERR_INVALID_PARAM;
1555+
}
1556+
} else {
1557+
if (!params) {
1558+
ucc_error("params cannot be NULL");
1559+
return UCC_ERR_INVALID_PARAM;
1560+
}
1561+
if (params->n_segments != 1) {
1562+
ucc_error("UCC only supports one mapping per call");
1563+
return UCC_ERR_INVALID_PARAM;
1564+
}
1565+
if (!params->segments) {
1566+
ucc_error("params->segments cannot be NULL");
1567+
return UCC_ERR_INVALID_PARAM;
1568+
}
15171569
}
15181570
return ucc_mem_map_export(context, mode, params, memh_size, memh);
15191571
}
@@ -1538,34 +1590,35 @@ ucc_status_t ucc_mem_unmap(ucc_mem_map_mem_h *memh)
15381590
return UCC_ERR_INVALID_PARAM;
15391591
}
15401592

1541-
lmemh = *memh;
1542-
ctx = (ucc_context_t *)lmemh->context;
1543-
tls = &ctx->all_tls;
1593+
lmemh = *memh;
1594+
ctx = (ucc_context_t *)lmemh->context;
1595+
tls = &ctx->all_tls;
1596+
status = UCC_OK;
15441597
for (i = 0; i < ctx->n_tl_ctx; i++) {
15451598
for (j = 0; j < lmemh->num_tls; j++) {
15461599
if (strcmp(lmemh->tl_h[j].tl_name, tls->names[i]) == 0) {
1600+
ucc_status_t s;
15471601
tl_lib = ucc_derived_of(ctx->tl_ctx[i]->super.lib, ucc_tl_lib_t);
1548-
status = tl_lib->iface->context.mem_unmap(
1602+
s = tl_lib->iface->context.mem_unmap(
15491603
(const ucc_base_context_t *)ctx->tl_ctx[i], lmemh->mode,
15501604
&lmemh->tl_h[j]);
1551-
if (status < UCC_ERR_NOT_IMPLEMENTED) {
1605+
if (s < UCC_ERR_NOT_IMPLEMENTED) {
15521606
ucc_error("error during unmap operation on TL %s",
15531607
lmemh->tl_h[j].tl_name);
1554-
return status;
1608+
status = s;
15551609
}
15561610
}
15571611
}
15581612
}
15591613

15601614
/* Free the TL handles array if it was allocated separately */
1561-
if (lmemh->tl_h && (lmemh->mode == UCC_MEM_MAP_MODE_EXPORT ||
1562-
lmemh->mode == UCC_MEM_MAP_MODE_IMPORT)) {
1615+
if (lmemh->tl_h) {
15631616
ucc_free(lmemh->tl_h);
15641617
}
15651618

15661619
/* Free the memory handle structure */
15671620
ucc_free(lmemh);
15681621
*memh = NULL;
15691622

1570-
return UCC_OK;
1623+
return status;
15711624
}

0 commit comments

Comments
 (0)