oshmem: Add actual symmetric remote keys deduplication calls

Invoke the lookup call, but only for ranks from different node where, scale is
actually problematic.

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

Author: tvegas1

oshmem/mca/spml/ucx/spml_ucx.c

@@ -373,6 +373,7 @@ int mca_spml_ucx_ctx_mkey_add(mca_spml_ucx_ctx_t *ucx_ctx, int pe, uint32_t segn
 {
     int rc;
     ucs_status_t err;
+    ucp_rkey_h rkey;
 
     rc = mca_spml_ucx_ctx_mkey_new(ucx_ctx, pe, segno, ucx_mkey);
     if (OSHMEM_SUCCESS != rc) {
@@ -381,11 +382,18 @@ int mca_spml_ucx_ctx_mkey_add(mca_spml_ucx_ctx_t *ucx_ctx, int pe, uint32_t segn
     }
 
     if (mkey->u.data) {
-        err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[pe].ucp_conn, mkey->u.data, &((*ucx_mkey)->rkey));
+        err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[pe].ucp_conn, mkey->u.data, &rkey);
         if (UCS_OK != err) {
             SPML_UCX_ERROR("failed to unpack rkey: %s", ucs_status_string(err));
             return OSHMEM_ERROR;
         }
+
+        if (!oshmem_proc_on_local_node(pe)) {
+            rkey = mca_spml_ucx_rkey_store_get(&mca_spml_ucx.rkey_store, ucx_ctx->ucp_worker[0], rkey);
+        }
+
+        (*ucx_mkey)->rkey = rkey;
+
         rc = mca_spml_ucx_ctx_mkey_cache(ucx_ctx, mkey, segno, pe);
         if (OSHMEM_SUCCESS != rc) {
             SPML_UCX_ERROR("mca_spml_ucx_ctx_mkey_cache failed");
@@ -400,7 +408,11 @@ int mca_spml_ucx_ctx_mkey_del(mca_spml_ucx_ctx_t *ucx_ctx, int pe, uint32_t segn
     ucp_peer_t *ucp_peer;
     int rc;
     ucp_peer = &(ucx_ctx->ucp_peers[pe]);
-    ucp_rkey_destroy(ucx_mkey->rkey);
+    if (!oshmem_proc_on_local_node(pe)) {
+        mca_spml_ucx_rkey_store_put(&mca_spml_ucx.rkey_store, ucx_mkey->rkey);
+    } else {
+        ucp_rkey_destroy(ucx_mkey->rkey);
+    }
     ucx_mkey->rkey = NULL;
     rc = mca_spml_ucx_peer_mkey_cache_del(ucp_peer, segno);
     if(OSHMEM_SUCCESS != rc){
@@ -1809,5 +1821,3 @@ int mca_spml_ucx_team_reduce(shmem_team_t team, void
 {
     return OSHMEM_ERR_NOT_IMPLEMENTED;
 }
-
-