Skip to content

Commit d038b0b

Browse files
wfaderhold21wfaderhold21
committed
TL/CUDA: add support for mem map
TOOLS: add global and local memory type TL/CUDA: add push algorithm TL/CUDA: add alltoallv memh TEST: add tl/cuda push algorithm gtest TEST: fix tl/cuda push gtest teardown UAF and 64-bit count flags TEST: MPI add user memh reg. (local + global) TEST/MPI: add dispatch hook for test_mem_map.cc so mem_map tests execute TEST/MPI: fix global memh and mem-map test lifecycle TL/CUDA: unify alltoall/alltoallv push algorithms TEST: remove tl cuda push gtest TEST/MPI: fix mem_map test issues
1 parent 7569893 commit d038b0b

35 files changed

Lines changed: 1526 additions & 100 deletions

‎src/components/tl/cuda/Makefile.am‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,17 @@ allgatherv = \
1919
allgatherv/allgatherv_ring.c \
2020
allgatherv/allgatherv_linear.c
2121

22-
alltoall = \
23-
alltoall/alltoall.h \
24-
alltoall/alltoall.c \
25-
alltoall/alltoall_ce.c
26-
27-
alltoallv = \
28-
alltoallv/alltoallv.h \
29-
alltoallv/alltoallv.c \
30-
alltoallv/alltoallv_ce.c
22+
alltoall = \
23+
alltoall/alltoall.h \
24+
alltoall/alltoall.c \
25+
alltoall/alltoall_ce.c \
26+
alltoall/alltoall_push.c
27+
28+
alltoallv = \
29+
alltoallv/alltoallv.h \
30+
alltoallv/alltoallv.c \
31+
alltoallv/alltoallv_ce.c \
32+
alltoallv/alltoallv_push.c
3133

3234
bcast = \
3335
bcast/bcast.h \

‎src/components/tl/cuda/alltoall/alltoall.c‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,25 @@
11
/**
2-
* Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
* Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
*
44
* See file LICENSE for terms.
55
*/
66

77
#include "alltoall.h"
88
#include "components/mc/ucc_mc.h"
99

10+
ucc_base_coll_alg_info_t
11+
ucc_tl_cuda_alltoall_algs[UCC_TL_CUDA_ALLTOALL_ALG_LAST + 1] = {
12+
[UCC_TL_CUDA_ALLTOALL_ALG_CE] =
13+
{.id = UCC_TL_CUDA_ALLTOALL_ALG_CE,
14+
.name = "ce",
15+
.desc = "copy-engine alltoall algorithm"},
16+
[UCC_TL_CUDA_ALLTOALL_ALG_PUSH] =
17+
{.id = UCC_TL_CUDA_ALLTOALL_ALG_PUSH,
18+
.name = "push",
19+
.desc = "push-based alltoall using pre-registered dst handles"},
20+
[UCC_TL_CUDA_ALLTOALL_ALG_LAST] = {
21+
.id = 0, .name = NULL, .desc = NULL}};
22+
1023
ucc_status_t ucc_tl_cuda_alltoall_ce_init(ucc_tl_cuda_task_t *task);
1124

1225
ucc_status_t ucc_tl_cuda_alltoall_ce_start(ucc_coll_task_t *task);

‎src/components/tl/cuda/alltoall/alltoall.h‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/**
2-
* Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
* Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
*
44
* See file LICENSE for terms.
55
*/
@@ -9,9 +9,33 @@
99

1010
#include "tl_cuda.h"
1111
#include "tl_cuda_coll.h"
12+
#include "components/base/ucc_base_iface.h"
13+
14+
enum {
15+
UCC_TL_CUDA_ALLTOALL_ALG_CE = 0,
16+
UCC_TL_CUDA_ALLTOALL_ALG_PUSH = 1,
17+
UCC_TL_CUDA_ALLTOALL_ALG_LAST
18+
};
19+
20+
extern ucc_base_coll_alg_info_t
21+
ucc_tl_cuda_alltoall_algs[UCC_TL_CUDA_ALLTOALL_ALG_LAST + 1];
22+
23+
static inline int ucc_tl_cuda_alltoall_alg_from_str(const char *str)
24+
{
25+
int i;
26+
for (i = 0; i < UCC_TL_CUDA_ALLTOALL_ALG_LAST; i++) {
27+
if (0 == strcasecmp(str, ucc_tl_cuda_alltoall_algs[i].name))
28+
return i;
29+
}
30+
return -1;
31+
}
1232

