Skip to content

Commit 579323b

Browse files
leoda1leoda1
authored
Add proper handling for self-communication in intra-node operations (flagos-ai#294)
1 parent ef484ca commit 579323b

2 files changed

Lines changed: 281 additions & 158 deletions

File tree

‎flagcx/core/group.cc‎

Lines changed: 211 additions & 144 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include <pthread.h>
1818
#include <queue>
1919
#include <stdio.h>
20+
#include <vector>
2021

2122
__thread int flagcxGroupDepth = 0;
2223
__thread bool flagcxGroupJobAbortFlag = false;
@@ -159,158 +160,224 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
159160
flagcxTasks *tasks = &comm->tasks;
160161
for (int i = 0; i < tasks->p2pOrderSteps; i++) {
161162
int peer = tasks->p2pOrder[i];
162-
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].sendQueue)) {
163-
flagcxTaskP2p *p2p =
164-
flagcxIntruQueueDequeue(&tasks->peers[peer].sendQueue);
165-
flagcxProxyOp *op;
166-
FLAGCXCHECK(flagcxCalloc(&op, 1));
167-
op->pattern = flagcxPatternSend;
168-
op->nbytes = p2p->bytes;
169-
op->recvbuff = (uint8_t *)p2p->buff;
170-
op->channelId = 0;
171-
op->root = peer;
172-
op->connection = comm->channels[op->channelId]
173-
.peers[peer]
174-
->send[0]
175-
.proxyConn.connection;
176-
op->args.chunkSize = CHUNKSIZE;
177-
op->args.chunkSteps = (p2p->bytes + CHUNKSIZE - 1) / (CHUNKSIZE);
178-
op->args.sendStepMask = MAXSTEPS - 1;
179-
op->args.deviceFuncRelaxedOrdering = deviceFuncRelaxedOrdering;
180-
op->stream = p2p->stream;
181-
if (op->connection->transport == TRANSPORT_P2P) {
182-
setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 0,
183-
&op->args.p2pOpHash, &op->args.p2pSlotIdx);
184-
setP2pSlotInfo(peer, comm->rank, p2p->bytes, p2p->dtype, 1,
185-
&op->args.p2pPeerOpHash, &op->args.p2pPeerSlotIdx);
186-
TRACE_CALL("Sender: [rank(%d), peerRank(%d)] -> [slotIdx(%ld), "
187-
"opHash(%d)]",
188-
comm->rank, peer, op->args.p2pSlotIdx,
189-
op->args.p2pOpHash);
190-
TRACE_CALL("Sender: [peerRank(%d), rank(%d)] -> [peerSlotIdx(%ld), "
191-
"peerOpHash(%d)]",
192-
peer, comm->rank, op->args.p2pPeerSlotIdx,
193-
op->args.p2pPeerOpHash);
194-
}
163+
if (peer != comm->rank) {
164+
// Handle cross-process send/recv: use proxy
165+
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].sendQueue)) {
166+
flagcxTaskP2p *p2p =
167+
flagcxIntruQueueDequeue(&tasks->peers[peer].sendQueue);
168+
flagcxProxyOp *op;
169+
FLAGCXCHECK(flagcxCalloc(&op, 1));
170+
op->pattern = flagcxPatternSend;
171+
op->nbytes = p2p->bytes;
172+
op->recvbuff = (uint8_t *)p2p->buff;
173+
op->channelId = 0;
174+
op->root = peer;
175+
op->connection = comm->channels[op->channelId]
176+
.peers[peer]
177+
->send[0]
178+
.proxyConn.connection;
179+
op->args.chunkSize = CHUNKSIZE;
180+
op->args.chunkSteps = (p2p->bytes + CHUNKSIZE - 1) / (CHUNKSIZE);
181+
op->args.sendStepMask = MAXSTEPS - 1;
182+
op->args.deviceFuncRelaxedOrdering = deviceFuncRelaxedOrdering;
183+
op->stream = p2p->stream;
184+
if (op->connection->transport == TRANSPORT_P2P) {
185+
setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 0,
186+
&op->args.p2pOpHash, &op->args.p2pSlotIdx);
187+
setP2pSlotInfo(peer, comm->rank, p2p->bytes, p2p->dtype, 1,
188+
&op->args.p2pPeerOpHash, &op->args.p2pPeerSlotIdx);
189+
TRACE_CALL("Sender: [rank(%d), peerRank(%d)] -> [slotIdx(%ld), "
190+
"opHash(%d)]",
191+
comm->rank, peer, op->args.p2pSlotIdx,
192+
op->args.p2pOpHash);
193+
TRACE_CALL(
194+
"Sender: [peerRank(%d), rank(%d)] -> [peerSlotIdx(%ld), "
195+
"peerOpHash(%d)]",
196+
peer, comm->rank, op->args.p2pPeerSlotIdx,
197+
op->args.p2pPeerOpHash);
198+
}
195199

