oshmem: sshmem: add UCX allocator

Signed-off-by: Alex Mikheev <[email protected]>

oshmem/mca/spml/ucx/spml_ucx.c

@@ -425,7 +425,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;
@@ -434,26 +435,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, 
@@ -477,9 +484,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:
@@ -493,6 +500,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)
@@ -501,8 +509,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/ucx/Makefile.am

@@ -0,0 +1,40 @@
+# Copyright (c) 2014      Mellanox Technologies, Inc.
+#                         All rights reserved.
+# $COPYRIGHT$
+#
+# Additional copyrights may follow
+#
+# $HEADER$
+#
+
+#dist_oshmemdata_DATA = help-oshmem-sshmem-mmap.txt
+
+AM_CPPFLAGS = $(sshmem_ucx_CPPFLAGS)
+
+sources = \
+        sshmem_ucx.h \
+        sshmem_ucx_component.c \
+        sshmem_ucx_module.c
+
+# Make the output library in this directory, and name it either
+# mca_<type>_<name>.la (for DSO builds) or libmca_<type>_<name>.la
+# (for static builds).
+
+if MCA_BUILD_oshmem_sshmem_ucx_DSO
+component_noinst =
+component_install = mca_sshmem_ucx.la
+else
+component_noinst = libmca_sshmem_ucx.la
+component_install =
+endif
+
+mcacomponentdir = $(oshmemlibdir)
+mcacomponent_LTLIBRARIES = $(component_install)
+mca_sshmem_ucx_la_SOURCES = $(sources)
+mca_sshmem_ucx_la_LDFLAGS = -module -avoid-version $(sshmem_ucx_LDFLAGS)
+mca_sshmem_ucx_la_LIBADD = $(sshmem_ucx_LIBS)
+
+noinst_LTLIBRARIES = $(component_noinst)
+libmca_sshmem_ucx_la_SOURCES =$(sources)
+libmca_sshmem_ucx_la_LDFLAGS = -module -avoid-version $(sshmem_ucx_LDFLAGS)
+

oshmem/mca/sshmem/ucx/configure.m4

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2017      Mellanox Technologies, Inc.
+ *                         All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ */
+
+# MCA_oshmem_sshmem_ucx_CONFIG([action-if-can-compile],
+#                    [action-if-cant-compile])
+# ------------------------------------------------
+AC_DEFUN([MCA_oshmem_sshmem_ucx_CONFIG],[
+    AC_CONFIG_FILES([oshmem/mca/sshmem/ucx/Makefile])
+
+    OMPI_CHECK_UCX([sshmem_ucx],
+                   [sshmem_ucx_happy="yes"],
+                   [sshmem_ucx_happy="no"])
+
+    AS_IF([test "$sshmem_ucx_happy" = "yes"],
+          [$1],
+          [$2])
+
+
+    # substitute in the things needed to build ucx 
+    AC_SUBST([sshmem_ucx_CFLAGS])
+    AC_SUBST([sshmem_ucx_CPPFLAGS])
+    AC_SUBST([sshmem_ucx_LDFLAGS])
+    AC_SUBST([sshmem_ucx_LIBS])
+])dnl
+

oshmem/mca/sshmem/ucx/sshmem_ucx.h

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2017      Mellanox Technologies, Inc.
+ *                         All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ */
+
+#ifndef MCA_SSHMEM_UCX_EXPORT_H
+#define MCA_SSHMEM_UCX_EXPORT_H
+
+#include "oshmem_config.h"
+
+#include "oshmem/mca/sshmem/sshmem.h"
+
+BEGIN_C_DECLS
+
+/**
+ * globally exported variable to hold the ucx component.
+ */
+typedef struct mca_sshmem_ucx_component_t {
+    /* base component struct */
+    mca_sshmem_base_component_t super;
+    /* priority for ucx component */
+    int priority;
+} mca_sshmem_ucx_component_t;
+
+OSHMEM_MODULE_DECLSPEC extern mca_sshmem_ucx_component_t
+mca_sshmem_ucx_component;
+
+typedef struct mca_sshmem_ucx_module_t {
+    mca_sshmem_base_module_t super;
+} mca_sshmem_ucx_module_t;
+extern mca_sshmem_ucx_module_t mca_sshmem_ucx_module;
+
+END_C_DECLS
+
+#endif /* MCA_SHMEM_UCX_EXPORT_H */

