Enhance ref counters

Author: Mamzi Bayatpour [email protected] ()

ompi/mca/osc/ucx/osc_ucx.h

@@ -135,7 +135,8 @@ typedef struct ompi_osc_ucx_module {
     bool lock_all_is_nocheck;
     bool no_locks;
     bool acc_single_intrinsic;
-    uint64_t acc_lock_refcnt;
+    uint64_t *acc_lock_refcnt;
+    uint64_t *pending_acc_finalize_refcnt;
     opal_common_ucx_ctx_t *ctx;
     opal_common_ucx_wpmem_t *mem;
     opal_common_ucx_wpmem_t *state_mem;

ompi/mca/osc/ucx/osc_ucx_comm.c

@@ -1523,6 +1523,8 @@ void ompi_osc_ucx_req_completion(void *request) {
                 if (temp_dt != NULL && !ompi_datatype_is_predefined(temp_dt)) {
                     OBJ_RELEASE(temp_dt);
                 }
+                module->pending_acc_finalize_refcnt[target]--;
+                assert(module->pending_acc_finalize_refcnt[target] >= 0);
                 break;
             }
             case ACC_GET_RESULTS_DATA:

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -840,7 +840,8 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, in
     OBJ_CONSTRUCT(&module->pending_posts, opal_list_t);
     module->start_grp_ranks = NULL;
     module->lock_all_is_nocheck = false;
-    module->acc_lock_refcnt = 0;
+    module->acc_lock_refcnt = calloc(comm_size, sizeof(uint64_t));
+    module->pending_acc_finalize_refcnt = calloc(comm_size, sizeof(uint64_t));
 
     if (!module->no_locks) {
         OBJ_CONSTRUCT(&module->outstanding_locks, opal_hash_table_t);
@@ -930,8 +931,13 @@ inline int ompi_osc_ucx_state_lock(
     OSC_UCX_GET_DEFAULT_EP(ep, module, target);
     int ret = OMPI_SUCCESS;
 
+    assert(module->pending_acc_finalize_refcnt[target] >= 0);
+    while (module->pending_acc_finalize_refcnt[target] > 0) {
+        opal_common_ucx_wpool_progress(mca_osc_ucx_component.wpool);
+    }
+
     if (force_lock || ompi_osc_need_acc_lock(module, target)) {
-        if (module->acc_lock_refcnt == 0) {
+        if (module->acc_lock_refcnt[target] == 0) {
             for (;;) {
                 ret = opal_common_ucx_wpmem_cmpswp(module->state_mem,
                                                 TARGET_LOCK_UNLOCKED, TARGET_LOCK_EXCLUSIVE,
@@ -948,13 +954,11 @@ inline int ompi_osc_ucx_state_lock(
                 opal_common_ucx_wpool_progress(mca_osc_ucx_component.wpool);
             }
         }
-
+        module->acc_lock_refcnt[target]++;
         *lock_acquired = true;
-        module->acc_lock_refcnt++;
     } else {
         *lock_acquired = false;
     }
-
     return OMPI_SUCCESS;
 }
 
@@ -977,16 +981,16 @@ inline int ompi_osc_ucx_state_unlock(
             return OMPI_ERROR;
         }
 
-        if (module->acc_lock_refcnt == 1) {
+        if (module->acc_lock_refcnt[target] == 1) {
             ret = opal_common_ucx_wpmem_fetch(module->state_mem,
                                             UCP_ATOMIC_FETCH_OP_SWAP, TARGET_LOCK_UNLOCKED,
                                             target, &result_value, sizeof(result_value),
                                             remote_addr, ep);
             assert(result_value == TARGET_LOCK_EXCLUSIVE);
         }
-        module->acc_lock_refcnt--;
-        assert(module->acc_lock_refcnt >= 0);
-    } else if (NULL != free_ptr){
+        module->acc_lock_refcnt[target]--;
+        assert(module->acc_lock_refcnt[target] >= 0);
+    } else if (NULL != free_ptr) {
         /* flush before freeing the buffer */
         ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target);
     }
@@ -1016,6 +1020,8 @@ inline int ompi_osc_ucx_nonblocking_ops_finalize(ompi_osc_ucx_module_t *module,
     ucx_req->phase = ACC_FINALIZE;
     ucx_req->acc_type = ANY;
     ucx_req->super.module = module;
+    ucx_req->lock_acquired = lock_acquired;
+    ucx_req->target = target;
 
     /* Fence any still active operations */
     ret = opal_common_ucx_wpmem_fence(module->mem);
@@ -1024,7 +1030,8 @@ inline int ompi_osc_ucx_nonblocking_ops_finalize(ompi_osc_ucx_module_t *module,
         return OMPI_ERROR;
     }
 
-    if (lock_acquired) {
+    module->pending_acc_finalize_refcnt[target]++;
+    if (lock_acquired && module->acc_lock_refcnt[target] == 1) {
         OSC_UCX_INCREMENT_OUTSTANDING_NB_OPS(module);
         ret = opal_common_ucx_wpmem_fetch_nb(module->state_mem,
                                         UCP_ATOMIC_FETCH_OP_SWAP, TARGET_LOCK_UNLOCKED,
@@ -1036,7 +1043,7 @@ inline int ompi_osc_ucx_nonblocking_ops_finalize(ompi_osc_ucx_module_t *module,
             return ret;
         }
     } else {
-        /* Lock is not acquired, but still, we need to know when the 
+        /* Lock is not acquired/ or ref count is not 1, but still, we need to know when the
          * acc is finalized so that we can free the temp buffers */
         OSC_UCX_INCREMENT_OUTSTANDING_NB_OPS(module);
         ret = opal_common_ucx_wpmem_flush_ep_nb(module->mem, target, ompi_osc_ucx_req_completion, ucx_req, ep);
@@ -1054,6 +1061,11 @@ inline int ompi_osc_ucx_nonblocking_ops_finalize(ompi_osc_ucx_module_t *module,
         }
     }
 
+    if (lock_acquired) {
+        module->acc_lock_refcnt[target]--;
+        assert(module->acc_lock_refcnt[target] >= 0);
+    }
+
     return ret;
 }