oshmem: sshmem: adds UCX allocator

Also some code cleanup

Author: alex-mikheev

oshmem/mca/spml/ucx/spml_ucx.c

@@ -396,7 +396,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
     size_t len;
     int my_pe = oshmem_my_proc_id();
     ucp_mem_map_params_t mem_map_params;
-    int seg;
+    int segno;
+    map_segment_t *mem_seg;
     unsigned flags;
 
     *count = 0;
@@ -405,26 +406,32 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
         return NULL;
     }
 
-    seg = memheap_find_segnum(addr);
+    segno   = memheap_find_segnum(addr);
+    mem_seg = memheap_find_seg(segno);
 
-    ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[seg].key;
+    ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[segno].key;
     mkeys[0].spml_context = ucx_mkey;
 
-    flags = 0;
-    if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
-        flags = UCP_MEM_MAP_NONBLOCK;
-    }
+    /* if possible use mem handle already created by ucx allocator */
+    if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
+        flags = 0;
+        if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
+            flags = UCP_MEM_MAP_NONBLOCK;
+        }
 
-    mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
-                                UCP_MEM_MAP_PARAM_FIELD_LENGTH |
-                                UCP_MEM_MAP_PARAM_FIELD_FLAGS;
-    mem_map_params.address    = addr;
-    mem_map_params.length     = size;
-    mem_map_params.flags      = flags;
+        mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
+            UCP_MEM_MAP_PARAM_FIELD_LENGTH |
+            UCP_MEM_MAP_PARAM_FIELD_FLAGS;
+        mem_map_params.address    = addr;
+        mem_map_params.length     = size;
+        mem_map_params.flags      = flags;
 
-    err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
-    if (UCS_OK != err) {
-        goto error_out;
+        err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
+        if (UCS_OK != err) {
+            goto error_out;
+        }
+    } else {
+        ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
     }
 
     err = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h, 
@@ -448,9 +455,9 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
     }
 
     mkeys[0].len     = len;
-    mkeys[0].va_base = mem_map_params.address;
+    mkeys[0].va_base = addr;
     *count = 1;
-    mca_spml_ucx_cache_mkey(&mkeys[0], seg, my_pe);
+    mca_spml_ucx_cache_mkey(&mkeys[0], segno, my_pe);
     return mkeys;
 
 error_unmap:
@@ -464,6 +471,7 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
 int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
 {
     spml_ucx_mkey_t   *ucx_mkey;
+    map_segment_t *mem_seg;
 
     MCA_SPML_CALL(fence());
     if (!mkeys)
@@ -472,8 +480,12 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
     if (!mkeys[0].spml_context) 
         return OSHMEM_SUCCESS;
 
-    ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
-    ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
+    mem_seg = memheap_find_va(mkeys[0].va_base);
+    
+    if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
+        ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
+        ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
+    }
 
     if (0 < mkeys[0].len) {
         ucp_rkey_buffer_release(mkeys[0].u.data);

oshmem/mca/sshmem/base/base.h

@@ -33,10 +33,6 @@ mca_sshmem_segment_create(map_segment_t *ds_buf,
                           const char *file_name,
                           size_t size);
 
-OSHMEM_DECLSPEC int
-mca_sshmem_ds_copy(const map_segment_t *from,
-                   map_segment_t *to);
-
 OSHMEM_DECLSPEC void *
 mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey);
 
@@ -148,23 +144,12 @@ OSHMEM_DECLSPEC extern mca_base_framework_t oshmem_sshmem_base_framework;
         "Warning %s:%d - %s()", __SSHMEM_FILE__, __LINE__, __func__, __VA_ARGS__)
 
 
+OSHMEM_DECLSPEC extern void shmem_ds_reset(map_segment_t *ds_buf);
+
 /*
  * Get unique file name
  */
-static inline char * oshmem_get_unique_file_name(uint64_t pe)
-{
-    char *file_name = NULL;
-
-    assert(mca_sshmem_base_backing_file_dir);
-
-    if (NULL == (file_name = calloc(OPAL_PATH_MAX, sizeof(char)))) {
-        return NULL;
-    }
-
-    snprintf(file_name, OPAL_PATH_MAX, "%s/shmem_job_%u_pe_%llu", mca_sshmem_base_backing_file_dir, ORTE_PROC_MY_NAME->jobid, (unsigned long long)pe);
-
-    return file_name;
-}
+OSHMEM_DECLSPEC extern char * oshmem_get_unique_file_name(uint64_t pe);
 
 END_C_DECLS
 

oshmem/mca/sshmem/base/sshmem_base_wrappers.c

@@ -15,7 +15,6 @@
 #include "oshmem/mca/sshmem/sshmem.h"
 #include "oshmem/mca/sshmem/base/base.h"
 
