UCX osc: allow to asynchronously compare-and-swap

Signed-off-by: Joseph Schuchart <[email protected]>

Author: devreal

ompi/mca/osc/ucx/osc_ucx_comm.c

@@ -66,6 +66,19 @@ static inline int check_sync_state(ompi_osc_ucx_module_t *module, int target,
     return OMPI_SUCCESS;
 }
 
+typedef struct {
+  size_t   dt_bytes;
+  uint64_t result_val;
+  void    *result_addr;
+} cas_result_handle_t;
+
+static void cas_result_cb(void *data) {
+  cas_result_handle_t *cas_result = (cas_result_handle_t*)data;
+  memcpy(cas_result->result_addr, &cas_result->result_val, cas_result->dt_bytes);
+  free(cas_result);
+}
+
+
 static inline int create_iov_list(const void *addr, int count, ompi_datatype_t *datatype,
                                   ucx_iovec_t **ucx_iov, uint32_t *ucx_iov_count) {
     int ret = OMPI_SUCCESS;
@@ -808,14 +821,35 @@ int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_a
     }
 
     ompi_datatype_type_size(dt, &dt_bytes);
-    ret = opal_common_ucx_wpmem_cmpswp(module->mem,*(uint64_t *)compare_addr,
-                                     *(uint64_t *)origin_addr, target,
-                                     result_addr, dt_bytes, remote_addr);
+    if (sizeof(uint64_t) < dt_bytes) {
+        return OMPI_ERR_NOT_SUPPORTED;
+    }
+
+    uint64_t compare_val;
+    memcpy(&compare_val, compare_addr, dt_bytes);
+
+    // the completion handler copies the returned value into the result address
+    // and free's the memory
+    cas_result_handle_t *cas_result = malloc(sizeof(*cas_result));
+    cas_result->dt_bytes = dt_bytes;
+    cas_result->result_addr = result_addr;
+    memcpy(&cas_result->result_val, origin_addr, dt_bytes);
+    ret = opal_common_ucx_wpmem_fetch_nb(module->mem, UCP_ATOMIC_FETCH_OP_CSWAP,
+                                         compare_val, target,
+                                         &cas_result->result_val, dt_bytes, remote_addr,
+                                         &cas_result_cb, cas_result);
 
     if (module->acc_single_intrinsic) {
         return ret;
     }
 
+    // fence before releasing the accumulate lock
+    ret = opal_common_ucx_wpmem_fence(module->mem);
+    if (ret != OMPI_SUCCESS) {
+        OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret);
+        // don't return error, try to release the accumulate lock
+    }
+
     return end_atomicity(module, target);
 }