OSHMEM: Add support for shmemx_malloc_with_hint()

oshmem/include/pshmemx.h

@@ -16,6 +16,11 @@
 extern "C" {
 #endif
 
+/*
+ * Symmetric heap routines
+ */
+OSHMEM_DECLSPEC  void* pshmemx_malloc_with_hint(size_t size, long hint);
+
 
 /*
  * Legacy API

oshmem/include/shmemx.h

@@ -18,11 +18,29 @@
 extern "C" {
 #endif
 
+enum {
+    SHMEM_HINT_NONE           = 0,
+    SHMEM_HINT_LOW_LAT_MEM    = 1 << 0,
+    SHMEM_HINT_HIGH_BW_MEM    = 1 << 1,
+    SHMEM_HINT_NEAR_NIC_MEM   = 1 << 2,
+    SHMEM_HINT_DEVICE_GPU_MEM = 1 << 3,
+    SHMEM_HINT_DEVICE_NIC_MEM = 1 << 4,
+
+    SHMEM_HINT_PSYNC          = 1 << 16,
+    SHMEM_HINT_PWORK          = 1 << 17,
+    SHMEM_HINT_ATOMICS        = 1 << 18
+};
+
 /*
  * All OpenSHMEM extension APIs that are not part of this specification must be defined in the shmemx.h include
  * file. These extensions shall use the shmemx_ prefix for all routine, variable, and constant names.
  */
 
+/*
+ * Symmetric heap routines
+ */
+OSHMEM_DECLSPEC  void* shmemx_malloc_with_hint(size_t size, long hint);
+
 /*
  * Elemental put routines
  */

oshmem/mca/memheap/base/base.h

@@ -41,23 +41,27 @@ OSHMEM_DECLSPEC int mca_memheap_base_select(void);
 extern int mca_memheap_base_already_opened;
 extern int mca_memheap_base_key_exchange;
 
-#define MCA_MEMHEAP_MAX_SEGMENTS    4 
-#define HEAP_SEG_INDEX  0
-#define SYMB_SEG_INDEX  1
-#define MCA_MEMHEAP_SEG_COUNT  (SYMB_SEG_INDEX+1)
+#define MCA_MEMHEAP_MAX_SEGMENTS    8
+#define HEAP_SEG_INDEX              0
 
 #define MEMHEAP_SEG_INVALID  0xFFFF
 
 
+typedef struct mca_memheap_base_config {
+    long            device_nic_mem_seg_size; /* Used for SHMEM_HINT_DEVICE_NIC_MEM */
+} mca_memheap_base_config_t;
+
+
 typedef struct mca_memheap_map {
     map_segment_t   mem_segs[MCA_MEMHEAP_MAX_SEGMENTS]; /* TODO: change into pointer array */
     int             n_segments;
     int             num_transports;
 } mca_memheap_map_t;
 
 extern mca_memheap_map_t mca_memheap_base_map;
+extern mca_memheap_base_config_t mca_memheap_base_config;
 
-int mca_memheap_base_alloc_init(mca_memheap_map_t *, size_t);
+int mca_memheap_base_alloc_init(mca_memheap_map_t *, size_t, long);
 void mca_memheap_base_alloc_exit(mca_memheap_map_t *);
 int mca_memheap_base_static_init(mca_memheap_map_t *);
 void mca_memheap_base_static_exit(mca_memheap_map_t *);
@@ -173,10 +177,12 @@ static inline int memheap_is_va_in_segment(void *va, int segno)
 
 static inline int memheap_find_segnum(void *va)
 {
-    if (OPAL_LIKELY(memheap_is_va_in_segment(va, SYMB_SEG_INDEX))) {
-        return SYMB_SEG_INDEX;
-    } else if (memheap_is_va_in_segment(va, HEAP_SEG_INDEX)) {
-        return HEAP_SEG_INDEX;
+    int i;
+
+    for (i = 0; i < mca_memheap_base_map.n_segments; i++) {
+        if (memheap_is_va_in_segment(va, i)) {
+            return i;
+        }
     }
     return MEMHEAP_SEG_INVALID;
 }
@@ -193,18 +199,17 @@ static inline void *map_segment_va2rva(mkey_segment_t *seg, void *va)
     return memheap_va2rva(va, seg->super.va_base, seg->rva_base);
 }
 
-static inline map_base_segment_t *map_segment_find_va(map_base_segment_t *segs, size_t elem_size, void *va) 
+static inline map_base_segment_t *map_segment_find_va(map_base_segment_t *segs,
+                                                      size_t elem_size, void *va)
 {
     map_base_segment_t *rseg;
+    int i;
 
-    rseg = (map_base_segment_t *)((char *)segs + elem_size * HEAP_SEG_INDEX);
-    if (OPAL_LIKELY(map_segment_is_va_in(rseg, va))) {
-        return rseg;
-    } 
-
-    rseg = (map_base_segment_t *)((char *)segs + elem_size * SYMB_SEG_INDEX);
-    if (OPAL_LIKELY(map_segment_is_va_in(rseg, va))) {
-        return rseg;
+    for (i = 0; i < MCA_MEMHEAP_MAX_SEGMENTS; i++) {
+        rseg = (map_base_segment_t *)((char *)segs + elem_size * i);
+        if (OPAL_LIKELY(map_segment_is_va_in(rseg, va))) {
+            return rseg;
+        }
     }
 
     return NULL;
@@ -214,21 +219,14 @@ void mkey_segment_init(mkey_segment_t *seg, sshmem_mkey_t *mkey, uint32_t segno)
 
 static inline map_segment_t *memheap_find_va(void* va)
 {
-    map_segment_t *s;
-
-    /* most probably there will be only two segments: heap and global data */
-    if (OPAL_LIKELY(memheap_is_va_in_segment(va, SYMB_SEG_INDEX))) {
-        s = &memheap_map->mem_segs[SYMB_SEG_INDEX];
-    } else if (memheap_is_va_in_segment(va, HEAP_SEG_INDEX)) {
-        s = &memheap_map->mem_segs[HEAP_SEG_INDEX];
-    } else if (memheap_map->n_segments - 2 > 0) {
-        s = bsearch(va,
-                    &memheap_map->mem_segs[SYMB_SEG_INDEX+1],
-                    memheap_map->n_segments - 2,
-                    sizeof(*s),
-                    mca_memheap_seg_cmp);
-    } else {
-        s = NULL;
+    map_segment_t *s = NULL;
+    int i;
+
+    for (i = 0; i < memheap_map->n_segments; i++) {
+        if (memheap_is_va_in_segment(va, i)) {
+            s = &memheap_map->mem_segs[i];
+            break;
+        }
     }
 
 #if MEMHEAP_BASE_DEBUG == 1

oshmem/mca/memheap/base/memheap_base_alloc.c

@@ -19,17 +19,21 @@
 #include "oshmem/mca/memheap/base/base.h"
 
 
-int mca_memheap_base_alloc_init(mca_memheap_map_t *map, size_t size)
+int mca_memheap_base_alloc_init(mca_memheap_map_t *map, size_t size, long hint)
 {
     int ret = OSHMEM_SUCCESS;
     char * seg_filename = NULL;
 
     assert(map);
-    assert(HEAP_SEG_INDEX == map->n_segments);
+    if (hint == 0) {
+        assert(HEAP_SEG_INDEX == map->n_segments);
+    } else {
+        assert(HEAP_SEG_INDEX < map->n_segments);
+    }
 
     map_segment_t *s = &map->mem_segs[map->n_segments];
     seg_filename = oshmem_get_unique_file_name(oshmem_my_proc_id());
-    ret = mca_sshmem_segment_create(s, seg_filename, size);
+    ret = mca_sshmem_segment_create(s, seg_filename, size, hint);
 
     if (OSHMEM_SUCCESS == ret) {
         map->n_segments++;
@@ -45,12 +49,34 @@ int mca_memheap_base_alloc_init(mca_memheap_map_t *map, size_t size)
 
 void mca_memheap_base_alloc_exit(mca_memheap_map_t *map)
 {
-    if (map) {
-        map_segment_t *s = &map->mem_segs[HEAP_SEG_INDEX];
+    int i;
+
+    if (!map) {
+        return;
+    }
+
+    for (i = 0; i < map->n_segments; ++i) {
+        map_segment_t *s = &map->mem_segs[i];
+        if (s->type != MAP_SEGMENT_STATIC) {
+            mca_sshmem_segment_detach(s, NULL);
+            mca_sshmem_unlink(s);
+        }
+    }
+}
 
-        assert(s);
+int mca_memheap_alloc_with_hint(size_t size, long hint, void** ptr)
+{
+    int i;
 
-        mca_sshmem_segment_detach(s, NULL);
-        mca_sshmem_unlink(s);
+    for (i = 0; i < mca_memheap_base_map.n_segments; i++) {
+        map_segment_t *s = &mca_memheap_base_map.mem_segs[i];
+        if (s->allocator && (hint && s->alloc_hints)) {
+            /* Do not fall back to default allocator since it will break the
+             * symmetry between PEs
+             */
+            return s->allocator->realloc(s, size, NULL, ptr);
+        }
     }
+
+    return MCA_MEMHEAP_CALL(alloc(size, ptr));
 }

oshmem/mca/memheap/base/memheap_base_frame.c

@@ -52,6 +52,12 @@ static int mca_memheap_base_register(mca_base_register_flag_t flags)
                                  MCA_BASE_VAR_SCOPE_READONLY,
                                  &mca_memheap_base_key_exchange);
 
+    mca_base_var_register("oshmem", "memheap", "base", "device_nic_mem_seg_size",
+                          "Size of memory block used for allocations with hint SHMEM_HINT_DEVICE_NIC_MEM",
+                          MCA_BASE_VAR_TYPE_LONG, NULL, 0,
+                          MCA_BASE_VAR_FLAG_SETTABLE, OPAL_INFO_LVL_3,
+                          MCA_BASE_VAR_SCOPE_LOCAL,
+                          &mca_memheap_base_config.device_nic_mem_seg_size);
 
     return OSHMEM_SUCCESS;
 }

oshmem/mca/memheap/base/memheap_base_mkey.c

@@ -749,7 +749,7 @@ void mkey_segment_init(mkey_segment_t *seg, sshmem_mkey_t *mkey, uint32_t segno)
 {
     map_segment_t *s;
 
-    if (segno >= MCA_MEMHEAP_SEG_COUNT) {
+    if (segno >= MCA_MEMHEAP_MAX_SEGMENTS) {
         return;
     }
 

oshmem/mca/memheap/base/memheap_base_select.c

@@ -21,6 +21,14 @@
 #include "oshmem/util/oshmem_util.h"
 #include "oshmem/mca/memheap/memheap.h"
 #include "oshmem/mca/memheap/base/base.h"
+#include "orte/mca/errmgr/errmgr.h"
+#include "oshmem/include/shmemx.h"
+#include "oshmem/mca/sshmem/base/base.h"
+
+
+mca_memheap_base_config_t mca_memheap_base_config = {
+    .device_nic_mem_seg_size = 0
+};
 
 mca_memheap_base_module_t mca_memheap = {0};
 
@@ -94,7 +102,7 @@ static memheap_context_t* _memheap_create(void)
 {
     int rc = OSHMEM_SUCCESS;
     static memheap_context_t context;
-    size_t user_size;
+    size_t user_size, size;
 
     user_size = _memheap_size();
     if (user_size < MEMHEAP_BASE_MIN_SIZE) {
@@ -105,7 +113,18 @@ static memheap_context_t* _memheap_create(void)
     /* Inititialize symmetric area */
     if (OSHMEM_SUCCESS == rc) {
         rc = mca_memheap_base_alloc_init(&mca_memheap_base_map,
-                                         user_size + MEMHEAP_BASE_PRIVATE_SIZE);
+                                         user_size + MEMHEAP_BASE_PRIVATE_SIZE, 0);
+    }
+
+    /* Initialize atomic symmetric area */
+    size = mca_memheap_base_config.device_nic_mem_seg_size;
+    if ((OSHMEM_SUCCESS == rc) && (size > 0)) {
+        rc = mca_memheap_base_alloc_init(&mca_memheap_base_map, size,
+                                         SHMEM_HINT_DEVICE_NIC_MEM);
+        if (rc == OSHMEM_ERR_NOT_IMPLEMENTED) {
+            /* do not treat NOT_IMPLEMENTED as error */
+            rc = OSHMEM_SUCCESS;
+        }
     }
 
     /* Inititialize static/global variables area */

oshmem/mca/memheap/base/memheap_base_static.c

@@ -49,7 +49,7 @@ int mca_memheap_base_static_init(mca_memheap_map_t *map)
     int ret = OSHMEM_SUCCESS;
 
     assert(map);
-    assert(SYMB_SEG_INDEX <= map->n_segments);
+    assert(HEAP_SEG_INDEX < map->n_segments);
 
     ret = _load_segments();
 

oshmem/mca/memheap/memheap.h

@@ -138,6 +138,8 @@ typedef struct mca_memheap_base_module_t mca_memheap_base_module_t;
 
 OSHMEM_DECLSPEC extern mca_memheap_base_module_t mca_memheap;
 
+int mca_memheap_alloc_with_hint(size_t size, long hint, void**);
+
 static inline int mca_memheap_base_mkey_is_shm(sshmem_mkey_t *mkey)
 {
     return (0 == mkey->len) && (MAP_SEGMENT_SHM_INVALID != (int)mkey->u.key);

oshmem/mca/spml/ucx/spml_ucx.c

@@ -36,6 +36,7 @@
 #include "oshmem/runtime/runtime.h"
 
 #include "oshmem/mca/spml/ucx/spml_ucx_component.h"
+#include "oshmem/mca/sshmem/ucx/sshmem_ucx.h"
 
 /* Turn ON/OFF debug output from build (default 0) */
 #ifndef SPML_UCX_PUT_DEBUG
@@ -267,7 +268,7 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
         OSHMEM_PROC_DATA(procs[i])->num_transports = 1;
         OSHMEM_PROC_DATA(procs[i])->transport_ids = spml_ucx_transport_ids;
 
-        for (j = 0; j < MCA_MEMHEAP_SEG_COUNT; j++) {
+        for (j = 0; j < MCA_MEMHEAP_MAX_SEGMENTS; j++) {
             mca_spml_ucx_ctx_default.ucp_peers[i].mkeys[j].key.rkey = NULL;
         }
 
@@ -438,7 +439,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
         }
 
     } else {
-        ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
+        mca_sshmem_ucx_segment_context_t *ctx = mem_seg->context;
+        ucx_mkey->mem_h = ctx->ucp_memh;
     }
 
     status = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
@@ -589,17 +591,19 @@ static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx
             goto error2;
         }
 
-        for (j = 0; j < MCA_MEMHEAP_SEG_COUNT; j++) {
+        for (j = 0; j < memheap_map->n_segments; j++) {
             mkey = &memheap_map->mem_segs[j].mkeys_cache[i][0];
             ucx_mkey = &ucx_ctx->ucp_peers[i].mkeys[j].key;
-            err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[i].ucp_conn,
-                                     mkey->u.data,
-                                     &ucx_mkey->rkey);
-            if (UCS_OK != err) {
-                SPML_UCX_ERROR("failed to unpack rkey");
-                goto error2;
+            if (mkey->u.data) {
+                err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[i].ucp_conn,
+                                         mkey->u.data,
+                                         &ucx_mkey->rkey);
+                if (UCS_OK != err) {
+                    SPML_UCX_ERROR("failed to unpack rkey");
+                    goto error2;
+                }
+                mca_spml_ucx_cache_mkey(ucx_ctx, mkey, j, i);
             }
-            mca_spml_ucx_cache_mkey(ucx_ctx, mkey, j, i);
         }
     }
 
@@ -747,6 +751,8 @@ int mca_spml_ucx_fence(shmem_ctx_t ctx)
     ucs_status_t err;
     mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
 
+    opal_atomic_wmb();
+
     err = ucp_worker_fence(ucx_ctx->ucp_worker);
     if (UCS_OK != err) {
          SPML_UCX_ERROR("fence failed: %s", ucs_status_string(err));
@@ -761,6 +767,8 @@ int mca_spml_ucx_quiet(shmem_ctx_t ctx)
     int ret;
     mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
 
+    opal_atomic_wmb();
+
     ret = opal_common_ucx_worker_flush(ucx_ctx->ucp_worker);
     if (OMPI_SUCCESS != ret) {
          oshmem_shmem_abort(-1);

oshmem/mca/spml/ucx/spml_ucx.h

@@ -61,7 +61,7 @@ typedef struct spml_ucx_cached_mkey spml_ucx_cached_mkey_t;
 
 struct ucp_peer {
     ucp_ep_h                 ucp_conn;
-    spml_ucx_cached_mkey_t   mkeys[MCA_MEMHEAP_SEG_COUNT];
+    spml_ucx_cached_mkey_t   mkeys[MCA_MEMHEAP_MAX_SEGMENTS];
 };
 typedef struct ucp_peer ucp_peer_t;
 

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -314,7 +314,7 @@ static void _ctx_cleanup(mca_spml_ucx_ctx_t *ctx)
     del_procs = malloc(sizeof(*del_procs) * nprocs);
 
     for (i = 0; i < nprocs; ++i) {
-        for (j = 0; j < MCA_MEMHEAP_SEG_COUNT; j++) {
+        for (j = 0; j < memheap_map->n_segments; j++) {
             if (ctx->ucp_peers[i].mkeys[j].key.rkey != NULL) {
                 ucp_rkey_destroy(ctx->ucp_peers[i].mkeys[j].key.rkey);
             }

oshmem/mca/sshmem/base/base.h

@@ -31,7 +31,7 @@ extern char* mca_sshmem_base_backing_file_dir;
 OSHMEM_DECLSPEC int
 mca_sshmem_segment_create(map_segment_t *ds_buf,
                           const char *file_name,
-                          size_t size);
+                          size_t size, long hint);
 
 OSHMEM_DECLSPEC void *
 mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey);