SPML/UCX: Add shmemx_alltoall_global_nb routine to shmemx.h

The new routine transfers the data asynchronously from the source PE to all
PEs in the OpenSHMEM job. The routine returns immediately. The source and
target buffers are reusable only after the completion of the routine.
After the data is transferred to the target buffers, the counter object
is updated atomically. The counter object can be read either using atomic
operations such as shmem_atomic_fetch or can use point-to-point synchronization
routines such as shmem_wait_until and shmem_test.

Signed-off-by: Mikhail Brinskii <mikhailb@mellanox.com>

Author: brminich

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -20,8 +20,6 @@
 #include "osc_ucx.h"
 #include "osc_ucx_request.h"
 
-#define UCX_VERSION(_major, _minor, _build) (((_major) * 100) + (_minor))
-
 #define memcpy_off(_dst, _src, _len, _off)        \
     memcpy(((char*)(_dst)) + (_off), _src, _len); \
     (_off) += (_len);

opal/mca/common/ucx/common_ucx.h

@@ -39,6 +39,9 @@ BEGIN_C_DECLS
 #define MCA_COMMON_UCX_PER_TARGET_OPS_THRESHOLD 1000
 #define MCA_COMMON_UCX_GLOBAL_OPS_THRESHOLD 1000
 
+#define UCX_VERSION(_major, _minor, _build) (((_major) * 100) + (_minor))
+
+
 #define _MCA_COMMON_UCX_QUOTE(_x) \
     # _x
 #define MCA_COMMON_UCX_QUOTE(_x) \

oshmem/include/shmemx.h

