Skip to content

Commit 1a95633

Browse files
yosefeyosefe
authored
Merge pull request #2717 from alex-mikheev/topic/sshmem_ucx
oshmem: sshmem: adds UCX allocator
2 parents 16ee880 + c9b5b12 commit 1a95633

File tree

13 files changed

+527
-276
lines changed

13 files changed

+527
-276
lines changed

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

Lines changed: 42 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -420,45 +420,53 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
420420
int *count)
421421
{
422422
sshmem_mkey_t *mkeys;
423-
ucs_status_t err;
423+
ucs_status_t status;
424424
spml_ucx_mkey_t *ucx_mkey;
425425
size_t len;
426-
int my_pe = oshmem_my_proc_id();
427426
ucp_mem_map_params_t mem_map_params;
428-
int seg;
427+
int segno;
428+
map_segment_t *mem_seg;
429429
unsigned flags;
430+
int my_pe = oshmem_my_proc_id();
430431

431432
*count = 0;
432433
mkeys = (sshmem_mkey_t *) calloc(1, sizeof(*mkeys));
433434
if (!mkeys) {
434435
return NULL;
435436
}
436437

437-
seg = memheap_find_segnum(addr);
438+
segno = memheap_find_segnum(addr);
439+
mem_seg = memheap_find_seg(segno);
438440

439-
ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[seg].key;
441+
ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[segno].key;
440442
mkeys[0].spml_context = ucx_mkey;
441443

442-
flags = 0;
443-
if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
444-
flags = UCP_MEM_MAP_NONBLOCK;
445-
}
444+
/* if possible use mem handle already created by ucx allocator */
445+
if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
446+
flags = 0;
447+
if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
448+
flags = UCP_MEM_MAP_NONBLOCK;
449+
}
446450

447-
mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
448-
UCP_MEM_MAP_PARAM_FIELD_LENGTH |
449-
UCP_MEM_MAP_PARAM_FIELD_FLAGS;
450-
mem_map_params.address = addr;
451-
mem_map_params.length = size;
452-
mem_map_params.flags = flags;
451+
mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
452+
UCP_MEM_MAP_PARAM_FIELD_LENGTH |
453+
UCP_MEM_MAP_PARAM_FIELD_FLAGS;
454+
mem_map_params.address = addr;
455+
mem_map_params.length = size;
456+
mem_map_params.flags = flags;
453457

454-
err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
455-
if (UCS_OK != err) {
456-
goto error_out;
458+
status = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
459+
if (UCS_OK != status) {
460+
goto error_out;
461+
}
462+
463+
} else {
464+
ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
457465
}
458466

459-
err = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
460-
&mkeys[0].u.data, &len);
461-
if (UCS_OK != err) {
467+
status = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
468+
&mkeys[0].u.data, &len);
469+
if (UCS_OK != status) {
462470
goto error_unmap;
463471
}
464472
if (len >= 0xffff) {
@@ -468,18 +476,18 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
468476
oshmem_shmem_abort(-1);
469477
}
470478

