static
mxm_mem_key_t *mca_spml_ikrit_get_mkey_slow(int pe, void *va, int ptl_id, void **rva)
{
sshmem_mkey_t *mkey;
retry:
mkey = mca_memheap_base_get_cached_mkey(pe, va, ptl_id, rva);
if (NULL == mkey) {
SPML_ERROR("pe=%d: %p is not address of shared variable", pe, va);
oshmem_shmem_abort(-1);
return NULL;
}
if (MXM_PTL_SHM == ptl_id) {
if (mca_memheap_base_can_local_copy(mkey, va)) {
return NULL;
}
/* if dst addr is on memheap and local copy is not allowed
* disable direct shm transport
*/
if (memheap_is_va_in_segment(va, HEAP_SEG_INDEX)) {
mca_spml_ikrit.mxm_peers[pe].ptl_id = MXM_PTL_RDMA;
}
/* going via mxm must always work */
ptl_id = MXM_PTL_RDMA;
goto retry;
}
return to_mxm_mkey(mkey);
}
void mca_spml_ucx_memuse_hook(void *addr, size_t length)
{
int my_pe;
spml_ucx_mkey_t *ucx_mkey;
ucp_mem_advise_params_t params;
ucs_status_t status;
if (!(mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX))) {
return;
}
my_pe = oshmem_my_proc_id();
ucx_mkey = &mca_spml_ucx_ctx_default.ucp_peers[my_pe].mkeys[HEAP_SEG_INDEX].key;
params.field_mask = UCP_MEM_ADVISE_PARAM_FIELD_ADDRESS |
UCP_MEM_ADVISE_PARAM_FIELD_LENGTH |
UCP_MEM_ADVISE_PARAM_FIELD_ADVICE;
params.address = addr;
params.length = length;
params.advice = UCP_MADV_WILLNEED;
status = ucp_mem_advise(mca_spml_ucx.ucp_context, ucx_mkey->mem_h, ¶ms);
if (UCS_OK != status) {
SPML_UCX_ERROR("ucp_mem_advise failed addr %p len %llu : %s",
addr, (unsigned long long)length, ucs_status_string(status));
}
}
sshmem_mkey_t *mca_spml_ucx_register(void* addr,
size_t size,
uint64_t shmid,
int *count)
{
sshmem_mkey_t *mkeys;
ucs_status_t status;
spml_ucx_mkey_t *ucx_mkey;
size_t len;
ucp_mem_map_params_t mem_map_params;
int segno;
map_segment_t *mem_seg;
unsigned flags;
int my_pe = oshmem_my_proc_id();
*count = 0;
mkeys = (sshmem_mkey_t *) calloc(1, sizeof(*mkeys));
if (!mkeys) {
return NULL;
}
segno = memheap_find_segnum(addr);
mem_seg = memheap_find_seg(segno);
ucx_mkey = &mca_spml_ucx_ctx_default.ucp_peers[my_pe].mkeys[segno].key;
mkeys[0].spml_context = ucx_mkey;
/* if possible use mem handle already created by ucx allocator */
if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
flags = 0;
if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
flags = UCP_MEM_MAP_NONBLOCK;
}
mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
UCP_MEM_MAP_PARAM_FIELD_LENGTH |
UCP_MEM_MAP_PARAM_FIELD_FLAGS;
mem_map_params.address = addr;
mem_map_params.length = size;
mem_map_params.flags = flags;
status = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
if (UCS_OK != status) {
goto error_out;
}
} else {
ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
}
status = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
&mkeys[0].u.data, &len);
if (UCS_OK != status) {
goto error_unmap;
}
if (len >= 0xffff) {
SPML_UCX_ERROR("packed rkey is too long: %llu >= %d",
(unsigned long long)len,
0xffff);
oshmem_shmem_abort(-1);
}
status = ucp_ep_rkey_unpack(mca_spml_ucx_ctx_default.ucp_peers[oshmem_group_self->my_pe].ucp_conn,
mkeys[0].u.data,
&ucx_mkey->rkey);
if (UCS_OK != status) {
SPML_UCX_ERROR("failed to unpack rkey");
goto error_unmap;
}
mkeys[0].len = len;
mkeys[0].va_base = addr;
*count = 1;
mca_spml_ucx_cache_mkey(&mca_spml_ucx_ctx_default, &mkeys[0], segno, my_pe);
return mkeys;
error_unmap:
ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
error_out:
free(mkeys);
return NULL ;
}
