Skip to content

Commit f1b095c

Browse files
xinzhao3xinzhao3
committed
OMPI/OSHMEM: bug-fix: store mkeys for each oshmem ctx.
Signed-off-by: Xin Zhao <[email protected]>
1 parent 352b667 commit f1b095c

File tree

13 files changed

+90
-54
lines changed

13 files changed

+90
-54
lines changed

‎oshmem/mca/atomic/ucx/atomic_ucx_cswap.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ int mca_atomic_ucx_cswap(shmem_ctx_t ctx,
4040
assert(NULL != prev);
4141

4242
*prev = value;
43-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, pe, target, (void *)&rva, mca_spml_self);
43+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
4444
status_ptr = ucp_atomic_fetch_nb(ucx_ctx->ucp_peers[pe].ucp_conn,
4545
UCP_ATOMIC_FETCH_OP_CSWAP, cond, prev, size,
4646
rva, ucx_mkey->rkey,

‎oshmem/mca/atomic/ucx/atomic_ucx_module.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ int mca_atomic_ucx_op(shmem_ctx_t ctx,
4747

4848
assert((8 == size) || (4 == size));
4949

50-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, pe, target, (void *)&rva, mca_spml_self);
50+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
5151
status = ucp_atomic_post(ucx_ctx->ucp_peers[pe].ucp_conn,
5252
op, value, size, rva,
5353
ucx_mkey->rkey);
@@ -70,7 +70,7 @@ int mca_atomic_ucx_fop(shmem_ctx_t ctx,
7070

7171
assert((8 == size) || (4 == size));
7272

73-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, pe, target, (void *)&rva, mca_spml_self);
73+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, pe, target, (void *)&rva, mca_spml_self);
7474
status_ptr = ucp_atomic_fetch_nb(ucx_ctx->ucp_peers[pe].ucp_conn,
7575
op, value, prev, size,
7676
rva, ucx_mkey->rkey,

‎oshmem/mca/memheap/base/base.h‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ void memheap_oob_destruct(void);
6969
OSHMEM_DECLSPEC int mca_memheap_base_is_symmetric_addr(const void* va);
7070
OSHMEM_DECLSPEC sshmem_mkey_t *mca_memheap_base_get_mkey(void* va,
7171
int tr_id);
72-
OSHMEM_DECLSPEC sshmem_mkey_t * mca_memheap_base_get_cached_mkey_slow(map_segment_t *s,
72+
OSHMEM_DECLSPEC sshmem_mkey_t * mca_memheap_base_get_cached_mkey_slow(shmem_ctx_t ctx,
73+
map_segment_t *s,
7374
int pe,
7475
void* va,
7576
int btl_id,
@@ -243,7 +244,8 @@ static inline map_segment_t *memheap_find_va(void* va)
243244
return s;
244245
}
245246

246-
static inline sshmem_mkey_t *mca_memheap_base_get_cached_mkey(int pe,
247+
static inline sshmem_mkey_t *mca_memheap_base_get_cached_mkey(shmem_ctx_t ctx,
248+
int pe,
247249
void* va,
248250
int btl_id,
249251
void** rva)
@@ -273,7 +275,7 @@ static inline sshmem_mkey_t *mca_memheap_base_get_cached_mkey(int pe,
273275
return mkey;
274276
}
275277

276-
return mca_memheap_base_get_cached_mkey_slow(s, pe, va, btl_id, rva);
278+
return mca_memheap_base_get_cached_mkey_slow(ctx, s, pe, va, btl_id, rva);
277279
}
278280

279281
static inline int mca_memheap_base_num_transports(void)

‎oshmem/mca/memheap/base/memheap_base_mkey.c‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ struct oob_comm {
5555
oob_comm_request_t req_pool[MEMHEAP_RECV_REQS_MAX];
5656
opal_list_t req_list;
5757
int is_inited;
58+
shmem_ctx_t ctx;
5859
};
5960

6061
mca_memheap_map_t* memheap_map = NULL;
@@ -66,7 +67,7 @@ static int send_buffer(int pe, opal_buffer_t *msg);
6667
static int oshmem_mkey_recv_cb(void);
6768

6869
/* pickup list of rkeys and remote va */
69-
static int memheap_oob_get_mkeys(int pe,
70+
static int memheap_oob_get_mkeys(shmem_ctx_t ctx, int pe,
7071
uint32_t va_seg_num,
7172
sshmem_mkey_t *mkey);
7273

@@ -142,7 +143,7 @@ static void memheap_attach_segment(sshmem_mkey_t *mkey, int tr_id)
142143
}
143144

144145

145-
static void unpack_remote_mkeys(opal_buffer_t *msg, int remote_pe)
146+
static void unpack_remote_mkeys(shmem_ctx_t ctx, opal_buffer_t *msg, int remote_pe)
146147
{
147148
int32_t cnt;
148149
int32_t n;
@@ -182,7 +183,7 @@ static void unpack_remote_mkeys(opal_buffer_t *msg, int remote_pe)
182183
} else {
183184
memheap_oob.mkeys[tr_id].u.key = MAP_SEGMENT_SHM_INVALID;
184185
}
185-
MCA_SPML_CALL(rmkey_unpack(&memheap_oob.mkeys[tr_id], memheap_oob.segno, remote_pe, tr_id));
186+
MCA_SPML_CALL(rmkey_unpack(ctx, &memheap_oob.mkeys[tr_id], memheap_oob.segno, remote_pe, tr_id));
186187
}
187188

188189
MEMHEAP_VERBOSE(5,
@@ -242,7 +243,7 @@ static void do_recv(int source_pe, opal_buffer_t* buffer)
242243
case MEMHEAP_RKEY_RESP:
243244
MEMHEAP_VERBOSE(5, "*** RKEY RESP");
244245
OPAL_THREAD_LOCK(&memheap_oob.lck);
245-
unpack_remote_mkeys(buffer, source_pe);
246+
unpack_remote_mkeys(memheap_oob.ctx, buffer, source_pe);
246247
memheap_oob.mkeys_rcvd = MEMHEAP_RKEY_RESP;
247248
opal_condition_broadcast(&memheap_oob.cond);
248249
OPAL_THREAD_UNLOCK(&memheap_oob.lck);
@@ -455,14 +456,14 @@ static int send_buffer(int pe, opal_buffer_t *msg)
455456
return rc;
456457
}
457458

458-
static int memheap_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
459+
static int memheap_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_mkey_t *mkeys)
459460
{
460461
opal_buffer_t *msg;
461462
uint8_t cmd;
462463
int i;
463464
int rc;
464465

465-
if (OSHMEM_SUCCESS == MCA_SPML_CALL(oob_get_mkeys(pe, seg, mkeys))) {
466+
if (OSHMEM_SUCCESS == MCA_SPML_CALL(oob_get_mkeys(ctx, pe, seg, mkeys))) {
466467
for (i = 0; i < memheap_map->num_transports; i++) {
467468
MEMHEAP_VERBOSE(5,
468469
"MKEY CALCULATED BY LOCAL SPML: pe: %d tr_id: %d %s",
@@ -478,6 +479,7 @@ static int memheap_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
478479
memheap_oob.mkeys = mkeys;
479480
memheap_oob.segno = seg;
480481
memheap_oob.mkeys_rcvd = 0;
482+
memheap_oob.ctx = ctx;
481483

482484
msg = OBJ_NEW(opal_buffer_t);
483485
if (!msg) {
@@ -645,7 +647,7 @@ void mca_memheap_modex_recv_all(void)
645647
}
646648
memheap_oob.mkeys = s->mkeys_cache[i];
647649
memheap_oob.segno = j;
648-
unpack_remote_mkeys(msg, i);
650+
unpack_remote_mkeys(oshmem_ctx_default, msg, i);
649651
}
650652
}
651653

@@ -674,7 +676,8 @@ void mca_memheap_modex_recv_all(void)
674676
}
675677
}
676678

677-
sshmem_mkey_t * mca_memheap_base_get_cached_mkey_slow(map_segment_t *s,
679+
sshmem_mkey_t * mca_memheap_base_get_cached_mkey_slow(shmem_ctx_t ctx,
680+
map_segment_t *s,
678681
int pe,
679682
void* va,
680683
int btl_id,
@@ -692,7 +695,7 @@ sshmem_mkey_t * mca_memheap_base_get_cached_mkey_slow(map_segment_t *s,
692695
if (!s->mkeys_cache[pe])
693696
return NULL ;
694697

695-
rc = memheap_oob_get_mkeys(pe,
698+
rc = memheap_oob_get_mkeys(ctx, pe,
696699
s - memheap_map->mem_segs,
697700
s->mkeys_cache[pe]);
698701
if (OSHMEM_SUCCESS != rc)

‎oshmem/mca/spml/base/base.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,11 +72,12 @@ OSHMEM_DECLSPEC int mca_spml_base_test(void* addr,
7272
void* value,
7373
int datatype,
7474
int *out_value);
75-
OSHMEM_DECLSPEC int mca_spml_base_oob_get_mkeys(int pe,
75+
OSHMEM_DECLSPEC int mca_spml_base_oob_get_mkeys(shmem_ctx_t ctx,
76+
int pe,
7677
uint32_t seg,
7778
sshmem_mkey_t *mkeys);
7879

79-
OSHMEM_DECLSPEC void mca_spml_base_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t seg, int pe, int tr_id);
80+
OSHMEM_DECLSPEC void mca_spml_base_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int pe, int tr_id);
8081
OSHMEM_DECLSPEC void mca_spml_base_rmkey_free(sshmem_mkey_t *mkey);
8182
OSHMEM_DECLSPEC void *mca_spml_base_rmkey_ptr(const void *dst_addr, sshmem_mkey_t *mkey, int pe);
8283

‎oshmem/mca/spml/base/spml_base.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,12 +247,12 @@ int mca_spml_base_wait_nb(void* handle)
247247
return OSHMEM_SUCCESS;
248248
}
249249

250-
int mca_spml_base_oob_get_mkeys(int pe, uint32_t segno, sshmem_mkey_t *mkeys)
250+
int mca_spml_base_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t segno, sshmem_mkey_t *mkeys)
251251
{
252252
return OSHMEM_ERROR;
253253
}
254254

255-
void mca_spml_base_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
255+
void mca_spml_base_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
256256
{
257257
}
258258

‎oshmem/mca/spml/ikrit/spml_ikrit.c‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ int mca_spml_ikrit_put_simple(void* dst_addr,
151151
void* src_addr,
152152
int dst);
153153

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

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

@@ -187,7 +187,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
187187
mca_spml_ikrit_get_mkey_slow
188188
};
189189

190-
static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
190+
static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
191191
{
192192
mxm_peer_t *peer;
193193

@@ -506,7 +506,7 @@ sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
506506
my_rank, i, addr, (unsigned long long)size,
507507
mca_spml_base_mkey2str(&mkeys[i]));
508508

509-
mca_spml_ikrit_cache_mkeys(&mkeys[i], memheap_find_segnum(addr), my_rank, i);
509+
mca_spml_ikrit_cache_mkeys(NULL, &mkeys[i], memheap_find_segnum(addr), my_rank, i);
510510
}
511511
*count = MXM_PTL_LAST;
512512