471-
err = ucp_ep_rkey_unpack(mca_spml_ucx.ucp_peers[oshmem_group_self->my_pe].ucp_conn,
472-
mkeys[0].u.data,
473-
&ucx_mkey->rkey);
474-
if (UCS_OK != err) {
479+
status = ucp_ep_rkey_unpack(mca_spml_ucx.ucp_peers[oshmem_group_self->my_pe].ucp_conn,
480+
mkeys[0].u.data,
481+
&ucx_mkey->rkey);
482+
if (UCS_OK != status) {
475483
SPML_ERROR("failed to unpack rkey");
476484
goto error_unmap;
477485
}
478486

479487
mkeys[0].len = len;
480-
mkeys[0].va_base = mem_map_params.address;
488+
mkeys[0].va_base = addr;
481489
*count = 1;
482-
mca_spml_ucx_cache_mkey(&mkeys[0], seg, my_pe);
490+
mca_spml_ucx_cache_mkey(&mkeys[0], segno, my_pe);
483491
return mkeys;
484492

485493
error_unmap:
@@ -493,6 +501,7 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
493501
int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
494502
{
495503
spml_ucx_mkey_t *ucx_mkey;
504+
map_segment_t *mem_seg;
496505

497506
MCA_SPML_CALL(fence());
498507
if (!mkeys)
@@ -501,8 +510,12 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
501510
if (!mkeys[0].spml_context)
502511
return OSHMEM_SUCCESS;
503512

504-
ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
505-
ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
513+
mem_seg = memheap_find_va(mkeys[0].va_base);
514+
515+
if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
516+
ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
517+
ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
518+
}
506519

507520
if (0 < mkeys[0].len) {
508521
ucp_rkey_buffer_release(mkeys[0].u.data);

‎oshmem/mca/sshmem/base/base.h‎

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,6 @@ mca_sshmem_segment_create(map_segment_t *ds_buf,
3333
const char *file_name,
3434
size_t size);
3535

36-
OSHMEM_DECLSPEC int
37-
mca_sshmem_ds_copy(const map_segment_t *from,
38-
map_segment_t *to);
39-
4036
OSHMEM_DECLSPEC void *
4137
mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey);
4238

@@ -148,23 +144,12 @@ OSHMEM_DECLSPEC extern mca_base_framework_t oshmem_sshmem_base_framework;
148144
"Warning %s:%d - %s()", __SSHMEM_FILE__, __LINE__, __func__, __VA_ARGS__)
149145

150146

147+
OSHMEM_DECLSPEC extern void shmem_ds_reset(map_segment_t *ds_buf);
148+
151149
/*
152150
* Get unique file name
153151
*/
154-
static inline char * oshmem_get_unique_file_name(uint64_t pe)
155-
{
156-
char *file_name = NULL;
157-
158-
assert(mca_sshmem_base_backing_file_dir);
159-
160-
if (NULL == (file_name = calloc(OPAL_PATH_MAX, sizeof(char)))) {
161-
return NULL;
162-
}
163-
164-
snprintf(file_name, OPAL_PATH_MAX, "%s/shmem_job_%u_pe_%llu", mca_sshmem_base_backing_file_dir, ORTE_PROC_MY_NAME->jobid, (unsigned long long)pe);
165-
166-
return file_name;
167-
}
152+
OSHMEM_DECLSPEC extern char * oshmem_get_unique_file_name(uint64_t pe);
168153

169154
END_C_DECLS
170155

‎oshmem/mca/sshmem/base/sshmem_base_wrappers.c‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
#include "oshmem/mca/sshmem/sshmem.h"
1616
#include "oshmem/mca/sshmem/base/base.h"
1717

