opal/common/ucx: Add fetch primitives to wpool

Signed-off-by: Artem Polyakov <artpol84@gmail.com>

opal/mca/common/ucx/common_ucx_int.h

@@ -166,6 +166,16 @@ int opal_common_ucx_atomic_fetch(ucp_ep_h ep, ucp_atomic_fetch_op_t opcode,
     return opal_common_ucx_wait_request(request, worker, "ucp_atomic_fetch_nb");
 }
 
+static inline
+ucs_status_ptr_t opal_common_ucx_atomic_fetch_nb(ucp_ep_h ep, ucp_atomic_fetch_op_t opcode,
+                                                 uint64_t value, void *result, size_t op_size,
+                                                 uint64_t remote_addr, ucp_rkey_h rkey,
+                                                 ucp_worker_h worker)
+{
+    return ucp_atomic_fetch_nb(ep, opcode, value, result, op_size,
+                               remote_addr, rkey, opal_common_ucx_empty_complete_cb);
+}
+
 static inline
 int opal_common_ucx_atomic_cswap(ucp_ep_h ep, uint64_t compare,
                                  uint64_t value, void *result, size_t op_size,

opal/mca/common/ucx/common_ucx_wpool.h

@@ -344,6 +344,70 @@ opal_common_ucx_wpmem_post(opal_common_ucx_wpmem_t *mem, ucp_atomic_post_op_t op
     return rc;
 }
 
+static inline int
+opal_common_ucx_wpmem_fetch(opal_common_ucx_wpmem_t *mem,
+                            ucp_atomic_fetch_op_t opcode, uint64_t value,
+                            int target, void *buffer, size_t len,
+                            uint64_t rem_addr)
+{
+    ucp_ep_h ep = NULL;
+    ucp_rkey_h rkey = NULL;
+    opal_common_ucx_winfo_t *winfo = NULL;
+    ucs_status_t status;
+    int rc = OPAL_SUCCESS;
+
+    rc = opal_common_ucx_tlocal_fetch(mem, target, &ep, &rkey, &winfo);
+    if(OPAL_SUCCESS != rc){
+        MCA_COMMON_UCX_VERBOSE(1, "tlocal_fetch failed: %d", rc);
+        return rc;
+    }
+    DBG_OUT("opal_common_ucx_mem_fetch(after _tlocal_fetch): mem = %p, ep = %p, rkey = %p, winfo = %p\n",
+            (void *)mem, (void *)ep, (void *)rkey, (void *)winfo);
+
+    /* Perform the operation */
+    opal_mutex_lock(&winfo->mutex);
+    status = opal_common_ucx_atomic_fetch(ep, opcode, value,
+                                          buffer, len,
+                                          rem_addr, rkey,
+                                          winfo->worker);
+    if (status != UCS_OK) {
+        MCA_COMMON_UCX_VERBOSE(1, "ucp_atomic_cswap64 failed: %d", status);
+        rc = OPAL_ERROR;
+    } else {
+        DBG_OUT("opal_common_ucx_mem_fetch(after opal_common_ucx_atomic_fetch): ep = %p, rkey = %p\n",
+                (void *)ep, (void *)rkey);
+    }
+    opal_mutex_unlock(&winfo->mutex);
+
+    return rc;
+}
+
+static inline int
+opal_common_ucx_wpmem_fetch_nb(opal_common_ucx_wpmem_t *mem,
+                                 ucp_atomic_fetch_op_t opcode,
+                                 uint64_t value,
+                                 int target, void *buffer, size_t len,
+                                 uint64_t rem_addr, ucs_status_ptr_t *ptr)
+{
+    ucp_ep_h ep = NULL;
+    ucp_rkey_h rkey = NULL;
+    opal_common_ucx_winfo_t *winfo = NULL;
+    int rc = OPAL_SUCCESS;
+    rc = opal_common_ucx_tlocal_fetch(mem, target, &ep, &rkey, &winfo);
+    if(OPAL_SUCCESS != rc){
+        MCA_COMMON_UCX_VERBOSE(1, "tlocal_fetch failed: %d", rc);
+        return rc;
+    }
+    /* Perform the operation */
+    opal_mutex_lock(&winfo->mutex);
+    (*ptr) = opal_common_ucx_atomic_fetch_nb(ep, opcode, value,
+                                             buffer, len,
+                                             rem_addr, rkey,
+                                             winfo->worker);
+    opal_mutex_unlock(&winfo->mutex);
+    return rc;
+}
+
 END_C_DECLS
 
 #endif // COMMON_UCX_WPOOL_H