@@ -168,6 +168,9 @@ OSHMEM_DECLSPEC void shmemx_int16_prod_to_all(int16_t *target, const int16_t *so
 OSHMEM_DECLSPEC void shmemx_int32_prod_to_all(int32_t *target, const int32_t *source, int nreduce, int PE_start, int logPE_stride, int PE_size, int32_t *pWrk, long *pSync);
 OSHMEM_DECLSPEC void shmemx_int64_prod_to_all(int64_t *target, const int64_t *source, int nreduce, int PE_start, int logPE_stride, int PE_size, int64_t *pWrk, long *pSync);
 
+/* Alltoall put with atomic counter increase */
+OSHMEM_DECLSPEC void shmemx_put_with_long_inc_all(void *target, const void *source, size_t size, long *counter);
+
 /*
  * Backward compatibility section
  */

oshmem/mca/spml/base/base.h

@@ -93,6 +93,10 @@ OSHMEM_DECLSPEC int mca_spml_base_get_nb(void *dst_addr,
                                          void **handle);
 
 OSHMEM_DECLSPEC void mca_spml_base_memuse_hook(void *addr, size_t length);
+
+OSHMEM_DECLSPEC int mca_spml_base_put_all_nb(void *target, const void *source,
+                                             size_t size, long *counter);
+
 /*
  * MCA framework
  */

oshmem/mca/spml/base/spml_base.c

@@ -280,3 +280,9 @@ int mca_spml_base_get_nb(void *dst_addr, size_t size,
 void mca_spml_base_memuse_hook(void *addr, size_t length)
 {
 }
+
+int mca_spml_base_put_all_nb(void *target, const void *source,
+                             size_t size, long *counter)
+{
+    return OSHMEM_ERR_NOT_IMPLEMENTED;
+}

oshmem/mca/spml/ikrit/spml_ikrit.c

@@ -179,6 +179,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
         mca_spml_base_rmkey_free,
         mca_spml_base_rmkey_ptr,
         mca_spml_base_memuse_hook,
+        mca_spml_base_put_all_nb,
 
         (void*)&mca_spml_ikrit
     },

oshmem/mca/spml/spml.h

@@ -314,6 +314,35 @@ typedef int (*mca_spml_base_module_send_fn_t)(void *buf,
                                               int dst,
                                               mca_spml_base_put_mode_t mode);
 
+/**
+ *  The routine transfers the data asynchronously from the source PE to all
+ *  PEs in the OpenSHMEM job. The routine returns immediately. The source and
+ *  target buffers are reusable only after the completion of the routine.
+ *  After the data is transferred to the target buffers, the counter object
+ *  is updated atomically. The counter object can be read either using atomic
+ *  operations such as shmem_atomic_fetch or can use point-to-point synchronization
+ *  routines such as shmem_wait_until and shmem_test.
+ *
+ *  Shmem_quiet may be used for completing the operation, but not required for
+ *  progress or completion. In a multithreaded OpenSHMEM program, the user
+ *  (the OpenSHMEM program) should ensure the correct ordering of
+ *  shmemx_alltoall_global calls.
+ *
+ *  @param dest        A symmetric data object that is large enough to receive
+ *                     “size” bytes of data.
+ *  @param source      A symmetric data object that contains “size” bytes of data
+ *                     for each PE in the OpenSHMEM job.
+ *  @param size        The number of bytes to be sent to each PE in the job.
+ *  @param counter     A symmetric data object to be atomically incremented after
+ *                     the target buffer is updated.
+ *
+ *  @return            OSHMEM_SUCCESS or failure status.
+ */
+typedef int (*mca_spml_base_module_put_all_nb_fn_t)(void *dest,
+                                                    const void *source,
+                                                    size_t size,
+                                                    long *counter);
+
 /**
  * Assures ordering of delivery of put() requests
  *
@@ -381,6 +410,7 @@ struct mca_spml_base_module_1_0_0_t {
     mca_spml_base_module_mkey_ptr_fn_t    spml_rmkey_ptr;
 
     mca_spml_base_module_memuse_hook_fn_t spml_memuse_hook;
+    mca_spml_base_module_put_all_nb_fn_t  spml_put_all_nb;
     void *self;
 };
 

oshmem/mca/spml/ucx/spml_ucx.c

@@ -32,6 +32,7 @@
 #include "oshmem/proc/proc.h"
 #include "oshmem/mca/spml/base/base.h"
 #include "oshmem/mca/spml/base/spml_base_putreq.h"
+#include "oshmem/mca/atomic/atomic.h"
 #include "oshmem/runtime/runtime.h"
 
 #include "oshmem/mca/spml/ucx/spml_ucx_component.h"
@@ -67,6 +68,7 @@ mca_spml_ucx_t mca_spml_ucx = {
         .spml_rmkey_free    = mca_spml_ucx_rmkey_free,
         .spml_rmkey_ptr     = mca_spml_ucx_rmkey_ptr,
         .spml_memuse_hook   = mca_spml_ucx_memuse_hook,
+        .spml_put_all_nb    = mca_spml_ucx_put_all_nb,
         .self               = (void*)&mca_spml_ucx
     },
 
@@ -439,8 +441,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
         ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
     }
 
-    status = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h, 
-                           &mkeys[0].u.data, &len); 
+    status = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
+                           &mkeys[0].u.data, &len);
     if (UCS_OK != status) {
         goto error_unmap;
     }
@@ -477,8 +479,6 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
 {
     spml_ucx_mkey_t   *ucx_mkey;
     map_segment_t *mem_seg;
-    int segno;
-    int my_pe = oshmem_my_proc_id();
 
     MCA_SPML_CALL(quiet(oshmem_ctx_default));
     if (!mkeys)
@@ -493,7 +493,7 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
     if (OPAL_UNLIKELY(NULL == mem_seg)) {
         return OSHMEM_ERROR;
     }
-    
+
     if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
         ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
     }
@@ -545,17 +545,15 @@ static inline void _ctx_remove(mca_spml_ucx_ctx_array_t *array, mca_spml_ucx_ctx
     opal_atomic_wmb ();
 }
 
-int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
+static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx_ctx_p)
 {
-    mca_spml_ucx_ctx_t *ucx_ctx;
     ucp_worker_params_t params;
     ucp_ep_params_t ep_params;
     size_t i, j, nprocs = oshmem_num_procs();
     ucs_status_t err;
-    int my_pe = oshmem_my_proc_id();
-    size_t len;
     spml_ucx_mkey_t *ucx_mkey;
     sshmem_mkey_t *mkey;
+    mca_spml_ucx_ctx_t *ucx_ctx;
     int rc = OSHMEM_ERROR;
 
     ucx_ctx = malloc(sizeof(mca_spml_ucx_ctx_t));
@@ -580,10 +578,6 @@ int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
         goto error;
     }
 
-    if (mca_spml_ucx.active_array.ctxs_count == 0) {
-        opal_progress_register(spml_ucx_ctx_progress);
-    }
-
     for (i = 0; i < nprocs; i++) {
         ep_params.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS;
         ep_params.address    = (ucp_address_t *)(mca_spml_ucx.remote_addrs_tbl[i]);
@@ -609,11 +603,8 @@ int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
         }
     }
 
-    SHMEM_MUTEX_LOCK(mca_spml_ucx.internal_mutex);
-    _ctx_add(&mca_spml_ucx.active_array, ucx_ctx);
-    SHMEM_MUTEX_UNLOCK(mca_spml_ucx.internal_mutex);
+    *ucx_ctx_p = ucx_ctx;
 
-    (*ctx) = (shmem_ctx_t)ucx_ctx;
     return OSHMEM_SUCCESS;
 
  error2:
@@ -634,6 +625,33 @@ int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
     return rc;
 }
 
+int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
+{
+    mca_spml_ucx_ctx_t *ucx_ctx;
+    int rc;
+
+    /* Take a lock controlling aux context. AUX context may set specific
+     * UCX parameters affecting worker creation, which are not needed for
+     * regular contexts. */
+    mca_spml_ucx_aux_lock();
+    rc = mca_spml_ucx_ctx_create_common(options, &ucx_ctx);
+    mca_spml_ucx_aux_unlock();
+    if (rc != OSHMEM_SUCCESS) {
+        return rc;
+    }
+
+    if (mca_spml_ucx.active_array.ctxs_count == 0) {
+        opal_progress_register(spml_ucx_ctx_progress);
+    }
+
+    SHMEM_MUTEX_LOCK(mca_spml_ucx.internal_mutex);
+    _ctx_add(&mca_spml_ucx.active_array, ucx_ctx);
+    SHMEM_MUTEX_UNLOCK(mca_spml_ucx.internal_mutex);
+
+    (*ctx) = (shmem_ctx_t)ucx_ctx;
+    return OSHMEM_SUCCESS;
+}
+
 void mca_spml_ucx_ctx_destroy(shmem_ctx_t ctx)
 {
     MCA_SPML_CALL(quiet(ctx));
@@ -748,6 +766,15 @@ int mca_spml_ucx_quiet(shmem_ctx_t ctx)
          oshmem_shmem_abort(-1);
          return ret;
     }