18-
/* ////////////////////////////////////////////////////////////////////////// */
1918
int
2019
mca_sshmem_segment_create(map_segment_t *ds_buf,
2120
const char *file_name,
@@ -28,19 +27,6 @@ mca_sshmem_segment_create(map_segment_t *ds_buf,
2827
return mca_sshmem_base_module->segment_create(ds_buf, file_name, size);
2928
}
3029

31-
/* ////////////////////////////////////////////////////////////////////////// */
32-
int
33-
mca_sshmem_ds_copy(const map_segment_t *from,
34-
map_segment_t *to)
35-
{
36-
if (!mca_sshmem_base_selected) {
37-
return OSHMEM_ERROR;
38-
}
39-
40-
return mca_sshmem_base_module->ds_copy(from, to);
41-
}
42-
43-
/* ////////////////////////////////////////////////////////////////////////// */
4430
void *
4531
mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
4632
{
@@ -51,7 +37,6 @@ mca_sshmem_segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
5137
return mca_sshmem_base_module->segment_attach(ds_buf, mkey);
5238
}
5339

54-
/* ////////////////////////////////////////////////////////////////////////// */
5540
int
5641
mca_sshmem_segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
5742
{
@@ -62,7 +47,6 @@ mca_sshmem_segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
6247
return mca_sshmem_base_module->segment_detach(ds_buf, mkey);
6348
}
6449

65-
/* ////////////////////////////////////////////////////////////////////////// */
6650
int
6751
mca_sshmem_unlink(map_segment_t *ds_buf)
6852
{
@@ -73,3 +57,31 @@ mca_sshmem_unlink(map_segment_t *ds_buf)
7357
return mca_sshmem_base_module->unlink(ds_buf);
7458
}
7559

60+
61+
char * oshmem_get_unique_file_name(uint64_t pe)
62+
{
63+
char *file_name = NULL;
64+
65+
assert(mca_sshmem_base_backing_file_dir);
66+
67+
if (NULL == (file_name = calloc(OPAL_PATH_MAX, sizeof(char)))) {
68+
return NULL;
69+
}
70+
71+
snprintf(file_name, OPAL_PATH_MAX, "%s/shmem_job_%u_pe_%llu", mca_sshmem_base_backing_file_dir, ORTE_PROC_MY_NAME->jobid, (unsigned long long)pe);
72+
73+
return file_name;
74+
}
75+
76+
77+
void
78+
shmem_ds_reset(map_segment_t *ds_buf)
79+
{
80+
MAP_SEGMENT_RESET_FLAGS(ds_buf);
81+
ds_buf->seg_id = MAP_SEGMENT_SHM_INVALID;
82+
ds_buf->super.va_base = 0;
83+
ds_buf->super.va_end = 0;
84+
ds_buf->seg_size = 0;
85+
ds_buf->type = MAP_SEGMENT_UNKNOWN;
86+
}
87+

‎oshmem/mca/sshmem/mmap/sshmem_mmap_module.c‎

Lines changed: 8 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -64,10 +64,6 @@ segment_create(map_segment_t *ds_buf,
6464
const char *file_name,
6565
size_t size);
6666

67-
static int
68-
ds_copy(const map_segment_t *from,
69-
map_segment_t *to);
70-
7167
static void *
7268
segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey);
7369

@@ -88,44 +84,13 @@ mca_sshmem_mmap_module_t mca_sshmem_mmap_module = {
8884
{
8985
module_init,
9086
segment_create,
91-
ds_copy,
9287
segment_attach,
9388
segment_detach,
9489
segment_unlink,
9590
module_finalize
9691
}
9792
};
9893

99-
/* ////////////////////////////////////////////////////////////////////////// */
100-
/* private utility functions */
101-
/* ////////////////////////////////////////////////////////////////////////// */
102-
103-
/* ////////////////////////////////////////////////////////////////////////// */
104-
/**
105-
* completely resets the contents of *ds_buf
106-
*/
107-
static inline void
108-
shmem_ds_reset(map_segment_t *ds_buf)
109-
{
110-
OPAL_OUTPUT_VERBOSE(
111-
(70, oshmem_sshmem_base_framework.framework_output,
112-
"%s: %s: shmem_ds_resetting "
113-
"(id: %d, size: %lu, name: %s)\n",
114-
mca_sshmem_mmap_component.super.base_version.mca_type_name,
115-
mca_sshmem_mmap_component.super.base_version.mca_component_name,
116-
ds_buf->seg_id, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
117-
);
118-
119-
MAP_SEGMENT_RESET_FLAGS(ds_buf);
120-
ds_buf->seg_id = MAP_SEGMENT_SHM_INVALID;
121-
ds_buf->super.va_base = 0;
122-
ds_buf->super.va_end = 0;
123-
ds_buf->seg_size = 0;
124-
ds_buf->type = MAP_SEGMENT_UNKNOWN;
125-
unlink(ds_buf->seg_name);
126-
memset(ds_buf->seg_name, '\0', sizeof(ds_buf->seg_name));
127-
}
128-
12994
/* ////////////////////////////////////////////////////////////////////////// */
13095
static int
13196
module_init(void)
@@ -142,30 +107,6 @@ module_finalize(void)
142107
return OSHMEM_SUCCESS;
143108
}
144109

