TL/UCP: Enable multiple packed buffers

openucx · Jan 15, 2025 · e66d36b · e66d36b
1 parent 6f13191
commit e66d36b
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git 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
@@ -246,16 +246,27 @@ 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;
 
     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
@@ -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);