Merge pull request #6605 from brminich/topic/shmem_all2all_put

SPML/UCX: Add shmemx_alltoall_global_nb routine to shmemx.h

Author: yosefe

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,21 @@ 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);
 
+/* shmemx_alltoall_global_nb is a nonblocking collective routine, where each PE
+ * exchanges “size” bytes of data with all other PEs in the OpenSHMEM job.
+
+ *  @param dest        A symmetric data object that is large enough to receive
+ *                     “size” bytes of data from each PE in the OpenSHMEM job.
+ *  @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.
+ */
+OSHMEM_DECLSPEC void shmemx_alltoall_global_nb(void *dest, 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 from each PE in the OpenSHMEM job.
+ *  @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 context creation. AUX context may set specific
+     * UCX parameters affecting worker creation, which are not needed for
+     * regular contexts. */
+    pthread_mutex_lock(&mca_spml_ucx.ctx_create_mutex);
+    rc = mca_spml_ucx_ctx_create_common(options, &ucx_ctx);
+    pthread_mutex_unlock(&mca_spml_ucx.ctx_create_mutex);
+    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;
 }
+
+/* this can be called with request==NULL in case of immediate completion */
+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) {
+        pthread_mutex_lock(&mca_spml_ucx.ctx_create_mutex);
+        opal_setenv("UCX_DC_MLX5_TX_POLICY", "rand", 0, &environ);
+    }
+
+    rc = mca_spml_ucx_ctx_create_common(SHMEM_CTX_PRIVATE, &mca_spml_ucx.aux_ctx);
+
+    if (rand_dci_supp) {
+        opal_unsetenv("UCX_DC_MLX5_TX_POLICY", &environ);
+        pthread_mutex_unlock(&mca_spml_ucx.ctx_create_mutex);
+    }
+
+    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++ == 0) {
+            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;
+    } 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;
+}