oshmem: Add symmetric remote key storage structure

After unpacking, we try to find an already existing equivalent remote key. When
found, we destroy the new one and use the old one instead.

Signed-off-by: Thomas Vegas <[email protected]>

Author: tvegas1

config/ompi_check_ucx.m4

@@ -125,7 +125,8 @@ AC_DEFUN([OMPI_CHECK_UCX],[
                           [#include <ucp/api/ucp.h>])
            AC_CHECK_DECLS([ucp_tag_send_nbx,
                            ucp_tag_send_sync_nbx,
-                           ucp_tag_recv_nbx],
+                           ucp_tag_recv_nbx,
+                           ucp_rkey_compare],
                           [], [],
                           [#include <ucp/api/ucp.h>])
            AC_CHECK_TYPES([ucp_request_param_t],

oshmem/mca/spml/ucx/spml_ucx.c

@@ -22,6 +22,7 @@
 #include "opal/datatype/opal_convertor.h"
 #include "opal/mca/common/ucx/common_ucx.h"
 #include "opal/util/opal_environ.h"
+#include "opal/util/minmax.h"
 #include "ompi/datatype/ompi_datatype.h"
 #include "ompi/mca/pml/pml.h"
 
@@ -138,6 +139,159 @@ mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx)
     return 0;
 }
 
+void mca_spml_ucx_rkey_store_init(mca_spml_ucx_rkey_store_t *store)
+{
+    store->array = NULL;
+    store->count = 0;
+    store->size  = 0;
+}
+
+void mca_spml_ucx_rkey_store_cleanup(mca_spml_ucx_rkey_store_t *store)
+{
+    int i;
+
+    for (i = 0; i < store->count; i++) {
+        if (store->array[i].refcnt != 0) {
+            SPML_UCX_ERROR("rkey store destroy: %d/%d has refcnt %d > 0",
+                           i, store->count, store->array[i].refcnt);
+        }
+
+        ucp_rkey_destroy(store->array[i].rkey);
+    }
+
+    free(store->array);
+}
+
+/**
+ * Find position in sorted array for existing or future entry
+ *
+ * @param[in]  store  Store of the entries
+ * @param[in]  worker Common worker for rkeys used
+ * @param[in]  rkey   Remote key to search for
+ * @param[out] index  Index of entry
+ *
+ * @return
+ *   OSHMEM_ERR_NOT_FOUND: index contains the position where future element
+ *                         should be inserted to keep array sorted
+ *   OSHMEM_SUCCESS      : index contains the position of the element
+ *   Other error         : index is not valid
+ */
+static int mca_spml_ucx_rkey_store_find(const mca_spml_ucx_rkey_store_t *store,
+                                        const ucp_worker_h worker,
+                                        const ucp_rkey_h rkey,
+                                        int *index)
+{
+#if HAVE_DECL_UCP_RKEY_COMPARE
+    ucp_rkey_compare_params_t params;
+    int i, result, m, end;
+    ucs_status_t status;
+
+    for (i = 0, end = store->count; i < end;) {
+        m = (i + end) / 2;
+
+        params.field_mask = 0;
+        status = ucp_rkey_compare(worker, store->array[m].rkey,
+                                  rkey, &params, &result);
+        if (status != UCS_OK) {
+            return OSHMEM_ERROR;
+        } else if (result == 0) {
+            *index = m;
+            return OSHMEM_SUCCESS;
+        } else if (result > 0) {
+            end = m;
+        } else {
+            i = m + 1;
+        }
+    }
+
+    *index = i;
+    return OSHMEM_ERR_NOT_FOUND;
+#else
+    return OSHMEM_ERROR;
+#endif
+}
+
+static void mca_spml_ucx_rkey_store_insert(mca_spml_ucx_rkey_store_t *store,
+                                           int i, ucp_rkey_h rkey)
+{
+    int size;
+    mca_spml_ucx_rkey_t *tmp;
+
+    if (store->count >= mca_spml_ucx.symmetric_rkey_max_count) {
+        return;
+    }
+
+    if (store->count >= store->size) {
+        size = opal_min(opal_max(store->size, 8) * 2,
+                        mca_spml_ucx.symmetric_rkey_max_count);
+        tmp  = realloc(store->array, size * sizeof(*store->array));
+        if (tmp == NULL) {
+            return;
+        }
+
+        store->array = tmp;
+        store->size  = size;
+    }
+
+    memmove(&store->array[i + 1], &store->array[i],
+            (store->count - i) * sizeof(*store->array));
+    store->array[i].rkey   = rkey;
+    store->array[i].refcnt = 1;
+    store->count++;
+    return;
+}
+
+/* Takes ownership of input ucp remote key */
+static ucp_rkey_h mca_spml_ucx_rkey_store_get(mca_spml_ucx_rkey_store_t *store,
+                                              ucp_worker_h worker,
+                                              ucp_rkey_h rkey)
+{
+    int ret, i;
+
+    if (mca_spml_ucx.symmetric_rkey_max_count == 0) {
+        return rkey;
+    }
+
+    ret = mca_spml_ucx_rkey_store_find(store, worker, rkey, &i);
+    if (ret == OSHMEM_SUCCESS) {
+        ucp_rkey_destroy(rkey);
+        store->array[i].refcnt++;
+        return store->array[i].rkey;
+    }
+
+    if (ret == OSHMEM_ERR_NOT_FOUND) {
+        mca_spml_ucx_rkey_store_insert(store, i, rkey);
+    }
+
+    return rkey;
+}
+
+static void mca_spml_ucx_rkey_store_put(mca_spml_ucx_rkey_store_t *store,
+                                        ucp_worker_h worker,
+                                        ucp_rkey_h rkey)
+{
+    mca_spml_ucx_rkey_t *entry;
+    int ret, i;
+
+    ret = mca_spml_ucx_rkey_store_find(store, worker, rkey, &i);
+    if (ret != OSHMEM_SUCCESS) {
+        goto out;
+    }
+
+    entry = &store->array[i];
+    assert(entry->rkey == rkey);
+    if (--entry->refcnt > 0) {
+        return;
+    }
+
+    memmove(&store->array[i], &store->array[i + 1],
+            (store->count - (i + 1)) * sizeof(*store->array));
+    store->count--;
+
+out:
+    ucp_rkey_destroy(rkey);
+}
+
 int mca_spml_ucx_enable(bool enable)
 {
     SPML_UCX_VERBOSE(50, "*** ucx ENABLED ****");
@@ -930,6 +1084,8 @@ static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx
         }
     }
 
+    mca_spml_ucx_rkey_store_init(&ucx_ctx->rkey_store);
+
     *ucx_ctx_p = ucx_ctx;
 
     return OSHMEM_SUCCESS;

oshmem/mca/spml/ucx/spml_ucx.h

@@ -76,18 +76,31 @@ struct ucp_peer {
     size_t                   mkeys_cnt;
 };
 typedef struct ucp_peer ucp_peer_t;
- 
+
+/* An rkey_store entry */
+typedef struct mca_spml_ucx_rkey {
+    ucp_rkey_h rkey;
+    int        refcnt;
+} mca_spml_ucx_rkey_t;
+
+typedef struct mca_spml_ucx_rkey_store {
+    mca_spml_ucx_rkey_t *array;
+    int                  size;
+    int                  count;
+} mca_spml_ucx_rkey_store_t;
+
 struct mca_spml_ucx_ctx {
-    ucp_worker_h            *ucp_worker;
-    ucp_peer_t              *ucp_peers;
-    long                     options;
-    opal_bitmap_t            put_op_bitmap;
-    unsigned long            nb_progress_cnt;
-    unsigned int             ucp_workers;
-    int                     *put_proc_indexes;
-    unsigned                 put_proc_count;
-    bool                     synchronized_quiet;
-    int                      strong_sync;
+    ucp_worker_h             *ucp_worker;
+    ucp_peer_t               *ucp_peers;
+    long                      options;
+    opal_bitmap_t             put_op_bitmap;
+    unsigned long             nb_progress_cnt;
+    unsigned int              ucp_workers;
+    int                      *put_proc_indexes;
+    unsigned                  put_proc_count;
+    bool                      synchronized_quiet;
+    int                       strong_sync;
+    mca_spml_ucx_rkey_store_t rkey_store;
 };
 typedef struct mca_spml_ucx_ctx mca_spml_ucx_ctx_t;
 
@@ -284,6 +297,9 @@ extern int mca_spml_ucx_team_reduce(shmem_team_t team, void
 extern unsigned
 mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx);
 
+extern void mca_spml_ucx_rkey_store_init(mca_spml_ucx_rkey_store_t *store);
+extern void mca_spml_ucx_rkey_store_cleanup(mca_spml_ucx_rkey_store_t *store);
+
 static inline int
 mca_spml_ucx_peer_mkey_get(ucp_peer_t *ucp_peer, int index, spml_ucx_cached_mkey_t **out_rmkey)
 {

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -260,7 +260,7 @@ static int mca_spml_ucx_component_close(void)
 
 static int spml_ucx_init(void)
 {
-    unsigned int i;
+    unsigned int i, ret;
     ucs_status_t err;
     ucp_config_t *ucp_config;
     ucp_params_t params;
@@ -336,6 +336,8 @@ static int spml_ucx_init(void)
         mca_spml_ucx_ctx_default.ucp_workers++;
     }
 
+    mca_spml_ucx_rkey_store_init(&mca_spml_ucx_ctx_default.rkey_store);
+
     wrk_attr.field_mask = UCP_WORKER_ATTR_FIELD_THREAD_MODE;
     err = ucp_worker_query(mca_spml_ucx_ctx_default.ucp_worker[0], &wrk_attr);
 
@@ -440,10 +442,25 @@ static void _ctx_cleanup(mca_spml_ucx_ctx_t *ctx)
     free(ctx->ucp_peers);
 }
 
+static void mca_spml_ucx_ctx_fini(mca_spml_ucx_ctx_t *ctx)
+{
+    unsigned int i;
+
+    mca_spml_ucx_rkey_store_cleanup(&ctx->rkey_store);
+    for (i = 0; i < ctx->ucp_workers; i++) {
+        ucp_worker_destroy(ctx->ucp_worker[i]);
+    }
+    free(ctx->ucp_worker);
+    if (ctx != &mca_spml_ucx_ctx_default) {
+        free(ctx);
+    }
+}
+
 static int mca_spml_ucx_component_fini(void)
 {
     int fenced = 0, i;
     int ret = OSHMEM_SUCCESS;
+    mca_spml_ucx_ctx_t *ctx;
 
     opal_progress_unregister(spml_ucx_default_progress);
     if (mca_spml_ucx.active_array.ctxs_count) {
@@ -496,36 +513,26 @@ static int mca_spml_ucx_component_fini(void)
         }
     }
 
-    /* delete all workers */
     for (i = 0; i < mca_spml_ucx.active_array.ctxs_count; i++) {
-        ucp_worker_destroy(mca_spml_ucx.active_array.ctxs[i]->ucp_worker[0]);
-        free(mca_spml_ucx.active_array.ctxs[i]->ucp_worker);
-        free(mca_spml_ucx.active_array.ctxs[i]);
+        mca_spml_ucx_ctx_fini(mca_spml_ucx.active_array.ctxs[i]);
     }
 
     for (i = 0; i < mca_spml_ucx.idle_array.ctxs_count; i++) {
-        ucp_worker_destroy(mca_spml_ucx.idle_array.ctxs[i]->ucp_worker[0]);
-        free(mca_spml_ucx.idle_array.ctxs[i]->ucp_worker);
-        free(mca_spml_ucx.idle_array.ctxs[i]);
+        mca_spml_ucx_ctx_fini(mca_spml_ucx.idle_array.ctxs[i]);
     }
 
     if (mca_spml_ucx_ctx_default.ucp_worker) {
-        for (i = 0; i < (signed int)mca_spml_ucx.ucp_workers; i++) {
-            ucp_worker_destroy(mca_spml_ucx_ctx_default.ucp_worker[i]);
-        }
-        free(mca_spml_ucx_ctx_default.ucp_worker);
+        mca_spml_ucx_ctx_fini(&mca_spml_ucx_ctx_default);
     }
 
     if (mca_spml_ucx.aux_ctx != NULL) {
-        ucp_worker_destroy(mca_spml_ucx.aux_ctx->ucp_worker[0]);
-        free(mca_spml_ucx.aux_ctx->ucp_worker);
+        mca_spml_ucx_ctx_fini(mca_spml_ucx.aux_ctx);
     }
 
     mca_spml_ucx.enabled = false;  /* not anymore */
 
     free(mca_spml_ucx.active_array.ctxs);
     free(mca_spml_ucx.idle_array.ctxs);
-    free(mca_spml_ucx.aux_ctx);
 
     SHMEM_MUTEX_DESTROY(mca_spml_ucx.internal_mutex);
     pthread_mutex_destroy(&mca_spml_ucx.ctx_create_mutex);