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

oshmem/mca/spml/ucx/spml_ucx.c

@@ -138,6 +138,128 @@ mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx)
     return 0;
 }
 
+int mca_spml_ucx_rkey_store_create(mca_spml_ucx_rkey_store_t *store,
+                                   ucp_worker_h worker,
+                                   int size)
+{
+    store->array = calloc(size, sizeof(*store->array));
+    if (store->array == NULL) {
+        return OSHMEM_ERR_OUT_OF_RESOURCE;
+    }
+
+    store->size   = size;
+    store->count  = 0;
+    store->worker = worker;
+    return OSHMEM_SUCCESS;
+}
+
+void mca_spml_ucx_rkey_store_destroy(mca_spml_ucx_rkey_store_t *store)
+{
+    int i;
+
+    for (i = 0; i < store->count; i++) {
+        assert(store->array[i].refcnt == 0);
+        ucp_rkey_destroy(store->array[i].rkey);
+    }
+
+    free(store->array);
+    store->array = NULL;
+    store->count = 0;
+    store->size  = 0;
+}
+
+static int mca_spml_ucx_rkey_store_find(const mca_spml_ucx_rkey_store_t *store,
+                                        const ucp_rkey_h target,
+                                        int *index)
+{
+#if HAVE_DECL_UCP_MEM_MAP_SYMMETRIC_RKEY
+    ucp_rkey_compare_params_t params = {};
+    int i, end, result, m;
+    mca_spml_ucx_rkey_t *entry;
+    ucs_status_t status;
+
+    for (i = 0, end = store->count; i < end;) {
+        m = (i + end) / 2;
+        entry = &store->array[m];
+
+        status = ucp_rkey_compare(store->worker, entry->rkey, target,
+                                  &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
+}
+
+/* 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 target)
+{
+    int ret, i;
+
+    if ((store->size == 0) || (worker != store->worker)) {
+        return target;
+    }
+
+    ret = mca_spml_ucx_rkey_store_find(store, target, &i);
+    if ((ret == OSHMEM_ERR_NOT_FOUND) && (store->count < store->size)) {
+        memmove(&store->array[i + 1], &store->array[i],
+                (store->count - i) * sizeof(*store->array));
+        store->array[i].rkey   = target;
+        store->array[i].refcnt = 0;
+        store->count++;
+    } else if (ret != OSHMEM_SUCCESS) {
+        assert((ret == OSHMEM_ERR_NOT_FOUND) || (ret == OSHMEM_ERROR));
+        return target; // internal error or max capacity reached
+    } else {
+        ucp_rkey_destroy(target);
+    }
+
+    store->array[i].refcnt++;
+    return store->array[i].rkey;
+}
+
+static void mca_spml_ucx_rkey_store_put(mca_spml_ucx_rkey_store_t *store,
+                                        ucp_rkey_h target)
+{
+    mca_spml_ucx_rkey_t *entry;
+    int ret, i;
+
+    if (store->count == 0) {
+        ucp_rkey_destroy(target);
+        return;
+    }
+
+    ret = mca_spml_ucx_rkey_store_find(store, target, &i);
+    if (ret != OSHMEM_SUCCESS) {
+        ucp_rkey_destroy(target);
+        return;
+    }
+
+    entry = &store->array[i];
+    entry->refcnt--;
+    if (entry->refcnt <= 0) {
+        ucp_rkey_destroy(entry->rkey);
+
+        memmove(&store->array[i], &store->array[i + 1],
+                (store->count - (i + 1)) * sizeof(*store->array));
+        store->count--;
+    }
+}
+
 int mca_spml_ucx_enable(bool enable)
 {
     SPML_UCX_VERBOSE(50, "*** ucx ENABLED ****");

oshmem/mca/spml/ucx/spml_ucx.h

@@ -101,34 +101,48 @@ typedef struct mca_spml_ucx_ctx_array {
     mca_spml_ucx_ctx_t       **ctxs;
 } mca_spml_ucx_ctx_array_t;
 
+typedef struct mca_spml_ucx_rkey {
+    ucp_rkey_h rkey;
+    uint16_t refcnt;
+} mca_spml_ucx_rkey_t;
+
+typedef struct mca_spml_ucx_rkey_store {
+    int                 size;
+    int                 count;
+    mca_spml_ucx_rkey_t *array;
+    ucp_worker_h        worker;
+} mca_spml_ucx_rkey_store_t;
+
 struct mca_spml_ucx {
-    mca_spml_base_module_t   super;
-    ucp_context_h            ucp_context;
-    int                      num_disconnect;
-    int                      heap_reg_nb;
-    bool                     enabled;
+    mca_spml_base_module_t    super;
+    ucp_context_h             ucp_context;
+    int                       num_disconnect;
+    int                       heap_reg_nb;
+    bool                      enabled;
     mca_spml_ucx_get_mkey_slow_fn_t get_mkey_slow;
-    char                     ***remote_addrs_tbl;
-    mca_spml_ucx_ctx_array_t active_array;
-    mca_spml_ucx_ctx_array_t idle_array;
-    int                      priority; /* component priority */
-    shmem_internal_mutex_t   internal_mutex;
-    pthread_mutex_t          ctx_create_mutex;
+    char                      ***remote_addrs_tbl;
+    mca_spml_ucx_ctx_array_t  active_array;
+    mca_spml_ucx_ctx_array_t  idle_array;
+    int                       priority; /* component priority */
+    shmem_internal_mutex_t    internal_mutex;
+    pthread_mutex_t           ctx_create_mutex;
     /* Fields controlling aux context for put_all_nb SPML routine */
-    bool                     async_progress;
-    int                      async_tick;
-    opal_event_base_t        *async_event_base;
-    opal_event_t             *tick_event;
-    mca_spml_ucx_ctx_t       *aux_ctx;
-    pthread_spinlock_t       async_lock;
-    int                      aux_refcnt;
-    unsigned long            nb_progress_thresh_global;
-    unsigned long            nb_put_progress_thresh;
-    unsigned long            nb_get_progress_thresh;
-    unsigned long            nb_ucp_worker_progress;
-    unsigned int             ucp_workers;
-    unsigned int             ucp_worker_cnt;
-    bool                     symmetric_rkey;
+    bool                      async_progress;
+    int                       async_tick;
+    opal_event_base_t         *async_event_base;
+    opal_event_t              *tick_event;
+    mca_spml_ucx_ctx_t        *aux_ctx;
+    pthread_spinlock_t        async_lock;
+    int                       aux_refcnt;
+    unsigned long             nb_progress_thresh_global;
+    unsigned long             nb_put_progress_thresh;
+    unsigned long             nb_get_progress_thresh;
+    unsigned long             nb_ucp_worker_progress;
+    unsigned int              ucp_workers;
+    unsigned int              ucp_worker_cnt;
+    mca_spml_ucx_rkey_store_t rkey_store;
+    bool                      symmetric_rkey;
+    int                       symmetric_rkey_size;
 };
 typedef struct mca_spml_ucx mca_spml_ucx_t;
 
@@ -284,6 +298,10 @@ 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 int mca_spml_ucx_rkey_store_create(mca_spml_ucx_rkey_store_t *store,
+                                          ucp_worker_h worker, int size);
+extern void mca_spml_ucx_rkey_store_destroy(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

@@ -157,6 +157,10 @@ static int mca_spml_ucx_component_register(void)
                                      "Enable remote keys deduplication to save memory at scale (best-effort)",
                                      &mca_spml_ucx.symmetric_rkey);
 
+    mca_spml_ucx_param_register_int("symmetric_rkey_size", 1000,
+                                    "Size of the symmetric key store. Default (1000)",
+                                    &mca_spml_ucx.symmetric_rkey_size);
+
     mca_spml_ucx_param_register_int("async_tick_usec", 3000,
                                     "Asynchronous progress tick granularity (in usec)",
                                     &mca_spml_ucx.async_tick);
@@ -260,7 +264,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;
@@ -357,6 +361,21 @@ static int spml_ucx_init(void)
                        -1, EV_PERSIST, mca_spml_ucx_async_cb, NULL);
     }
 
+    if (mca_spml_ucx.symmetric_rkey) {
+        if (mca_spml_ucx.symmetric_rkey_size <= 0) {
+            SPML_UCX_ERROR("failed to create symmetric rkey store with size %d",
+                           mca_spml_ucx.symmetric_rkey_size);
+            return OSHMEM_ERROR;
+        }
+
+        ret = mca_spml_ucx_rkey_store_create(&mca_spml_ucx.rkey_store,
+                                             mca_spml_ucx_ctx_default.ucp_worker[0],
+                                             mca_spml_ucx.symmetric_rkey_size);
+        if (OSHMEM_SUCCESS != ret) {
+            SPML_UCX_ERROR("failed to initialize symmetric rkey store");
+        }
+    }
+
     mca_spml_ucx.aux_ctx    = NULL;
     mca_spml_ucx.aux_refcnt = 0;
 
@@ -496,6 +515,8 @@ static int mca_spml_ucx_component_fini(void)
         }
     }
 
+    mca_spml_ucx_rkey_store_destroy(&mca_spml_ucx.rkey_store);
+
     /* 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]);