Skip to content
22 changes: 15 additions & 7 deletions opal/mca/common/ucx/common_ucx.c
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,11 @@ void opal_common_ucx_mca_proc_added(void)
#endif
}

OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced)
{
return opal_pmix.fence_nb(NULL, 0, opal_common_ucx_mca_fence_complete_cb, (void *)fenced);
}

OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker)
{
volatile int fenced = 0;
Expand Down Expand Up @@ -182,9 +187,8 @@ static void opal_common_ucx_wait_all_requests(void **reqs, int count, ucp_worker
}
}

OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
size_t my_rank, size_t max_disconnect, ucp_worker_h worker)
{
OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *procs, size_t count,
size_t my_rank, size_t max_disconnect, ucp_worker_h worker) {
size_t num_reqs;
size_t max_reqs;
void *dreq, **dreqs;
Expand Down Expand Up @@ -232,10 +236,14 @@ OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, s
opal_common_ucx_wait_all_requests(dreqs, num_reqs, worker);
free(dreqs);

if (OPAL_SUCCESS != (ret = opal_common_ucx_mca_pmix_fence(worker))) {
return ret;
}

return OPAL_SUCCESS;
}

OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
size_t my_rank, size_t max_disconnect, ucp_worker_h worker)
{
opal_common_ucx_del_procs_nofence(procs, count, my_rank, max_disconnect, worker);

return opal_common_ucx_mca_pmix_fence(worker);
}

3 changes: 3 additions & 0 deletions opal/mca/common/ucx/common_ucx.h
Original file line number Diff line number Diff line change
Expand Up @@ -105,8 +105,11 @@ OPAL_DECLSPEC void opal_common_ucx_mca_deregister(void);
OPAL_DECLSPEC void opal_common_ucx_mca_proc_added(void);
OPAL_DECLSPEC void opal_common_ucx_empty_complete_cb(void *request, ucs_status_t status);
OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker);
OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence_nb(int *fenced);
OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, size_t count,
size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *procs, size_t count,
size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
OPAL_DECLSPEC void opal_common_ucx_mca_var_register(const mca_base_component_t *component);

static inline
Expand Down
10 changes: 5 additions & 5 deletions oshmem/mca/spml/ikrit/spml_ikrit.c
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ int mca_spml_ikrit_put_simple(void* dst_addr,
void* src_addr,
int dst);

static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *, uint32_t seg, int remote_pe, int tr_id);
static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *, uint32_t seg, int remote_pe, int tr_id);

static mxm_mem_key_t *mca_spml_ikrit_get_mkey_slow(int pe, void *va, int ptl_id, void **rva);

Expand Down Expand Up @@ -185,7 +185,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
mca_spml_ikrit_get_mkey_slow
};

static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
{
mxm_peer_t *peer;

Expand Down Expand Up @@ -504,7 +504,7 @@ sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
my_rank, i, addr, (unsigned long long)size,
mca_spml_base_mkey2str(&mkeys[i]));

mca_spml_ikrit_cache_mkeys(NULL, &mkeys[i], memheap_find_segnum(addr), my_rank, i);
mca_spml_ikrit_cache_mkeys(&mkeys[i], memheap_find_segnum(addr), my_rank, i);
}
*count = MXM_PTL_LAST;

Expand Down Expand Up @@ -548,7 +548,7 @@ int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys)

}

int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_mkey_t *mkeys)
int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
{
int ptl;

Expand All @@ -567,7 +567,7 @@ int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_m
mkeys[ptl].len = 0;
mkeys[ptl].va_base = mca_memheap_seg2base_va(seg);
mkeys[ptl].u.key = MAP_SEGMENT_SHM_INVALID;
mca_spml_ikrit_cache_mkeys(NULL, &mkeys[ptl], seg, pe, ptl);
mca_spml_ikrit_cache_mkeys(&mkeys[ptl], seg, pe, ptl);
return OSHMEM_SUCCESS;
}

Expand Down
2 changes: 1 addition & 1 deletion oshmem/mca/spml/ikrit/spml_ikrit.h
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ extern sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
uint64_t shmid,
int *count);
extern int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys);
extern int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe,
extern int mca_spml_ikrit_oob_get_mkeys(int pe,
uint32_t segno,
sshmem_mkey_t *mkeys);

Expand Down
Loading