oshmem: spml: add memory allocation hook

The hook is called from memheap when memory range
is going to be allocated by smalloc(), realloc() and others.

ucx spml uses this hook to call ucp_mem_advise in order to speedup
non blocking memory mapping.

Signed-off-by: Alex Mikheev <alexm@mellanox.com>

oshmem/mca/memheap/buddy/memheap_buddy.c

@@ -469,7 +469,7 @@ static int _do_alloc(uint32_t order,
     }
 
     *p_buff = (void*) addr;
-    /* no barrier because it is not required by spec! */
+    MCA_SPML_CALL(memuse_hook(addr, 1<<order));
     return OSHMEM_SUCCESS;
 
     alloc_error: _buddy_free(&memheap_buddy, offset, order, heap);

oshmem/mca/memheap/ptmalloc/memheap_ptmalloc.c

@@ -85,6 +85,7 @@ int mca_memheap_ptmalloc_alloc(size_t size, void** p_buff)
     if (NULL == *p_buff)
         return OSHMEM_ERROR;
 
+    MCA_SPML_CALL(memuse_hook(*p_buff, size));
     return OSHMEM_SUCCESS;
 }
 
@@ -113,6 +114,7 @@ int mca_memheap_ptmalloc_align(size_t align, size_t size, void **p_buff)
     if (NULL == *p_buff)
         return OSHMEM_ERROR;
 
+    MCA_SPML_CALL(memuse_hook(*p_buff, size));
     return OSHMEM_SUCCESS;
 }
 
@@ -132,6 +134,7 @@ int mca_memheap_ptmalloc_realloc(size_t new_size,
     if (!*p_new_buff)
         return OSHMEM_ERR_OUT_OF_RESOURCE;
 
+    MCA_SPML_CALL(memuse_hook(*p_new_buff, new_size));
     return OSHMEM_SUCCESS;
 }
 

oshmem/mca/spml/base/base.h

@@ -84,6 +84,7 @@ OSHMEM_DECLSPEC int mca_spml_base_get_nb(void *dst_addr,
                                          int src,
                                          void **handle);
 
+OSHMEM_DECLSPEC void mca_spml_base_memuse_hook(void *addr, size_t length);
 /*
  * MCA framework
  */

oshmem/mca/spml/base/spml_base.c

@@ -177,3 +177,7 @@ int mca_spml_base_get_nb(void *dst_addr, size_t size,
 {
     return OSHMEM_ERROR;
 }
+
+void mca_spml_base_memuse_hook(void *addr, size_t length)
+{
+}

oshmem/mca/spml/ikrit/spml_ikrit.c

@@ -171,6 +171,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
         mca_spml_ikrit_fence,
         mca_spml_ikrit_cache_mkeys,
         mca_spml_base_rmkey_free,
+        mca_spml_base_memuse_hook,
 
         (void*)&mca_spml_ikrit
     }

oshmem/mca/spml/spml.h

@@ -277,6 +277,15 @@ typedef int (*mca_spml_base_module_fence_fn_t)(void);
  */
 typedef int (*mca_spml_base_module_wait_nb_fn_t)(void *);
 
+/**
+ * Called by memheap when memory is allocated by shmalloc(),
+ * shcalloc(), shmemalign() or shrealloc()
+ *
+ * @param addr   base address of the registered buffer.
+ * @param size   the size of the buffer to be registered.
+ */
+typedef void (*mca_spml_base_module_memuse_hook_fn_t)(void *, size_t);
+
 /**
  *  SPML instance.
  */
@@ -304,6 +313,8 @@ struct mca_spml_base_module_1_0_0_t {
 
     mca_spml_base_module_mkey_unpack_fn_t spml_rmkey_unpack;
     mca_spml_base_module_mkey_free_fn_t   spml_rmkey_free;
+
+    mca_spml_base_module_memuse_hook_fn_t spml_memuse_hook;
     void *self;
 };
 

oshmem/mca/spml/ucx/spml_ucx.c

@@ -43,7 +43,6 @@
 #define SPML_UCX_PUT_DEBUG    0
 #endif
 
-
 mca_spml_ucx_t mca_spml_ucx = {
     {
         /* Init mca_spml_base_module_t */
@@ -65,6 +64,7 @@ mca_spml_ucx_t mca_spml_ucx = {
                                every spml */
         mca_spml_ucx_rmkey_unpack,
         mca_spml_ucx_rmkey_free,
+        mca_spml_ucx_memuse_hook,
         (void*)&mca_spml_ucx
     },
 
@@ -385,6 +385,34 @@ error_fatal:
     return;
 }
 
+void mca_spml_ucx_memuse_hook(void *addr, size_t length)
+{
+    int my_pe = oshmem_my_proc_id();
+    spml_ucx_mkey_t   *ucx_mkey;
+    ucp_mem_advise_params_t params;
+    ucs_status_t status;
+
+    if (!(mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX))) {
+        return;
+    }
+
+    ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[HEAP_SEG_INDEX].key;
+
+    params.field_mask = UCP_MEM_ADVISE_PARAM_FIELD_ADDRESS |
+                        UCP_MEM_ADVISE_PARAM_FIELD_LENGTH |
+                        UCP_MEM_ADVISE_PARAM_FIELD_ADVICE;
+
+    params.address = addr;
+    params.length  = length;
+    params.advice  = UCP_MADV_WILLNEED;
+
+    status = ucp_mem_advise(mca_spml_ucx.ucp_context, ucx_mkey->mem_h, &params);
+    if (UCS_OK != status) {
+        SPML_ERROR("ucp_mem_advise failed addr %p len %llu",
+                   addr, (unsigned long long)length);
+    }
+}
+
 sshmem_mkey_t *mca_spml_ucx_register(void* addr,
                                          size_t size,
                                          uint64_t shmid,

oshmem/mca/spml/ucx/spml_ucx.h

@@ -108,6 +108,8 @@ extern sshmem_mkey_t *mca_spml_ucx_register(void* addr,
                                                 int *count);
 extern int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys);
 
+extern void mca_spml_ucx_memuse_hook(void *addr, size_t length);
+
 extern void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id);
 extern void mca_spml_ucx_rmkey_free(sshmem_mkey_t *mkey);
 

oshmem/mca/spml/yoda/spml_yoda.c

@@ -65,6 +65,7 @@ mca_spml_yoda_module_t mca_spml_yoda = {
         mca_spml_yoda_fence,
         mca_spml_base_rmkey_unpack,
         mca_spml_base_rmkey_free,
+        mca_spml_base_memuse_hook,
 
         (void *)&mca_spml_yoda
     }