UCX osc: atomic fetch-and-op only on 32 and 64bit values

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

Author: devreal

ompi/mca/osc/ucx/osc_ucx_comm.c

@@ -24,6 +24,10 @@
         return OMPI_ERROR;                                                       \
     }
 
+/* macro to check whether UCX supports atomic operation on the size the operands */
+#define ATOMIC_SIZE_SUPPORTED(_size) (sizeof(uint32_t) == _size || \
+                                      sizeof(uint64_t) == _size)
+
 typedef struct ucx_iovec {
     void *addr;
     size_t len;
@@ -384,9 +388,8 @@ bool use_atomic_op(
         ompi_datatype_type_size(origin_dt, &origin_dt_bytes);
         ompi_datatype_type_size(target_dt, &target_dt_bytes);
         /* UCX only supports 32 and 64-bit operands atm */
-        if (sizeof(uint64_t) >= origin_dt_bytes &&
-            sizeof(uint32_t) <= origin_dt_bytes &&
-            origin_dt_bytes  == target_dt_bytes &&
+        if (ATOMIC_SIZE_SUPPORTED(origin_dt_bytes) &&
+            origin_dt_bytes  == target_dt_bytes    &&
             origin_count     == target_count) {
             return true;
         }
@@ -775,7 +778,7 @@ int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_a
     }
 
     ompi_datatype_type_size(dt, &dt_bytes);
-    if (4 == dt_bytes || 8 == dt_bytes) {
+    if (ATOMIC_SIZE_SUPPORTED(dt_bytes)) {
         // fast path using UCX atomic operations
         return do_atomic_compare_and_swap(origin_addr, compare_addr,
                                           result_addr, dt, target,
@@ -818,6 +821,7 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
                               struct ompi_datatype_t *dt, int target,
                               ptrdiff_t target_disp, struct ompi_op_t *op,
                               struct ompi_win_t *win) {
+    size_t dt_bytes;
     ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module;
     int ret = OMPI_SUCCESS;
 
@@ -826,12 +830,15 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
         return ret;
     }
 
-    if (op == &ompi_mpi_op_no_op.op || op == &ompi_mpi_op_replace.op ||
-        op == &ompi_mpi_op_sum.op) {
+    ompi_datatype_type_size(dt, &dt_bytes);
+
+    /* UCX atomics are only supported on 32 and 64 bit values */
+    if (ATOMIC_SIZE_SUPPORTED(dt_bytes) &&
+        (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);
         ucp_atomic_fetch_op_t opcode;
-        size_t dt_bytes;
         bool lock_acquired = false;
 
         if (!module->acc_single_intrinsic) {