SPML/UCX: added strong sync for fence

- added infra for strong sync for fence calls
- added [non]blocking get and flush implementations to strong fence
- refactoring for strong sync parameters: used global config
  instead of per-context config

Signed-off-by: Sergey Oblomov <sergeyo@nvidia.com>
(cherry picked from commit 78185f9)

Author: Sergey Oblomov

config/ompi_check_ucx.m4

@@ -123,7 +123,8 @@ AC_DEFUN([OMPI_CHECK_UCX],[
                                         [#include <ucp/api/ucp.h>])
                          AC_CHECK_DECLS([ucp_ep_flush_nb, ucp_worker_flush_nb,
                                          ucp_request_check_status, ucp_put_nb, ucp_get_nb,
-                                         ucp_put_nbx, ucp_get_nbx, ucp_atomic_op_nbx],
+                                         ucp_put_nbx, ucp_get_nbx, ucp_atomic_op_nbx,
+                                         ucp_ep_flush_nbx],
                                         [], [],
                                         [#include <ucp/api/ucp.h>])
                          AC_CHECK_DECLS([ucm_test_events,

oshmem/mca/spml/ucx/spml_ucx.c

@@ -78,14 +78,15 @@ mca_spml_ucx_t mca_spml_ucx = {
     .num_disconnect         = 1,
     .heap_reg_nb            = 0,
     .enabled                = 0,
-    .get_mkey_slow          = NULL
+    .get_mkey_slow          = NULL,
+    .synchronized_quiet     = false,
+    .strong_sync            = SPML_UCX_STRONG_ORDERING_NONE
 };
 
 mca_spml_ucx_ctx_t mca_spml_ucx_ctx_default = {
     .ucp_worker         = NULL,
     .ucp_peers          = NULL,
-    .options            = 0,
-    .synchronized_quiet = false
+    .options            = 0
 };
 
 #if HAVE_DECL_UCP_ATOMIC_OP_NBX
@@ -404,7 +405,7 @@ int mca_spml_ucx_init_put_op_mask(mca_spml_ucx_ctx_t *ctx, size_t nprocs)
 {
     int res;
 
-    if (mca_spml_ucx.synchronized_quiet) {
+    if (mca_spml_ucx_is_strong_ordering()) {
         ctx->put_proc_indexes = malloc(nprocs * sizeof(*ctx->put_proc_indexes));
         if (NULL == ctx->put_proc_indexes) {
             return OSHMEM_ERR_OUT_OF_RESOURCE;
@@ -426,7 +427,7 @@ int mca_spml_ucx_init_put_op_mask(mca_spml_ucx_ctx_t *ctx, size_t nprocs)
 
 int mca_spml_ucx_clear_put_op_mask(mca_spml_ucx_ctx_t *ctx)
 {
-    if (mca_spml_ucx.synchronized_quiet && ctx->put_proc_indexes) {
+    if (mca_spml_ucx_is_strong_ordering() && ctx->put_proc_indexes) {
         OBJ_DESTRUCT(&ctx->put_op_bitmap);
         free(ctx->put_proc_indexes);
     }
@@ -843,7 +844,6 @@ static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx
     ucx_ctx->options = options;
     ucx_ctx->ucp_worker = calloc(1, sizeof(ucp_worker_h));
     ucx_ctx->ucp_workers = 1;
-    ucx_ctx->synchronized_quiet = mca_spml_ucx_ctx_default.synchronized_quiet;
 
     params.field_mask  = UCP_WORKER_PARAM_FIELD_THREAD_MODE;
     if (oshmem_mpi_thread_provided == SHMEM_THREAD_SINGLE || options & SHMEM_CTX_PRIVATE || options & SHMEM_CTX_SERIALIZED) {
@@ -1178,14 +1178,81 @@ int mca_spml_ucx_put_nb_wprogress(shmem_ctx_t ctx, void* dst_addr, size_t size,
     return ucx_status_to_oshmem_nb(status);
 }
 
+static int mca_spml_ucx_strong_sync(shmem_ctx_t ctx)
+{
+    mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
+    ucs_status_ptr_t request;
+    static int flush_get_data;
+    unsigned i;
+    int ret;
+    int idx;
+#if !(HAVE_DECL_UCP_EP_FLUSH_NBX || HAVE_DECL_UCP_EP_FLUSH_NB)
+    ucs_status_t status;
+#endif
+
+    for (i = 0; i < ucx_ctx->put_proc_count; i++) {
+        idx = ucx_ctx->put_proc_indexes[i];
+
+        switch (mca_spml_ucx.strong_sync) {
+        case SPML_UCX_STRONG_ORDERING_NONE:
+        case SPML_UCX_STRONG_ORDERING_GETNB:
+            ret = mca_spml_ucx_get_nb(ctx,
+                                      ucx_ctx->ucp_peers[idx].mkeys[SPML_UCX_SERVICE_SEG]->super.super.va_base,
+                                      sizeof(flush_get_data), &flush_get_data, idx, NULL);
+            break;
+        case SPML_UCX_STRONG_ORDERING_GET:
+            ret = mca_spml_ucx_get(ctx,
+                                   ucx_ctx->ucp_peers[idx].mkeys[SPML_UCX_SERVICE_SEG]->super.super.va_base,
+                                   sizeof(flush_get_data), &flush_get_data, idx);
+            break;
+#if HAVE_DECL_UCP_EP_FLUSH_NBX
+        case SPML_UCX_STRONG_ORDERING_FLUSH:
+            request = ucp_ep_flush_nbx(ucx_ctx->ucp_peers[idx].ucp_conn,
+                                       &mca_spml_ucx_request_param_b);
+            ret = opal_common_ucx_wait_request(request, ucx_ctx->ucp_worker[0], "ucp_flush_nbx");
+#elif HAVE_DECL_UCP_EP_FLUSH_NB
+            request = ucp_ep_flush_nb(ucx_ctx->ucp_peers[idx].ucp_conn, 0, opal_common_ucx_empty_complete_cb);
+            ret = opal_common_ucx_wait_request(request, ucx_ctx->ucp_worker[0], "ucp_flush_nb");
+#else
+            status = ucp_ep_flush(ucx_ctx->ucp_peers[idx].ucp_conn);
+            ret = (status == UCS_OK) ? OPAL_SUCCESS : OPAL_ERROR;
+#endif
+            break;
+        default:
+            /* unknown mode */
+            ret = OMPI_SUCCESS;
+            break;
+        }
+
+        if (OPAL_UNLIKELY(OMPI_SUCCESS != ret)) {
+            oshmem_shmem_abort(-1);
+            return ret;
+        }
+
+        opal_bitmap_clear_bit(&ucx_ctx->put_op_bitmap, idx);
+    }
+
+    ucx_ctx->put_proc_count = 0;
+    return OSHMEM_SUCCESS;
+}
+
 int mca_spml_ucx_fence(shmem_ctx_t ctx)
 {
+    mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
     ucs_status_t err;
+    int ret;
     unsigned int i = 0;
-    mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
 
     opal_atomic_wmb();
 
+    if (mca_spml_ucx.strong_sync != SPML_UCX_STRONG_ORDERING_NONE) {
+        ret = mca_spml_ucx_strong_sync(ctx);
+        if (ret != OSHMEM_SUCCESS) {
+            oshmem_shmem_abort(-1);
+            return ret;
+        }
+    }
+
     for (i=0; i < ucx_ctx->ucp_workers; i++) {
         if (ucx_ctx->ucp_worker[i] != NULL) {
             err = ucp_worker_fence(ucx_ctx->ucp_worker[i]);
@@ -1201,26 +1268,16 @@ int mca_spml_ucx_fence(shmem_ctx_t ctx)
 
 int mca_spml_ucx_quiet(shmem_ctx_t ctx)
 {
-    int flush_get_data;
     int ret;
     unsigned i;
-    int idx;
     mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
 
     if (mca_spml_ucx.synchronized_quiet) {
-        for (i = 0; i < ucx_ctx->put_proc_count; i++) {
-            idx = ucx_ctx->put_proc_indexes[i];
-            ret = mca_spml_ucx_get_nb(ctx,
-                                      ucx_ctx->ucp_peers[idx].mkeys[SPML_UCX_SERVICE_SEG]->super.super.va_base,
-                                      sizeof(flush_get_data), &flush_get_data, idx, NULL);
-            if (OMPI_SUCCESS != ret) {
-                oshmem_shmem_abort(-1);
-                return ret;
-            }
-
-            opal_bitmap_clear_bit(&ucx_ctx->put_op_bitmap, idx);
+        ret = mca_spml_ucx_strong_sync(ctx);
+        if (ret != OSHMEM_SUCCESS) {
+            oshmem_shmem_abort(-1);
+            return ret;
         }
-        ucx_ctx->put_proc_count = 0;
     }
 
     opal_atomic_wmb();

oshmem/mca/spml/ucx/spml_ucx.h

@@ -48,6 +48,13 @@ BEGIN_C_DECLS
 #define SPML_UCX_TRANSP_CNT 1
 #define SPML_UCX_SERVICE_SEG 0
 
+enum {
+    SPML_UCX_STRONG_ORDERING_NONE  = 0, /* don't use strong ordering */
+    SPML_UCX_STRONG_ORDERING_GETNB = 1, /* use non-blocking read to provide ordering */
+    SPML_UCX_STRONG_ORDERING_GET   = 2, /* use blocking read to provide ordering*/
+    SPML_UCX_STRONG_ORDERING_FLUSH = 3  /* flush EP to provide ordering */
+};
+
 /**
  * UCX SPML module
  */
@@ -79,7 +86,6 @@ struct mca_spml_ucx_ctx {
     unsigned int             ucp_workers;
     int                     *put_proc_indexes;
     unsigned                 put_proc_count;
-    bool                     synchronized_quiet;
 };
 typedef struct mca_spml_ucx_ctx mca_spml_ucx_ctx_t;
 
@@ -115,6 +121,7 @@ struct mca_spml_ucx {
     pthread_spinlock_t       async_lock;
     int                      aux_refcnt;
     bool                     synchronized_quiet;
+    int                      strong_sync;
     unsigned long            nb_progress_thresh_global;
     unsigned long            nb_put_progress_thresh;
     unsigned long            nb_get_progress_thresh;
@@ -294,9 +301,15 @@ static inline int ucx_status_to_oshmem_nb(ucs_status_t status)
 #endif
 }
 
+static inline int mca_spml_ucx_is_strong_ordering(void)
+{
+    return (mca_spml_ucx.strong_sync != SPML_UCX_STRONG_ORDERING_NONE) ||
+           mca_spml_ucx.synchronized_quiet;
+}
+
 static inline void mca_spml_ucx_remote_op_posted(mca_spml_ucx_ctx_t *ctx, int dst)
 {
-    if (OPAL_UNLIKELY(ctx->synchronized_quiet)) {
+    if (OPAL_UNLIKELY(mca_spml_ucx_is_strong_ordering())) {
         if (!opal_bitmap_is_set_bit(&ctx->put_op_bitmap, dst)) {
             ctx->put_proc_indexes[ctx->put_proc_count++] = dst;
             opal_bitmap_set_bit(&ctx->put_op_bitmap, dst);

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -159,7 +159,12 @@ static int mca_spml_ucx_component_register(void)
 
     mca_spml_ucx_param_register_bool("synchronized_quiet", 0,
                                      "Use synchronized quiet on shmem_quiet or shmem_barrier_all operations",
-                                     &mca_spml_ucx_ctx_default.synchronized_quiet);
+                                     &mca_spml_ucx.synchronized_quiet);
+
+    mca_spml_ucx_param_register_int("strong_sync", 0,
+                                    "Use strong synchronization on shmem_quiet, shmem_fence or shmem_barrier_all operations: "
+                                    "0 - don't do strong synchronization, 1 - use non blocking get, 2 - use blocking get, 3 - use flush operation",
+                                    &mca_spml_ucx.strong_sync);
 
     mca_spml_ucx_param_register_ulong("nb_progress_thresh_global", 0,
                                     "Number of nb_put or nb_get operations before ucx progress is triggered. Disabled by default (0). Setting this value will override nb_put/get_progress_thresh.",
@@ -384,6 +389,13 @@ mca_spml_ucx_component_init(int* priority,
         return NULL ;
 
     SPML_UCX_VERBOSE(50, "*** ucx initialized ****");
+
+    if ((mca_spml_ucx.strong_sync < SPML_UCX_STRONG_ORDERING_NONE) ||
+        (mca_spml_ucx.strong_sync > SPML_UCX_STRONG_ORDERING_FLUSH)) {
+        SPML_UCX_ERROR("incorrect value of strong_sync parameter: %d",
+                       mca_spml_ucx.strong_sync);
+    }
+
     return &mca_spml_ucx.super;
 }