+
+    /* If put_all_nb op/s is/are being executed asynchronously, need to wait its
+     * completion as well. */
+    if (ctx == oshmem_ctx_default) {
+        while (mca_spml_ucx.aux_refcnt) {
+            opal_progress();
+        }
+    }
+
     return OSHMEM_SUCCESS;
 }
 
@@ -785,3 +812,99 @@ int mca_spml_ucx_send(void* buf,
 
     return rc;
 }
+
+static void mca_spml_ucx_put_all_complete_cb(void *request, ucs_status_t status)
+{
+    if (mca_spml_ucx.async_progress && (--mca_spml_ucx.aux_refcnt == 0)) {
+        opal_event_evtimer_del(mca_spml_ucx.tick_event);
+        opal_progress_unregister(spml_ucx_progress_aux_ctx);
+    }
+
+    if (request != NULL) {
+        ucp_request_free(request);
+    }
+}
+
+/* Should be called with AUX lock taken */
+static int mca_spml_ucx_create_aux_ctx(void)
+{
+    unsigned major      = 0;
+    unsigned minor      = 0;
+    unsigned rel_number = 0;
+    int rc;
+    bool rand_dci_supp;
+
+    ucp_get_version(&major, &minor, &rel_number);
+    rand_dci_supp = UCX_VERSION(major, minor, rel_number) >= UCX_VERSION(1, 6, 0);
+
+    if (rand_dci_supp) {
+        opal_setenv("UCX_DC_TX_POLICY", "rand", 1, &environ);
+        opal_setenv("UCX_DC_MLX5_TX_POLICY", "rand", 1, &environ);
+    }
+
+    rc = mca_spml_ucx_ctx_create_common(SHMEM_CTX_PRIVATE, &mca_spml_ucx.aux_ctx);
+
+    if (rand_dci_supp) {
+        opal_unsetenv("UCX_DC_TX_POLICY", &environ);
+        opal_unsetenv("UCX_DC_MLX5_TX_POLICY", &environ);
+    }
+
+    return rc;
+}
+
+int mca_spml_ucx_put_all_nb(void *dest, const void *source, size_t size, long *counter)
+{
+    int my_pe = oshmem_my_proc_id();
+    long val  = 1;
+    int peer, dst_pe, rc;
+    shmem_ctx_t ctx;
+    struct timeval tv;
+    void *request;
+
+    mca_spml_ucx_aux_lock();
+    if (mca_spml_ucx.async_progress) {
+        if (mca_spml_ucx.aux_ctx == NULL) {
+            rc = mca_spml_ucx_create_aux_ctx();
+            if (rc != OMPI_SUCCESS) {
+                mca_spml_ucx_aux_unlock();
+                oshmem_shmem_abort(-1);
+            }
+        }
+
+        if (!mca_spml_ucx.aux_refcnt) {
+            tv.tv_sec  = 0;
+            tv.tv_usec = mca_spml_ucx.async_tick;
+            opal_event_evtimer_add(mca_spml_ucx.tick_event, &tv);
+            opal_progress_register(spml_ucx_progress_aux_ctx);
+        }
+        ctx = (shmem_ctx_t)mca_spml_ucx.aux_ctx;
+        ++mca_spml_ucx.aux_refcnt;
+    } else {
+        ctx = oshmem_ctx_default;
+    }
+
+    for (peer = 0; peer < oshmem_num_procs(); peer++) {
+        dst_pe = (peer + my_pe) % oshmem_group_all->proc_count;
+        rc = mca_spml_ucx_put_nb(ctx,
+                                 (void*)((uintptr_t)dest + my_pe * size),
+                                 size,
+                                 (void*)((uintptr_t)source + dst_pe * size),
+                                 dst_pe, NULL);
+        RUNTIME_CHECK_RC(rc);
+
+        mca_spml_ucx_fence(ctx);
+
+        rc = MCA_ATOMIC_CALL(add(ctx, (void*)counter, val, sizeof(val), dst_pe));
+        RUNTIME_CHECK_RC(rc);
+    }
+
+    request = ucp_worker_flush_nb(((mca_spml_ucx_ctx_t*)ctx)->ucp_worker, 0,
+                                  mca_spml_ucx_put_all_complete_cb);
+    if (!UCS_PTR_IS_PTR(request)) {
+        mca_spml_ucx_put_all_complete_cb(NULL, UCS_PTR_STATUS(request));
+    }
+
+    mca_spml_ucx_aux_unlock();
+
+    return OSHMEM_SUCCESS;
+}