Skip to content

Commit 132fbd9

Browse files
alex-mikheevalex-mikheev
committed
oshmem: sshmem: add UCX allocator
Signed-off-by: Alex Mikheev <[email protected]>
1 parent e038e3f commit 132fbd9

File tree

6 files changed

+464
-20
lines changed

6 files changed

+464
-20
lines changed

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

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -425,7 +425,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
425425
size_t len;
426426
int my_pe = oshmem_my_proc_id();
427427
ucp_mem_map_params_t mem_map_params;
428-
int seg;
428+
int segno;
429+
map_segment_t *mem_seg;
429430
unsigned flags;
430431

431432
*count = 0;
@@ -434,26 +435,32 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
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+
err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
459+
if (UCS_OK != err) {
460+
goto error_out;
461+
}
462+
} else {
463+
ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
457464
}
458465

459466
err = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
@@ -477,9 +484,9 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
477484
}
478485

479486
mkeys[0].len = len;
480-
mkeys[0].va_base = mem_map_params.address;
487+
mkeys[0].va_base = addr;
481488
*count = 1;
482-
mca_spml_ucx_cache_mkey(&mkeys[0], seg, my_pe);
489+
mca_spml_ucx_cache_mkey(&mkeys[0], segno, my_pe);
483490
return mkeys;
484491

