From 450ea6684c9369153d8f71580eb0e929ceb89fc7 Mon Sep 17 00:00:00 2001
From: Igor Ivanov <Igor.Ivanov@itseez.com>
Date: Mon, 14 Mar 2016 14:10:28 +0200
Subject: [PATCH] oshmem: Align OSHMEM API with spec v1.3 (update spml/ucx)

---
 oshmem/mca/spml/ucx/spml_ucx.c | 30 ++++++++++++++++++++++++++++--
 oshmem/mca/spml/ucx/spml_ucx.h |  5 +++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/oshmem/mca/spml/ucx/spml_ucx.c b/oshmem/mca/spml/ucx/spml_ucx.c
index 2ccbfc108d..028d95b7fb 100644
--- a/oshmem/mca/spml/ucx/spml_ucx.c
+++ b/oshmem/mca/spml/ucx/spml_ucx.c
@@ -53,9 +53,9 @@ mca_spml_ucx_t mca_spml_ucx = {
         mca_spml_ucx_deregister,
         mca_spml_base_oob_get_mkeys,
         mca_spml_ucx_put,
-        mca_spml_base_put_nb, /* todo: mca_spml_ucx_put_nb, */
+        mca_spml_ucx_put_nb,
         mca_spml_ucx_get,
-        mca_spml_base_get_nb, /* todo: mca_spml_ucx_get_nb, */
+        mca_spml_ucx_get_nb,
         mca_spml_ucx_recv,
         mca_spml_ucx_send,
         mca_spml_base_wait,
@@ -391,6 +391,19 @@ int mca_spml_ucx_get(void *src_addr, size_t size, void *dst_addr, int src)
     return ucx_status_to_oshmem(status);
 }
 
+int mca_spml_ucx_get_nb(void *src_addr, size_t size, void *dst_addr, int src, void **handle)
+{
+    void *rva;
+    ucs_status_t status;
+    spml_ucx_mkey_t *ucx_mkey;
+
+    ucx_mkey = mca_spml_ucx_get_mkey(src, src_addr, &rva);
+    status = ucp_get_nbi(mca_spml_ucx.ucp_peers[src].ucp_conn, dst_addr, size,
+                     (uint64_t)rva, ucx_mkey->rkey);
+
+    return ucx_status_to_oshmem(status);
+}
+
 int mca_spml_ucx_put(void* dst_addr, size_t size, void* src_addr, int dst)
 {
     void *rva;
@@ -404,6 +417,19 @@ int mca_spml_ucx_put(void* dst_addr, size_t size, void* src_addr, int dst)
     return ucx_status_to_oshmem(status);
 }
 
+int mca_spml_ucx_put_nb(void* dst_addr, size_t size, void* src_addr, int dst, void **handle)
+{
+    void *rva;
+    ucs_status_t status;
+    spml_ucx_mkey_t *ucx_mkey;
+
+    ucx_mkey = mca_spml_ucx_get_mkey(dst, dst_addr, &rva);
+    status = ucp_put_nbi(mca_spml_ucx.ucp_peers[dst].ucp_conn, src_addr, size,
+                     (uint64_t)rva, ucx_mkey->rkey);
+
+    return ucx_status_to_oshmem(status);
+}
+
 int mca_spml_ucx_fence(void)
 {
     ucs_status_t err;
diff --git a/oshmem/mca/spml/ucx/spml_ucx.h b/oshmem/mca/spml/ucx/spml_ucx.h
index b7fe35a1e2..424347b3e8 100644
--- a/oshmem/mca/spml/ucx/spml_ucx.h
+++ b/oshmem/mca/spml/ucx/spml_ucx.h
@@ -70,6 +70,11 @@ extern int mca_spml_ucx_get(void* dst_addr,
                               size_t size,
                               void* src_addr,
                               int src);
+extern int mca_spml_ucx_get_nb(void* dst_addr,
+                              size_t size,
+                              void* src_addr,
+                              int src,
+                              void **handle);
 
 extern int mca_spml_ucx_put(void* dst_addr,
                               size_t size,