From 986ca000f8398a00be07ccaf1ab4f7a0c84e69d9 Mon Sep 17 00:00:00 2001
From: Alex Mikheev <alexm@mellanox.com>
Date: Tue, 10 Jan 2017 14:21:36 +0200
Subject: [PATCH] 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      |  2 +-
 .../mca/memheap/ptmalloc/memheap_ptmalloc.c   |  3 ++
 oshmem/mca/spml/base/base.h                   |  1 +
 oshmem/mca/spml/base/spml_base.c              |  4 +++
 oshmem/mca/spml/ikrit/spml_ikrit.c            |  1 +
 oshmem/mca/spml/spml.h                        | 11 +++++++
 oshmem/mca/spml/ucx/spml_ucx.c                | 30 ++++++++++++++++++-
 oshmem/mca/spml/ucx/spml_ucx.h                |  2 ++
 oshmem/mca/spml/yoda/spml_yoda.c              |  1 +
 9 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/oshmem/mca/memheap/buddy/memheap_buddy.c b/oshmem/mca/memheap/buddy/memheap_buddy.c
index c7715f4dd1..b3f4ca1df9 100644
--- a/oshmem/mca/memheap/buddy/memheap_buddy.c
+++ b/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);
diff --git a/oshmem/mca/memheap/ptmalloc/memheap_ptmalloc.c b/oshmem/mca/memheap/ptmalloc/memheap_ptmalloc.c
index e416f66bc5..f6f5b48bb7 100644
--- a/oshmem/mca/memheap/ptmalloc/memheap_ptmalloc.c
+++ b/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;
 }
 
diff --git a/oshmem/mca/spml/base/base.h b/oshmem/mca/spml/base/base.h
index a0fd613d4d..4a0eb3e735 100644
--- a/oshmem/mca/spml/base/base.h
+++ b/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
  */
diff --git a/oshmem/mca/spml/base/spml_base.c b/oshmem/mca/spml/base/spml_base.c
index bdaf013438..75ccda7393 100644
--- a/oshmem/mca/spml/base/spml_base.c
+++ b/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)
+{
+}
diff --git a/oshmem/mca/spml/ikrit/spml_ikrit.c b/oshmem/mca/spml/ikrit/spml_ikrit.c
index 1374ceb455..e9f89ce90b 100644
--- a/oshmem/mca/spml/ikrit/spml_ikrit.c
+++ b/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
     }
diff --git a/oshmem/mca/spml/spml.h b/oshmem/mca/spml/spml.h
index f081b8b7b1..16b372b2dc 100644
--- a/oshmem/mca/spml/spml.h
+++ b/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;
 };
 
diff --git a/oshmem/mca/spml/ucx/spml_ucx.c b/oshmem/mca/spml/ucx/spml_ucx.c
index c0d67deef2..fbdc13f9f5 100644
--- a/oshmem/mca/spml/ucx/spml_ucx.c
+++ b/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,
diff --git a/oshmem/mca/spml/ucx/spml_ucx.h b/oshmem/mca/spml/ucx/spml_ucx.h
index 4b248e026a..b524031d3f 100644
--- a/oshmem/mca/spml/ucx/spml_ucx.h
+++ b/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);
 
diff --git a/oshmem/mca/spml/yoda/spml_yoda.c b/oshmem/mca/spml/yoda/spml_yoda.c
index 8e87c58df7..ebdceab8c9 100644
--- a/oshmem/mca/spml/yoda/spml_yoda.c
+++ b/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
     }