1333
ucc_status_t ucc_tl_cuda_alltoall_init(ucc_base_coll_args_t *coll_args,
1434
ucc_base_team_t *tl_team,
1535
ucc_coll_task_t **task_p);
1636

37+
ucc_status_t ucc_tl_cuda_alltoall_push_init(ucc_base_coll_args_t *coll_args,
38+
ucc_base_team_t *tl_team,
39+
ucc_coll_task_t **task_p);
40+
1741
#endif

‎src/components/tl/cuda/alltoall/alltoall_push.c‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/**
2+
* Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3+
*
4+
* See file LICENSE for terms.
5+
*/
6+
7+
/*
8+
* Push-based alltoall with one CPU SHM barrier.
9+
*
10+
* Each rank pushes its chunks directly to peer destination buffers. Once all
11+
* pushes complete (detected via cudaEventQuery on the rank's own stream), the
12+
* rank enters a single CPU SHM barrier. When all ranks have arrived, every
13+
* rbuf is fully populated and UCC_OK can be returned.
14+
*
15+
* This replaces the pull-based CE algorithm's TWO barriers with ONE:
16+
* - No SETUP barrier: we push from our own sbuf so no cudaStreamWaitEvent on
17+
* peer events is needed, eliminating the stale-event problem entirely.
18+
* - ONE FINAL barrier: ensures all peers have completed their pushes into our
19+
* rbuf before we signal completion to the user.
20+
*
21+
* Counts and displacements are uniform, so the destination offset is
22+
* rank * chunk and no SETUP exchange is required (needs_setup == 0). The
23+
* shared stage machine lives in alltoallv/alltoallv_push.c; this file only
24+
* does the alltoall-specific argument parsing.
25+
*
26+
* Requirements:
27+
* - global_memh_dst: peer destination buffer handles pre-exchanged.
28+
* - No proxies (push writes directly to peer rbuf, not via a proxy rank).
29+
*/
30+
31+
#include "alltoall.h"
32+
#include "../alltoallv/alltoallv.h"
33+
34+
ucc_status_t ucc_tl_cuda_alltoall_push_init(ucc_base_coll_args_t *coll_args,
35+
ucc_base_team_t *tl_team,
36+
ucc_coll_task_t **task_p)
37+
{
38+
ucc_tl_cuda_team_t *team = ucc_derived_of(tl_team, ucc_tl_cuda_team_t);
39+
ucc_tl_cuda_task_t *task;
40+
ucc_coll_args_t *args;
41+
ucc_status_t status;
42+
43+
if (UCC_IS_INPLACE(coll_args->args)) {
44+
return UCC_ERR_NOT_SUPPORTED;
45+
}
46+
47+
status = ucc_tl_cuda_task_init(coll_args, team, &task);
48+
if (ucc_unlikely(status != UCC_OK)) {
49+
return status;
50+
}
51+
52+
args = &TASK_ARGS(task);
53+
54+
/* Uniform counts: chunk = count/nranks, dst offset = rank*chunk — both
55+
* computable locally, so the SETUP barrier is skipped. */
56+
task->alltoallv_push.sbuf = args->src.info.buffer;
57+
task->alltoallv_push.rbuf = args->dst.info.buffer;
58+
task->alltoallv_push.sdt = args->src.info.datatype;
59+
task->alltoallv_push.rdt = args->dst.info.datatype;
60+
task->alltoallv_push.scnts = NULL;
61+
task->alltoallv_push.rcnts = NULL;
62+
task->alltoallv_push.sdispl = NULL;
63+
task->alltoallv_push.rdispl = NULL;
64+
task->alltoallv_push.needs_setup = 0;
65+
task->alltoallv_push.global_memh_dst = args->dst_memh.global_memh;
66+
67+
status = ucc_tl_cuda_alltoallv_push_setup(task);
68+
if (ucc_unlikely(status != UCC_OK)) {
69+
goto err;
70+
}
71+
72+
*task_p = &task->super;
73+
return UCC_OK;
74+
err:
75+
ucc_tl_cuda_task_put(task);
76+
return status;
77+
}