-/* ////////////////////////////////////////////////////////////////////////// */
 int
 mca_sshmem_segment_create(map_segment_t *ds_buf,
                           const char *file_name,
@@ -28,19 +27,6 @@ mca_sshmem_segment_create(map_segment_t *ds_buf,
     return mca_sshmem_base_module->segment_create(ds_buf, file_name, size);
 }
 
-/* ////////////////////////////////////////////////////////////////////////// */
-int
-mca_sshmem_ds_copy(const map_segment_t *from,
-                   map_segment_t *to)
-{
-    if (!mca_sshmem_base_selected) {
-        return OSHMEM_ERROR;
-    }
-
-    return mca_sshmem_base_module->ds_copy(from, to);
-}
-
-/* ////////////////////////////////////////////////////////////////////////// */
 void *
 mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
 {
@@ -51,7 +37,6 @@ mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
     return mca_sshmem_base_module->segment_attach(ds_buf, mkey);
 }
 
-/* ////////////////////////////////////////////////////////////////////////// */
 int
 mca_sshmem_segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
 {
@@ -62,7 +47,6 @@ mca_sshmem_segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
     return mca_sshmem_base_module->segment_detach(ds_buf, mkey);
 }
 
-/* ////////////////////////////////////////////////////////////////////////// */
 int
 mca_sshmem_unlink(map_segment_t *ds_buf)
 {
@@ -73,3 +57,31 @@ mca_sshmem_unlink(map_segment_t *ds_buf)
     return mca_sshmem_base_module->unlink(ds_buf);
 }
 
+
+char * oshmem_get_unique_file_name(uint64_t pe)
+{
+    char *file_name = NULL;
+
+    assert(mca_sshmem_base_backing_file_dir);
+
+    if (NULL == (file_name = calloc(OPAL_PATH_MAX, sizeof(char)))) {
+        return NULL;
+    }
+
+    snprintf(file_name, OPAL_PATH_MAX, "%s/shmem_job_%u_pe_%llu", mca_sshmem_base_backing_file_dir, ORTE_PROC_MY_NAME->jobid, (unsigned long long)pe);
+
+    return file_name;
+}
+
+
+void
+shmem_ds_reset(map_segment_t *ds_buf)
+{
+    MAP_SEGMENT_RESET_FLAGS(ds_buf);
+    ds_buf->seg_id = MAP_SEGMENT_SHM_INVALID;
+    ds_buf->super.va_base = 0;
+    ds_buf->super.va_end = 0;
+    ds_buf->seg_size = 0;
+    ds_buf->type = MAP_SEGMENT_UNKNOWN;
+}
+

oshmem/mca/sshmem/mmap/sshmem_mmap_module.c

@@ -64,10 +64,6 @@ segment_create(map_segment_t *ds_buf,
                const char *file_name,
                size_t size);
 
-static int
-ds_copy(const map_segment_t *from,
-        map_segment_t *to);
-
 static void *
 segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey);
 
@@ -88,44 +84,13 @@ mca_sshmem_mmap_module_t mca_sshmem_mmap_module = {
     {
         module_init,
         segment_create,
-        ds_copy,
         segment_attach,
         segment_detach,
         segment_unlink,
         module_finalize
     }
 };
 
-/* ////////////////////////////////////////////////////////////////////////// */
-/* private utility functions */
-/* ////////////////////////////////////////////////////////////////////////// */
-
-/* ////////////////////////////////////////////////////////////////////////// */
-/**
- * completely resets the contents of *ds_buf
- */
-static inline void
-shmem_ds_reset(map_segment_t *ds_buf)
-{
-    OPAL_OUTPUT_VERBOSE(
-        (70, oshmem_sshmem_base_framework.framework_output,
-         "%s: %s: shmem_ds_resetting "
-         "(id: %d, size: %lu, name: %s)\n",
-         mca_sshmem_mmap_component.super.base_version.mca_type_name,
-         mca_sshmem_mmap_component.super.base_version.mca_component_name,
-         ds_buf->seg_id, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
-    );
-
-    MAP_SEGMENT_RESET_FLAGS(ds_buf);
-    ds_buf->seg_id = MAP_SEGMENT_SHM_INVALID;
-    ds_buf->super.va_base = 0;
-    ds_buf->super.va_end = 0;
-    ds_buf->seg_size = 0;
-    ds_buf->type = MAP_SEGMENT_UNKNOWN;
-    unlink(ds_buf->seg_name);
-    memset(ds_buf->seg_name, '\0', sizeof(ds_buf->seg_name));
-}
-
 /* ////////////////////////////////////////////////////////////////////////// */
 static int
 module_init(void)
@@ -142,30 +107,6 @@ module_finalize(void)
     return OSHMEM_SUCCESS;
 }
 
