Skip to content

Commit e57e18f

Browse files
jladd-mlnxjladd-mlnx
authored
Merge pull request #6290 from xinzhao3/topic/oshmem_mkeys
OMPI/OSHMEM: bug-fix: store mkeys for each oshmem ctx.
2 parents 9b3a9c2 + f1b095c commit e57e18f

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
@@ -149,7 +149,7 @@ int mca_spml_ikrit_put_simple(void* dst_addr,
149149
void* src_addr,
150150
int dst);
151151

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

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

@@ -185,7 +185,7 @@ mca_spml_ikrit_t mca_spml_ikrit = {
185185
mca_spml_ikrit_get_mkey_slow
186186
};
187187

188-
static void mca_spml_ikrit_cache_mkeys(sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
188+
static void mca_spml_ikrit_cache_mkeys(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t seg, int dst_pe, int tr_id)
189189
{
190190
mxm_peer_t *peer;
191191

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

507-
mca_spml_ikrit_cache_mkeys(&mkeys[i], memheap_find_segnum(addr), my_rank, i);
507+
mca_spml_ikrit_cache_mkeys(NULL, &mkeys[i], memheap_find_segnum(addr), my_rank, i);
508508
}
509509
*count = MXM_PTL_LAST;
510510

@@ -548,7 +548,7 @@ int mca_spml_ikrit_deregister(sshmem_mkey_t *mkeys)
548548

549549
}
550550

551-
int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
551+
int mca_spml_ikrit_oob_get_mkeys(shmem_ctx_t ctx, int pe, uint32_t seg, sshmem_mkey_t *mkeys)
552552
{
553553
int ptl;
554554

@@ -567,7 +567,7 @@ int mca_spml_ikrit_oob_get_mkeys(int pe, uint32_t seg, sshmem_mkey_t *mkeys)
567567
mkeys[ptl].len = 0;
568568
mkeys[ptl].va_base = mca_memheap_seg2base_va(seg);
569569
mkeys[ptl].u.key = MAP_SEGMENT_SHM_INVALID;
570-
mca_spml_ikrit_cache_mkeys(&mkeys[ptl], seg, pe, ptl);
570+
mca_spml_ikrit_cache_mkeys(NULL, &mkeys[ptl], seg, pe, ptl);
571571
return OSHMEM_SUCCESS;
572572
}
573573

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

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

‎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
@@ -42,7 +42,7 @@
4242
#endif
4343

4444
static
45-
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(int pe, void *va, void **rva);
45+
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(shmem_ctx_t ctx, int pe, void *va, void **rva);
4646