196-
// launch proxyRegister op if not yet registered
197-
flagcxConnector *peerConns[] = {
198-
comm->channels[op->channelId].peers[peer]->send};
199-
FLAGCXCHECK(flagcxNetRegisterBuffer(
200-
comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
201-
&op->args.regHandle));
202-
// we don't use semaphore tracking for device func for the moment
203-
if (deviceAsyncLoad && deviceAsyncStore) {
204-
FLAGCXCHECK(deviceAdaptor->eventCreate(&op->event,
205-
flagcxEventDisableTiming));
206-
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
207-
std::vector<void *> argList;
208-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
209-
(void **)&op->args.dlArgs, sizeof(bool), flagcxMemDevice,
210-
op->stream));
211-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
212-
(void **)&op->args.dEventReady, sizeof(bool), flagcxMemDevice,
213-
op->stream));
214-
FLAGCXCHECK(deviceAdaptor->launchDeviceFunc(
215-
op->stream, deviceAsyncStore, op->args.dEventReady));
216-
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
217-
(void *)&op->args.hEventReady, (void *)op->args.dEventReady,
218-
sizeof(bool), flagcxMemcpyDeviceToHost, op->stream, NULL));
219-
argList = {(void *)&op->args.eventRecorded,
220-
(void *)&op->args.hlArgs, (void *)op->args.dlArgs};
221-
funcQueue.push({op->stream, op->event, argList.data()});
222-
argsQueue.push(std::move(argList));
223-
} else {
224-
op->args.semaphore = semaphore;
225-
op->event = semaphore->getEvent();
226-
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
227-
semaphore->counter++;
228-
if (semaphore->counter == 1) {
229-
launchStream = op->stream;
200+
// launch proxyRegister op if not yet registered
201+
flagcxConnector *peerConns[] = {
202+
comm->channels[op->channelId].peers[peer]->send};
203+
FLAGCXCHECK(flagcxNetRegisterBuffer(
204+
comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
205+
&op->args.regHandle));
206+
// we don't use semaphore tracking for device func for the moment
207+
if (deviceAsyncLoad && deviceAsyncStore) {
208+
FLAGCXCHECK(deviceAdaptor->eventCreate(&op->event,
209+
flagcxEventDisableTiming));
210+
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
211+
std::vector<void *> argList;
212+
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
213+
(void **)&op->args.dlArgs, sizeof(bool), flagcxMemDevice,
214+
op->stream));
215+
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
216+
(void **)&op->args.dEventReady, sizeof(bool), flagcxMemDevice,
217+
op->stream));
218+
FLAGCXCHECK(deviceAdaptor->launchDeviceFunc(
219+
op->stream, deviceAsyncStore, op->args.dEventReady));
220+
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
221+
(void *)&op->args.hEventReady, (void *)op->args.dEventReady,
222+
sizeof(bool), flagcxMemcpyDeviceToHost, op->stream, NULL));
223+
argList = {(void *)&op->args.eventRecorded,
224+
(void *)&op->args.hlArgs, (void *)op->args.dlArgs};
225+
funcQueue.push({op->stream, op->event, argList.data()});
226+
argsQueue.push(std::move(argList));
227+
} else {
228+
op->args.semaphore = semaphore;
229+
op->event = semaphore->getEvent();
230+
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
231+
if (launchStream == nullptr) {
232+
launchStream = op->stream;
233+
semaphore->counter++;
234+
} else {
235+
FLAGCXCHECK(
236+
deviceAdaptor->streamWaitEvent(launchStream, op->event));
237+
}
230238
}
231-
FLAGCXCHECK(
232-
deviceAdaptor->streamWaitEvent(launchStream, op->event));
233-
}
234-
FLAGCXCHECK(flagcxProxySaveOp(comm, op));
235-
free(p2p);
236-
}
237-
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].recvQueue)) {
238-
flagcxTaskP2p *p2p =
239-
flagcxIntruQueueDequeue(&tasks->peers[peer].recvQueue);
240-
flagcxProxyOp *op;
241-
FLAGCXCHECK(flagcxCalloc(&op, 1));
242-
op->pattern = flagcxPatternRecv;
243-
op->nbytes = p2p->bytes;
244-
op->recvbuff = (uint8_t *)p2p->buff;
245-
op->channelId = 0;
246-
op->root = peer;
247-
op->connection = comm->channels[op->channelId]
248-
.peers[peer]
249-
->recv[0]
250-
.proxyConn.connection;
251-
op->args.chunkSize = CHUNKSIZE;
252-
op->args.chunkSteps = (p2p->bytes + CHUNKSIZE - 1) / (CHUNKSIZE);
253-
op->args.sendStepMask = MAXSTEPS - 1;
254-
op->args.deviceFuncRelaxedOrdering = deviceFuncRelaxedOrdering;
255-
op->stream = p2p->stream;
256-
if (op->connection->transport == TRANSPORT_P2P) {
257-
setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 1,
258-
&op->args.p2pOpHash, &op->args.p2pSlotIdx);
259-
setP2pSlotInfo(peer, comm->rank, p2p->bytes, p2p->dtype, 0,
260-
&op->args.p2pPeerOpHash, &op->args.p2pPeerSlotIdx);
261-
TRACE_CALL("Receiver: [rank(%d), peerRank(%d)] -> [slotIdx(%ld), "
262-
"opHash(%d)]",
263-
comm->rank, peer, op->args.p2pSlotIdx,
264-
op->args.p2pOpHash);
265-
TRACE_CALL("Receiver: [peerRank(%d), rank(%d)] -> "
266-
"[peerSlotIdx(%ld), peerOpHash(%d)]",
267-
peer, comm->rank, op->args.p2pPeerSlotIdx,
268-
op->args.p2pPeerOpHash);
239+
FLAGCXCHECK(flagcxProxySaveOp(comm, op));
240+
free(p2p);
269241
}
270242

