4242#endif
4343
4444static
45- spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (shmem_ctx_t ctx , int pe , void * va , void * * rva );
45+ spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (int pe , void * va , void * * rva );
4646
4747mca_spml_ucx_t mca_spml_ucx = {
4848 .super = {
@@ -308,11 +308,11 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
308308
309309
310310static
311- spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (shmem_ctx_t ctx , int pe , void * va , void * * rva )
311+ spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (int pe , void * va , void * * rva )
312312{
313313 sshmem_mkey_t * r_mkey ;
314314
315- r_mkey = mca_memheap_base_get_cached_mkey (ctx , pe , va , 0 , rva );
315+ r_mkey = mca_memheap_base_get_cached_mkey (pe , va , 0 , rva );
316316 if (OPAL_UNLIKELY (!r_mkey )) {
317317 SPML_UCX_ERROR ("pe=%d: %p is not address of symmetric variable" ,
318318 pe , va );
@@ -350,24 +350,31 @@ void *mca_spml_ucx_rmkey_ptr(const void *dst_addr, sshmem_mkey_t *mkey, int pe)
350350#endif
351351}
352352
353- void mca_spml_ucx_rmkey_unpack (shmem_ctx_t ctx , sshmem_mkey_t * mkey , uint32_t segno , int pe , int tr_id )
353+ static void mca_spml_ucx_cache_mkey (mca_spml_ucx_ctx_t * ucx_ctx , sshmem_mkey_t * mkey , uint32_t segno , int dst_pe )
354+ {
355+ ucp_peer_t * peer ;
356+
357+ peer = & (ucx_ctx -> ucp_peers [dst_pe ]);
358+ mkey_segment_init (& peer -> mkeys [segno ].super , mkey , segno );
359+ }
360+
361+ void mca_spml_ucx_rmkey_unpack (sshmem_mkey_t * mkey , uint32_t segno , int pe , int tr_id )
354362{
355363 spml_ucx_mkey_t * ucx_mkey ;
356- mca_spml_ucx_ctx_t * ucx_ctx = (mca_spml_ucx_ctx_t * )ctx ;
357364 ucs_status_t err ;
358365
359- ucx_mkey = & ucx_ctx -> ucp_peers [pe ].mkeys [segno ].key ;
366+ ucx_mkey = & mca_spml_ucx_ctx_default . ucp_peers [pe ].mkeys [segno ].key ;
360367
361- err = ucp_ep_rkey_unpack (ucx_ctx -> ucp_peers [pe ].ucp_conn ,
362- mkey -> u .data ,
368+ err = ucp_ep_rkey_unpack (mca_spml_ucx_ctx_default . ucp_peers [pe ].ucp_conn ,
369+ mkey -> u .data ,
363370 & ucx_mkey -> rkey );
364371 if (UCS_OK != err ) {
365372 SPML_UCX_ERROR ("failed to unpack rkey: %s" , ucs_status_string (err ));
366373 goto error_fatal ;
367374 }
368375
369376 mkey -> spml_context = ucx_mkey ;
370- mca_spml_ucx_cache_mkey (ucx_ctx , mkey , segno , pe );
377+ mca_spml_ucx_cache_mkey (& mca_spml_ucx_ctx_default , mkey , segno , pe );
371378 return ;
372379
373380error_fatal :
@@ -628,7 +635,7 @@ int mca_spml_ucx_get(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_add
628635 ucs_status_t status ;
629636#endif
630637
631- ucx_mkey = mca_spml_ucx_get_mkey (ctx , src , src_addr , & rva , & mca_spml_ucx );
638+ ucx_mkey = mca_spml_ucx_get_mkey (ucx_ctx , src , src_addr , & rva , & mca_spml_ucx );
632639#if HAVE_DECL_UCP_GET_NB
633640 request = ucp_get_nb (ucx_ctx -> ucp_peers [src ].ucp_conn , dst_addr , size ,
634641 (uint64_t )rva , ucx_mkey -> rkey , opal_common_ucx_empty_complete_cb );
@@ -647,7 +654,7 @@ int mca_spml_ucx_get_nb(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_
647654 spml_ucx_mkey_t * ucx_mkey ;
648655 mca_spml_ucx_ctx_t * ucx_ctx = (mca_spml_ucx_ctx_t * )ctx ;
649656
650- ucx_mkey = mca_spml_ucx_get_mkey (ctx , src , src_addr , & rva , & mca_spml_ucx );
657+ ucx_mkey = mca_spml_ucx_get_mkey (ucx_ctx , src , src_addr , & rva , & mca_spml_ucx );
651658 status = ucp_get_nbi (ucx_ctx -> ucp_peers [src ].ucp_conn , dst_addr , size ,
652659 (uint64_t )rva , ucx_mkey -> rkey );
653660
@@ -665,7 +672,7 @@ int mca_spml_ucx_put(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_add
665672 ucs_status_t status ;
666673#endif
667674
668- ucx_mkey = mca_spml_ucx_get_mkey (ctx , dst , dst_addr , & rva , & mca_spml_ucx );
675+ ucx_mkey = mca_spml_ucx_get_mkey (ucx_ctx , dst , dst_addr , & rva , & mca_spml_ucx );
669676#if HAVE_DECL_UCP_PUT_NB
670677 request = ucp_put_nb (ucx_ctx -> ucp_peers [dst ].ucp_conn , src_addr , size ,
671678 (uint64_t )rva , ucx_mkey -> rkey , opal_common_ucx_empty_complete_cb );
@@ -684,7 +691,7 @@ int mca_spml_ucx_put_nb(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_
684691 spml_ucx_mkey_t * ucx_mkey ;
685692 mca_spml_ucx_ctx_t * ucx_ctx = (mca_spml_ucx_ctx_t * )ctx ;
686693
687- ucx_mkey = mca_spml_ucx_get_mkey (ctx , dst , dst_addr , & rva , & mca_spml_ucx );
694+ ucx_mkey = mca_spml_ucx_get_mkey (ucx_ctx , dst , dst_addr , & rva , & mca_spml_ucx );
688695 status = ucp_put_nbi (ucx_ctx -> ucp_peers [dst ].ucp_conn , src_addr , size ,
689696 (uint64_t )rva , ucx_mkey -> rkey );
690697
0 commit comments