Skip to content

Commit 185e99d

Browse files
MC952-archMC952-arch
authored
[CRL] Use multiple semaphores tracking for p2p ops in a group (flagos-ai#318)
1 parent ea34ccf commit 185e99d

12 files changed

Lines changed: 142 additions & 100 deletions

File tree

flagcx/core/collectives.cc renamed to flagcx/core/flagcx_hetero.cc

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
1-
#include "collectives.h"
1+
#include "flagcx_hetero.h"
22
#include "group.h"
33
#include "net.h"
44
#include "transport.h"
55
#include "type.h"
66

77
flagcxResult_t flagcxHeteroSend(const void *sendbuff, size_t count,
88
flagcxDataType_t datatype, int peer,
9-
flagcxHeteroComm_t comm,
10-
flagcxStream_t stream) {
9+
flagcxHeteroComm_t comm, flagcxStream_t stream,
10+
int groupIdx) {
11+
assert(groupIdx >= 0 && groupIdx < FLAGCX_MAX_SUBGROUPS);
1112
flagcxHeteroGroupStart();
1213
int channelId = 0;
1314
if (comm->channels[channelId].peers[peer]->send[0].connected == 0) {
@@ -22,6 +23,7 @@ flagcxResult_t flagcxHeteroSend(const void *sendbuff, size_t count,
2223
p2p->chunk = 0;
2324
p2p->dtype = datatype;
2425
p2p->stream = stream;
26+
p2p->groupIdx = groupIdx;
2527
if (flagcxIntruQueueEmpty(&tasks->peers[peer].sendQueue))
2628
tasks->p2pOrder[tasks->p2pOrderSteps++] = peer;
2729
flagcxIntruQueueEnqueue(&tasks->peers[peer].sendQueue, p2p);
@@ -33,11 +35,11 @@ flagcxResult_t flagcxHeteroSend(const void *sendbuff, size_t count,
3335

3436
flagcxResult_t flagcxHeteroRecv(void *recvbuff, size_t count,
3537
flagcxDataType_t datatype, int peer,
36-
flagcxHeteroComm_t comm,
37-
flagcxStream_t stream) {
38+
flagcxHeteroComm_t comm, flagcxStream_t stream,
39+
int groupIdx) {
40+
assert(groupIdx >= 0 && groupIdx < FLAGCX_MAX_SUBGROUPS);
3841
flagcxHeteroGroupStart();
3942
int channelId = 0;
40-
4143
if (comm->channels[channelId].peers[peer]->recv[0].connected == 0) {
4244
comm->connectRecv[peer] |= (1UL << channelId);
4345
flagcxGroupCommPreconnect(comm);
@@ -50,6 +52,7 @@ flagcxResult_t flagcxHeteroRecv(void *recvbuff, size_t count,
5052
p2p->chunk = 0;
5153
p2p->dtype = datatype;
5254
p2p->stream = stream;
55+
p2p->groupIdx = groupIdx;
5356
if (flagcxIntruQueueEmpty(&tasks->peers[peer].recvQueue))
5457
tasks->p2pOrder[tasks->p2pOrderSteps++] = peer;
5558
flagcxIntruQueueEnqueue(&tasks->peers[peer].recvQueue, p2p);

‎flagcx/core/group.cc‎

Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
#include "group.h"
88
#include "adaptor.h"
99
#include "assert.h"
10-
#include "collectives.h"
1110
#include "debug.h"
11+
#include "flagcx_hetero.h"
1212
#include "launch_kernel.h"
1313
#include "net.h"
1414
#include "p2p.h"
@@ -75,6 +75,26 @@ void *flagcxAsyncJobMain(void *arg) {
7575
return arg;
7676
}
7777

78+
flagcxResult_t createAndLookupSemaphore(
79+
std::map<int, std::shared_ptr<flagcxSemaphore>> &semaphoreMap,
80+
std::shared_ptr<flagcxSemaphore> &semaphore, int &subGroupCount,
81+
int roundIdx) {
82+
if (semaphoreMap.find(roundIdx) != semaphoreMap.end()) {
83+
semaphore = semaphoreMap[roundIdx];
84+
} else {
85+
if (deviceAsyncKernel) {
86+
semaphore = std::make_shared<flagcxDeviceSemaphore>();
87+
} else {
88+
semaphore = std::make_shared<flagcxHostSemaphore>();
89+
}
90+
if (semaphoreMap.empty()) {
91+
subGroupCount++;
92+
}
93+
semaphoreMap[roundIdx] = semaphore;
94+
}
95+
return flagcxSuccess;
96+
}
97+
7898
static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
7999
flagcxResult_t ret = flagcxSuccess;
80100
// bool errorJobAbortFlag = false;
@@ -88,14 +108,11 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
88108
*asyncJobsMain = gjob->asyncJobsPtr;
89109
// volatile bool *groupAbortFlag = gjob->abortFlagPtr;
90110

91-
// Each groupLaunch we create a semaphore to track the p2p ops
92-
// and a stream to launch host or device func
93-
std::shared_ptr<flagcxSemaphore> semaphore;
94-
if (deviceAsyncKernel) {
95-
semaphore = std::make_shared<flagcxDeviceSemaphore>();
96-
} else {
97-
semaphore = std::make_shared<flagcxHostSemaphore>();
98-
}
111+
// Each groupLaunch we create a set of sub-groups of semaphores to track the
112+
// p2p ops and a stream to launch host or device func
113+
int subGroupCount = 0;
114+
std::map<int, std::shared_ptr<flagcxSemaphore>>
115+
semaphoreMapList[FLAGCX_MAX_SUBGROUPS];
99116
flagcxStream_t launchStream = nullptr;
100117
flagcxEvent_t launchEvent = nullptr;
101118

@@ -143,6 +160,7 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
143160
do {
144161
flagcxTasks *tasks = &comm->tasks;
145162
int nRanks = comm->nRanks;
163+
int localRanks = comm->localRanks;
146164

147165
// Round 0: handle self send/recv (local copy)
148166
{
@@ -162,6 +180,10 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
162180
if (sendTasks[i]->bytes == recvTasks[j]->bytes &&
163181
sendTasks[i]->dtype == recvTasks[j]->dtype) {
164182
if (sendTasks[i]->buff != recvTasks[j]->buff) {
183+
std::shared_ptr<flagcxSemaphore> semaphore;
184+
createAndLookupSemaphore(
185+
semaphoreMapList[sendTasks[i]->groupIdx], semaphore,
186+
subGroupCount, 0);
165187
flagcxProxyOp *op;
166188
FLAGCXCHECK(flagcxCalloc(&op, 1));
167189
op->pattern = flagcxPatternSend;
@@ -209,6 +231,7 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
209231
// Round 1..nRanks-1: use p2pSchedule to pair recv/send with different
210232
// peers
211233
for (int round = 1; round < nRanks; round++) {
234+
int roundIdx = round % localRanks;
212235
int recvPeer = comm->p2pSchedule[round].recvRank;
213236
int sendPeer = comm->p2pSchedule[round].sendRank;
214237

@@ -220,6 +243,9 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
220243
flagcxTaskP2p *p2p =
221244
flagcxIntruQueueDequeue(&tasks->peers[recvPeer].recvQueue);
222245
int peer = recvPeer;
246+
std::shared_ptr<flagcxSemaphore> semaphore;
247+
createAndLookupSemaphore(semaphoreMapList[p2p->groupIdx], semaphore,
248+
subGroupCount, roundIdx);
223249
flagcxProxyOp *op;
224250
FLAGCXCHECK(flagcxCalloc(&op, 1));
225251
op->pattern = flagcxPatternRecv;
@@ -299,6 +325,9 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
299325
flagcxTaskP2p *p2p =
300326
flagcxIntruQueueDequeue(&tasks->peers[sendPeer].sendQueue);
301327
int peer = sendPeer;
328+
std::shared_ptr<flagcxSemaphore> semaphore;
329+
createAndLookupSemaphore(semaphoreMapList[p2p->groupIdx], semaphore,
330+
subGroupCount, roundIdx);
302331
flagcxProxyOp *op;
303332
FLAGCXCHECK(flagcxCalloc(&op, 1));
304333
op->pattern = flagcxPatternSend;
@@ -376,15 +405,24 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
376405
}
377406

378407
if (launchStream != nullptr && launchEvent != nullptr) {
408+
for (int i = 0; i < subGroupCount; i++) {
409+
auto &semaphoreMap = semaphoreMapList[i];
410+
for (auto it = semaphoreMap.begin(); it != semaphoreMap.end(); ++it) {
411+
auto &semaphore = it->second;
412+
if (deviceAsyncKernel) {
413+
FLAGCXCHECK(
414+
deviceAdaptor->launchDeviceFunc(launchStream, deviceAsyncKernel,
415+
(void *)semaphore->getSignals()));
416+
} else {
417+
FLAGCXCHECK(deviceAdaptor->launchHostFunc(
418+
launchStream, cpuAsyncKernel, (void *)semaphore.get()));
419+
}
420+
}
421+
}
379422
if (deviceAsyncKernel) {
380-
FLAGCXCHECK(deviceAdaptor->launchDeviceFunc(
381-
launchStream, deviceAsyncKernel, (void *)semaphore->getSignals()));
382-
} else {
383-
FLAGCXCHECK(deviceAdaptor->launchHostFunc(launchStream, cpuAsyncKernel,
384-
(void *)semaphore.get()));
423+
// device semaphore need this event to signal completion
424+
FLAGCXCHECK(deviceAdaptor->eventRecord(launchEvent, launchStream));
385425
}
386-
// device semaphore need this event to signal completion
387-
FLAGCXCHECK(deviceAdaptor->eventRecord(launchEvent, launchStream));
388426
}
389427

390428
while (!flagcxIntruQueueEmpty(asyncJobsMain)) {

‎flagcx/core/include/collectives.h‎

Lines changed: 0 additions & 15 deletions
This file was deleted.

‎flagcx/core/include/flagcx_hetero.h‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,24 @@
1+
#ifndef FLAGCX_HETERO_H_
2+
#define FLAGCX_HETERO_H_
3+
14
#include "flagcx.h"
25
#include "type.h"
36

47
typedef struct flagcxHeteroComm *flagcxHeteroComm_t;
58

69
flagcxResult_t flagcxHeteroGetVersion(int *version);
710

11+
/* C++ style */
812
flagcxResult_t flagcxHeteroSend(const void *sendbuff, size_t count,
913
flagcxDataType_t datatype, int peer,
10-
flagcxHeteroComm_t comm, flagcxStream_t stream);
14+
flagcxHeteroComm_t comm, flagcxStream_t stream,
15+
int groupIdx = 0);
1116

17+
/* C++ style */
1218
flagcxResult_t flagcxHeteroRecv(void *recvbuff, size_t count,
1319
flagcxDataType_t datatype, int peer,
14-
flagcxHeteroComm_t comm, flagcxStream_t stream);
20+
flagcxHeteroComm_t comm, flagcxStream_t stream,
21+
int groupIdx = 0);
1522

1623
flagcxResult_t flagcxHeteroGroupStart();
1724

@@ -28,3 +35,5 @@ flagcxResult_t flagcxHeteroCommUserRank(const flagcxHeteroComm_t comm,
2835
int *rank);
2936

3037
flagcxResult_t flagcxHeteroCommDestroy(flagcxHeteroComm_t comm);
38+
39+
#endif

‎flagcx/core/include/group.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,17 @@
1010
#include "assert.h"
1111
#include "comm.h"
1212

13+
/**
14+
* Maximum number of subgroups allowed in a group operation.
15+
* The value 4 was chosen as a reasonable upper bound based on expected usage
16+
* patterns and to limit resource consumption (e.g., semaphores). If more than 4
17+
* subgroups are needed, the operation will need to be refactored to support a
18+
* higher limit, or FLAGCX_MAX_SUBGROUPS should be increased accordingly.
19+
* Exceeding this limit may result in errors or undefined behavior.
20+
* Consider making this value configurable if requirements change.
21+
*/
22+
#define FLAGCX_MAX_SUBGROUPS 4
23+
1324
typedef flagcxResult_t (*flagcxInitFunc_t)(flagcxHeteroComm_t *newcomm,
1425
int ndev, flagcxUniqueId commId,
1526
int myrank, int cudaDev);

‎flagcx/core/include/info.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@ struct flagcxTaskP2p {
120120
// Stateful chunk index. If a p2p gets "cut" over two plans this keeps track
121121
// of where it left off.
122122
int chunk;
123+
int groupIdx;
123124
flagcxDataType_t dtype;
124125
flagcxStream_t stream;
125126
};

‎flagcx/core/include/proxy.h‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,6 @@ struct flagcxProxyArgs {
100100
int posted = 0;
101101
int copied = 0;
102102
int postFlush = 0;
103-
int flushed = 0;
104103
int transmitted = 0;
105104
int sendStepMask;
106105
size_t totalCopySize;

‎flagcx/core/init.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
#include "adaptor.h"
88
#include "bootstrap.h"
99
#include "check.h"
10-
#include "collectives.h"
1110
#include "flagcx.h"
11+
#include "flagcx_hetero.h"
1212
#include "group.h"
1313
#include "net.h"
1414
#include "p2p.h"

0 commit comments

Comments
 (0)