diff --git a/oshmem/mca/spml/ucx/spml_ucx.c b/oshmem/mca/spml/ucx/spml_ucx.c index 3f49a5ea9a..2ce55cb21e 100644 --- a/oshmem/mca/spml/ucx/spml_ucx.c +++ b/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; } diff --git a/oshmem/mca/spml/ucx/spml_ucx.h b/oshmem/mca/spml/ucx/spml_ucx.h index 0e0e01b7f1..4b248e026a 100644 --- a/oshmem/mca/spml/ucx/spml_ucx.h +++ b/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; diff --git a/oshmem/mca/spml/ucx/spml_ucx_component.c b/oshmem/mca/spml/ucx/spml_ucx_component.c index e77818f840..d76d6b1b12 100644 --- a/oshmem/mca/spml/ucx/spml_ucx_component.c +++ b/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; }