opal/common/ucx: introduce internal UCX request in wpool.

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

Author: xinzhao3

opal/mca/common/ucx/Makefile.am

@@ -16,13 +16,15 @@ headers =                         \
         common_ucx.h              \
         common_ucx_int.h          \
         common_ucx_wpool.h        \
-        common_ucx_wpool_int.h
+        common_ucx_wpool_int.h    \
+        common_ucx_request.h
 
 # Source files
 
 sources = \
         common_ucx.c              \
-        common_ucx_wpool.c
+        common_ucx_wpool.c        \
+        common_ucx_request.c
 
 # Help file
 

opal/mca/common/ucx/common_ucx.h

@@ -15,5 +15,6 @@
 
 #include "common_ucx_int.h"
 #include "common_ucx_wpool.h"
+#include "common_ucx_request.h"
 
 #endif

opal/mca/common/ucx/common_ucx_int.h

@@ -2,6 +2,7 @@
 #define COMMON_UCX_INT_H
 
 #include "opal_config.h"
+#include "common_ucx_request.h"
 
 #include <stdint.h>
 
@@ -170,10 +171,11 @@ 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_send_callback_t req_handler,
                                                  ucp_worker_h worker)
 {
     return ucp_atomic_fetch_nb(ep, opcode, value, result, op_size,
-                               remote_addr, rkey, opal_common_ucx_empty_complete_cb);
+                               remote_addr, rkey, req_handler);
 }
 
 static inline

opal/mca/common/ucx/common_ucx_request.c

@@ -0,0 +1,17 @@
+#include "common_ucx_request.h"
+
+OPAL_DECLSPEC void
+opal_common_ucx_req_init(void *request) {
+    opal_common_ucx_request_t *req = (opal_common_ucx_request_t *)request;
+    req->ext_req = NULL;
+    req->ext_cb = NULL;
+}
+
+OPAL_DECLSPEC void
+opal_common_ucx_req_completion(void *request, ucs_status_t status) {
+    opal_common_ucx_request_t *req = (opal_common_ucx_request_t *)request;
+    if (req->ext_cb != NULL) {
+        (*req->ext_cb)(req->ext_req);
+    }
+    ucp_request_release(req);
+}

opal/mca/common/ucx/common_ucx_request.h

@@ -0,0 +1,17 @@
+#ifndef COMMON_UCX_REQUEST_H
+#define COMMON_UCX_REQUEST_H
+
+#include "opal_config.h"
+#include <ucp/api/ucp.h>
+
+typedef void (*opal_common_ucx_user_req_handler_t)(void *request);
+
+typedef struct {
+    void *ext_req;
+    opal_common_ucx_user_req_handler_t ext_cb;
+} opal_common_ucx_request_t;
+
+OPAL_DECLSPEC void opal_common_ucx_req_init(void *request);
+OPAL_DECLSPEC void opal_common_ucx_req_completion(void *request, ucs_status_t status);
+
+#endif // COMMON_UCX_REQUEST_H

opal/mca/common/ucx/common_ucx_wpool.c

@@ -126,9 +126,7 @@ opal_common_ucx_wpool_free(opal_common_ucx_wpool_t *wpool)
 
 OPAL_DECLSPEC int
 opal_common_ucx_wpool_init(opal_common_ucx_wpool_t *wpool,
-                               int proc_world_size,
-                               ucp_request_init_callback_t req_init_ptr,
-                               size_t req_size, bool enable_mt)
+                           int proc_world_size, bool enable_mt)
 {
     ucp_config_t *config = NULL;
     ucp_params_t context_params;
@@ -164,8 +162,8 @@ opal_common_ucx_wpool_init(opal_common_ucx_wpool_t *wpool,
                               UCP_FEATURE_AMO64;
     context_params.mt_workers_shared = (enable_mt ? 1 : 0);
     context_params.estimated_num_eps = proc_world_size;
-    context_params.request_init = req_init_ptr;
-    context_params.request_size = req_size;
+    context_params.request_init = opal_common_ucx_req_init;
+    context_params.request_size = sizeof(opal_common_ucx_request_t);
 
     status = ucp_init(&context_params, config, &wpool->ucp_ctx);
     ucp_config_release(config);
@@ -1272,4 +1270,3 @@ opal_common_ucx_wpmem_fence(opal_common_ucx_wpmem_t *mem) {
     /* TODO */
     return OPAL_SUCCESS;
 }
-

opal/mca/common/ucx/common_ucx_wpool.h

@@ -5,6 +5,7 @@
 #include "opal_config.h"
 
 #include "common_ucx_int.h"
+#include "common_ucx_request.h"
 #include <stdint.h>
 
 #include <ucp/api/ucp.h>
@@ -176,9 +177,7 @@ static inline void opal_common_ucx_wpool_dbg_init(void)
 OPAL_DECLSPEC opal_common_ucx_wpool_t * opal_common_ucx_wpool_allocate(void);
 OPAL_DECLSPEC void opal_common_ucx_wpool_free(opal_common_ucx_wpool_t *wpool);
 OPAL_DECLSPEC int opal_common_ucx_wpool_init(opal_common_ucx_wpool_t *wpool,
-                                             int proc_world_size,
-                                             ucp_request_init_callback_t req_init_ptr,
-                                             size_t req_size, bool enable_mt);
+                                             int proc_world_size, bool enable_mt);
 OPAL_DECLSPEC void opal_common_ucx_wpool_finalize(opal_common_ucx_wpool_t *wpool);
 OPAL_DECLSPEC void opal_common_ucx_wpool_progress(opal_common_ucx_wpool_t *wpool);
 
@@ -394,27 +393,40 @@ opal_common_ucx_wpmem_fetch(opal_common_ucx_wpmem_t *mem,
 
 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_atomic_fetch_op_t opcode,
+                               uint64_t value,
+                               int target, void *buffer, size_t len,
+                               uint64_t rem_addr,
+                               opal_common_ucx_user_req_handler_t user_req_cb,
+                               void *user_req_ptr)
 {
     ucp_ep_h ep = NULL;
     ucp_rkey_h rkey = NULL;
     opal_common_ucx_winfo_t *winfo = NULL;
     int rc = OPAL_SUCCESS;
+    opal_common_ucx_request_t *req;
+
     rc = opal_common_ucx_tlocal_fetch(mem, target, &ep, &rkey, &winfo);
     if(OPAL_UNLIKELY(OPAL_SUCCESS != rc)){
         MCA_COMMON_UCX_ERROR("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);
+    req = opal_common_ucx_atomic_fetch_nb(ep, opcode, value, buffer, len,
+                                          rem_addr, rkey, opal_common_ucx_req_completion,
+                                          winfo->worker);
     opal_mutex_unlock(&winfo->mutex);
+
+    if (UCS_PTR_IS_PTR(req)) {
+        req->ext_req = user_req_ptr;
+        req->ext_cb = user_req_cb;
+    } else {
+        if (user_req_cb != NULL) {
+            (*user_req_cb)(user_req_ptr);
+        }
+    }
+
     return rc;
 }