Skip to content
This repository was archived by the owner on Sep 30, 2022. It is now read-only.

Commit 908cf4e

Browse files
alex-mikheevalex-mikheev
committed
OSHMEM/UCX: implements atomic support
ucx atomic component has a real code now. fixes bug in spml ucx add_procs removes redundant parameter checks from atomic components. (cherry picked from commit f627608)
1 parent 1529f93 commit 908cf4e

File tree

8 files changed

+138
-88
lines changed

8 files changed

+138
-88
lines changed

‎oshmem/mca/atomic/basic/atomic_basic_fadd.c

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -27,32 +27,25 @@ int mca_atomic_basic_fadd(void *target,
2727
struct oshmem_op_t *op)
2828
{
2929
int rc = OSHMEM_SUCCESS;
30+
long long temp_value = 0;
3031

31-
if (!target || !value) {
32-
rc = OSHMEM_ERROR;
33-
}
34-
35-
if (rc == OSHMEM_SUCCESS) {
36-
long long temp_value = 0;
37-
38-
atomic_basic_lock(pe);
32+
atomic_basic_lock(pe);
3933

40-
rc = MCA_SPML_CALL(get(target, nlong, (void*)&temp_value, pe));
34+
rc = MCA_SPML_CALL(get(target, nlong, (void*)&temp_value, pe));
4135

42-
if (prev)
43-
memcpy(prev, (void*) &temp_value, nlong);
36+
if (prev)
37+
memcpy(prev, (void*) &temp_value, nlong);
4438

45-
op->o_func.c_fn((void*) value,
46-
(void*) &temp_value,
47-
nlong / op->dt_size);
39+
op->o_func.c_fn((void*) value,
40+
(void*) &temp_value,
41+
nlong / op->dt_size);
4842

49-
if (rc == OSHMEM_SUCCESS) {
50-
rc = MCA_SPML_CALL(put(target, nlong, (void*)&temp_value, pe));
51-
shmem_quiet();
52-
}
53-
54-
atomic_basic_unlock(pe);
43+
if (rc == OSHMEM_SUCCESS) {
44+
rc = MCA_SPML_CALL(put(target, nlong, (void*)&temp_value, pe));
45+
shmem_quiet();
5546
}
5647

48+
atomic_basic_unlock(pe);
49+
5750
return rc;
5851
}

‎oshmem/mca/atomic/mxm/atomic_mxm_cswap.c

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -43,18 +43,6 @@ int mca_atomic_mxm_cswap(void *target,
4343
ptl_id = -1;
4444
mxm_err = MXM_OK;
4545

46-
if (!prev || !target || !value) {
47-
ATOMIC_ERROR("[#%d] Whether target, value or prev are not defined",
48-
my_pe);
49-
oshmem_shmem_abort(-1);
50-
return OSHMEM_ERR_BAD_PARAM;
51-
}
52-
if ((pe < 0) || (pe >= oshmem_num_procs())) {
53-
ATOMIC_ERROR("[#%d] PE=%d not valid", my_pe, pe);
54-
oshmem_shmem_abort(-1);
55-
return OSHMEM_ERR_BAD_PARAM;
56-
}
57-
5846
switch (nlong) {
5947
case 1:
6048
nlong_order = 0;

‎oshmem/mca/atomic/mxm/atomic_mxm_fadd.c

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -45,18 +45,6 @@ int mca_atomic_mxm_fadd(void *target,
4545
ptl_id = -1;
4646
mxm_err = MXM_OK;
4747

48-
if (!target || !value) {
49-
ATOMIC_ERROR("[#%d] target or value are not defined", my_pe);
50-
oshmem_shmem_abort(-1);
51-
return OSHMEM_ERR_BAD_PARAM;
52-
}
53-
54-
if ((pe < 0) || (pe >= oshmem_num_procs())) {
55-
ATOMIC_ERROR("[#%d] PE=%d not valid", my_pe, pe);
56-
oshmem_shmem_abort(-1);
57-
return OSHMEM_ERR_BAD_PARAM;
58-
}
59-
6048
switch (nlong) {
6149
case 1:
6250
nlong_order = 0;

‎oshmem/mca/atomic/ucx/atomic_ucx_component.c

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ static int ucx_open(void)
9797
*/
9898
if (strcmp(mca_spml_base_selected_component.spmlm_version.mca_component_name, "ucx")) {
9999
ATOMIC_VERBOSE(5,
100-
"Can not use atomic/ucx because spml ikrit component disabled");
100+
"Can not use atomic/ucx because spml ucx component disabled");
101101
return OSHMEM_ERR_NOT_AVAILABLE;
102102
}
103103
mca_spml_self = (mca_spml_ucx_t *)mca_spml.self;

‎oshmem/mca/atomic/ucx/atomic_ucx_cswap.c

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,8 @@
1313
#include <stdlib.h>
1414

1515
#include "oshmem/constants.h"
16-
#include "oshmem/mca/spml/spml.h"
1716
#include "oshmem/mca/atomic/atomic.h"
1817
#include "oshmem/mca/atomic/base/base.h"
19-
#include "oshmem/mca/memheap/memheap.h"
20-
#include "oshmem/mca/memheap/base/base.h"
2118
#include "oshmem/runtime/runtime.h"
2219

2320
#include "atomic_ucx.h"
@@ -29,6 +26,45 @@ int mca_atomic_ucx_cswap(void *target,
2926
size_t nlong,
3027
int pe)
3128
{
32-
return OSHMEM_SUCCESS;
29+
ucs_status_t status;
30+
spml_ucx_mkey_t *ucx_mkey;
31+
uint64_t rva;
32+
33+
ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
34+
if (NULL == cond) {
35+
switch (nlong) {
36+
case 4:
37+
status = ucp_atomic_swap32(mca_spml_self->ucp_peers[pe].ucp_conn,
38+
*(uint32_t *)value, rva, ucx_mkey->rkey, prev);
39+
break;
40+
case 8:
41+
status = ucp_atomic_swap64(mca_spml_self->ucp_peers[pe].ucp_conn,
42+
*(uint64_t *)value, rva, ucx_mkey->rkey, prev);
43+
break;
44+
default:
45+
goto err_size;
46+
}
47+
}
48+
else {
49+
switch (nlong) {
50+
case 4:
51+
status = ucp_atomic_cswap32(mca_spml_self->ucp_peers[pe].ucp_conn,
52+
*(uint32_t *)cond, *(uint32_t *)value, rva, ucx_mkey->rkey, prev);
53+
break;
54+
case 8:
55+
status = ucp_atomic_cswap64(mca_spml_self->ucp_peers[pe].ucp_conn,
56+
*(uint64_t *)cond, *(uint64_t *)value, rva, ucx_mkey->rkey, prev);
57+
break;
58+
default:
59+
goto err_size;
60+
}
61+
}
62+
63+
return ucx_status_to_oshmem(status);
64+
65+
err_size:
66+
ATOMIC_ERROR("[#%d] Type size must be 1/2/4 or 8 bytes.", my_pe);
67+
return OSHMEM_ERROR;
3368
}
3469

70+

‎oshmem/mca/atomic/ucx/atomic_ucx_fadd.c

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,8 @@
1313
#include <stdlib.h>
1414

1515
#include "oshmem/constants.h"
16-
#include "oshmem/op/op.h"
17-
#include "oshmem/mca/spml/spml.h"
1816
#include "oshmem/mca/atomic/atomic.h"
1917
#include "oshmem/mca/atomic/base/base.h"
20-
#include "oshmem/mca/memheap/memheap.h"
21-
#include "oshmem/mca/memheap/base/base.h"
22-
#include "oshmem/runtime/runtime.h"
2318

2419
#include "atomic_ucx.h"
2520

@@ -30,6 +25,44 @@ int mca_atomic_ucx_fadd(void *target,
3025
int pe,
3126
struct oshmem_op_t *op)
3227
{
33-
/* TODO: actual code */
34-
return OSHMEM_SUCCESS;
28+
ucs_status_t status;
29+
spml_ucx_mkey_t *ucx_mkey;
30+
uint64_t rva;
31+
32+
ucx_mkey = mca_spml_ucx_get_mkey(pe, target, (void *)&rva);
33+
34+
if (NULL == prev) {
35+
switch (nlong) {
36+
case 4:
37+
status = ucp_atomic_add32(mca_spml_self->ucp_peers[pe].ucp_conn,
38+
*(uint32_t *)value, rva, ucx_mkey->rkey);
39+
break;
40+
case 8:
41+
status = ucp_atomic_add64(mca_spml_self->ucp_peers[pe].ucp_conn,
42+
*(uint64_t *)value, rva, ucx_mkey->rkey);
43+
break;
44+
default:
45+
goto err_size;
46+
}
47+
}
48+
else {
49+
switch (nlong) {
50+
case 4:
51+
status = ucp_atomic_fadd32(mca_spml_self->ucp_peers[pe].ucp_conn,
52+
*(uint32_t *)value, rva, ucx_mkey->rkey, prev);
53+
break;
54+
case 8:
55+
status = ucp_atomic_fadd64(mca_spml_self->ucp_peers[pe].ucp_conn,
56+
*(uint64_t *)value, rva, ucx_mkey->rkey, prev);
57+
break;
58+
default:
59+
goto err_size;
60+
}
61+
}
62+
63+
return ucx_status_to_oshmem(status);
64+
65+
err_size:
66+
ATOMIC_ERROR("[#%d] Type size must be 1/2/4 or 8 bytes.", my_pe);
67+
return OSHMEM_ERROR;
3568
}

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

Lines changed: 16 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ mca_spml_ucx_t mca_spml_ucx = {
5353
mca_spml_ucx_deregister,
5454
mca_spml_base_oob_get_mkeys,
5555
mca_spml_ucx_put,
56-
NULL, //mca_spml_ucx_put_nb,
56+
NULL, /* todo: mca_spml_ucx_put_nb, */
5757
mca_spml_ucx_get,
5858
mca_spml_ucx_recv,
5959
mca_spml_ucx_send,
@@ -174,7 +174,9 @@ static void dump_address(int pe, char *addr, size_t len)
174174
#endif
175175
}
176176

177-
int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
177+
static char spml_ucx_transport_ids[1] = { 0 };
178+
179+
int mca_spml_ucx_add_procs(ompi_proc_t **procs, size_t nprocs)
178180
{
179181
size_t i, n;
180182
int rc = OSHMEM_ERROR;
@@ -208,7 +210,6 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
208210
/* Get the EP connection requests for all the processes from modex */
209211
for (n = 0; n < nprocs; ++n) {
210212
i = (my_rank + n) % nprocs;
211-
//if (i == my_rank) continue;
212213
dump_address(i, (char *)(wk_raddrs + wk_roffs[i]), wk_rsizes[i]);
213214
err = ucp_ep_create(mca_spml_ucx.ucp_worker,
214215
(ucp_address_t *)(wk_raddrs + wk_roffs[i]),
@@ -217,6 +218,8 @@ int mca_spml_ucx_add_procs(ompi_proc_t** procs, size_t nprocs)
217218
SPML_ERROR("ucp_ep_create failed!!!\n");
218219
goto error2;
219220
}
221+
OSHMEM_PROC_DATA(procs[i])->num_transports = 1;
222+
OSHMEM_PROC_DATA(procs[i])->transport_ids = spml_ucx_transport_ids;
220223
}
221224

222225
ucp_worker_release_address(mca_spml_ucx.ucp_worker, wk_local_addr);
@@ -377,46 +380,27 @@ int mca_spml_ucx_deregister(sshmem_mkey_t *mkeys)
377380
int mca_spml_ucx_get(void *src_addr, size_t size, void *dst_addr, int src)
378381
{
379382
void *rva;
380-
sshmem_mkey_t *r_mkey;
381-
ucs_status_t err;
383+
ucs_status_t status;
382384
spml_ucx_mkey_t *ucx_mkey;
383385

384-
r_mkey = mca_memheap_base_get_cached_mkey(src, src_addr, 0, &rva);
385-
if (OPAL_UNLIKELY(!r_mkey)) {
386-
SPML_ERROR("pe=%d: %p is not address of shared variable",
387-
src, src_addr);
388-
oshmem_shmem_abort(-1);
389-
return OSHMEM_ERROR;
390-
}
391-
392-
ucx_mkey = (spml_ucx_mkey_t *)(r_mkey->spml_context);
393-
err = ucp_get(mca_spml_ucx.ucp_peers[src].ucp_conn, dst_addr, size,
394-
(uint64_t)rva, ucx_mkey->rkey);
386+
ucx_mkey = mca_spml_ucx_get_mkey(src, src_addr, &rva);
387+
status = ucp_get(mca_spml_ucx.ucp_peers[src].ucp_conn, dst_addr, size,
388+
(uint64_t)rva, ucx_mkey->rkey);
395389

396-
return OPAL_LIKELY(UCS_OK == err) ? OSHMEM_SUCCESS : OSHMEM_ERROR;
390+
return ucx_status_to_oshmem(status);
397391
}
398392

399393
int mca_spml_ucx_put(void* dst_addr, size_t size, void* src_addr, int dst)
400394
{
401395
void *rva;
402-
sshmem_mkey_t *r_mkey;
403-
ucs_status_t err;
396+
ucs_status_t status;
404397
spml_ucx_mkey_t *ucx_mkey;
405398

406-
r_mkey = mca_memheap_base_get_cached_mkey(dst, dst_addr, 0, &rva);
407-
if (OPAL_UNLIKELY(!r_mkey)) {
408-
SPML_ERROR("pe=%d: %p is not address of shared variable",
409-
dst, dst_addr);
410-
oshmem_shmem_abort(-1);
411-
return OSHMEM_ERROR;
412-
}
413-
414-
ucx_mkey = (spml_ucx_mkey_t *)(r_mkey->spml_context);
415-
416-
err = ucp_put(mca_spml_ucx.ucp_peers[dst].ucp_conn, src_addr, size,
417-
(uint64_t)rva, ucx_mkey->rkey);
399+
ucx_mkey = mca_spml_ucx_get_mkey(dst, dst_addr, &rva);
400+
status = ucp_put(mca_spml_ucx.ucp_peers[dst].ucp_conn, src_addr, size,
401+
(uint64_t)rva, ucx_mkey->rkey);
418402

419-
return OPAL_LIKELY(UCS_OK == err) ? OSHMEM_SUCCESS : OSHMEM_ERROR;
403+
return ucx_status_to_oshmem(status);
420404
}
421405

422406
int mca_spml_ucx_fence(void)

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,17 @@
1818

1919
#include "oshmem_config.h"
2020
#include "oshmem/request/request.h"
21+
#include "oshmem/mca/spml/base/base.h"
2122
#include "oshmem/mca/spml/spml.h"
2223
#include "oshmem/util/oshmem_util.h"
2324
#include "oshmem/mca/spml/base/spml_base_putreq.h"
2425
#include "oshmem/proc/proc.h"
2526
#include "oshmem/mca/spml/base/spml_base_request.h"
2627
#include "oshmem/mca/spml/base/spml_base_getreq.h"
28+
#include "oshmem/runtime/runtime.h"
29+
30+
#include "oshmem/mca/memheap/memheap.h"
31+
#include "oshmem/mca/memheap/base/base.h"
2732

2833
#include "orte/runtime/orte_globals.h"
2934

@@ -100,6 +105,29 @@ extern int mca_spml_ucx_fence(void);
100105
extern int mca_spml_ucx_quiet(void);
101106
extern int spml_ucx_progress(void);
102107

108+
109+
110+
static inline spml_ucx_mkey_t *
111+
mca_spml_ucx_get_mkey(int pe, void *va, void **rva)
112+
{
113+
sshmem_mkey_t *r_mkey;
114+
115+
r_mkey = mca_memheap_base_get_cached_mkey(pe, va, 0, rva);
116+
if (OPAL_UNLIKELY(!r_mkey)) {
117+
SPML_ERROR("pe=%d: %p is not address of symmetric variable",
118+
pe, va);
119+
oshmem_shmem_abort(-1);
120+
return NULL;
121+
}
122+
return (spml_ucx_mkey_t *)(r_mkey->spml_context);
123+
}
124+
125+
static inline int ucx_status_to_oshmem(ucs_status_t status)
126+
{
127+
return OPAL_LIKELY(UCS_OK == status) ? OSHMEM_SUCCESS : OSHMEM_ERROR;
128+
}
129+
130+
103131
END_C_DECLS
104132

105133
#endif

0 commit comments

Comments
 (0)