4343#define SPML_UCX_PUT_DEBUG 0
4444#endif
4545
46+ static
47+ spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (int pe , void * va , void * * rva );
48+
4649mca_spml_ucx_t mca_spml_ucx = {
4750 {
4851 /* Init mca_spml_base_module_t */
@@ -73,7 +76,9 @@ mca_spml_ucx_t mca_spml_ucx = {
7376 NULL , /* ucp_peers */
7477 0 , /* using_mem_hooks */
7578 1 , /* num_disconnect */
76- 0 /* heap_reg_nb */
79+ 0 , /* heap_reg_nb */
80+ 0 , /* enabled */
81+ mca_spml_ucx_get_mkey_slow
7782};
7883
7984int mca_spml_ucx_enable (bool enable )
@@ -328,6 +333,7 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
328333}
329334
330335
336+ static
331337spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow (int pe , void * va , void * * rva )
332338{
333339 sshmem_mkey_t * r_mkey ;
@@ -530,7 +536,7 @@ int mca_spml_ucx_get(void *src_addr, size_t size, void *dst_addr, int src)
530536 ucs_status_t status ;
531537 spml_ucx_mkey_t * ucx_mkey ;
532538
533- ucx_mkey = mca_spml_ucx_get_mkey (src , src_addr , & rva );
539+ ucx_mkey = mca_spml_ucx_get_mkey (src , src_addr , & rva , & mca_spml_ucx );
534540 status = ucp_get (mca_spml_ucx .ucp_peers [src ].ucp_conn , dst_addr , size ,
535541 (uint64_t )rva , ucx_mkey -> rkey );
536542
@@ -543,7 +549,7 @@ int mca_spml_ucx_get_nb(void *src_addr, size_t size, void *dst_addr, int src, vo
543549 ucs_status_t status ;
544550 spml_ucx_mkey_t * ucx_mkey ;
545551
546- ucx_mkey = mca_spml_ucx_get_mkey (src , src_addr , & rva );
552+ ucx_mkey = mca_spml_ucx_get_mkey (src , src_addr , & rva , & mca_spml_ucx );
547553 status = ucp_get_nbi (mca_spml_ucx .ucp_peers [src ].ucp_conn , dst_addr , size ,
548554 (uint64_t )rva , ucx_mkey -> rkey );
549555
@@ -556,7 +562,7 @@ int mca_spml_ucx_put(void* dst_addr, size_t size, void* src_addr, int dst)
556562 ucs_status_t status ;
557563 spml_ucx_mkey_t * ucx_mkey ;
558564
559- ucx_mkey = mca_spml_ucx_get_mkey (dst , dst_addr , & rva );
565+ ucx_mkey = mca_spml_ucx_get_mkey (dst , dst_addr , & rva , & mca_spml_ucx );
560566 status = ucp_put (mca_spml_ucx .ucp_peers [dst ].ucp_conn , src_addr , size ,
561567 (uint64_t )rva , ucx_mkey -> rkey );
562568
@@ -569,7 +575,7 @@ int mca_spml_ucx_put_nb(void* dst_addr, size_t size, void* src_addr, int dst, vo
569575 ucs_status_t status ;
570576 spml_ucx_mkey_t * ucx_mkey ;
571577
572- ucx_mkey = mca_spml_ucx_get_mkey (dst , dst_addr , & rva );
578+ ucx_mkey = mca_spml_ucx_get_mkey (dst , dst_addr , & rva , & mca_spml_ucx );
573579 status = ucp_put_nbi (mca_spml_ucx .ucp_peers [dst ].ucp_conn , src_addr , size ,
574580 (uint64_t )rva , ucx_mkey -> rkey );
575581
0 commit comments