Skip to content

Commit bf263bd

Browse files
leoda1leoda1
authored
[CRL] Optimize rmaProxyProgress and Adds batched one-sided PUT operations to improve RDMA throughput (flagos-ai#461)
1 parent e75e482 commit bf263bd

9 files changed

Lines changed: 844 additions & 240 deletions

File tree

‎docs/environment_variables.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ This document provides a comprehensive reference for all environment variables u
7676
| `FLAGCX_P2P_DISABLE` | 0 | When set to 1, disables P2P transport |
7777
| `FLAGCX_P2P_SCHEDULE_DISABLE` | 0 | When set to 1, disables P2P scheduling optimization |
7878
| `FLAGCX_DEVICE_FUNC_PATH` | None | Path to device function library for async kernel loading |
79+
| `FLAGCX_RMA_QUEUE_SIZE` | 256 | Per-peer circular buffer depth for the RMA proxy thread. Must be a power of two. Increasing this allows more in-flight RDMA descriptors per peer before the producer blocks |
80+
| `FLAGCX_RMA_BATCH_MAX` | 256 | Maximum number of RDMA PUT descriptors batched into a single `iputBatch` call by the RMA proxy thread. Capped internally at 256. Set to 1 to disable batching and fall back to one-at-a-time `iput` |
7981

8082
---
8183

‎flagcx/adaptor/include/flagcx_net_adaptor.h‎

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#define FLAGCX_NET_ADAPTOR_H_
77

88
#include "flagcx.h"
9+
#include "string.h"
910

1011
#ifdef __cplusplus
1112
extern "C" {
@@ -23,6 +24,8 @@ typedef enum {
2324
// listen, connect, accept, closeSend, closeRecv, closeListen,
2425
// regMr, regMrDmaBuf, deregMr, isend, irecv, iflush, test,
2526
// iput, iget, iputSignal, getDevFromName
27+
// v2 — adds iputBatch (optional one-sided batch WRITE)
28+
2629
struct flagcxNetAdaptor_v1 {
2730
// Basic functions
2831
const char *name;
@@ -56,6 +59,58 @@ struct flagcxNetAdaptor_v1 {
5659
void **mhandles, void **request);
5760
flagcxResult_t (*test)(void *request, int *done, int *sizes);
5861

62+
// One-sided (per-window MR: separate src/dst handles for independent buffers)
63+
flagcxResult_t (*iput)(void *sendComm, uint64_t srcOff, uint64_t dstOff,
64+
size_t size, int srcRank, int dstRank,
65+
void **srcHandles, void **dstHandles, void **request);
66+
// RDMA READ: pull data from remote srcRank into local dstRank buffer
67+
flagcxResult_t (*iget)(void *sendComm, uint64_t srcOff, uint64_t dstOff,
68+
size_t size, int srcRank, int dstRank,
69+
void **srcHandles, void **dstHandles, void **request);
70+
// Data + signal combined (NCCL GIN-aligned: enables chained WRITE + ATOMIC)
71+
// When size == 0, only signal ATOMIC is posted (signal-only mode)
72+
flagcxResult_t (*iputSignal)(void *sendComm, uint64_t srcOff, uint64_t dstOff,
73+
size_t size, int srcRank, int dstRank,
74+
void **srcHandles, void **dstHandles,
75+
uint64_t signalOff, void **signalHandles,
76+
uint64_t signalValue, void **request);
77+
// Device name lookup
78+
flagcxResult_t (*getDevFromName)(char *name, int *dev);
79+
};
80+
81+
struct flagcxNetAdaptor_latest {
82+
// Basic functions
83+
const char *name;
84+
flagcxResult_t (*init)();
85+
flagcxResult_t (*devices)(int *ndev);
86+
flagcxResult_t (*getProperties)(int dev, void *props);
87+
88+
// Setup functions
89+
flagcxResult_t (*listen)(int dev, void *handle, void **listenComm);
90+
flagcxResult_t (*connect)(int dev, void *handle, void **sendComm);
91+
flagcxResult_t (*accept)(void *listenComm, void **recvComm);
92+
flagcxResult_t (*closeSend)(void *sendComm);
93+
flagcxResult_t (*closeRecv)(void *recvComm);
94+
flagcxResult_t (*closeListen)(void *listenComm);
95+
96+
// Memory region functions
97+
flagcxResult_t (*regMr)(void *comm, void *data, size_t size, int type,
98+
int mrFlags, void **mhandle);
99+
flagcxResult_t (*regMrDmaBuf)(void *comm, void *data, size_t size, int type,
100+
uint64_t offset, int fd, int mrFlags,
101+
void **mhandle);
102+
flagcxResult_t (*deregMr)(void *comm, void *mhandle);
103+
104+
// Two-sided functions
105+
flagcxResult_t (*isend)(void *sendComm, void *data, size_t size, int tag,
106+
void *mhandle, void *phandle, void **request);
107+
flagcxResult_t (*irecv)(void *recvComm, int n, void **data, size_t *sizes,
108+
int *tags, void **mhandles, void **phandles,
109+
void **request);
110+
flagcxResult_t (*iflush)(void *recvComm, int n, void **data, int *sizes,
111+
void **mhandles, void **request);
112+
flagcxResult_t (*test)(void *request, int *done, int *sizes);
113+
59114
// One-sided (per-window MR: separate src/dst handles for independent buffers)
60115
flagcxResult_t (*iput)(void *sendComm, uint64_t srcOff, uint64_t dstOff,
61116
size_t size, int srcRank, int dstRank,
@@ -74,8 +129,22 @@ struct flagcxNetAdaptor_v1 {
74129

75130
// Device name lookup
76131
flagcxResult_t (*getDevFromName)(char *name, int *dev);
132+
// Optional one-side batch WRITE.
133+
flagcxResult_t (*iputBatch)(void *sendComm, int count,
134+
const uint64_t *srcOffs, const uint64_t *dstOffs,
135+
const size_t *sizes, int srcRank, int dstRank,
136+
void **srcHandles, void **dstHandles,
137+
void **requests, int *posted);
77138
};
78-
#define flagcxNetAdaptor flagcxNetAdaptor_v1
139+
140+
#define flagcxNetAdaptor flagcxNetAdaptor_latest
141+
142+
static inline void
143+
flagcxNetAdaptorUpgrade(const struct flagcxNetAdaptor_v1 *src,
144+
struct flagcxNetAdaptor_latest *dst) {
145+
memset(dst, 0, sizeof(*dst));
146+
memcpy(dst, src, sizeof(struct flagcxNetAdaptor_v1));
147+
}
79148

80149
// Versioned export symbol name
81150
#define FLAGCX_NET_ADAPTOR_PLUGIN_SYMBOL_V1 flagcxNetAdaptorPlugin_v1

‎flagcx/adaptor/net/ibrc_adaptor.cc‎

Lines changed: 114 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2476,6 +2476,116 @@ flagcxResult_t flagcxIbIput(void *sendComm, uint64_t srcOff, uint64_t dstOff,
24762476
return flagcxSuccess;
24772477
}
24782478

2479+
flagcxResult_t flagcxIbIputBatch(void *sendComm, int count,
2480+
const uint64_t *srcOffs,
2481+
const uint64_t *dstOffs, const size_t *sizes,
2482+
int srcRank, int dstRank, void **srcHandles,
2483+
void **dstHandles, void **requests,
2484+
int *posted) {
2485+
if (posted == NULL || requests == NULL)
2486+
return flagcxInvalidArgument;
2487+
*posted = 0;
2488+
if (count <= 0)
2489+
return flagcxSuccess;
2490+
if (count > MAX_REQUESTS)
2491+
return flagcxInvalidArgument;
2492+
2493+
struct flagcxIbSendComm *comm = (struct flagcxIbSendComm *)sendComm;
2494+
struct flagcxOneSideHandleInfo *srcInfo =
2495+
(struct flagcxOneSideHandleInfo *)srcHandles;
2496+
struct flagcxOneSideHandleInfo *dstInfo =
2497+
(struct flagcxOneSideHandleInfo *)dstHandles;
2498+
if (comm == NULL || srcInfo == NULL || dstInfo == NULL || srcOffs == NULL ||
2499+
dstOffs == NULL || sizes == NULL) {
2500+
return flagcxInvalidArgument;
2501+
}
2502+
2503+
int qpIdx = comm->base.qpIndex;
2504+
comm->base.qpIndex = (qpIdx + 1) % comm->base.nqps;
2505+
struct flagcxIbQp *qp = &comm->base.qps[qpIdx];
2506+
int devIndex = qp->devIndex;
2507+
int lkey = srcInfo->lkeys[srcRank];
2508+
int rkey = dstInfo->rkeys[dstRank];
2509+
2510+
struct ibv_send_wr wrs[MAX_REQUESTS];
2511+
struct ibv_sge sges[MAX_REQUESTS];
2512+
struct flagcxIbRequest *reqs[MAX_REQUESTS];
2513+
memset(wrs, 0, count * sizeof(struct ibv_send_wr));
2514+
memset(sges, 0, count * sizeof(struct ibv_sge));
2515+
memset(reqs, 0, count * sizeof(struct flagcxIbRequest *));
2516+
2517+
flagcxResult_t res = flagcxSuccess;
2518+
struct ibv_send_wr *bad_wr = NULL;
2519+
for (int i = 0; i < count; i++) {
2520+
struct flagcxIbRequest *req = NULL;
2521+
res = flagcxIbGetRequest(&comm->base, &req);
2522+
if (res != flagcxSuccess) {
2523+
goto fail_before_post;
2524+
}
2525+
reqs[i] = req;
2526+
req->type = FLAGCX_NET_IB_REQ_IPUT;
2527+
req->sock = &comm->base.sock;
2528+
for (int d = 0; d < comm->base.ndevs; d++) {
2529+
req->devBases[d] = &comm->devs[d].base;
2530+
}
2531+
2532+
void *srcPtr = (void *)(srcInfo->baseVas[srcRank] + srcOffs[i]);
2533+
void *dstPtr = (void *)(dstInfo->baseVas[dstRank] + dstOffs[i]);
2534+
2535+
wrs[i].opcode = IBV_WR_RDMA_WRITE;
2536+
wrs[i].send_flags = IBV_SEND_SIGNALED;
2537+
wrs[i].wr_id = req - comm->base.reqs;
2538+
wrs[i].next = (i + 1 == count) ? NULL : &wrs[i + 1];
2539+
wrs[i].wr.rdma.remote_addr = (uint64_t)dstPtr;
2540+
wrs[i].wr.rdma.rkey = rkey;
2541+
wrs[i].sg_list = &sges[i];
2542+
wrs[i].num_sge = 1;
2543+
2544+
sges[i].addr = (uintptr_t)srcPtr;
2545+
sges[i].length = (uint32_t)sizes[i];
2546+
if ((size_t)sges[i].length != sizes[i]) {
2547+
WARN("flagcxIbIputBatch: transfer size %zu exceeds ibv_sge 32-bit limit",
2548+
sizes[i]);
2549+
res = flagcxInvalidArgument;
2550+
goto fail_before_post;
2551+
}
2552+
sges[i].lkey = lkey;
2553+
}
2554+
2555+
res = flagcxWrapIbvPostSend(qp->qp, wrs, &bad_wr);
2556+
if (res != flagcxSuccess) {
2557+
int first_failed = bad_wr ? (int)(bad_wr - wrs) : 0;
2558+
if (first_failed < 0 || first_failed > count)
2559+
first_failed = 0;
2560+
for (int i = first_failed; i < count; i++) {
2561+
if (reqs[i] != NULL) {
2562+
flagcxIbFreeRequest(reqs[i]);
2563+
reqs[i] = NULL;
2564+
}
2565+
}
2566+
for (int i = 0; i < first_failed; i++) {
2567+
flagcxIbAddEvent(reqs[i], devIndex, &comm->devs[devIndex].base);
2568+
requests[i] = reqs[i];
2569+
}
2570+
*posted = first_failed;
2571+
return res;
2572+
}
2573+
2574+
for (int i = 0; i < count; i++) {
2575+
flagcxIbAddEvent(reqs[i], devIndex, &comm->devs[devIndex].base);
2576+
requests[i] = reqs[i];
2577+
}
2578+
*posted = count;
2579+
return flagcxSuccess;
2580+
2581+
fail_before_post:
2582+
for (int i = 0; i < count; i++) {
2583+
if (reqs[i] != NULL)
2584+
flagcxIbFreeRequest(reqs[i]);
2585+
}
2586+
return res;
2587+
}
2588+
24792589
flagcxResult_t flagcxIbIget(void *sendComm, uint64_t srcOff, uint64_t dstOff,
24802590
size_t size, int srcRank, int dstRank,
24812591
void **srcHandles, void **dstHandles,
@@ -2641,4 +2751,7 @@ struct flagcxNetAdaptor flagcxNetIb = {
26412751
flagcxIbIput, flagcxIbIget, flagcxIbIputSignal,
26422752

26432753
// Device name lookup
2644-
flagcxIbGetDevFromName};
2754+
flagcxIbGetDevFromName,
2755+
2756+
// Optional one-sided batch WRITE
2757+
flagcxIbIputBatch};

‎flagcx/adaptor/net/net_plugin_load.cc‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
************************************************************************/
44

55
#include "adaptor_plugin_load.h"
6+
#include "alloc.h"
67
#include "core.h"
78
#include "flagcx_net_adaptor.h"
89
#include "net.h"
@@ -15,6 +16,7 @@
1516
static void *netPluginDlHandle = NULL;
1617
static int netPluginRefCount = 0;
1718
static std::mutex netPluginMutex;
19+
static struct flagcxNetAdaptor *upgradedNetPluginAdaptor = NULL;
1820

1921
extern struct flagcxNetAdaptor *flagcxNetAdaptors[3];
2022

@@ -35,12 +37,9 @@ static flagcxResult_t flagcxNetAdaptorPluginLoad() {
3537
return flagcxSuccess;
3638
}
3739

38-
// Future: When v2 is introduced, try dlsym("flagcxNetAdaptorPlugin_v2")
39-
// first, then fall back to "flagcxNetAdaptorPlugin_v1" and wrap in a v1→v2
40-
// shim.
41-
struct flagcxNetAdaptor *plugin = (struct flagcxNetAdaptor *)dlsym(
40+
struct flagcxNetAdaptor_v1 *pluginV1 = (struct flagcxNetAdaptor_v1 *)dlsym(
4241
netPluginDlHandle, "flagcxNetAdaptorPlugin_v1");
43-
if (plugin == NULL) {
42+
if (pluginV1 == NULL) {
4443
WARN("ADAPTOR/Plugin: Failed to find symbol 'flagcxNetAdaptorPlugin_v1' in "
4544
"'%s': %s",
4645
envValue, dlerror());
@@ -49,6 +48,15 @@ static flagcxResult_t flagcxNetAdaptorPluginLoad() {
4948
return flagcxSuccess;
5049
}
5150

51+
if (flagcxCalloc(&upgradedNetPluginAdaptor, 1) != flagcxSuccess) {
52+
WARN("ADAPTOR/Plugin: Failed to allocate upgraded net adaptor struct");
53+
flagcxAdaptorClosePluginLib(netPluginDlHandle);
54+
netPluginDlHandle = NULL;
55+
return flagcxSystemError;
56+
}
57+
flagcxNetAdaptorUpgrade(pluginV1, upgradedNetPluginAdaptor);
58+
struct flagcxNetAdaptor *plugin = upgradedNetPluginAdaptor;
59+
5260
// Validate function pointers that all built-in net adaptors implement.
5361
// Fields left NULL in some adaptors (regMrDmaBuf, iput, iget, iputSignal,
5462
// getDevFromName) are intentionally not checked here.
@@ -62,6 +70,8 @@ static flagcxResult_t flagcxNetAdaptorPluginLoad() {
6270
WARN("ADAPTOR/Plugin: Net adaptor plugin '%s' is missing required function "
6371
"pointers",
6472
envValue);
73+
free(upgradedNetPluginAdaptor);
74+
upgradedNetPluginAdaptor = NULL;
6575
flagcxAdaptorClosePluginLib(netPluginDlHandle);
6676
netPluginDlHandle = NULL;
6777
return flagcxSuccess;
@@ -76,6 +86,10 @@ static flagcxResult_t flagcxNetAdaptorPluginLoad() {
7686
static flagcxResult_t flagcxNetAdaptorPluginUnload() {
7787
flagcxNetAdaptors[0] = nullptr;
7888
flagcxNetStates[0] = flagcxNetStateInit;
89+
if (upgradedNetPluginAdaptor != NULL) {
90+
free(upgradedNetPluginAdaptor);
91+
upgradedNetPluginAdaptor = NULL;
92+
}
7993
flagcxAdaptorClosePluginLib(netPluginDlHandle);
8094
netPluginDlHandle = NULL;
8195
return flagcxSuccess;

0 commit comments

Comments
 (0)