Skip to content

Commit 289595e

Browse files
xinzhao3Tomislav Janjusic
xinzhao3
authored and
Tomislav Janjusic
committed
OMPI/OSHMEM: bug-fix: store mkeys for each oshmem ctx.
Signed-off-by: Xin Zhao <xinz@mellanox.com> Signed-off-by: Tomislav Janjusic <tomislavj@mellanox.com>
1 parent 79ba752 commit 289595e

File tree

13 files changed

+133
-69
lines changed

13 files changed

+133
-69
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/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

0 commit comments

Comments
 (0)