ATOMIC/UCX: fixed atomic swap/cswap issues

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

oshmem/mca/atomic/atomic.h

@@ -33,6 +33,8 @@
 
 BEGIN_C_DECLS
 
+#define ATOMIC_PTR_2_INT(ptr, size) ((size) == 8 ? *(uint64_t*)(ptr) : *(uint32_t*)(ptr))
+
 /* ******************************************************************** */
 
 struct oshmem_op_t;

oshmem/mca/atomic/mxm/atomic_mxm_cswap.c

@@ -55,7 +55,7 @@ int mca_atomic_mxm_cswap(void *target,
     mca_atomic_mxm_req_init(&sreq, pe, target, nlong);
 
     *prev                     = value;
-    sreq.op.atomic.value      = value;
+    sreq.op.atomic.value      = cond;
     sreq.base.data.buffer.ptr = prev;
     sreq.opcode               = MXM_REQ_OP_ATOMIC_CSWAP;
 

oshmem/mca/atomic/ucx/atomic_ucx_cswap.c

@@ -26,7 +26,6 @@ int mca_atomic_ucx_cswap(void *target,
                          size_t size,
                          int pe)
 {
-    int status;
     ucs_status_ptr_t status_ptr;
     spml_ucx_mkey_t *ucx_mkey;
     uint64_t rva;
@@ -39,7 +38,7 @@ int mca_atomic_ucx_cswap(void *target,
     assert(NULL != prev);
 
     *prev      = value;
-    ucx_mkey   = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
+    ucx_mkey   = mca_spml_ucx_get_mkey(pe, target, (void *)&rva, &mca_spml_ucx);
     status_ptr = ucp_atomic_fetch_nb(mca_spml_self->ucp_peers[pe].ucp_conn,
                                      UCP_ATOMIC_FETCH_OP_CSWAP, cond, prev, size,
                                      rva, ucx_mkey->rkey,

oshmem/mca/atomic/ucx/atomic_ucx_module.c

@@ -45,12 +45,9 @@ int mca_atomic_ucx_op(void *target,
     uint64_t rva;
     uint64_t val;
 
-    if ((8 != size) && (4 != size)) {
-        ATOMIC_ERROR("[#%d] Type size must be 4 or 8 bytes.", my_pe);
-        return OSHMEM_ERROR;
-    }
+    assert((8 == size) || (4 == size));
 
-    ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
+    ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva, &mca_spml_ucx);
     status = ucp_atomic_post(mca_spml_self->ucp_peers[pe].ucp_conn,
                              op, value, size, rva,
                              ucx_mkey->rkey);
@@ -65,17 +62,13 @@ int mca_atomic_ucx_fop(void *target,
                        int pe,
                        ucp_atomic_fetch_op_t op)
 {
-    ucs_status_t status;
     ucs_status_ptr_t status_ptr;
     spml_ucx_mkey_t *ucx_mkey;
     uint64_t rva;
 
-    if ((8 != size) && (4 != size)) {
-        ATOMIC_ERROR("[#%d] Type size must be 4 or 8 bytes.", my_pe);
-        return OSHMEM_ERROR;
-    }
+    assert((8 == size) || (4 == size));
 
-    ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
+    ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva, &mca_spml_ucx);
     status_ptr = ucp_atomic_fetch_nb(mca_spml_self->ucp_peers[pe].ucp_conn,
                                      op, value, prev, size,
                                      rva, ucx_mkey->rkey,

oshmem/shmem/c/shmem_cswap.c

@@ -38,9 +38,9 @@
         size = sizeof(value);                                       \
         rc = MCA_ATOMIC_CALL(cswap(                                 \
             (void*)target,                                          \
-            (void*)&out_value,                                      \
-            cond,                                                   \
-            value,                                                  \
+            &out_value,                                             \
+            ATOMIC_PTR_2_INT(&cond, sizeof(cond)),                  \
+            ATOMIC_PTR_2_INT(&value, sizeof(value)),                \
             size,                                                   \
             pe));                                                   \
         RUNTIME_CHECK_RC(rc);                                       \

oshmem/shmem/c/shmem_swap.c

@@ -37,7 +37,7 @@
         rc = MCA_ATOMIC_CALL(swap(                                  \
             (void*)target,                                          \
             (void*)&out_value,                                      \
-            value,                                                  \
+            ATOMIC_PTR_2_INT(&value, sizeof(value)),                \
             size,                                                   \
             pe));                                                   \
         RUNTIME_CHECK_RC(rc);                                       \