diff --git a/src/components/tl/ucp/tl_ucp_context.c b/src/components/tl/ucp/tl_ucp_context.c index b9dd6b8c29..bf2784d7b0 100644 --- a/src/components/tl/ucp/tl_ucp_context.c +++ b/src/components/tl/ucp/tl_ucp_context.c @@ -570,7 +570,6 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, ucp_mem_h mh = NULL; ucs_status_t status; ucp_memh_pack_params_t pack_params; - //size_t pack_len; // basically an import here if (type == UCC_MEM_MAP_TYPE_GLOBAL) { diff --git a/src/components/tl/ucp/tl_ucp_sendrecv.h b/src/components/tl/ucp/tl_ucp_sendrecv.h index 38786cd2ec..fa02790abe 100644 --- a/src/components/tl/ucp/tl_ucp_sendrecv.h +++ b/src/components/tl/ucp/tl_ucp_sendrecv.h @@ -246,6 +246,9 @@ static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_ ucc_mem_map_memh_t *memh = map_memh; ucc_tl_ucp_memh_data_t *tl_data = (ucc_tl_ucp_memh_data_t *)memh->tl_h[tl_index].tl_data; uint64_t base, end; + ucs_status_t ucs_status; + int i; + size_t offset; base = (uint64_t)memh->address; end = base + memh->len; @@ -253,9 +256,17 @@ static inline ucc_status_t ucc_tl_ucp_check_memh(ucp_ep_h *ep, void *va, uint64_ if ((uint64_t)va >= base && (uint64_t)va < end) { *rva = (uint64_t)PTR_OFFSET(memh->address, ((uint64_t)va - (uint64_t)memh->address)); if (NULL == tl_data->rkey) { - // this is incorrect as pack buffer will have size, then packed data - ucs_status_t ucs_status = - ucp_ep_rkey_unpack(*ep, memh->pack_buffer + sizeof(size_t) * 3, + offset = 0; + /* find pack location for tl */ + for (i = 0; i < tl_index; i++) { + size_t *p = PTR_OFFSET(memh->pack_buffer, offset); + if (p[0] == tl_index) { + break; + } + offset += p[1]; + } + ucs_status = + ucp_ep_rkey_unpack(*ep, PTR_OFFSET(memh->pack_buffer, offset + sizeof(size_t) * 4), &tl_data->rkey); if (UCS_OK != ucs_status) { return ucs_status_to_ucc_status(ucs_status); diff --git a/src/core/ucc_context.c b/src/core/ucc_context.c index c7b5685181..5586a85fbd 100644 --- a/src/core/ucc_context.c +++ b/src/core/ucc_context.c @@ -1171,6 +1171,9 @@ ucc_status_t ucc_mem_map_import(ucc_context_h context, } strncpy(local_memh->tl_h[i].tl_name, tls->names[i], 8); } + local_memh->type = UCC_MEM_MAP_TYPE_GLOBAL; + /* fix context as it will be incorrect on a different system */ + local_memh->context = ctx; return status; } @@ -1241,7 +1244,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, /* allocate exported memh, copy items over */ exported_memh = (ucc_mem_map_memh_t *)ucc_calloc( - 1, sizeof(ucc_mem_map_memh_t) + total_pack_size + sizeof(size_t) * ctx->n_tl_ctx, "exported memh"); + 1, sizeof(ucc_mem_map_memh_t) + total_pack_size + 2 * sizeof(size_t) * ctx->n_tl_ctx, "exported memh"); if (!exported_memh) { ucc_error("failed to allocate handle for exported buffers"); return UCC_ERR_NO_MEMORY; @@ -1250,9 +1253,13 @@ ucc_status_t ucc_mem_map_export(ucc_context_h context, /* copying */ exported_memh->tl_h = local_memh->tl_h; for (i = 0, offset = 0; i < ctx->n_tl_ctx; i++) { + uint64_t tl_index = i; if (local_memh->tl_h[i].packed_size == 0) { continue; } + memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), + &tl_index, sizeof(size_t)); + offset += sizeof(size_t); memcpy(PTR_OFFSET(exported_memh->pack_buffer, offset), &local_memh->tl_h[i].packed_size, sizeof(size_t)); offset += sizeof(size_t);