4747
mca_spml_ucx_t mca_spml_ucx = {
4848
.super = {
@@ -306,11 +306,11 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
306306

307307

308308
static
309-
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(int pe, void *va, void **rva)
309+
spml_ucx_mkey_t * mca_spml_ucx_get_mkey_slow(shmem_ctx_t ctx, int pe, void *va, void **rva)
310310
{
311311
sshmem_mkey_t *r_mkey;
312312

313-
r_mkey = mca_memheap_base_get_cached_mkey(pe, va, 0, rva);
313+
r_mkey = mca_memheap_base_get_cached_mkey(ctx, pe, va, 0, rva);
314314
if (OPAL_UNLIKELY(!r_mkey)) {
315315
SPML_UCX_ERROR("pe=%d: %p is not address of symmetric variable",
316316
pe, va);
@@ -348,31 +348,24 @@ void *mca_spml_ucx_rmkey_ptr(const void *dst_addr, sshmem_mkey_t *mkey, int pe)
348348
#endif
349349
}
350350

351-
static void mca_spml_ucx_cache_mkey(mca_spml_ucx_ctx_t *ucx_ctx, sshmem_mkey_t *mkey, uint32_t segno, int dst_pe)
352-
{
353-
ucp_peer_t *peer;
354-
355-
peer = &(ucx_ctx->ucp_peers[dst_pe]);
356-
mkey_segment_init(&peer->mkeys[segno].super, mkey, segno);
357-
}
358-
359-
void mca_spml_ucx_rmkey_unpack(sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
351+
void mca_spml_ucx_rmkey_unpack(shmem_ctx_t ctx, sshmem_mkey_t *mkey, uint32_t segno, int pe, int tr_id)
360352
{
361353
spml_ucx_mkey_t *ucx_mkey;
354+
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
362355
ucs_status_t err;
363356

364-
ucx_mkey = &mca_spml_ucx_ctx_default.ucp_peers[pe].mkeys[segno].key;
357+
ucx_mkey = &ucx_ctx->ucp_peers[pe].mkeys[segno].key;
365358

366-
err = ucp_ep_rkey_unpack(mca_spml_ucx_ctx_default.ucp_peers[pe].ucp_conn,
367-
mkey->u.data,
359+
err = ucp_ep_rkey_unpack(ucx_ctx->ucp_peers[pe].ucp_conn,
360+
mkey->u.data,
368361
&ucx_mkey->rkey);
369362
if (UCS_OK != err) {
370363
SPML_UCX_ERROR("failed to unpack rkey: %s", ucs_status_string(err));
371364
goto error_fatal;
372365
}
373366

374367
mkey->spml_context = ucx_mkey;
375-
mca_spml_ucx_cache_mkey(&mca_spml_ucx_ctx_default, mkey, segno, pe);
368+
mca_spml_ucx_cache_mkey(ucx_ctx, mkey, segno, pe);
376369
return;
377370

378371
error_fatal:
@@ -633,7 +626,7 @@ int mca_spml_ucx_get(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_add
633626
ucs_status_t status;
634627
#endif
635628

636-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, src, src_addr, &rva, &mca_spml_ucx);
629+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, src, src_addr, &rva, &mca_spml_ucx);
637630
#if HAVE_DECL_UCP_GET_NB
638631
request = ucp_get_nb(ucx_ctx->ucp_peers[src].ucp_conn, dst_addr, size,
639632
(uint64_t)rva, ucx_mkey->rkey, opal_common_ucx_empty_complete_cb);
@@ -652,7 +645,7 @@ int mca_spml_ucx_get_nb(shmem_ctx_t ctx, void *src_addr, size_t size, void *dst_
652645
spml_ucx_mkey_t *ucx_mkey;
653646
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
654647

655-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, src, src_addr, &rva, &mca_spml_ucx);
648+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, src, src_addr, &rva, &mca_spml_ucx);
656649
status = ucp_get_nbi(ucx_ctx->ucp_peers[src].ucp_conn, dst_addr, size,
657650
(uint64_t)rva, ucx_mkey->rkey);
658651

@@ -670,7 +663,7 @@ int mca_spml_ucx_put(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_add
670663
ucs_status_t status;
671664
#endif
672665

673-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, dst, dst_addr, &rva, &mca_spml_ucx);
666+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, dst, dst_addr, &rva, &mca_spml_ucx);
674667
#if HAVE_DECL_UCP_PUT_NB
675668
request = ucp_put_nb(ucx_ctx->ucp_peers[dst].ucp_conn, src_addr, size,
676669
(uint64_t)rva, ucx_mkey->rkey, opal_common_ucx_empty_complete_cb);
@@ -689,7 +682,7 @@ int mca_spml_ucx_put_nb(shmem_ctx_t ctx, void* dst_addr, size_t size, void* src_
689682
spml_ucx_mkey_t *ucx_mkey;
690683
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
691684

692-
ucx_mkey = mca_spml_ucx_get_mkey(ucx_ctx, dst, dst_addr, &rva, &mca_spml_ucx);
685+
ucx_mkey = mca_spml_ucx_get_mkey(ctx, dst, dst_addr, &rva, &mca_spml_ucx);
693686
status = ucp_put_nbi(ucx_ctx->ucp_peers[dst].ucp_conn, src_addr, size,
694687
(uint64_t)rva, ucx_mkey->rkey);
695688

0 commit comments

Comments
 (0)