Merge pull request #2845 from alex-mikheev/topic/oshmem_mem_prefetch

oshmem: spml: add memory allocation hook

Author: yosefe

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
     },
 
@@ -289,7 +289,7 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
                             &ep_params,
                             &mca_spml_ucx.ucp_peers[i].ucp_conn);
         if (UCS_OK != err) {
-            SPML_ERROR("ucp_ep_create failed!!!\n");
+            SPML_ERROR("ucp_ep_create failed: %s\n", ucs_status_string(err));
             goto error2;
         }
         OSHMEM_PROC_DATA(procs[i])->num_transports = 1;
@@ -372,7 +372,7 @@ void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int
             mkey->u.data, 
             &ucx_mkey->rkey); 
     if (UCS_OK != err) {
-        SPML_ERROR("failed to unpack rkey");
+        SPML_ERROR("failed to unpack rkey: %s", ucs_status_string(err));
         goto error_fatal;
     }
 
@@ -385,6 +385,34 @@ void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int
     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 : %s",
+                   addr, (unsigned long long)length, ucs_status_string(status));
+    }
+}
+
 sshmem_mkey_t *mca_spml_ucx_register(void* addr,
                                          size_t size,
                                          uint64_t shmid,
@@ -540,7 +568,7 @@ int mca_spml_ucx_fence(void)
 
     err = ucp_worker_flush(mca_spml_ucx.ucp_worker);
     if (UCS_OK != err) {
-        SPML_ERROR("fence failed");
+        SPML_ERROR("fence failed: %s", ucs_status_string(err));
          oshmem_shmem_abort(-1);
          return OSHMEM_ERROR;
     }
@@ -553,7 +581,7 @@ int mca_spml_ucx_quiet(void)
 
     err = ucp_worker_flush(mca_spml_ucx.ucp_worker);
     if (UCS_OK != err) {
-        SPML_ERROR("fence failed");
+        SPML_ERROR("fence failed: %s", ucs_status_string(err));
          oshmem_shmem_abort(-1);
          return OSHMEM_ERROR;
     }

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
     }