Merge pull request #2455 from yosefe/topic/ucp-uct-nonblock-mem-reg-api

spml_ucx: allow registering the heap in non-blocking mode.

oshmem/mca/spml/ucx/spml_ucx.c

@@ -72,7 +72,8 @@ mca_spml_ucx_t mca_spml_ucx = {
     NULL,   /* ucp_worker */
     NULL,   /* ucp_peers */
     0,      /* using_mem_hooks */
-    1       /* num_disconnect */
+    1,      /* num_disconnect */
+    0       /* heap_reg_nb */
 };
 
 int mca_spml_ucx_enable(bool enable)
@@ -390,6 +391,7 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
     size_t len;
     int my_pe = oshmem_my_proc_id();
     int seg;
+    unsigned flags;
 
     *count = 0;
     mkeys = (sshmem_mkey_t *) calloc(1, sizeof(*mkeys));
@@ -402,7 +404,11 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
     ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[seg].key;
     mkeys[0].spml_context = ucx_mkey;
 
-    err = ucp_mem_map(mca_spml_ucx.ucp_context, &addr, size, 0, &ucx_mkey->mem_h);
+    flags = 0;
+    if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
+        flags = UCP_MEM_MAP_NONBLOCK;
+    }
+    err = ucp_mem_map(mca_spml_ucx.ucp_context, &addr, size, flags, &ucx_mkey->mem_h);
     if (UCS_OK != err) {
         goto error_out;
     }

oshmem/mca/spml/ucx/spml_ucx.h

@@ -64,6 +64,7 @@ struct mca_spml_ucx {
     ucp_worker_h             ucp_worker;
     ucp_peer_t              *ucp_peers;
     int                      num_disconnect;
+    int                      heap_reg_nb;
 
     int                      priority; /* component priority */
     bool                     enabled;

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -101,6 +101,10 @@ static int mca_spml_ucx_component_register(void)
                                     "How may disconnects go in parallel",
                                     &mca_spml_ucx.num_disconnect);
 
+    mca_spml_ucx_param_register_int("heap_reg_nb", 0,
+                                    "Use non-blocking memory registration for shared heap",
+                                    &mca_spml_ucx.heap_reg_nb);
+
     return OSHMEM_SUCCESS;
 }