@@ -550,7 +550,7 @@ int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys)
550550

551551
}
552552

553-
int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
553+
int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_mkey_t *mkeys)
554554
{
555555
int ptl;
556556

@@ -569,7 +569,7 @@ int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
569569
mkeys[ptl].len = 0;
570570
mkeys[ptl].va_base = mca_memheap_seg2base_va(seg);
571571
mkeys[ptl].u.key = MAP_SEGMENT_SHM_INVALID;
572-
mca_spml_ikrit_cache_mkeys(&mkeys[ptl], seg, pe, ptl);
572+
mca_spml_ikrit_cache_mkeys(NULL, &mkeys[ptl], seg, pe, ptl);
573573
return OSHMEM_SUCCESS;
574574
}
575575

‎oshmem/mca/spml/ikrit/spml_ikrit.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ extern sshmem_mkey_t *mca_spml_ikrit_register(void* addr,
183183
uint64_t shmid,
184184
int *count);
185185
extern int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys);
186-
extern int mca_spml_ikrit_oob_get_mkeys(int pe,
186+
extern int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe,
187187
uint32_t segno,
188188
sshmem_mkey_t *mkeys);
189189

‎oshmem/mca/spml/spml.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ typedef int (*mca_spml_base_module_test_fn_t)(void* addr,
132132
*
133133
* @param mkey remote mkey
134134
*/
135-
typedef void (*mca_spml_base_module_mkey_unpack_fn_t)(sshmem_mkey_t *, uint32_t segno, int remote_pe, int tr_id);
135+
typedef void (*mca_spml_base_module_mkey_unpack_fn_t)(shmem_ctx_t ctx, sshmem_mkey_t *, uint32_t segno, int remote_pe, int tr_id);
136136

137137
/**
138138
* If possible, get a pointer to the remote memory described by the mkey
@@ -180,7 +180,7 @@ typedef int (*mca_spml_base_module_deregister_fn_t)(sshmem_mkey_t *mkeys);
180180
*
181181
* @return OSHMEM_SUCCSESS if keys are found
182182
*/
183-
typedef int (*mca_spml_base_module_oob_get_mkeys_fn_t)(int pe,
183+
typedef int (*mca_spml_base_module_oob_get_mkeys_fn_t)(shmem_ctx_t ctx, int pe,
184184
uint32_t seg,
185185
sshmem_mkey_t *mkeys);
186186

‎oshmem/mca/spml/ucx/spml_ucx.c‎

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
#endif
4646

4747
static
48-
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(int pe, void *va, void **rva);
48+
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(shmem_ctx_t ctx, int pe, void *va, void **rva);
4949

5050
mca_spml_ucx_t mca_spml_ucx = {
5151
.super = {
@@ -309,11 +309,11 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
309309

310310

311311
static
312-
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(int pe, void *va, void **rva)
312+
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(shmem_ctx_t ctx, int pe, void *va, void **rva)
313313
{
314314
sshmem_mkey_t *r_mkey;
315315

316-
r_mkey = mca_memheap_base_get_cached_mkey(pe, va, 0, rva);
316+
r_mkey = mca_memheap_base_get_cached_mkey(ctx, pe, va, 0, rva);
317317
if (OPAL_UNLIKELY(!r_mkey)) {
318318
SPML_UCX_ERROR("pe=%d: %p is not address of symmetric variable",
319319
pe, va);
@@ -351,31 +351,24 @@ void *mca_spml_ucx_rmkey_ptr(const void *dst_addr, sshmem_mkey_t *mkey, int pe)
351351
#endif
352352
}
353353

354-
static void mca_spml_ucx_cache_mkey(mca_spml_ucx_ctx_t *ucx_ctx, sshmem_mkey_t *mkey, uint32_t segno, int dst_pe)
355-
{
356-
ucp_peer_t *peer;
357-
358-
peer = &(ucx_ctx->ucp_peers[dst_pe]);
359-
mkey_segment_init(&peer->mkeys[segno].super, mkey, segno);
360-
}
361-
362-
void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
354+
void mca_spml_ucx_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
363355
{
364356
spml_ucx_mkey_t *ucx_mkey;
357+
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
365358
ucs_status_t err;
366359

367-
ucx_mkey = &mca_spml_ucx_ctx_default.ucp_peers[pe].mkeys[segno].key;
360+
ucx_mkey = &ucx_ctx->ucp_peers[pe].mkeys[segno].key;
368361

369-
err = ucp_ep_rkey_unpack(mca_spml_ucx_ctx_default.ucp_peers[pe].ucp_conn,
370-
mkey->u.data,
362+
err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[pe].ucp_conn,
363+
mkey->u.data,
371364
&ucx_mkey->rkey);
372365
if (UCS_OK != err) {
373366
SPML_UCX_ERROR("failed to unpack rkey: %s", ucs_status_string(err));
374367
goto error_fatal;
375368
}
376369

377370
mkey->spml_context = ucx_mkey;
378-
mca_spml_ucx_cache_mkey(&mca_spml_ucx_ctx_default, mkey, segno, pe);
371+
mca_spml_ucx_cache_mkey(ucx_ctx, mkey, segno, pe);
379372
return;
380373

381374
error_fatal:
@@ -636,7 +629,7 @@ int mca_spml_ucx_get(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_add
636629
ucs_status_t status;
637630
#endif
638631

639-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, src, src_addr, &rva, &mca_spml_ucx);
632+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, src, src_addr, &rva, &mca_spml_ucx);
640633
#if HAVE_DECL_UCP_GET_NB
641634
request = ucp_get_nb(ucx_ctx->ucp_peers[src].ucp_conn, dst_addr, size,
642635
(uint64_t)rva, ucx_mkey->rkey, opal_common_ucx_empty_complete_cb);
@@ -655,7 +648,7 @@ int mca_spml_ucx_get_nb(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_
655648
spml_ucx_mkey_t *ucx_mkey;
656649
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
657650

658-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, src, src_addr, &rva, &mca_spml_ucx);
651+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, src, src_addr, &rva, &mca_spml_ucx);
659652
status = ucp_get_nbi(ucx_ctx->ucp_peers[src].ucp_conn, dst_addr, size,
660653
(uint64_t)rva, ucx_mkey->rkey);
661654

@@ -673,7 +666,7 @@ int mca_spml_ucx_put(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_add
673666
ucs_status_t status;
674667
#endif
675668

676-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, dst, dst_addr, &rva, &mca_spml_ucx);
669+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, dst, dst_addr, &rva, &mca_spml_ucx);
677670
#if HAVE_DECL_UCP_PUT_NB
678671
request = ucp_put_nb(ucx_ctx->ucp_peers[dst].ucp_conn, src_addr, size,
679672
(uint64_t)rva, ucx_mkey->rkey, opal_common_ucx_empty_complete_cb);
@@ -692,7 +685,7 @@ int mca_spml_ucx_put_nb(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_
692685
spml_ucx_mkey_t *ucx_mkey;
693686
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
694687

695-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, dst, dst_addr, &rva, &mca_spml_ucx);
688+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, dst, dst_addr, &rva, &mca_spml_ucx);
696689
status = ucp_put_nbi(ucx_ctx->ucp_peers[dst].ucp_conn, src_addr, size,
697690
(uint64_t)rva, ucx_mkey->rkey);
698691

0 commit comments

Comments
 (0)