OSC/UCX: Adding the following optimzations: 1) Reuse the same worker/eps in

single threaded applications, this is helpful if an application
creates many windows, therefore, we avoid the unnecessary overheads  and 2) adding the truely nonblocking
MPI_Accumulate/Get_Accumulate.

Signed-off-by: Mamzi Bayatpour  <[email protected]>
Co-authored-by: Tomislav Janjusic <[email protected]>

Author: Mamzi Bayatpour [email protected] ()

ompi/mca/osc/ucx/osc_ucx.h

@@ -27,6 +27,7 @@
 #define OMPI_OSC_UCX_ATTACH_MAX    48
 #define OMPI_OSC_UCX_MEM_ADDR_MAX_LEN  1024
 
+
 typedef struct ompi_osc_ucx_component {
     ompi_osc_base_component_t super;
     opal_common_ucx_wpool_t *wpool;
@@ -125,6 +126,7 @@ typedef struct ompi_osc_ucx_module {
     opal_common_ucx_wpmem_t *mem;
     opal_common_ucx_wpmem_t *state_mem;
 
+    bool skip_sync_check;
     bool noncontig_shared_win;
     size_t *sizes;
     /* in shared windows, shmem_addrs can be used for direct load store to
@@ -150,6 +152,17 @@ typedef struct ompi_osc_ucx_lock {
 #define OSC_UCX_GET_EP(comm_, rank_) (ompi_comm_peer_lookup(comm_, rank_)->proc_endpoints[OMPI_PROC_ENDPOINT_TAG_UCX])
 #define OSC_UCX_GET_DISP(module_, rank_) ((module_->disp_unit < 0) ? module_->disp_units[rank_] : module_->disp_unit)
 
+extern bool mpi_thread_multiple_enabled;
+
+#define OSC_UCX_GET_DEFAULT_EP(_ep_ptr, _comm, _target)                 \
+    if (mpi_thread_multiple_enabled) {                                  \
+        _ep_ptr = NULL;                                                 \
+    } else {                                                            \
+        _ep_ptr = (ucp_ep_h *)&(OSC_UCX_GET_EP(_comm, _target));        \
+    }
+
+#define OSC_UCX_OUTSTANDING_OPS_FLUSH_THRESHOLD 64
+
 int ompi_osc_ucx_shared_query(struct ompi_win_t *win, int rank, size_t *size,
         int *disp_unit, void * baseptr);
 int ompi_osc_ucx_win_attach(struct ompi_win_t *win, void *base, size_t len);
@@ -169,6 +182,11 @@ int ompi_osc_ucx_accumulate(const void *origin_addr, int origin_count,
                             int target, ptrdiff_t target_disp, int target_count,
                             struct ompi_datatype_t *target_dt,
                             struct ompi_op_t *op, struct ompi_win_t *win);
+int ompi_osc_ucx_accumulate_nb(const void *origin_addr, int origin_count,
+                            struct ompi_datatype_t *origin_dt,
+                            int target, ptrdiff_t target_disp, int target_count,
+                            struct ompi_datatype_t *target_dt,
+                            struct ompi_op_t *op, struct ompi_win_t *win);
 int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_addr,
                                   void *result_addr, struct ompi_datatype_t *dt,
                                   int target, ptrdiff_t target_disp,
@@ -184,6 +202,13 @@ int ompi_osc_ucx_get_accumulate(const void *origin_addr, int origin_count,
                                 int target_rank, ptrdiff_t target_disp,
                                 int target_count, struct ompi_datatype_t *target_datatype,
                                 struct ompi_op_t *op, struct ompi_win_t *win);
+int ompi_osc_ucx_get_accumulate_nb(const void *origin_addr, int origin_count,
+                                struct ompi_datatype_t *origin_datatype,
+                                void *result_addr, int result_count,
+                                struct ompi_datatype_t *result_datatype,
+                                int target_rank, ptrdiff_t target_disp,
+                                int target_count, struct ompi_datatype_t *target_datatype,
+                                struct ompi_op_t *op, struct ompi_win_t *win);
 int ompi_osc_ucx_rput(const void *origin_addr, int origin_count,
                       struct ompi_datatype_t *origin_dt,
                       int target, ptrdiff_t target_disp, int target_count,
@@ -229,9 +254,11 @@ int ompi_osc_find_attached_region_position(ompi_osc_dynamic_win_info_t *dynamic_
                                            int min_index, int max_index,
                                            uint64_t base, size_t len, int *insert);
 extern inline bool ompi_osc_need_acc_lock(ompi_osc_ucx_module_t *module, int target);
-extern inline int ompi_osc_state_lock(ompi_osc_ucx_module_t *module, int target,
+extern inline int ompi_osc_ucx_state_lock(ompi_osc_ucx_module_t *module, int target,
         bool *lock_acquired, bool force_lock);
-extern inline int ompi_osc_state_unlock(ompi_osc_ucx_module_t *module, int target,
+extern inline int ompi_osc_ucx_state_unlock(ompi_osc_ucx_module_t *module, int target,
         bool lock_acquired, void *free_ptr);
+extern inline int ompi_osc_ucx_state_unlock_nb(ompi_osc_ucx_module_t *module, int target,
+        bool lock_acquired, struct ompi_win_t *win);
 
 #endif /* OMPI_OSC_UCX_H */

ompi/mca/osc/ucx/osc_ucx_active_target.c

@@ -165,25 +165,28 @@ 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;
     int i, size;
     int ret = OMPI_SUCCESS;
+    ucp_ep_h *ep;
 
     if (module->epoch_type.access != START_COMPLETE_EPOCH) {
         return OMPI_ERR_RMA_SYNC;
     }
 
-    module->epoch_type.access = NONE_EPOCH;
-
     ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_WORKER, 0/*ignore*/);
     if (ret != OMPI_SUCCESS) {
         return ret;
     }
 
+    module->epoch_type.access = NONE_EPOCH;
+
     size = ompi_group_size(module->start_group);
     for (i = 0; i < size; i++) {
         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
 
+        OSC_UCX_GET_DEFAULT_EP(ep, module->comm, module->start_grp_ranks[i]);
+
         ret = opal_common_ucx_wpmem_post(module->state_mem, UCP_ATOMIC_POST_OP_ADD,
                                        1, module->start_grp_ranks[i], sizeof(uint64_t),
-                                       remote_addr);
+                                       remote_addr, ep);
         if (ret != OMPI_SUCCESS) {
             OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_post failed: %d", ret);
         }
@@ -204,6 +207,7 @@ int ompi_osc_ucx_complete(struct ompi_win_t *win) {
 
 int ompi_osc_ucx_post(struct ompi_group_t *group, int mpi_assert, struct ompi_win_t *win) {
     ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module;
+    ucp_ep_h *ep;
     int ret = OMPI_SUCCESS;
 
     if (module->epoch_type.exposure != NONE_EPOCH) {
@@ -243,12 +247,12 @@ int ompi_osc_ucx_post(struct ompi_group_t *group, int mpi_assert, struct ompi_wi
             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;
 
-
+            OSC_UCX_GET_DEFAULT_EP(ep, module->comm, ranks_in_win_grp[i]);
 
             /* do fop first to get an post index */
             ret = opal_common_ucx_wpmem_fetch(module->state_mem, UCP_ATOMIC_FETCH_OP_FADD,
                                             1, ranks_in_win_grp[i], &result,
-                                            sizeof(result), remote_addr);
+                                            sizeof(result), remote_addr, ep);
 
             if (ret != OMPI_SUCCESS) {
                 ret = OMPI_ERROR;
@@ -265,7 +269,7 @@ int ompi_osc_ucx_post(struct ompi_group_t *group, int mpi_assert, struct ompi_wi
                 result =  myrank + 1;
                 ret = opal_common_ucx_wpmem_cmpswp(module->state_mem, 0, result,
                                                  ranks_in_win_grp[i], &result, sizeof(result),
-                                                 remote_addr);
+                                                 remote_addr, ep);
 
                 if (ret != OMPI_SUCCESS) {
                     ret = OMPI_ERROR;