oshmem/mca/sshmem/ucx/sshmem_ucx_component.c

@@ -0,0 +1,129 @@
+/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil -*- */
+/*
+ * Copyright (c) 2017      Mellanox Technologies, Inc.
+ *                         All rights reserved.
+ * $COPYRIGHT$
+ *
+ * Additional copyrights may follow
+ *
+ * $HEADER$
+ */
+
+#include "oshmem_config.h"
+
+#include "opal/constants.h"
+
+#include "oshmem/mca/sshmem/sshmem.h"
+#include "oshmem/mca/sshmem/base/base.h"
+#include "oshmem/mca/spml/base/base.h"
+
+#include "sshmem_ucx.h"
+
+/**
+ * public string showing the shmem ompi_ucx component version number
+ */
+const char *mca_sshmem_ucx_component_version_string =
+    "OSHMEM ucx sshmem MCA component version " OSHMEM_VERSION;
+
+
+/**
+ * local functions
+ */
+static int ucx_register(void);
+static int ucx_open(void);
+static int ucx_close(void);
+static int ucx_query(mca_base_module_t **module, int *priority);
+static int ucx_runtime_query(mca_base_module_t **module,
+                             int *priority,
+                             const char *hint);
+
+/**
+ * instantiate the public struct with all of our public information
+ * and pointers to our public functions in it
+ */
+mca_sshmem_ucx_component_t mca_sshmem_ucx_component = {
+    /* ////////////////////////////////////////////////////////////////////// */
+    /* super */
+    /* ////////////////////////////////////////////////////////////////////// */
+    {
+        /**
+         * common MCA component data
+         */
+        .base_version = {
+            MCA_SSHMEM_BASE_VERSION_2_0_0,
+
+            /* component name and version */
+            .mca_component_name = "ucx",
+            MCA_BASE_MAKE_VERSION(component, OSHMEM_MAJOR_VERSION, OSHMEM_MINOR_VERSION,
+                                  OSHMEM_RELEASE_VERSION),
+
+            .mca_open_component = ucx_open,
+            .mca_close_component = ucx_close,
+            .mca_query_component = ucx_query,
+            .mca_register_component_params = ucx_register,
+        },
+        /* MCA v2.0.0 component meta data */
+        .base_data = {
+            /* the component is checkpoint ready */
+            MCA_BASE_METADATA_PARAM_CHECKPOINT
+        },
+        .runtime_query = ucx_runtime_query,
+    },
+};
+
+static int
+ucx_runtime_query(mca_base_module_t **module,
+                   int *priority,
+                   const char *hint)
+{
+    /* check that spml ucx was selected. Otherwise disqualify */
+    if (strcmp(mca_spml_base_selected_component.spmlm_version.mca_component_name, "ucx")) {
+        *module = NULL;
+        return OSHMEM_ERR_NOT_AVAILABLE;
+    }
+
+    *priority = mca_sshmem_ucx_component.priority;
+    /* use lowest priority because UCX does not support
+     * fixed address mapping yet
+     */
+    *priority = 0;
+    *module = (mca_base_module_t *)&mca_sshmem_ucx_module.super;
+    return OPAL_SUCCESS;
+}
+
+static int
+ucx_register(void)
+{
+    /* (default) priority - set high to make ucx the default */
+    mca_sshmem_ucx_component.priority = 100;
+    mca_base_component_var_register (&mca_sshmem_ucx_component.super.base_version,
+                                     "priority", "Priority for sshmem ucx "
+                                     "component (default: 100)", MCA_BASE_VAR_TYPE_INT,
+                                     NULL, 0, MCA_BASE_VAR_FLAG_SETTABLE,
+                                     OPAL_INFO_LVL_3,
+                                     MCA_BASE_VAR_SCOPE_ALL_EQ,
+                                     &mca_sshmem_ucx_component.priority);
+
+    return OSHMEM_SUCCESS;
+}
+
+static int
+ucx_open(void)
+{
+    return OSHMEM_SUCCESS;
+}
+
+static int
+ucx_query(mca_base_module_t **module, int *priority)
+{
+    *priority = mca_sshmem_ucx_component.priority;
+    *module = (mca_base_module_t *)&mca_sshmem_ucx_module.super;
+    return OSHMEM_SUCCESS;
+}
+
+static int
+ucx_close(void)
+{
+    return OSHMEM_SUCCESS;
+}
+