TL/UCP: various memh fixes

Author: wfaderhold21

src/components/tl/ucp/tl_ucp_context.c

@@ -716,8 +716,7 @@ ucc_status_t ucc_tl_ucp_mem_map_offload_import(ucc_tl_ucp_context_t *ctx,
         if (strncmp(tl_h->tl_name, name, UCC_MEM_MAP_TL_NAME_LEN - 1) == 0) {
             break;
         }
-        /* this is not the index, skip this section of buffer if exists */
-        offset += *packed_size;
+        offset += UCC_MEM_MAP_TL_NAME_LEN + sizeof(size_t) + *packed_size;
     }
     if (i == l_memh->num_tls) {
         tl_error(ctx->super.super.lib, "unable to find TL UCP in memory handle");
@@ -767,6 +766,7 @@ ucc_status_t ucc_tl_ucp_mem_map(const ucc_base_context_t *context, ucc_mem_map_m
         if (UCC_OK != ucc_status) {
             tl_error(ctx->super.super.lib, "failed to import memory handle");
             ucc_free(m_data);
+            tl_h->tl_data = NULL;
             return ucc_status;
         }
     }
@@ -785,6 +785,9 @@ ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, ucc_mem_map
     }
 
     data = (ucc_tl_ucp_memh_data_t *)memh->tl_data;
