Skip to content

Commit 1b4d25e

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

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
@@ -396,7 +396,8 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
396396
size_t len;
397397
int my_pe = oshmem_my_proc_id();
398398
ucp_mem_map_params_t mem_map_params;
399-
int seg;
399+
int segno;
400+
map_segment_t *mem_seg;
400401
unsigned flags;
401402

402403
*count = 0;
@@ -405,26 +406,32 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
405406
return NULL;
406407
}
407408

408-
seg = memheap_find_segnum(addr);
409+
segno = memheap_find_segnum(addr);
410+
mem_seg = memheap_find_seg(segno);
409411

410-
ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[seg].key;
412+
ucx_mkey = &mca_spml_ucx.ucp_peers[my_pe].mkeys[segno].key;
411413
mkeys[0].spml_context = ucx_mkey;
412414

413-
flags = 0;
414-
if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
415-
flags = UCP_MEM_MAP_NONBLOCK;
416-
}
415+
/* if possible use mem handle already created by ucx allocator */
416+
if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
417+
flags = 0;
418+
if (mca_spml_ucx.heap_reg_nb && memheap_is_va_in_segment(addr, HEAP_SEG_INDEX)) {
419+
flags = UCP_MEM_MAP_NONBLOCK;
420+
}
417421

418-
mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
419-
UCP_MEM_MAP_PARAM_FIELD_LENGTH |
420-
UCP_MEM_MAP_PARAM_FIELD_FLAGS;
421-
mem_map_params.address = addr;
422-
mem_map_params.length = size;
423-
mem_map_params.flags = flags;
422+
mem_map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
423+
UCP_MEM_MAP_PARAM_FIELD_LENGTH |
424+
UCP_MEM_MAP_PARAM_FIELD_FLAGS;
425+
mem_map_params.address = addr;
426+
mem_map_params.length = size;
427+
mem_map_params.flags = flags;
424428

425-
err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
426-
if (UCS_OK != err) {
427-
goto error_out;
429+
err = ucp_mem_map(mca_spml_ucx.ucp_context, &mem_map_params, &ucx_mkey->mem_h);
430+
if (UCS_OK != err) {
431+
goto error_out;
432+
}
433+
} else {
434+
ucx_mkey->mem_h = (ucp_mem_h)mem_seg->context;
428435
}
429436

430437
err = ucp_rkey_pack(mca_spml_ucx.ucp_context, ucx_mkey->mem_h,
@@ -448,9 +455,9 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
448455
}
449456

450457
mkeys[0].len = len;
451-
mkeys[0].va_base = mem_map_params.address;
458+
mkeys[0].va_base = addr;
452459
*count = 1;
453-
mca_spml_ucx_cache_mkey(&mkeys[0], seg, my_pe);
460+
mca_spml_ucx_cache_mkey(&mkeys[0], segno, my_pe);
454461
return mkeys;
455462

456463
error_unmap:
@@ -464,6 +471,7 @@ sshmem_mkey_t *mca_spml_ucx_register(void* addr,
464471
int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
465472
{
466473
spml_ucx_mkey_t *ucx_mkey;
474+
map_segment_t *mem_seg;
467475

468476
MCA_SPML_CALL(fence());
469477
if (!mkeys)
@@ -472,8 +480,12 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
472480
if (!mkeys[0].spml_context)
473481
return OSHMEM_SUCCESS;
474482

475-
ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
476-
ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
483+
mem_seg = memheap_find_va(mkeys[0].va_base);
484+
485+
if (MAP_SEGMENT_ALLOC_UCX != mem_seg->type) {
486+
ucx_mkey = (spml_ucx_mkey_t *)mkeys[0].spml_context;
487+
ucp_mem_unmap(mca_spml_ucx.ucp_context, ucx_mkey->mem_h);
488+
}
477489

478490
if (0 < mkeys[0].len) {
479491
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)