oshmem: Add symmetric remote key environment variable

This is used to activate symmetric remote key registration and deduplication.

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

Author: tvegas1

config/ompi_check_ucx.m4

@@ -107,7 +107,8 @@ AC_DEFUN([OMPI_CHECK_UCX],[
                            UCP_ATOMIC_FETCH_OP_FXOR,
                            UCP_PARAM_FIELD_ESTIMATED_NUM_PPN,
                            UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK,
-                           UCP_OP_ATTR_FLAG_MULTI_SEND],
+                           UCP_OP_ATTR_FLAG_MULTI_SEND,
+                           UCP_MEM_MAP_SYMMETRIC_RKEY],
                           [], [],
                           [#include <ucp/api/ucp.h>])
            AC_CHECK_DECLS([UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS],

oshmem/mca/spml/ucx/spml_ucx.c

@@ -126,6 +126,18 @@ static ucp_request_param_t mca_spml_ucx_request_param_b = {
 };
 #endif
 
+unsigned
+mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx)
+{
+#if HAVE_DECL_UCP_MEM_MAP_SYMMETRIC_RKEY
+    if (spml_ucx->symmetric_rkey_max_count > 0) {
+        return UCP_MEM_MAP_SYMMETRIC_RKEY;
+    }
+#endif
+
+    return 0;
+}
+
 int mca_spml_ucx_enable(bool enable)
 {
     SPML_UCX_VERBOSE(50, "*** ucx ENABLED ****");
@@ -725,7 +737,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
                                     UCP_MEM_MAP_PARAM_FIELD_FLAGS;
         mem_map_params.address    = addr;
         mem_map_params.length     = size;
-        mem_map_params.flags      = flags;
+        mem_map_params.flags      = flags |
+            mca_spml_ucx_mem_map_flags_symmetric_rkey(&mca_spml_ucx);
 
         status = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &mem_h);
         if (UCS_OK != status) {

oshmem/mca/spml/ucx/spml_ucx.h

@@ -128,6 +128,7 @@ struct mca_spml_ucx {
     unsigned long            nb_ucp_worker_progress;
     unsigned int             ucp_workers;
     unsigned int             ucp_worker_cnt;
+    int                      symmetric_rkey_max_count;
 };
 typedef struct mca_spml_ucx mca_spml_ucx_t;
 
@@ -280,6 +281,8 @@ extern int mca_spml_ucx_team_fcollect(shmem_team_t team, void
 extern int mca_spml_ucx_team_reduce(shmem_team_t team, void
         *dest, const void *source, size_t nreduce, int operation, int datatype);
 
+extern unsigned
+mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx);
 
 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

@@ -153,6 +153,10 @@ static int mca_spml_ucx_component_register(void)
                                      "Enable asynchronous progress thread",
                                      &mca_spml_ucx.async_progress);
 
+    mca_spml_ucx_param_register_int("symmetric_rkey_max_count", 0,
+                                    "Size of the symmetric key store. Non-zero to enable, typical use 5000",
+                                    &mca_spml_ucx.symmetric_rkey_max_count);
+
     mca_spml_ucx_param_register_int("async_tick_usec", 3000,
                                     "Asynchronous progress tick granularity (in usec)",
                                     &mca_spml_ucx.async_tick);

oshmem/mca/sshmem/ucx/sshmem_ucx_module.c

@@ -124,7 +124,8 @@ segment_create_internal(map_segment_t *ds_buf, void *address, size_t size,
 
     mem_map_params.address    = address;
     mem_map_params.length     = size;
-    mem_map_params.flags      = flags;
+    mem_map_params.flags      = flags |
+        mca_spml_ucx_mem_map_flags_symmetric_rkey(spml);
 
     status = ucp_mem_map(spml->ucp_context, &mem_map_params, &mem_h);
     if (UCS_OK != status) {