UCX osc: add support for acc_single_intrinsic info key / mca param

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

Author: devreal

ompi/mca/osc/ucx/osc_ucx.h

@@ -34,6 +34,7 @@ typedef struct ompi_osc_ucx_component {
     int num_incomplete_req_ops;
     int num_modules;
     bool no_locks; /* Default value of the no_locks info key for new windows */
+    bool acc_single_intrinsic;
     unsigned int priority;
 } ompi_osc_ucx_component_t;
 
@@ -115,6 +116,7 @@ typedef struct ompi_osc_ucx_module {
     int *start_grp_ranks;
     bool lock_all_is_nocheck;
     bool no_locks;
+    bool acc_single_intrinsic;
     opal_common_ucx_ctx_t *ctx;
     opal_common_ucx_wpmem_t *mem;
     opal_common_ucx_wpmem_t *state_mem;

ompi/mca/osc/ucx/osc_ucx_comm.c

@@ -323,6 +323,149 @@ static inline int get_dynamic_win_info(uint64_t remote_addr, ompi_osc_ucx_module
     return ret;
 }
 
+static int atomic_op_replace_sum(
+    ompi_osc_ucx_module_t  *module,
+    struct ompi_op_t       *op,
+    int                     target,
+    const void             *origin_addr,
+    int                     origin_count,
+    struct ompi_datatype_t *origin_dt,
+    ptrdiff_t               target_disp,
+    int                     target_count,
+    struct ompi_datatype_t *target_dt,
+    void                   *result_addr)
+{
+    int ret = OMPI_SUCCESS;
+    size_t origin_dt_bytes;
+    size_t target_dt_bytes;
+    ompi_datatype_type_size(origin_dt, &origin_dt_bytes);
+    ompi_datatype_type_size(target_dt, &target_dt_bytes);
+
+    if (origin_dt_bytes  > sizeof(uint64_t) ||
+        origin_dt_bytes != target_dt_bytes  ||
+        target_count    != origin_count) {
+        return OMPI_ERR_NOT_SUPPORTED;
+    }
+
+    uint64_t remote_addr = (module->addrs[target]) + target_disp * OSC_UCX_GET_DISP(module, target);
+
+    if (module->flavor == MPI_WIN_FLAVOR_DYNAMIC) {
+        ret = get_dynamic_win_info(remote_addr, module, target);
+        if (ret != OMPI_SUCCESS) {
+            return ret;
+        }
+    }
+
+    ucp_atomic_fetch_op_t opcode;
+    if (op == &ompi_mpi_op_replace.op) {
+        opcode = UCP_ATOMIC_FETCH_OP_SWAP;
+    } else {
+        opcode = UCP_ATOMIC_FETCH_OP_FADD;
+    }
+
+    for (int i = 0; i < origin_count; ++i) {
+        uint64_t value = 0;
+        memcpy(&value, origin_addr, origin_dt_bytes);
+        ret = opal_common_ucx_wpmem_fetch_nb(module->mem, opcode, value, target,
+                                             result_addr ? result_addr : &(module->req_result),
+                                             origin_dt_bytes, remote_addr, NULL, NULL);
+
+        // advance origin and remote address
+        origin_addr  = (void*)((intptr_t)origin_addr + origin_dt_bytes);
+        remote_addr += origin_dt_bytes;
+        if (result_addr) {
+            result_addr = (void*)((intptr_t)result_addr + origin_dt_bytes);
+        }
+    }
+
+    return ret;
+}
+
+static int atomic_op_cswap(
+    ompi_osc_ucx_module_t  *module,
+    struct ompi_op_t       *op,
+    int                     target,
+    const void             *origin_addr,
+    int                     origin_count,
+    struct ompi_datatype_t *origin_dt,
+    ptrdiff_t               target_disp,
+    int                     target_count,
+    struct ompi_datatype_t *target_dt,
+    void                   *result_addr)
+{
+    int ret = OMPI_SUCCESS;
+    size_t origin_dt_bytes;
+    size_t target_dt_bytes;
+    ompi_datatype_type_size(origin_dt, &origin_dt_bytes);
+    ompi_datatype_type_size(target_dt, &target_dt_bytes);
+
+    if (origin_dt_bytes  > sizeof(uint64_t) ||
+        origin_dt_bytes != target_dt_bytes  ||
+        target_count    != origin_count) {
+        return OMPI_ERR_NOT_SUPPORTED;
+    }
+
+    uint64_t remote_addr = (module->addrs[target]) + target_disp * OSC_UCX_GET_DISP(module, target);
+
+    if (module->flavor == MPI_WIN_FLAVOR_DYNAMIC) {
+        ret = get_dynamic_win_info(remote_addr, module, target);
+        if (ret != OMPI_SUCCESS) {
+            return ret;
+        }
+    }
+
+    for (int i = 0; i < origin_count; ++i) {
+
+        uint64_t tmp_val;
+        do {
+            uint64_t target_val = 0;
+
+            // get the value from the origin
+            ret = opal_common_ucx_wpmem_putget(module->mem, OPAL_COMMON_UCX_GET,
+                                               target, &target_val, origin_dt_bytes,
+                                               remote_addr);
+            if (ret != OMPI_SUCCESS) {
+                return ret;
+            }
+
+            ret = opal_common_ucx_wpmem_flush(module->mem, OPAL_COMMON_UCX_SCOPE_EP, target);
+            if (ret != OMPI_SUCCESS) {
+                return ret;
+            }
+
+            tmp_val = target_val;
+            // compute the result value
+            ompi_op_reduce(op, (void *)origin_addr, &tmp_val, 1, origin_dt);
+
+            // compare-and-swap the resulting value
+            ret = opal_common_ucx_wpmem_cmpswp(module->mem, target_val, tmp_val,
+                                               target, &tmp_val, origin_dt_bytes,
+                                               remote_addr);
+            if (ret != OMPI_SUCCESS) {
+                return ret;
+            }
+
+            // check whether the conditional swap was successful
+            if (tmp_val == target_val) {
+              break;
+            }
+
+        } while (1);
+
+        // store the result if necessary
+        if (NULL != result_addr) {
+          memcpy(result_addr, &tmp_val, origin_dt_bytes);
+          result_addr = (void*)((intptr_t)result_addr + origin_dt_bytes);
+        }
+        // advance origin and remote address
+        origin_addr  = (void*)((intptr_t)origin_addr + origin_dt_bytes);
+        remote_addr += origin_dt_bytes;
+    }
+
+    return ret;
+}
+
+
 int ompi_osc_ucx_put(const void *origin_addr, int origin_count, struct ompi_datatype_t *origin_dt,
                      int target, ptrdiff_t target_disp, int target_count,
                      struct ompi_datatype_t *target_dt, struct ompi_win_t *win) {
@@ -449,6 +592,22 @@ int ompi_osc_ucx_accumulate(const void *origin_addr, int origin_count,
         return ret;
     }
 
+    if (module->acc_single_intrinsic) {
+        if (op == &ompi_mpi_op_replace.op || op == &ompi_mpi_op_sum.op) {
+            ret = atomic_op_replace_sum(module, op, target,
+                                        origin_addr, origin_count, origin_dt,
+                                        target_disp, target_count, target_dt,
+                                        &(module->req_result));
+        } else {
+            ret = atomic_op_cswap(module, op, target,
+                                  origin_addr, origin_count, origin_dt,
+                                  target_disp, target_count, target_dt,
+                                  &(module->req_result));
+        }
+        return ret;
+    }
+
+
     ret = start_atomicity(module, target);
     if (ret != OMPI_SUCCESS) {
         return ret;
@@ -569,9 +728,11 @@ int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_a
         return ret;
     }
 
-    ret = start_atomicity(module, target);
-    if (ret != OMPI_SUCCESS) {
-        return ret;
+    if (!module->acc_single_intrinsic) {
+        ret = start_atomicity(module, target);
+        if (ret != OMPI_SUCCESS) {
+            return ret;
+        }
     }
 
     if (module->flavor == MPI_WIN_FLAVOR_DYNAMIC) {
@@ -585,7 +746,8 @@ int ompi_osc_ucx_compare_and_swap(const void *origin_addr, const void *compare_a
     ret = opal_common_ucx_wpmem_cmpswp(module->mem,*(uint64_t *)compare_addr,
                                      *(uint64_t *)origin_addr, target,
                                      result_addr, dt_bytes, remote_addr);
-    if (ret != OMPI_SUCCESS) {
+
+    if (module->acc_single_intrinsic) {
         return ret;
     }
 
@@ -611,9 +773,11 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
         ucp_atomic_fetch_op_t opcode;
         size_t dt_bytes;
 
-        ret = start_atomicity(module, target);
-        if (ret != OMPI_SUCCESS) {
-            return ret;
+        if (!module->acc_single_intrinsic) {
+            ret = start_atomicity(module, target);
+            if (ret != OMPI_SUCCESS) {
+                return ret;
+            }
         }
 
         if (module->flavor == MPI_WIN_FLAVOR_DYNAMIC) {
@@ -636,7 +800,8 @@ int ompi_osc_ucx_fetch_and_op(const void *origin_addr, void *result_addr,
 
         ret = opal_common_ucx_wpmem_fetch(module->mem, opcode, value, target,
                                         (void *)result_addr, dt_bytes, remote_addr);
-        if (ret != OMPI_SUCCESS) {
+
+        if (module->acc_single_intrinsic) {
             return ret;
         }
 
@@ -662,6 +827,20 @@ int ompi_osc_ucx_get_accumulate(const void *origin_addr, int origin_count,
         return ret;
     }
 
+    if (module->acc_single_intrinsic) {
+        if (op == &ompi_mpi_op_replace.op || op == &ompi_mpi_op_sum.op) {
+            ret = atomic_op_replace_sum(module, op, target,
+                                        origin_addr, origin_count, origin_dt,
+                                        target_disp, target_count, target_dt, result_addr);
+        } else {
+            ret = atomic_op_cswap(module, op, target,
+                                  origin_addr, origin_count, origin_dt,
+                                  target_disp, target_count, target_dt, result_addr);
+        }
+        return ret;
+    }
+
+
     ret = start_atomicity(module, target);
     if (ret != OMPI_SUCCESS) {
         return ret;

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -72,7 +72,8 @@ ompi_osc_ucx_component_t mca_osc_ucx_component = {
     .wpool                  = NULL,
     .env_initialized        = false,
     .num_incomplete_req_ops = 0,
-    .num_modules            = 0
+    .num_modules            = 0,
+    .acc_single_intrinsic   = false
 };
 
 ompi_osc_ucx_module_t ompi_osc_ucx_module_template = {
@@ -167,6 +168,15 @@ static int component_register(void) {
                                            MCA_BASE_VAR_SCOPE_GROUP, &mca_osc_ucx_component.no_locks);
     free(description_str);
 
+    mca_osc_ucx_component.acc_single_intrinsic = false;
+    opal_asprintf(&description_str, "Enable optimizations for MPI_Fetch_and_op, MPI_Accumulate, etc for codes "
+             "that will not use anything more than a single predefined datatype (default: %s)",
+             mca_osc_ucx_component.acc_single_intrinsic  ? "true" : "false");
+    (void) mca_base_component_var_register(&mca_osc_ucx_component.super.osc_version, "acc_single_intrinsic",
+                                           description_str, MCA_BASE_VAR_TYPE_BOOL, NULL, 0, 0, OPAL_INFO_LVL_5,
+                                           MCA_BASE_VAR_SCOPE_GROUP, &mca_osc_ucx_component.acc_single_intrinsic);
+    free(description_str);
+
     opal_common_ucx_mca_var_register(&mca_osc_ucx_component.super.osc_version);
 
     return OMPI_SUCCESS;
@@ -389,6 +399,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, in
     module->flavor = flavor;
     module->size = size;
     module->no_locks = check_config_value_bool ("no_locks", info);
+    module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info);
 
     /* share everyone's displacement units. Only do an allgather if
        strictly necessary, since it requires O(p) state. */