485492
error_unmap:
@@ -493,6 +500,7 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
493500
int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
494501
{
495502
spml_ucx_mkey_t *ucx_mkey;
503+
map_segment_t *mem_seg;
496504

497505
MCA_SPML_CALL(fence());
498506
if (!mkeys)
@@ -501,8 +509,12 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
501509
if (!mkeys[0].spml_context)
502510
return OSHMEM_SUCCESS;
503511

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

507519
if (0 < mkeys[0].len) {
508520
ucp_rkey_buffer_release(mkeys[0].u.data);

‎oshmem/mca/sshmem/ucx/Makefile.am‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
# Copyright (c) 2014 Mellanox Technologies, Inc.
2+
# All rights reserved.
3+
# $COPYRIGHT$
4+
#
5+
# Additional copyrights may follow
6+
#
7+
# $HEADER$
8+
#
9+
10+
#dist_oshmemdata_DATA = help-oshmem-sshmem-mmap.txt
11+
12+
AM_CPPFLAGS = $(sshmem_ucx_CPPFLAGS)
13+
14+
sources = \
15+
sshmem_ucx.h \
16+
sshmem_ucx_component.c \
17+
sshmem_ucx_module.c
18+
19+
# Make the output library in this directory, and name it either
20+
# mca_<type>_<name>.la (for DSO builds) or libmca_<type>_<name>.la
21+
# (for static builds).
22+
23+
if MCA_BUILD_oshmem_sshmem_ucx_DSO
24+
component_noinst =
25+
component_install = mca_sshmem_ucx.la
26+
else
27+
component_noinst = libmca_sshmem_ucx.la
28+
component_install =
29+
endif
30+
31+
mcacomponentdir = $(oshmemlibdir)
32+
mcacomponent_LTLIBRARIES = $(component_install)
33+
mca_sshmem_ucx_la_SOURCES = $(sources)
34+
mca_sshmem_ucx_la_LDFLAGS = -module -avoid-version $(sshmem_ucx_LDFLAGS)
35+
mca_sshmem_ucx_la_LIBADD = $(sshmem_ucx_LIBS)
36+
37+
noinst_LTLIBRARIES = $(component_noinst)
38+
libmca_sshmem_ucx_la_SOURCES =$(sources)
39+
libmca_sshmem_ucx_la_LDFLAGS = -module -avoid-version $(sshmem_ucx_LDFLAGS)
40+

‎oshmem/mca/sshmem/ucx/configure.m4‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* Copyright (c) 2017 Mellanox Technologies, Inc.
3+
* All rights reserved.
4+
* $COPYRIGHT$
5+
*
6+
* Additional copyrights may follow
7+
*
8+
* $HEADER$
9+
*/
10+
11+
# MCA_oshmem_sshmem_ucx_CONFIG([action-if-can-compile],
12+
# [action-if-cant-compile])
13+
# ------------------------------------------------
14+
AC_DEFUN([MCA_oshmem_sshmem_ucx_CONFIG],[
15+
AC_CONFIG_FILES([oshmem/mca/sshmem/ucx/Makefile])
16+
17+
OMPI_CHECK_UCX([sshmem_ucx],
18+
[sshmem_ucx_happy="yes"],
19+
[sshmem_ucx_happy="no"])
20+
21+
AS_IF([test "$sshmem_ucx_happy" = "yes"],
22+
[$1],
23+
[$2])
24+
25+
26+
# substitute in the things needed to build ucx
27+
AC_SUBST([sshmem_ucx_CFLAGS])
28+
AC_SUBST([sshmem_ucx_CPPFLAGS])
29+
AC_SUBST([sshmem_ucx_LDFLAGS])
30+
AC_SUBST([sshmem_ucx_LIBS])
31+
])dnl
32+

‎oshmem/mca/sshmem/ucx/sshmem_ucx.h‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright (c) 2017 Mellanox Technologies, Inc.
3+
* All rights reserved.
4+
* $COPYRIGHT$
5+
*
6+
* Additional copyrights may follow
7+
*
8+
* $HEADER$
9+
*/
10+
11+
#ifndef MCA_SSHMEM_UCX_EXPORT_H
12+
#define MCA_SSHMEM_UCX_EXPORT_H
13+
14+
#include "oshmem_config.h"
15+
16+
#include "oshmem/mca/sshmem/sshmem.h"
17+
18+
BEGIN_C_DECLS
19+
20+
/**
21+
* globally exported variable to hold the ucx component.
22+
*/
23+
typedef struct mca_sshmem_ucx_component_t {
24+
/* base component struct */
25+
mca_sshmem_base_component_t super;
26+
/* priority for ucx component */
27+
int priority;
28+
} mca_sshmem_ucx_component_t;
29+
30+
OSHMEM_MODULE_DECLSPEC extern mca_sshmem_ucx_component_t
31+
mca_sshmem_ucx_component;
32+
33+
typedef struct mca_sshmem_ucx_module_t {
34+
mca_sshmem_base_module_t super;
35+
} mca_sshmem_ucx_module_t;
36+
extern mca_sshmem_ucx_module_t mca_sshmem_ucx_module;
37+
38+
END_C_DECLS
39+
40+
#endif /* MCA_SHMEM_UCX_EXPORT_H */

‎oshmem/mca/sshmem/ucx/sshmem_ucx_component.c‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil -*- */
2+
/*
3+
* Copyright (c) 2017 Mellanox Technologies, Inc.
4+
* All rights reserved.
5+
* $COPYRIGHT$
6+
*
7+
* Additional copyrights may follow
8+
*
9+
* $HEADER$
10+
*/
11+
12+
#include "oshmem_config.h"
13+
14+
#include "opal/constants.h"
15+
16+
#include "oshmem/mca/sshmem/sshmem.h"
17+
#include "oshmem/mca/sshmem/base/base.h"
18+
#include "oshmem/mca/spml/base/base.h"
19+
20+
#include "sshmem_ucx.h"
21+
22+
/**
23+
* public string showing the shmem ompi_ucx component version number
24+
*/
25+
const char *mca_sshmem_ucx_component_version_string =
26+
"OSHMEM ucx sshmem MCA component version " OSHMEM_VERSION;
27+
28+
29+
/**
30+
* local functions
31+
*/
32+
static int ucx_register(void);
33+
static int ucx_open(void);
34+
static int ucx_close(void);
35+
static int ucx_query(mca_base_module_t **module, int *priority);
36+
static int ucx_runtime_query(mca_base_module_t **module,
37+
int *priority,
38+
const char *hint);
39+
40+
/**
41+
* instantiate the public struct with all of our public information
42+
* and pointers to our public functions in it
43+
*/
44+
mca_sshmem_ucx_component_t mca_sshmem_ucx_component = {
45+
/* ////////////////////////////////////////////////////////////////////// */
46+
/* super */
47+
/* ////////////////////////////////////////////////////////////////////// */
48+
{
49+
/**
50+
* common MCA component data
51+
*/
52+
.base_version = {
53+
MCA_SSHMEM_BASE_VERSION_2_0_0,
54+
55+
/* component name and version */
56+
.mca_component_name = "ucx",
57+
MCA_BASE_MAKE_VERSION(component, OSHMEM_MAJOR_VERSION, OSHMEM_MINOR_VERSION,
58+
OSHMEM_RELEASE_VERSION),
59+
60+
.mca_open_component = ucx_open,
61+
.mca_close_component = ucx_close,
62+
.mca_query_component = ucx_query,
63+
.mca_register_component_params = ucx_register,
64+
},
65+
/* MCA v2.0.0 component meta data */
66+
.base_data = {
67+
/* the component is checkpoint ready */
68+
MCA_BASE_METADATA_PARAM_CHECKPOINT
69+
},
70+
.runtime_query = ucx_runtime_query,
71+
},
72+
};
73+
74+
static int
75+
ucx_runtime_query(mca_base_module_t **module,
76+
int *priority,
77+
const char *hint)
78+
{
79+
/* check that spml ucx was selected. Otherwise disqualify */
80+
if (strcmp(mca_spml_base_selected_component.spmlm_version.mca_component_name, "ucx")) {
81+
*module = NULL;
82+
return OSHMEM_ERR_NOT_AVAILABLE;
83+
}
84+
85+
*priority = mca_sshmem_ucx_component.priority;
86+
/* use lowest priority because UCX does not support
87+
* fixed address mapping yet
88+
*/
89+
*priority = 0;
90+
*module = (mca_base_module_t *)&mca_sshmem_ucx_module.super;
91+
return OPAL_SUCCESS;
92+
}
93+
94+
static int
95+
ucx_register(void)
96+
{
97+
/* (default) priority - set high to make ucx the default */
98+
mca_sshmem_ucx_component.priority = 100;
99+
mca_base_component_var_register (&mca_sshmem_ucx_component.super.base_version,
100+
"priority", "Priority for sshmem ucx "
101+
"component (default: 100)", MCA_BASE_VAR_TYPE_INT,
102+
NULL, 0, MCA_BASE_VAR_FLAG_SETTABLE,
103+
OPAL_INFO_LVL_3,
104+
MCA_BASE_VAR_SCOPE_ALL_EQ,
105+
&mca_sshmem_ucx_component.priority);
106+
107+
return OSHMEM_SUCCESS;
108+
}
109+
110+
static int
111+
ucx_open(void)
112+
{
113+
return OSHMEM_SUCCESS;
114+
}
115+
116+
static int
117+
ucx_query(mca_base_module_t **module, int *priority)
118+
{
119+
*priority = mca_sshmem_ucx_component.priority;
120+
*module = (mca_base_module_t *)&mca_sshmem_ucx_module.super;
121+
return OSHMEM_SUCCESS;
122+
}
123+
124+
static int
125+
ucx_close(void)
126+
{
127+
return OSHMEM_SUCCESS;
128+
}
129+

0 commit comments

Comments
 (0)