MCA/UCX: atomic add/swap are moved to new UCX atomic API

Signed-off-by: Sergey Oblomov <sergeyo@mellanox.com>

oshmem/mca/atomic/ucx/atomic_ucx.h

@@ -60,6 +60,10 @@ struct mca_atomic_ucx_module_t {
 typedef struct mca_atomic_ucx_module_t mca_atomic_ucx_module_t;
 OBJ_CLASS_DECLARATION(mca_atomic_ucx_module_t);
 
+
+void mca_atomic_ucx_complete_cb(void *request, ucs_status_t status);
+ucs_status_t mca_atomic_ucx_wait_request(ucs_status_ptr_t request);
+
 END_C_DECLS
 
 #endif /* MCA_ATOMIC_MXM_H */

oshmem/mca/atomic/ucx/atomic_ucx_cswap.c

@@ -27,23 +27,23 @@ int mca_atomic_ucx_cswap(void *target,
                          int pe)
 {
     ucs_status_t status;
+    ucs_status_ptr_t status_ptr;
     spml_ucx_mkey_t *ucx_mkey;
     uint64_t rva;
+    uint64_t val;
+
+    if ((8 != nlong) && (4 != nlong)) {
+        ATOMIC_ERROR("[#%d] Type size must be 4 or 8 bytes.", my_pe);
+        return OSHMEM_ERROR;
+    }
 
     ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva); 
+    val = (8 == nlong) ? *(uint64_t*)value : *(uint32_t*)value;
     if (NULL == cond) {
-        switch (nlong) {
-            case 4:
-                status = ucp_atomic_swap32(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint32_t *)value, rva, ucx_mkey->rkey, prev);
-                break;
-            case 8:
-                status = ucp_atomic_swap64(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint64_t *)value, rva, ucx_mkey->rkey, prev);
-                break;
-            default:
-                goto err_size;
-        }
+        status_ptr = ucp_atomic_fetch_nb(mca_spml_self->ucp_peers[pe].ucp_conn,
+                UCP_ATOMIC_FETCH_OP_SWAP, val, prev, nlong,
+                rva, ucx_mkey->rkey, mca_atomic_ucx_complete_cb);
+        status = mca_atomic_ucx_wait_request(status_ptr);
     }
     else {
         switch (nlong) {
@@ -56,15 +56,12 @@ int mca_atomic_ucx_cswap(void *target,
                         *(uint64_t *)cond, *(uint64_t *)value, rva, ucx_mkey->rkey, prev);
                 break;
             default:
-                goto err_size;
+                assert(0); /* should not be here */
+                break;
         }
     }
 
     return ucx_status_to_oshmem(status);
-
-err_size:
-    ATOMIC_ERROR("[#%d] Type size must be 1/2/4 or 8 bytes.", my_pe);
-    return OSHMEM_ERROR;
 }
 
 

oshmem/mca/atomic/ucx/atomic_ucx_fadd.c

@@ -26,43 +26,29 @@ int mca_atomic_ucx_fadd(void *target,
                         struct oshmem_op_t *op)
 {
     ucs_status_t status;
+    ucs_status_ptr_t status_ptr;
     spml_ucx_mkey_t *ucx_mkey;
     uint64_t rva;
+    uint64_t val;
+
+    if ((8 != nlong) && (4 != nlong)) {
+        ATOMIC_ERROR("[#%d] Type size must be 4 or 8 bytes.", my_pe);
+        return OSHMEM_ERROR;
+    }
 
     ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
+    val = (8 == nlong) ? *(uint64_t*)value : *(uint32_t*)value;
 
     if (NULL == prev) {
-        switch (nlong) {
-            case 4:
-                status = ucp_atomic_add32(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint32_t *)value, rva, ucx_mkey->rkey);
-                break;
-            case 8:
-                status = ucp_atomic_add64(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint64_t *)value, rva, ucx_mkey->rkey);
-                break;
-            default:
-                goto err_size;
-        }
+        status = ucp_atomic_post(mca_spml_self->ucp_peers[pe].ucp_conn, 
+                UCP_ATOMIC_POST_OP_ADD, val, nlong, rva, ucx_mkey->rkey);
     }
     else {
-        switch (nlong) {
-            case 4:
-                status = ucp_atomic_fadd32(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint32_t *)value, rva, ucx_mkey->rkey, prev);
-                break;
-            case 8:
-                status = ucp_atomic_fadd64(mca_spml_self->ucp_peers[pe].ucp_conn, 
-                        *(uint64_t *)value, rva, ucx_mkey->rkey, prev);
-                break;
-            default:
-                goto err_size;
-        }
+        status_ptr = ucp_atomic_fetch_nb(mca_spml_self->ucp_peers[pe].ucp_conn, 
+                UCP_ATOMIC_FETCH_OP_FADD, val, prev, nlong,
+                rva, ucx_mkey->rkey, mca_atomic_ucx_complete_cb);
+        status = mca_atomic_ucx_wait_request(status_ptr);
     }
 
     return ucx_status_to_oshmem(status);
-
-err_size:
-    ATOMIC_ERROR("[#%d] Type size must be 1/2/4 or 8 bytes.", my_pe);
-    return OSHMEM_ERROR;
 }

oshmem/mca/atomic/ucx/atomic_ucx_module.c

@@ -49,3 +49,24 @@ mca_atomic_ucx_query(int *priority)
     return NULL ;
 }
 
+void mca_atomic_ucx_complete_cb(void *request, ucs_status_t status)
+{
+}
+
+ucs_status_t mca_atomic_ucx_wait_request(ucs_status_ptr_t request)
+{
+    ucs_status_t status;
+
+    /* check for request completed or failed */
+    if (UCS_OK == request) {
+        return UCS_OK;
+    } else if (UCS_PTR_IS_ERR(request)) {
+        return UCS_PTR_STATUS(request);
+    }
+
+    while (UCS_INPROGRESS == (status = ucp_request_check_status(request))) {
+        ucp_worker_progress(mca_spml_self->ucp_worker);
+    }
+    ucp_request_free(request);
+    return status;
+}