void mkey_segment_init(mkey_segment_t *seg, sshmem_mkey_t *mkey, uint32_t segno)
{
map_segment_t *s;
if (segno >= MCA_MEMHEAP_SEG_COUNT) {
return;
}
s = memheap_find_seg(segno);
assert(NULL != s);
seg->super.va_base = s->super.va_base;
seg->super.va_end = s->super.va_end;
seg->rva_base = mkey->va_base;
}
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 ;
}