145-
/* ////////////////////////////////////////////////////////////////////////// */
146-
static int
147-
ds_copy(const map_segment_t *from,
148-
map_segment_t *to)
149-
{
150-
memcpy(to, from, sizeof(map_segment_t));
151-
152-
OPAL_OUTPUT_VERBOSE(
153-
(70, oshmem_sshmem_base_framework.framework_output,
154-
"%s: %s: ds_copy complete "
155-
"from: (id: %d, size: %lu, "
156-
"name: %s flags: 0x%02x) "
157-
"to: (id: %d, size: %lu, "
158-
"name: %s flags: 0x%02x)\n",
159-
mca_sshmem_mmap_component.super.base_version.mca_type_name,
160-
mca_sshmem_mmap_component.super.base_version.mca_component_name,
161-
from->seg_id, (unsigned long)from->seg_size, from->seg_name,
162-
from->flags, to->seg_id, (unsigned long)to->seg_size, to->seg_name,
163-
to->flags)
164-
);
165-
166-
return OSHMEM_SUCCESS;
167-
}
168-
/* ////////////////////////////////////////////////////////////////////////// */
169110

170111
static int
171112
segment_create(map_segment_t *ds_buf,
@@ -225,11 +166,11 @@ segment_create(map_segment_t *ds_buf,
225166
OPAL_OUTPUT_VERBOSE(
226167
(70, oshmem_sshmem_base_framework.framework_output,
227168
"%s: %s: create %s "
228-
"(id: %d, addr: %p size: %lu, name: %s)\n",
169+
"(id: %d, addr: %p size: %lu)\n",
229170
mca_sshmem_mmap_component.super.base_version.mca_type_name,
230171
mca_sshmem_mmap_component.super.base_version.mca_component_name,
231172
(rc ? "failure" : "successful"),
232-
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
173+
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
233174
);
234175

235176
return rc;
@@ -316,10 +257,10 @@ segment_attach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
316257
OPAL_OUTPUT_VERBOSE(
317258
(70, oshmem_sshmem_base_framework.framework_output,
318259
"%s: %s: attach successful "
319-
"(id: %d, addr: %p size: %lu, name: %s | va_base: 0x%p len: %d key %llx)\n",
260+
"(id: %d, addr: %p size: %lu | va_base: 0x%p len: %d key %llx)\n",
320261
mca_sshmem_mmap_component.super.base_version.mca_type_name,
321262
mca_sshmem_mmap_component.super.base_version.mca_component_name,
322-
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name,
263+
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size,
323264
mkey->va_base, mkey->len, (unsigned long long)mkey->u.key)
324265
);
325266

@@ -338,10 +279,10 @@ segment_detach(map_segment_t *ds_buf, sshmem_mkey_t *mkey)
338279
OPAL_OUTPUT_VERBOSE(
339280
(70, oshmem_sshmem_base_framework.framework_output,
340281
"%s: %s: detaching "
341-
"(id: %d, addr: %p size: %lu, name: %s)\n",
282+
"(id: %d, addr: %p size: %lu)\n",
342283
mca_sshmem_mmap_component.super.base_version.mca_type_name,
343284
mca_sshmem_mmap_component.super.base_version.mca_component_name,
344-
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
285+
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
345286
);
346287

347288
munmap((void *)ds_buf->super.va_base, ds_buf->seg_size);
@@ -363,10 +304,10 @@ segment_unlink(map_segment_t *ds_buf)
363304
OPAL_OUTPUT_VERBOSE(
364305
(70, oshmem_sshmem_base_framework.framework_output,
365306
"%s: %s: unlinking "
366-
"(id: %d, addr: %p size: %lu, name: %s)\n",
307+
"(id: %d, addr: %p size: %lu)\n",
367308
mca_sshmem_mmap_component.super.base_version.mca_type_name,
368309
mca_sshmem_mmap_component.super.base_version.mca_component_name,
369-
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size, ds_buf->seg_name)
310+
ds_buf->seg_id, ds_buf->super.va_base, (unsigned long)ds_buf->seg_size)
370311
);
371312

372313
/* don't completely reset. in particular, only reset

0 commit comments

Comments
 (0)