-/* ////////////////////////////////////////////////////////////////////////// */
-static int
-ds_copy(const map_segment_t *from,
-        map_segment_t *to)
-{
-    memcpy(to, from, sizeof(map_segment_t));
-
-    OPAL_OUTPUT_VERBOSE(
-        (70, oshmem_sshmem_base_framework.framework_output,
-         "%s: %s: ds_copy complete "
-         "from: (id: %d, size: %lu, "
-         "name: %s flags: 0x%02x) "
-         "to: (id: %d, size: %lu, "
-         "name: %s flags: 0x%02x)\n",
-         mca_sshmem_mmap_component.super.base_version.mca_type_name,
-         mca_sshmem_mmap_component.super.base_version.mca_component_name,
-         from->seg_id, (unsigned long)from->seg_size, from->seg_name,
-         from->flags, to->seg_id, (unsigned long)to->seg_size, to->seg_name,
-         to->flags)
-    );
-
-    return OSHMEM_SUCCESS;
-}
-/* ////////////////////////////////////////////////////////////////////////// */
 
 static int
 segment_create(map_segment_t *ds_buf,
@@ -225,11 +166,11 @@ segment_create(map_segment_t *ds_buf,
     OPAL_OUTPUT_VERBOSE(
           (70, oshmem_sshmem_base_framework.framework_output,
            "%s: %s: create %s "
-           "(id: %d, addr: %p size: %lu, name: %s)\n",
+           "(id: %d, addr: %p size: %lu)\n",
            mca_sshmem_mmap_component.super.base_version.mca_type_name,
            mca_sshmem_mmap_component.super.base_version.mca_component_name,
            (rc ? "failure" : "successful"),
-           ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
+           ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
       );
 
     return rc;
@@ -316,10 +257,10 @@ segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
     OPAL_OUTPUT_VERBOSE(
         (70, oshmem_sshmem_base_framework.framework_output,
          "%s: %s: attach successful "
-            "(id: %d, addr: %p size: %lu, name: %s | va_base: 0x%p len: %d key %llx)\n",
+            "(id: %d, addr: %p size: %lu | va_base: 0x%p len: %d key %llx)\n",
             mca_sshmem_mmap_component.super.base_version.mca_type_name,
             mca_sshmem_mmap_component.super.base_version.mca_component_name,
-            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name,
+            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size,
             mkey->va_base, mkey->len, (unsigned long long)mkey->u.key)
         );
 
@@ -338,10 +279,10 @@ segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
     OPAL_OUTPUT_VERBOSE(
         (70, oshmem_sshmem_base_framework.framework_output,
          "%s: %s: detaching "
-            "(id: %d, addr: %p size: %lu, name: %s)\n",
+            "(id: %d, addr: %p size: %lu)\n",
             mca_sshmem_mmap_component.super.base_version.mca_type_name,
             mca_sshmem_mmap_component.super.base_version.mca_component_name,
-            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
+            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
     );
 
     munmap((void *)ds_buf->super.va_base, ds_buf->seg_size);
@@ -363,10 +304,10 @@ segment_unlink(map_segment_t *ds_buf)
     OPAL_OUTPUT_VERBOSE(
         (70, oshmem_sshmem_base_framework.framework_output,
          "%s: %s: unlinking "
-            "(id: %d, addr: %p size: %lu, name: %s)\n",
+            "(id: %d, addr: %p size: %lu)\n",
             mca_sshmem_mmap_component.super.base_version.mca_type_name,
             mca_sshmem_mmap_component.super.base_version.mca_component_name,
-            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
+            ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
     );
 
     /* don't completely reset.  in particular, only reset

oshmem/mca/sshmem/sshmem.h

@@ -21,7 +21,6 @@
  *
  * - module_init
  * - segment_create
- * - ds_copy
  * - segment_attach
  * - segment_detach
  * - unlink
@@ -74,20 +73,6 @@ typedef struct mca_sshmem_base_component_2_0_0_t mca_sshmem_base_component_t;
 typedef int
 (*mca_sshmem_base_module_init_fn_t)(void);
 
-/**
- * copy shmem data structure information pointed to by from to the structure
- * pointed to by to.
- *
- * @param from  source pointer (IN).
- *
- * @param to    destination pointer (OUT).
- *
- * @return OSHMEM_SUCCESS on success.
- */
-typedef int
-(*mca_sshmem_base_ds_copy_fn_t)(const map_segment_t *from,
-                                map_segment_t *to);
-
 /**
  * create a new shared memory segment and initialize members in structure
  * pointed to by ds_buf.
@@ -153,7 +138,6 @@ typedef int (*mca_sshmem_base_module_finalize_fn_t)(void);
 struct mca_sshmem_base_module_2_0_0_t {
     mca_sshmem_base_module_init_fn_t            module_init;
     mca_sshmem_base_module_segment_create_fn_t  segment_create;
-    mca_sshmem_base_ds_copy_fn_t                ds_copy;
     mca_sshmem_base_module_segment_attach_fn_t  segment_attach;
     mca_sshmem_base_module_segment_detach_fn_t  segment_detach;
     mca_sshmem_base_module_unlink_fn_t          unlink;