271-
// launch proxyRegister op if not yet registered
272-
flagcxConnector *peerConns[] = {
273-
comm->channels[op->channelId].peers[peer]->recv};
274-
FLAGCXCHECK(flagcxNetRegisterBuffer(
275-
comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
276-
&op->args.regHandle));
277-
// we don't use semaphore tracking for device func for the moment
278-
if (deviceAsyncLoad && deviceAsyncStore) {
279-
std::vector<void *> argList;
280-
FLAGCXCHECK(deviceAdaptor->eventCreate(&op->event,
281-
flagcxEventDisableTiming));
282-
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
283-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
284-
(void **)&op->args.dlArgs, sizeof(bool), flagcxMemDevice,
285-
op->stream));
286-
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
287-
(void **)&op->args.dEventReady, sizeof(bool), flagcxMemDevice,
288-
op->stream));
289-
FLAGCXCHECK(deviceAdaptor->launchDeviceFunc(
290-
op->stream, deviceAsyncStore, op->args.dEventReady));
291-
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
292-
(void *)&op->args.hEventReady, (void *)op->args.dEventReady,
293-
sizeof(bool), flagcxMemcpyDeviceToHost, op->stream, NULL));
294-
argList = {(void *)&op->args.eventRecorded,
295-
(void *)&op->args.hlArgs, (void *)op->args.dlArgs};
296-
funcQueue.push({op->stream, op->event, argList.data()});
297-
argsQueue.push(std::move(argList));
298-
} else {
299-
op->args.semaphore = semaphore;
300-
op->event = semaphore->getEvent();
301-
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
302-
semaphore->counter++;
303-
if (semaphore->counter == 1) {
304-
launchStream = op->stream;
243+
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].recvQueue)) {
244+
flagcxTaskP2p *p2p =
245+
flagcxIntruQueueDequeue(&tasks->peers[peer].recvQueue);
246+
flagcxProxyOp *op;
247+
FLAGCXCHECK(flagcxCalloc(&op, 1));
248+
op->pattern = flagcxPatternRecv;
249+
op->nbytes = p2p->bytes;
250+
op->recvbuff = (uint8_t *)p2p->buff;
251+
op->channelId = 0;
252+
op->root = peer;
253+
op->connection = comm->channels[op->channelId]
254+
.peers[peer]
255+
->recv[0]
256+
.proxyConn.connection;
257+
op->args.chunkSize = CHUNKSIZE;
258+
op->args.chunkSteps = (p2p->bytes + CHUNKSIZE - 1) / (CHUNKSIZE);
259+
op->args.sendStepMask = MAXSTEPS - 1;
260+
op->args.deviceFuncRelaxedOrdering = deviceFuncRelaxedOrdering;
261+
op->stream = p2p->stream;
262+
if (op->connection->transport == TRANSPORT_P2P) {
263+
setP2pSlotInfo(comm->rank, peer, p2p->bytes, p2p->dtype, 1,
264+
&op->args.p2pOpHash, &op->args.p2pSlotIdx);
265+
setP2pSlotInfo(peer, comm->rank, p2p->bytes, p2p->dtype, 0,
266+
&op->args.p2pPeerOpHash, &op->args.p2pPeerSlotIdx);
267+
TRACE_CALL("Receiver: [rank(%d), peerRank(%d)] -> [slotIdx(%ld), "
268+
"opHash(%d)]",
269+
comm->rank, peer, op->args.p2pSlotIdx,
270+
op->args.p2pOpHash);
271+
TRACE_CALL("Receiver: [peerRank(%d), rank(%d)] -> "
272+
"[peerSlotIdx(%ld), peerOpHash(%d)]",
273+
peer, comm->rank, op->args.p2pPeerSlotIdx,
274+
op->args.p2pPeerOpHash);
305275
}
306-
FLAGCXCHECK(
307-
deviceAdaptor->streamWaitEvent(launchStream, op->event));
276+
277+
// launch proxyRegister op if not yet registered
278+
flagcxConnector *peerConns[] = {
279+
comm->channels[op->channelId].peers[peer]->recv};
280+
FLAGCXCHECK(flagcxNetRegisterBuffer(
281+
comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
282+
&op->args.regHandle));
283+
// we don't use semaphore tracking for device func for the moment
284+
if (deviceAsyncLoad && deviceAsyncStore) {
285+
std::vector<void *> argList;
286+
FLAGCXCHECK(deviceAdaptor->eventCreate(&op->event,
287+
flagcxEventDisableTiming));
288+
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
289+
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
290+
(void **)&op->args.dlArgs, sizeof(bool), flagcxMemDevice,
291+
op->stream));
292+
FLAGCXCHECK(deviceAdaptor->deviceMalloc(
293+
(void **)&op->args.dEventReady, sizeof(bool), flagcxMemDevice,
294+
op->stream));
295+
FLAGCXCHECK(deviceAdaptor->launchDeviceFunc(
296+
op->stream, deviceAsyncStore, op->args.dEventReady));
297+
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
298+
(void *)&op->args.hEventReady, (void *)op->args.dEventReady,
299+
sizeof(bool), flagcxMemcpyDeviceToHost, op->stream, NULL));
300+
argList = {(void *)&op->args.eventRecorded,
301+
(void *)&op->args.hlArgs, (void *)op->args.dlArgs};
302+
funcQueue.push({op->stream, op->event, argList.data()});
303+
argsQueue.push(std::move(argList));
304+
} else {
305+
op->args.semaphore = semaphore;
306+
op->event = semaphore->getEvent();
307+
FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
308+
if (launchStream == nullptr) {
309+
launchStream = op->stream;
310+
semaphore->counter++;
311+
} else {
312+
FLAGCXCHECK(
313+
deviceAdaptor->streamWaitEvent(launchStream, op->event));
314+
}
315+
}
316+
FLAGCXCHECK(flagcxProxySaveOp(comm, op));
317+
free(p2p);
308318
}
309-
FLAGCXCHECK(flagcxProxySaveOp(comm, op));
310-
free(p2p);
319+
} else {
320+
std::vector<flagcxTaskP2p *> sendTasks;
321+
std::vector<flagcxTaskP2p *> recvTasks;
322+
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].sendQueue))
323+
sendTasks.push_back(
324+
flagcxIntruQueueDequeue(&tasks->peers[peer].sendQueue));
325+
while (!flagcxIntruQueueEmpty(&tasks->peers[peer].recvQueue))
326+
recvTasks.push_back(
327+
flagcxIntruQueueDequeue(&tasks->peers[peer].recvQueue));
328+
329+
for (size_t i = 0; i < sendTasks.size();) {
330+
bool matched = false;
331+
for (size_t j = 0; j < recvTasks.size(); j++) {
332+
if (sendTasks[i]->bytes == recvTasks[j]->bytes &&
333+
sendTasks[i]->dtype == recvTasks[j]->dtype) {
334+
if (sendTasks[i]->buff != recvTasks[j]->buff) {
335+
flagcxEvent_t selfEvent = semaphore->getEvent();
336+
FLAGCXCHECK(deviceAdaptor->eventRecord(selfEvent,
337+
sendTasks[i]->stream));
338+
if (launchStream == nullptr) {
339+
launchStream = sendTasks[i]->stream;
340+
} else {
341+
FLAGCXCHECK(deviceAdaptor->streamWaitEvent(launchStream,
342+
selfEvent));
343+
}
344+
345+
FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
346+
recvTasks[j]->buff, sendTasks[i]->buff,
347+
sendTasks[i]->bytes, flagcxMemcpyDeviceToDevice,
348+
sendTasks[i]->stream, NULL));
349+
}
350+
free(sendTasks[i]);
351+
free(recvTasks[j]);
352+
353+
sendTasks.erase(sendTasks.begin() + i);
354+
recvTasks.erase(recvTasks.begin() + j);
355+
matched = true;
356+
break;
357+
}
358+
}
359+
if (!matched)
360+
i++;
361+
}
362+
for (auto *task : sendTasks)
363+
flagcxIntruQueueEnqueue(&tasks->peers[peer].sendQueue, task);
364+
for (auto *task : recvTasks)
365+
flagcxIntruQueueEnqueue(&tasks->peers[peer].recvQueue, task);
366+
}
367+
}
368+
// Clean up p2pOrder: remove peers with empty queues, keep peers with
369+
// pending operations
370+
int newOrderSteps = 0;
371+
for (int i = 0; i < tasks->p2pOrderSteps; i++) {
372+
int peer = tasks->p2pOrder[i];
373+
// Keep peer in order if it still has pending send or recv operations
374+
if (!flagcxIntruQueueEmpty(&tasks->peers[peer].sendQueue) ||
375+
!flagcxIntruQueueEmpty(&tasks->peers[peer].recvQueue)) {
376+
tasks->p2pOrder[newOrderSteps++] = peer;
311377
}
312378
}
313-
comm->tasks.p2pOrderSteps = 0;
379+
tasks->p2pOrderSteps = newOrderSteps;
380+
314381
comm = comm->groupNext;
315382
} while (comm != nullptr);
316383
}

0 commit comments

Comments
 (0)