UCX osc: safely load/store 64bit integer from variable size pointer

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

Author: devreal

ompi/mca/osc/ucx/osc_ucx_comm.c

@@ -454,7 +454,7 @@ static int do_atomic_op_intrinsic(
         if (is_no_op) {
             value = 0;
         } else {
-            memcpy(&value, origin_addr, origin_dt_bytes);
+            value = opal_common_ucx_load_uint64(origin_addr, origin_dt_bytes);
         }
         ret = opal_common_ucx_wpmem_fetch_nb(module->mem, opcode, value, target,
                                              output_addr, origin_dt_bytes, remote_addr,
@@ -756,13 +756,11 @@ int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_a
         }
     }
 
-    uint64_t compare_val;
-    memcpy(&compare_val, compare_addr, dt_bytes);
-    memcpy(result_addr, origin_addr, dt_bytes);
-    ret = opal_common_ucx_wpmem_fetch_nb(module->mem, UCP_ATOMIC_FETCH_OP_CSWAP,
-                                         compare_val, target,
-                                         result_addr, dt_bytes, remote_addr,
-                                         NULL, NULL);
+    uint64_t compare_val = opal_common_ucx_load_uint64(compare_addr, dt_bytes);
+    uint64_t value       = opal_common_ucx_load_uint64(origin_addr,  dt_bytes);
+    ret = opal_common_ucx_wpmem_cmpswp_nb(module->mem, compare_val, value, target,
+                                          result_addr, dt_bytes, remote_addr,
+                                          NULL, NULL);
 
     if (module->acc_single_intrinsic) {
         return ret;
@@ -785,8 +783,8 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
 
     if (op == &ompi_mpi_op_no_op.op || op == &ompi_mpi_op_replace.op ||
         op == &ompi_mpi_op_sum.op) {
+        uint64_t value;
         uint64_t remote_addr = (module->addrs[target]) + target_disp * OSC_UCX_GET_DISP(module, target);
-        uint64_t value = origin_addr ? *(uint64_t *)origin_addr : 0;
         ucp_atomic_fetch_op_t opcode;
         size_t dt_bytes;
         bool lock_acquired = false;
@@ -805,7 +803,7 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
             }
         }
 
-        ompi_datatype_type_size(dt, &dt_bytes);
+        value = origin_addr ? opal_common_ucx_load_uint64(origin_addr, dt_bytes) : 0;
 
         if (op == &ompi_mpi_op_replace.op) {
             opcode = UCP_ATOMIC_FETCH_OP_SWAP;

opal/mca/common/ucx/common_ucx.h

@@ -115,6 +115,42 @@ OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *
                                                size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
 OPAL_DECLSPEC void opal_common_ucx_mca_var_register(const mca_base_component_t *component);
 
+
+/**
+ * Load an integer value of \c size bytes from \c ptr and cast it to uint64_t.
+ */
+static inline
+uint64_t opal_common_ucx_load_uint64(void *ptr, size_t size)
+{
+    if (sizeof(uint8_t) == size) {
+      return *(uint8_t*)ptr;
+    } else if (sizeof(uint16_t) == size) {
+      return *(uint16_t*)ptr;
+    } else if (sizeof(uint32_t) == size) {
+      return *(uint32_t*)ptr;
+    } else {
+      return *(uint64_t*)ptr;
+    }
+}
+
+/**
+ * Cast and store a uint64_t value to a value of \c size bytes pointed to by \c ptr.
+ */
+static inline
+void opal_common_ucx_store_uint64(uint64_t value, void *ptr, size_t size)
+{
+    if (sizeof(uint8_t) == size) {
+      *(uint8_t*)ptr  = value;
+    } else if (sizeof(uint16_t) == size) {
+      *(uint16_t*)ptr = value;
+    } else if (sizeof(uint32_t) == size) {
+      *(uint32_t*)ptr = value;
+    } else {
+      *(uint64_t*)ptr = value;
+    }
+}
+
+
 static inline
 ucs_status_t opal_common_ucx_request_status(ucs_status_ptr_t request)
 {
@@ -206,13 +242,7 @@ int opal_common_ucx_atomic_cswap(ucp_ep_h ep, uint64_t compare,
                                  uint64_t remote_addr, ucp_rkey_h rkey,
                                  ucp_worker_h worker)
 {
-    if (op_size == sizeof(uint64_t)) {
-        *(uint64_t*)result = value;
-    } else {
-        assert(op_size == sizeof(uint32_t));
-        *(uint32_t*)result = value;
-    }
-
+    opal_common_ucx_store_uint64(value, result, op_size);
     return opal_common_ucx_atomic_fetch(ep, UCP_ATOMIC_FETCH_OP_CSWAP, compare, result,
                                         op_size, remote_addr, rkey, worker);
 }
@@ -224,13 +254,7 @@ ucs_status_ptr_t opal_common_ucx_atomic_cswap_nb(ucp_ep_h ep, uint64_t compare,
                                                  ucp_send_callback_t req_handler,
                                                  ucp_worker_h worker)
 {
-    if (op_size == sizeof(uint64_t)) {
-        *(uint64_t*)result = value;
-    } else {
-        assert(op_size == sizeof(uint32_t));
-        *(uint32_t*)result = value;
-    }
-
+    opal_common_ucx_store_uint64(value, result, op_size);
     return opal_common_ucx_atomic_fetch_nb(ep, UCP_ATOMIC_FETCH_OP_CSWAP, compare, result,
                                            op_size, remote_addr, rkey, req_handler, worker);
 }