modifying osc mt code

contrib/platform/mellanox/optimized

@@ -1,4 +1,4 @@
-enable_mca_no_build=coll-ml,btl-uct
+enable_mca_no_build=coll-ml
 enable_debug_symbols=yes
 enable_orterun_prefix_by_default=yes
 with_verbs=no

ompi/mca/osc/ucx/osc_ucx.h

@@ -24,16 +24,9 @@
 #define OMPI_OSC_UCX_ATTACH_MAX    32
 #define OMPI_OSC_UCX_RKEY_BUF_MAX  1024
 
-typedef struct ompi_osc_ucx_win_info {
-    ucp_rkey_h rkey;
-    uint64_t addr;
-    bool rkey_init;
-} ompi_osc_ucx_win_info_t;
-
 typedef struct ompi_osc_ucx_component {
     ompi_osc_base_component_t super;
-    ucp_context_h ucp_context;
-    ucp_worker_h ucp_worker;
+    opal_common_ucx_wpool_t *wpool;
     bool enable_mpi_threads;
     opal_free_list_t requests; /* request free list for the r* communication variants */
     bool env_initialized; /* UCX environment is initialized or not */
@@ -97,12 +90,10 @@ typedef struct ompi_osc_ucx_state {
 typedef struct ompi_osc_ucx_module {
     ompi_osc_base_module_t super;
     struct ompi_communicator_t *comm;
-    ucp_mem_h memh; /* remote accessible memory */
     int flavor;
     size_t size;
-    ucp_mem_h state_memh;
-    ompi_osc_ucx_win_info_t *win_info_array;
-    ompi_osc_ucx_win_info_t *state_info_array;
+    uint64_t *addrs;
+    uint64_t *state_addrs;
     int disp_unit; /* if disp_unit >= 0, then everyone has the same
                     * disp unit size; if disp_unit == -1, then we
                     * need to look at disp_units */
@@ -122,6 +113,9 @@ typedef struct ompi_osc_ucx_module {
     uint64_t req_result;
     int *start_grp_ranks;
     bool lock_all_is_nocheck;
+    opal_common_ucx_ctx_t *ctx;
+    opal_common_ucx_mem_t *mem;
+    opal_common_ucx_mem_t *state_mem;
 } ompi_osc_ucx_module_t;
 
 typedef enum locktype {

ompi/mca/osc/ucx/osc_ucx_active_target.c

@@ -60,7 +60,7 @@ static inline void ompi_osc_ucx_handle_incoming_post(ompi_osc_ucx_module_t *modu
 
 int ompi_osc_ucx_fence(int assert, struct ompi_win_t *win) {
     ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module;
-    int ret;
+    int ret = OMPI_SUCCESS;
 
     if (module->epoch_type.access != NONE_EPOCH &&
         module->epoch_type.access != FENCE_EPOCH) {
@@ -74,7 +74,7 @@ int ompi_osc_ucx_fence(int assert, struct ompi_win_t *win) {
     }
 
     if (!(assert & MPI_MODE_NOPRECEDE)) {
-        ret = opal_common_ucx_worker_flush(mca_osc_ucx_component.ucp_worker);
+        ret = opal_common_ucx_mem_flush(module->mem, OPAL_COMMON_UCX_SCOPE_WORKER, 0/*ignore*/);
         if (ret != OMPI_SUCCESS) {
             return ret;
         }
@@ -147,7 +147,7 @@ int ompi_osc_ucx_start(struct ompi_group_t *group, int assert, struct ompi_win_t
 
                 ompi_osc_ucx_handle_incoming_post(module, &(module->state.post_state[i]), ranks_in_win_grp, size);
             }
-            ucp_worker_progress(mca_osc_ucx_component.ucp_worker);
+            opal_common_ucx_workers_progress(mca_osc_ucx_component.wpool);
         }
 
         module->post_count = 0;
@@ -163,7 +163,6 @@ int ompi_osc_ucx_start(struct ompi_group_t *group, int assert, struct ompi_win_t
 
 int ompi_osc_ucx_complete(struct ompi_win_t *win) {
     ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module;
-    ucs_status_t status;
     int i, size;
     int ret = OMPI_SUCCESS;
 
@@ -173,29 +172,30 @@ int ompi_osc_ucx_complete(struct ompi_win_t *win) {
 
     module->epoch_type.access = NONE_EPOCH;
 
-    ret = opal_common_ucx_worker_flush(mca_osc_ucx_component.ucp_worker);
+    ret = opal_common_ucx_mem_flush(module->mem, OPAL_COMMON_UCX_SCOPE_WORKER, 0/*ignore*/);
     if (ret != OMPI_SUCCESS) {
         return ret;
     }
+
     module->global_ops_num = 0;
     memset(module->per_target_ops_nums, 0,
            sizeof(int) * ompi_comm_size(module->comm));
 
     size = ompi_group_size(module->start_group);
     for (i = 0; i < size; i++) {
-        uint64_t remote_addr = (module->state_info_array)[module->start_grp_ranks[i]].addr + OSC_UCX_STATE_COMPLETE_COUNT_OFFSET; /* write to state.complete_count on remote side */
-        ucp_rkey_h rkey = (module->state_info_array)[module->start_grp_ranks[i]].rkey;
-        ucp_ep_h ep = OSC_UCX_GET_EP(module->comm, module->start_grp_ranks[i]);
-
-        status = ucp_atomic_post(ep, UCP_ATOMIC_POST_OP_ADD, 1,
-                                 sizeof(uint64_t), remote_addr, rkey);
-        if (status != UCS_OK) {
-            OSC_UCX_VERBOSE(1, "ucp_atomic_post failed: %d", status);
+        uint64_t remote_addr = module->state_addrs[module->start_grp_ranks[i]] + OSC_UCX_STATE_COMPLETE_COUNT_OFFSET; // write to state.complete_count on remote side
+
+        ret = opal_common_ucx_mem_post(module->mem, UCP_ATOMIC_POST_OP_ADD,
+                                       1, module->start_grp_ranks[i], sizeof(uint64_t),
+                                       remote_addr);
+        if (ret != OMPI_SUCCESS) {
+            OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_post failed: %d", ret);
         }
 
-        ret = opal_common_ucx_ep_flush(ep, mca_osc_ucx_component.ucp_worker);
-        if (OMPI_SUCCESS != ret) {
-            OSC_UCX_VERBOSE(1, "opal_common_ucx_ep_flush failed: %d", ret);
+        ret = opal_common_ucx_mem_flush(module->mem, OPAL_COMMON_UCX_SCOPE_EP,
+                                        module->start_grp_ranks[i]);
+        if (ret != OMPI_SUCCESS) {
+            return ret;
         }
     }
 
@@ -243,25 +243,29 @@ int ompi_osc_ucx_post(struct ompi_group_t *group, int assert, struct ompi_win_t
         }
 
         for (i = 0; i < size; i++) {
-            uint64_t remote_addr = (module->state_info_array)[ranks_in_win_grp[i]].addr + OSC_UCX_STATE_POST_INDEX_OFFSET; /* write to state.post_index on remote side */
-            ucp_rkey_h rkey = (module->state_info_array)[ranks_in_win_grp[i]].rkey;
-            ucp_ep_h ep = OSC_UCX_GET_EP(module->comm, ranks_in_win_grp[i]);
+            uint64_t remote_addr = module->state_addrs[ranks_in_win_grp[i]] + OSC_UCX_STATE_POST_INDEX_OFFSET; // write to state.post_index on remote side
             uint64_t curr_idx = 0, result = 0;
 
             /* do fop first to get an post index */
-            opal_common_ucx_atomic_fetch(ep, UCP_ATOMIC_FETCH_OP_FADD, 1,
-                                         &result, sizeof(result),
-                                         remote_addr, rkey, mca_osc_ucx_component.ucp_worker);
+            ret = opal_common_ucx_mem_fetch(module->mem, UCP_ATOMIC_FETCH_OP_FADD,
+                                            1, ranks_in_win_grp[i], &result,
+                                            sizeof(result), remote_addr);
+            if (ret != OMPI_SUCCESS) {
+                return OMPI_ERROR;
+            }
 
             curr_idx = result & (OMPI_OSC_UCX_POST_PEER_MAX - 1);
 
-            remote_addr = (module->state_info_array)[ranks_in_win_grp[i]].addr + OSC_UCX_STATE_POST_STATE_OFFSET + sizeof(uint64_t) * curr_idx;
+            remote_addr = module->state_addrs[ranks_in_win_grp[i]] + OSC_UCX_STATE_POST_STATE_OFFSET + sizeof(uint64_t) * curr_idx;
 
             /* do cas to send post message */
             do {
-                opal_common_ucx_atomic_cswap(ep, 0, (uint64_t)myrank + 1, &result,
-                                             sizeof(result), remote_addr, rkey,
-                                             mca_osc_ucx_component.ucp_worker);
+                ret = opal_common_ucx_mem_cmpswp(module->mem, 0, result,
+                                                 myrank + 1, &result, sizeof(result),
+                                                 remote_addr);
+                if (ret != OMPI_SUCCESS) {
+                    return OMPI_ERROR;
+                }
 
                 if (result == 0)
                     break;
@@ -302,7 +306,7 @@ int ompi_osc_ucx_wait(struct ompi_win_t *win) {
 
     while (module->state.complete_count != (uint64_t)size) {
         /* not sure if this is required */
-        ucp_worker_progress(mca_osc_ucx_component.ucp_worker);
+        opal_common_ucx_workers_progress(mca_osc_ucx_component.wpool);
     }
 
     module->state.complete_count = 0;