1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
|
/*************************************************************************
* Copyright (c) 2016-2018, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#include "core.h"
#include "common_coll.h"
extern struct ncclTransport p2pTransport;
extern struct ncclTransport shmTransport;
extern struct ncclTransport netTransport;
struct ncclTransport ncclTransports[NTRANSPORTS] = {
p2pTransport,
shmTransport,
netTransport,
};
static void FifoPullArgs(struct transportProxyInfo* info, struct ncclProxyArgs *args) {
struct ncclProxyArgs *fifoArgs = info->argsFifo + (info->argsFifoHead % TRANSPORT_PROXY_FIFO_SIZE);
pthread_mutex_lock(&info->mutex);
while (fifoArgs->active == 0)
pthread_cond_wait(&info->cond, &info->mutex);
__sync_synchronize();
memcpy(args, fifoArgs, sizeof(struct ncclProxyArgs));
__sync_synchronize();
fifoArgs->active = 0;
pthread_cond_signal(&info->cond);
pthread_mutex_unlock(&info->mutex);
info->argsFifoHead++;
}
static struct ncclProxyArgs* FifoGetNextArgs(struct transportProxyInfo* info) {
if (info == NULL) return NULL;
struct ncclProxyArgs* fifoArgs = info->argsFifo + (info->argsFifoTail % TRANSPORT_PROXY_FIFO_SIZE);
pthread_mutex_lock(&info->mutex);
while (fifoArgs->active == 1)
pthread_cond_wait(&info->cond, &info->mutex);
pthread_mutex_unlock(&info->mutex);
info->argsFifoTail++;
return fifoArgs;
}
static void FifoPushArgs(struct transportProxyInfo* info) {
if (info == NULL) return;
struct ncclProxyArgs* fifoArgs = info->argsFifo + ((info->argsFifoTail-1) % TRANSPORT_PROXY_FIFO_SIZE);
if (fifoArgs->active == 0) return;
pthread_mutex_lock(&info->mutex);
pthread_cond_signal(&info->cond);
pthread_mutex_unlock(&info->mutex);
}
static void WaitProxyReady(struct transportProxyInfo* info) {
pthread_mutex_lock(&info->mutex);
while (info->proxyReady == 0)
pthread_cond_wait(&info->cond, &info->mutex);
pthread_mutex_unlock(&info->mutex);
}
static void SetProxyReady(struct transportProxyInfo* info) {
pthread_mutex_lock(&info->mutex);
info->proxyReady = 1;
pthread_cond_signal(&info->cond);
pthread_mutex_unlock(&info->mutex);
}
static void StopProxy(struct transportProxyInfo* info) {
struct ncclProxyArgs* fifoArgs = FifoGetNextArgs(info);
fifoArgs->active = -1;
FifoPushArgs(info);
}
#define RECV 0
#define SEND 1
static bool NeedProxy(int type, int pattern, struct ncclRing* ring, int nranks) {
enum proxyMode mode = proxyPatternMode(pattern);
if (mode == proxyRing) return true;
/* In chains, one rank does not need a proxy. Let's figure out which one it is */
int root = proxyPatternRoot(pattern);
// Which index in the reorganized rings should we compare root against */
const int myrank = 0, nextrank = 1, prevrank = nranks-1;
int index = mode == proxyFrom ?
/* no recv / no send if root = */
/* bcast */ (type == RECV ? myrank : nextrank ):
/* reduce */ (type == RECV ? prevrank : myrank );
int rank = ring->userRanks[index];
return (root != rank);
}
static void SaveProxy(struct ncclConnector* connector, struct ncclProxyArgs* args, int needProxy) {
struct transportProxyInfo* info = connector->proxyInfo;
if (info == NULL) return;
struct ncclProxyArgs* fifoArgs = FifoGetNextArgs(info);
args->needProxy = needProxy;
__sync_synchronize();
memcpy(fifoArgs, args, sizeof(struct ncclProxyArgs));
__sync_synchronize();
fifoArgs->active = 1;
}
ncclResult_t transportSaveProxies(int substeps, int subchunks, int nstepsPerRound, int nblocksPerRound, size_t nbytes, int pattern, struct ncclComm* comm) {
int llMode, nrings, nthreads;
ncclGetCollResource(comm, nbytes, &nrings, &nthreads, &llMode);
nbytes = llMode ? nbytes * 2 : nbytes;
substeps = llMode ? 1 : substeps;
subchunks = llMode ? NCCL_LL_CHUNKS : subchunks;
int buffSize = llMode ? NCCL_LL_BUFF_SIZE : comm->rings[0].buffSize;
int nrounds = (int)(DIVUP(nbytes, ((size_t)nrings * nblocksPerRound * (buffSize/subchunks)))); // Fixed 32-bit overflow
int nsteps = nstepsPerRound * nrounds * substeps;
TRACE(NET,"opCount %lx substeps %d subchunks %d nrounds %d nsteps %d comm %p", comm->opCount, subchunks, subchunks, nrounds, nsteps, comm);
TRACE(NET,"opCount %lx nbytes %zi nrings %d buffSize %d pattern %d comm %p", comm->opCount, nbytes, nrings, buffSize, pattern, comm);
for (int r=0; r<nrings; r++) {
struct ncclRing* ring = comm->rings+((comm->myParams->gridDim.x+r)%comm->nRings);
struct ncclProxyArgs args = { ring, substeps*subchunks, nsteps, comm->opCount, llMode, 0 };
SaveProxy(&ring->recv, &args, NeedProxy(RECV, pattern, ring, comm->nRanks));
SaveProxy(&ring->send, &args, NeedProxy(SEND, pattern, ring, comm->nRanks));
}
return ncclSuccess;
}
ncclResult_t transportStartProxies(ncclComm* comm) {
for (int r=0; r<comm->nRings; r++) {
FifoPushArgs(comm->rings[r].send.proxyInfo);
FifoPushArgs(comm->rings[r].recv.proxyInfo);
}
pthread_yield(); // Let other threads run
return ncclSuccess;
}
void* persistentThread(void *opaqueInfo) {
struct transportProxyInfo* info = (struct transportProxyInfo*)opaqueInfo;
// We need to initialize the context before launching any NCCL cuda kernel,
// otherwise we would create it during the first cudaMemcpyAsync inside the
// proxy function and that would cause a deadlock
cudaSetDevice(info->comm->cudaDev);
// Signal the main thread the context is created and it can proceed.
SetProxyReady(info);
while (1) {
struct ncclProxyArgs args;
FifoPullArgs(info, &args);
if (args.active == -1) {
// Main thread asked to stop
return NULL;
}
ncclResult_t res = info->func(&args);
if (res != ncclSuccess) {
WARN("%s:%d -> %d [Proxy thread error]", __FILE__, __LINE__, res);
}
}
}
ncclResult_t transportCreateProxy(int type, struct ncclRing* ring, struct ncclComm* comm) {
struct ncclConnector* connector = (type == RECV) ? &ring->recv : &ring->send;
threadFunc_t proxyfunc = (threadFunc_t) ((type == RECV) ? connector->transport->recv.proxy : connector->transport->send.proxy);
if (proxyfunc) {
TRACE(NET,"type %d ring %p proxyfunc %p comm %p", type, ring, proxyfunc, comm);
struct transportProxyInfo* info;
NCCLCHECK(ncclCalloc(&info, 1));
connector->proxyInfo = info;
info->comm = comm;
info->cond = PTHREAD_COND_INITIALIZER;
info->mutex = PTHREAD_MUTEX_INITIALIZER;
info->func = proxyfunc;
info->argsFifoHead = info->argsFifoTail = 0;
info->proxyReady = 0;
pthread_create(&connector->proxyInfo->thread, NULL, persistentThread, info);
// Wait for thread to initialize its CUDA context.
WaitProxyReady(info);
}
return ncclSuccess;
}
ncclResult_t transportDestroyProxy(struct ncclConnector* connector) {
if (connector->proxyInfo) {
StopProxy(connector->proxyInfo);
pthread_join(connector->proxyInfo->thread, NULL);
free(connector->proxyInfo);
connector->proxyInfo = NULL;
}
return ncclSuccess;
}
|