‎src/components/tl/cuda/alltoallv/alltoallv.c‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/**
2-
* Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
* Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
* Copyright (c) Meta Platforms, Inc. and affiliates. 2022.
44
*
55
* See file LICENSE for terms.
@@ -8,6 +8,19 @@
88
#include "alltoallv.h"
99
#include "components/mc/ucc_mc.h"
1010

11+
ucc_base_coll_alg_info_t
12+
ucc_tl_cuda_alltoallv_algs[UCC_TL_CUDA_ALLTOALLV_ALG_LAST + 1] = {
13+
[UCC_TL_CUDA_ALLTOALLV_ALG_CE] =
14+
{.id = UCC_TL_CUDA_ALLTOALLV_ALG_CE,
15+
.name = "ce",
16+
.desc = "copy-engine alltoallv algorithm"},
17+
[UCC_TL_CUDA_ALLTOALLV_ALG_PUSH] =
18+
{.id = UCC_TL_CUDA_ALLTOALLV_ALG_PUSH,
19+
.name = "push",
20+
.desc = "push-based alltoallv using pre-registered dst handles"},
21+
[UCC_TL_CUDA_ALLTOALLV_ALG_LAST] = {
22+
.id = 0, .name = NULL, .desc = NULL}};
23+
1124
ucc_status_t ucc_tl_cuda_alltoallv_ce_init(ucc_tl_cuda_task_t *task);
1225

1326
ucc_status_t ucc_tl_cuda_alltoallv_ce_start(ucc_coll_task_t *task);

‎src/components/tl/cuda/alltoallv/alltoallv.h‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/**
2-
* Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
* Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
* Copyright (c) Meta Platforms, Inc. and affiliates. 2022.
44
*
55
* See file LICENSE for terms.
@@ -10,6 +10,26 @@
1010

1111
#include "../tl_cuda.h"
1212
#include "../tl_cuda_coll.h"
13+
#include "components/base/ucc_base_iface.h"
14+
15+
enum {
16+
UCC_TL_CUDA_ALLTOALLV_ALG_CE = 0,
17+
UCC_TL_CUDA_ALLTOALLV_ALG_PUSH = 1,
18+
UCC_TL_CUDA_ALLTOALLV_ALG_LAST
19+
};
20+
21+
extern ucc_base_coll_alg_info_t
22+
ucc_tl_cuda_alltoallv_algs[UCC_TL_CUDA_ALLTOALLV_ALG_LAST + 1];
23+
24+
static inline int ucc_tl_cuda_alltoallv_alg_from_str(const char *str)
25+
{
26+
int i;
27+
for (i = 0; i < UCC_TL_CUDA_ALLTOALLV_ALG_LAST; i++) {
28+
if (0 == strcasecmp(str, ucc_tl_cuda_alltoallv_algs[i].name))
29+
return i;
30+
}
31+
return -1;
32+
}
1333

1434
ucc_status_t ucc_tl_cuda_alltoallv_ce_finalize(ucc_coll_task_t *coll_task);
1535

@@ -24,6 +44,12 @@ ucc_status_t ucc_tl_cuda_alltoallv_init(ucc_base_coll_args_t *coll_args,
2444
ucc_base_team_t *tl_team,
2545
ucc_coll_task_t **task_p);
2646

47+
ucc_status_t ucc_tl_cuda_alltoallv_push_init(ucc_base_coll_args_t *coll_args,
48+
ucc_base_team_t *tl_team,
49+
ucc_coll_task_t **task_p);
50+
51+
ucc_status_t ucc_tl_cuda_alltoallv_push_setup(ucc_tl_cuda_task_t *task);
52+
2753
/**
2854
* @brief Post a copy operation using CUDA copy engine
2955
*

0 commit comments

Comments
 (0)