+    if (!data) {
+        return UCC_OK;
+    }
     if (mode == UCC_MEM_MAP_MODE_EXPORT || mode == UCC_MEM_MAP_MODE_EXPORT_OFFLOAD) {
         status = ucp_mem_unmap(ctx->worker.ucp_context, data->rinfo.mem_h);
         if (status != UCS_OK) {
@@ -802,7 +805,16 @@ ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, ucc_mem_map
             data->rinfo.packed_key = NULL;
         }
     } else if (mode == UCC_MEM_MAP_MODE_IMPORT || mode == UCC_MEM_MAP_MODE_IMPORT_OFFLOAD) {
-        // need to free rkeys (data->rkey) , packed memh (data->packed_memh)
+        ucc_status_t ret = UCC_OK;
+        if (mode == UCC_MEM_MAP_MODE_IMPORT_OFFLOAD && data->rinfo.mem_h) {
+            ucs_status_t ucs_st = ucp_mem_unmap(ctx->worker.ucp_context,
+                                                 data->rinfo.mem_h);
+            if (ucs_st != UCS_OK) {
+                tl_error(ctx->super.super.lib,
+                         "ucp_mem_unmap (import offload) failed: %d", ucs_st);
+                ret = ucs_status_to_ucc_status(ucs_st);
+            }
+        }
         if (data->packed_memh) {
             ucp_memh_buffer_release(data->packed_memh, NULL);
         }
@@ -812,6 +824,11 @@ ucc_status_t ucc_tl_ucp_mem_unmap(const ucc_base_context_t *context, ucc_mem_map
         if (data->rkey) {
             ucp_rkey_destroy(data->rkey);
         }
+        if (ret != UCC_OK) {
+            ucc_free(data);
+            memh->tl_data = NULL;
+            return ret;
+        }
     } else {
         ucc_error("Unknown mem map mode entered: %d", mode);
         return UCC_ERR_INVALID_PARAM;
@@ -872,7 +889,7 @@ ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context,
      *
      * packed_key_size | packed_memh_size | packed_key | packed_memh
      */
-    packed_buffer = ucc_malloc(sizeof(size_t) * 2 + data->packed_memh_len +
+    packed_buffer = ucc_malloc(UCC_TL_UCP_MEMH_TL_HEADER_SIZE + data->packed_memh_len +
                                    data->rinfo.packed_key_len,
                                "packed buffer");
     if (!packed_buffer) {
@@ -899,9 +916,11 @@ ucc_status_t ucc_tl_ucp_memh_pack(const ucc_base_context_t *context,
 
 failed_alloc_buffer:
     ucp_rkey_buffer_release(data->rinfo.packed_key);
+    data->rinfo.packed_key = NULL;
 failed_rkey_pack:
     if (data->packed_memh) {
         ucp_memh_buffer_release(data->packed_memh, NULL);
+        data->packed_memh = NULL;
     }
     return ucc_status;
 }

src/core/ucc_context.c

@@ -11,15 +11,13 @@
 #include "components/tl/ucc_tl.h"
 #include "components/topo/ucc_sysinfo.h"
 #include "utils/ucc_malloc.h"
+#include "utils/ucc_math.h"
 #include "utils/ucc_log.h"
 #include "utils/ucc_list.h"
 #include "utils/ucc_string.h"
 #include "utils/ucc_debug.h"
 #include "ucc_progress_queue.h"
 
-_Static_assert(UCC_MEM_MAP_TL_NAME_LEN == sizeof(size_t),
-               "UCC_MEM_MAP_TL_NAME_LEN must equal sizeof(size_t) for pack layout");
-
 static uint32_t ucc_context_seq_num = 0;
 static ucc_config_field_t ucc_context_config_table[] = {
     {"ESTIMATED_NUM_EPS", "0",
@@ -1246,6 +1244,7 @@ ucc_status_t ucc_mem_map_import(ucc_context_h        context,
 
     if (ctx->n_tl_ctx == 0) {
         ucc_debug("No TL contexts available for import");
+        local_memh->tl_h    = NULL;
         local_memh->mode    = mode;
         local_memh->context = ctx;
         *memh_size          = 0;
@@ -1328,7 +1327,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h         context,
     int                       tls;
     int                       tlh_index;
 
-    if (!params) {
+    if (mode != UCC_MEM_MAP_MODE_EXPORT_OFFLOAD && !params) {
         ucc_error("params cannot be NULL");
         return UCC_ERR_INVALID_PARAM;
     }
@@ -1433,8 +1432,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h         context,
         if (mode == UCC_MEM_MAP_MODE_EXPORT) {
             ucc_free(local_memh->tl_h);
             ucc_free(local_memh);
+            *memh = NULL;
         }
-        *memh      = NULL;
         *memh_size = 0;
         return UCC_OK;
     }
@@ -1478,14 +1477,20 @@ ucc_status_t ucc_mem_map_export(ucc_context_h         context,
             }
         }
     } else {
-        exported_memh->tl_h = local_memh->tl_h;
+        exported_memh->tl_h = (ucc_mem_map_tl_t *)ucc_calloc(
+            packed_tls, sizeof(ucc_mem_map_tl_t), "packed tl memh");
+        if (!exported_memh->tl_h) {
+            ucc_error("failed to allocate handle for exported buffers' tl handles");
+            status = UCC_ERR_NO_MEMORY;
+            ucc_free(exported_memh);
+            goto failed_pack;
+        }
 
-        for (i = 0; i < local_memh->num_tls; i++) {
+        for (i = 0, offset = 0, tls = 0; i < (int)local_memh->num_tls; i++) {
             if (!local_memh->tl_h[i].packed_size) {
                 ucc_free(packed_buffers[i]);
                 continue;
             }
-            /* these are already packed in order */
             strncpy(PTR_OFFSET(exported_memh->pack_buffer, offset),
                    local_memh->tl_h[i].tl_name, UCC_MEM_MAP_TL_NAME_LEN);
             offset += UCC_MEM_MAP_TL_NAME_LEN;
@@ -1496,6 +1501,8 @@ ucc_status_t ucc_mem_map_export(ucc_context_h         context,
                    packed_buffers[i], local_memh->tl_h[i].packed_size);
             ucc_free(packed_buffers[i]);
             offset += local_memh->tl_h[i].packed_size;
+            memcpy(&exported_memh->tl_h[tls++], &local_memh->tl_h[i],
+                   sizeof(ucc_mem_map_tl_t));
         }
     }
 
@@ -1506,9 +1513,7 @@ ucc_status_t ucc_mem_map_export(ucc_context_h         context,
     exported_memh->num_tls     = packed_tls;
     *memh                      = exported_memh;
     *memh_size                 = sizeof(ucc_mem_map_memh_t) + offset;
-    if (mode == UCC_MEM_MAP_MODE_EXPORT) {
-        ucc_free(local_memh->tl_h);
-    }
+    ucc_free(local_memh->tl_h);
     ucc_free(local_memh);
     ucc_free(packed_buffers);
     return UCC_OK;