Skip to content

Commit 3e6db3e

Browse files
MC952-archMC952-arch
committed
fix(rma): remove memHandleInit dependency, separate STREAM_OPS/HOST_FUNC buffers
1 parent db260d3 commit 3e6db3e

7 files changed

Lines changed: 112 additions & 130 deletions

File tree

‎.github/workflows/test.yml‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,13 @@ jobs:
7777
export PATH=$MPI_HOME/bin:$PATH
7878
export LD_LIBRARY_PATH=/__w/FlagCX/FlagCX/build/lib:$LD_LIBRARY_PATH
7979
set -e
80-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_alltoall -b 128M -e 1G -f 2 -p 1
81-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_alltoallv -b 128M -e 1G -f 2 -p 1
82-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_sendrecv -b 128M -e 1G -f 2 -p 1
83-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_allgather -b 128M -e 1G -f 2 -p 1
84-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_broadcast -b 128M -e 1G -f 2 -r 0 -p 1
85-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_gather -b 128M -e 1G -f 2 -r 0 -p 1
86-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_scatter -b 128M -e 1G -f 2 -r 0 -p 1
80+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_alltoall -b 128M -e 1G -f 2 -p 1
81+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_alltoallv -b 128M -e 1G -f 2 -p 1
82+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_sendrecv -b 128M -e 1G -f 2 -p 1
83+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_allgather -b 128M -e 1G -f 2 -p 1
84+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_broadcast -b 128M -e 1G -f 2 -r 0 -p 1
85+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_gather -b 128M -e 1G -f 2 -r 0 -p 1
86+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_USE_HETERO_COMM=1 $PERF_BIN/perf_scatter -b 128M -e 1G -f 2 -r 0 -p 1
8787
8888
- name: "Registration -R 1 (homoRunner)"
8989
run: |

‎.github/workflows/unit-test.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ jobs:
9292
cd /__w/FlagCX/FlagCX/test/unittest/runner
9393
export MPI_HOME=/usr/local/mpi
9494
export PATH=$MPI_HOME/bin:$PATH
95-
mpirun -np 8 --allow-run-as-root -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_CLUSTER_SPLIT_LIST=2 ./build/bin/runner_mpi_tests
95+
mpirun -np 8 --allow-run-as-root -x FLAGCX_DEBUG=TRACE -x FLAGCX_DEBUG_SUBSYS=ALL -x FLAGCX_MEM_ENABLE=1 -x FLAGCX_CLUSTER_SPLIT_LIST=2 ./build/bin/runner_mpi_tests
9696
9797
- name: Run multi-cluster tests using NET transport with mpirun
9898
run: |

‎.github/workflows/unittest-rma.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,4 +46,6 @@ jobs:
4646
- name: Run RMA tests
4747
run: |
4848
cd /__w/FlagCX/FlagCX
49+
export FLAGCX_DEBUG=TRACE
50+
export FLAGCX_DEBUG_SUBSYS=ALL
4951
bash test/script/rma_test.sh

‎flagcx/core/flagcx_hetero.cc‎

Lines changed: 76 additions & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -567,74 +567,61 @@ flagcxResult_t flagcxHeteroRmaProxyStart(flagcxHeteroComm_t comm) {
567567
flagcxIntruQueueConstruct(&proxy->inProgressQueues[p]);
568568
}
569569

570-
// Allocate GPU-visible memory for doneSeqsDev (stream-based synchronization).
571-
// Use GDR to get CPU-mappable device memory.
570+
// Allocate device memory for stream-based synchronization.
572571
proxy->doneSeqsDev = NULL;
573572
proxy->doneSeqsCpu = NULL;
574-
proxy->doneSeqsGdrHandle = NULL;
575-
flagcxResult_t memRes =
576-
deviceAdaptor->memHandleInit(comm->cudaDev, &proxy->doneSeqsGdrHandle);
577-
if (memRes == flagcxSuccess && proxy->doneSeqsGdrHandle != NULL) {
578-
memRes = deviceAdaptor->gdrMemAlloc((void **)&proxy->doneSeqsDev,
579-
nRanks * sizeof(uint64_t),
580-
proxy->doneSeqsGdrHandle);
581-
if (memRes == flagcxSuccess) {
582-
// Map device memory to CPU address space for progress thread writes
583-
memRes = deviceAdaptor->gdrPtrMmap((void **)&proxy->doneSeqsCpu,
584-
proxy->doneSeqsDev,
585-
nRanks * sizeof(uint64_t));
586-
if (memRes == flagcxSuccess) {
587-
// Zero-initialize doneSeqsDev
588-
memset((void *)proxy->doneSeqsCpu, 0, nRanks * sizeof(uint64_t));
589-
} else {
590-
WARN("flagcxHeteroRmaProxyStart: gdrPtrMmap failed, stream sync "
591-
"unavailable");
592-
deviceAdaptor->gdrMemFree(proxy->doneSeqsDev, proxy->doneSeqsGdrHandle);
593-
proxy->doneSeqsDev = NULL;
594-
}
595-
} else {
596-
WARN("flagcxHeteroRmaProxyStart: gdrMemAlloc failed, stream sync "
597-
"unavailable");
598-
}
599-
} else {
600-
WARN("flagcxHeteroRmaProxyStart: memHandleInit failed, stream sync "
601-
"unavailable");
602-
}
603-
604-
// Allocate GPU-visible memory for readySeqsDev (stream→proxy handshake).
605-
// Same pattern as doneSeqs: GDR-mapped device memory for CPU polling.
606573
proxy->readySeqsDev = NULL;
607574
proxy->readySeqsCpu = NULL;
608-
proxy->readySeqsGdrHandle = NULL;
609-
memRes =
610-
deviceAdaptor->memHandleInit(comm->cudaDev, &proxy->readySeqsGdrHandle);
611-
if (memRes == flagcxSuccess && proxy->readySeqsGdrHandle != NULL) {
575+
if (deviceAdaptor->gdrMemAlloc != NULL) {
576+
flagcxResult_t memRes = deviceAdaptor->gdrMemAlloc(
577+
(void **)&proxy->doneSeqsDev, nRanks * sizeof(uint64_t), NULL);
578+
if (memRes != flagcxSuccess)
579+
proxy->doneSeqsDev = NULL;
612580
memRes = deviceAdaptor->gdrMemAlloc((void **)&proxy->readySeqsDev,
613-
nRanks * sizeof(uint64_t),
614-
proxy->readySeqsGdrHandle);
615-
if (memRes == flagcxSuccess) {
616-
memRes = deviceAdaptor->gdrPtrMmap((void **)&proxy->readySeqsCpu,
617-
proxy->readySeqsDev,
618-
nRanks * sizeof(uint64_t));
619-
if (memRes == flagcxSuccess) {
620-
memset((void *)proxy->readySeqsCpu, 0, nRanks * sizeof(uint64_t));
621-
} else {
622-
WARN("flagcxHeteroRmaProxyStart: gdrPtrMmap for readySeqs failed");
623-
deviceAdaptor->gdrMemFree(proxy->readySeqsDev,
624-
proxy->readySeqsGdrHandle);
625-
proxy->readySeqsDev = NULL;
626-
}
627-
} else {
628-
WARN("flagcxHeteroRmaProxyStart: gdrMemAlloc for readySeqs failed");
581+
nRanks * sizeof(uint64_t), NULL);
582+
if (memRes != flagcxSuccess)
583+
proxy->readySeqsDev = NULL;
584+
}
585+
586+
// Map device memory to CPU if adaptor supports gdrPtrMmap (e.g. kunlunxin).
587+
// This gives the proxy thread direct CPU access to the device buffers.
588+
if (deviceAdaptor->gdrPtrMmap != NULL) {
589+
if (proxy->doneSeqsDev != NULL) {
590+
if (deviceAdaptor->gdrPtrMmap((void **)&proxy->doneSeqsCpu,
591+
proxy->doneSeqsDev,
592+
nRanks * sizeof(uint64_t)) != flagcxSuccess)
593+
proxy->doneSeqsCpu = NULL;
594+
}
595+
if (proxy->readySeqsDev != NULL) {
596+
if (deviceAdaptor->gdrPtrMmap((void **)&proxy->readySeqsCpu,
597+
proxy->readySeqsDev,
598+
nRanks * sizeof(uint64_t)) != flagcxSuccess)
599+
proxy->readySeqsCpu = NULL;
629600
}
630-
} else {
631-
WARN("flagcxHeteroRmaProxyStart: memHandleInit for readySeqs failed");
632601
}
633602

634-
// Determine sync method: STREAM_OPS if env enabled AND GDR buffers available
603+
// If no CPU mapping available, allocate separate host memory for HOST_FUNC.
604+
if (proxy->doneSeqsCpu == NULL)
605+
proxy->doneSeqsCpu = (volatile uint64_t *)calloc(nRanks, sizeof(uint64_t));
606+
if (proxy->readySeqsCpu == NULL)
607+
proxy->readySeqsCpu = (volatile uint64_t *)calloc(nRanks, sizeof(uint64_t));
608+
609+
// Zero-init device buffers
610+
if (proxy->doneSeqsDev != NULL)
611+
deviceAdaptor->deviceMemset(proxy->doneSeqsDev, 0,
612+
nRanks * sizeof(uint64_t), flagcxMemDevice,
613+
NULL);
614+
if (proxy->readySeqsDev != NULL)
615+
deviceAdaptor->deviceMemset(proxy->readySeqsDev, 0,
616+
nRanks * sizeof(uint64_t), flagcxMemDevice,
617+
NULL);
618+
619+
// STREAM_OPS requires: device buffers AND CPU mapping of those same buffers
620+
// (so proxy thread can write doneSeqsDev from CPU). Only possible when
621+
// gdrPtrMmap is available (doneSeqsCpu points into doneSeqsDev).
635622
proxy->useStreamOps =
636-
(flagcxParamRmaStreamOps() == 1 && proxy->readySeqsDev != NULL &&
637-
proxy->doneSeqsDev != NULL)
623+
(flagcxParamRmaStreamOps() == 1 && proxy->doneSeqsDev != NULL &&
624+
proxy->readySeqsDev != NULL && deviceAdaptor->gdrPtrMmap != NULL)
638625
? 1
639626
: 0;
640627
INFO(FLAGCX_INIT, "RMA proxy sync method: %s",
@@ -646,28 +633,21 @@ flagcxResult_t flagcxHeteroRmaProxyStart(flagcxHeteroComm_t comm) {
646633
if (pthread_create(&proxy->thread, NULL, flagcxRmaProxyProgressThread,
647634
proxy) != 0) {
648635
WARN("flagcxHeteroRmaProxyStart: pthread_create failed");
649-
// Free GDR memory for stream-based sync (doneSeqs)
650-
if (proxy->doneSeqsCpu != NULL) {
651-
deviceAdaptor->gdrPtrMunmap((void *)proxy->doneSeqsCpu,
652-
nRanks * sizeof(uint64_t));
653-
}
654-
if (proxy->doneSeqsDev != NULL && proxy->doneSeqsGdrHandle != NULL) {
655-
deviceAdaptor->gdrMemFree(proxy->doneSeqsDev, proxy->doneSeqsGdrHandle);
656-
}
657-
if (proxy->doneSeqsGdrHandle != NULL) {
658-
deviceAdaptor->memHandleDestroy(comm->cudaDev, proxy->doneSeqsGdrHandle);
659-
}
660-
// Free GDR memory for readySeqs
661-
if (proxy->readySeqsCpu != NULL) {
662-
deviceAdaptor->gdrPtrMunmap((void *)proxy->readySeqsCpu,
663-
nRanks * sizeof(uint64_t));
664-
}
665-
if (proxy->readySeqsDev != NULL && proxy->readySeqsGdrHandle != NULL) {
666-
deviceAdaptor->gdrMemFree(proxy->readySeqsDev, proxy->readySeqsGdrHandle);
667-
}
668-
if (proxy->readySeqsGdrHandle != NULL) {
669-
deviceAdaptor->memHandleDestroy(comm->cudaDev, proxy->readySeqsGdrHandle);
636+
if (deviceAdaptor->gdrPtrMunmap != NULL) {
637+
if (proxy->doneSeqsCpu != NULL)
638+
deviceAdaptor->gdrPtrMunmap((void *)proxy->doneSeqsCpu,
639+
nRanks * sizeof(uint64_t));
640+
if (proxy->readySeqsCpu != NULL)
641+
deviceAdaptor->gdrPtrMunmap((void *)proxy->readySeqsCpu,
642+
nRanks * sizeof(uint64_t));
643+
} else {
644+
free((void *)proxy->doneSeqsCpu);
645+
free((void *)proxy->readySeqsCpu);
670646
}
647+
if (proxy->doneSeqsDev != NULL)
648+
deviceAdaptor->gdrMemFree(proxy->doneSeqsDev, NULL);
649+
if (proxy->readySeqsDev != NULL)
650+
deviceAdaptor->gdrMemFree(proxy->readySeqsDev, NULL);
671651
for (int p = 0; p < nRanks; p++)
672652
pthread_mutex_destroy(&proxy->peerProducerMutexes[p]);
673653
free(proxy->circularBuffers);
@@ -702,29 +682,24 @@ flagcxResult_t flagcxHeteroRmaProxyStop(flagcxHeteroComm_t comm) {
702682
for (int p = 0; p < proxy->nRanks; p++)
703683
pthread_mutex_destroy(&proxy->peerProducerMutexes[p]);
704684

705-
// Free GDR memory for stream-based sync (doneSeqs)
706-
if (proxy->doneSeqsCpu != NULL) {
707-
deviceAdaptor->gdrPtrMunmap((void *)proxy->doneSeqsCpu,
708-
proxy->nRanks * sizeof(uint64_t));
709-
}
710-
if (proxy->doneSeqsDev != NULL && proxy->doneSeqsGdrHandle != NULL) {
711-
deviceAdaptor->gdrMemFree(proxy->doneSeqsDev, proxy->doneSeqsGdrHandle);
712-
}
713-
if (proxy->doneSeqsGdrHandle != NULL) {
714-
deviceAdaptor->memHandleDestroy(comm->cudaDev, proxy->doneSeqsGdrHandle);
685+
// Free CPU mappings / host memory
686+
if (deviceAdaptor->gdrPtrMunmap != NULL) {
687+
if (proxy->doneSeqsCpu != NULL)
688+
deviceAdaptor->gdrPtrMunmap((void *)proxy->doneSeqsCpu,
689+
proxy->nRanks * sizeof(uint64_t));
690+
if (proxy->readySeqsCpu != NULL)
691+
deviceAdaptor->gdrPtrMunmap((void *)proxy->readySeqsCpu,
692+
proxy->nRanks * sizeof(uint64_t));
693+
} else {
694+
free((void *)proxy->doneSeqsCpu);
695+
free((void *)proxy->readySeqsCpu);
715696
}
716697

717-
// Free GDR memory for readySeqs
718-
if (proxy->readySeqsCpu != NULL) {
719-
deviceAdaptor->gdrPtrMunmap((void *)proxy->readySeqsCpu,
720-
proxy->nRanks * sizeof(uint64_t));
721-
}
722-
if (proxy->readySeqsDev != NULL && proxy->readySeqsGdrHandle != NULL) {
723-
deviceAdaptor->gdrMemFree(proxy->readySeqsDev, proxy->readySeqsGdrHandle);
724-
}
725-
if (proxy->readySeqsGdrHandle != NULL) {
726-
deviceAdaptor->memHandleDestroy(comm->cudaDev, proxy->readySeqsGdrHandle);
727-
}
698+
// Free device memory
699+
if (proxy->doneSeqsDev != NULL)
700+
deviceAdaptor->gdrMemFree(proxy->doneSeqsDev, NULL);
701+
if (proxy->readySeqsDev != NULL)
702+
deviceAdaptor->gdrMemFree(proxy->readySeqsDev, NULL);
728703

729704
free(proxy->circularBuffers);
730705
free((void *)proxy->pis);

‎flagcx/core/include/flagcx_hetero.h‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -62,19 +62,19 @@ struct flagcxRmaProxyState {
6262
volatile uint64_t *doneSeqs; // [nRanks]
6363
volatile uint32_t *inFlights; // [nRanks]
6464

65-
// GPU-visible done sequence counters for stream-based synchronization.
66-
// Written by progress thread (via GDR-mapped CPU pointer), waited on by
67-
// GPU stream via streamWaitValue64.
68-
uint64_t *doneSeqsDev; // [nRanks] device pointer (GPU-visible)
69-
volatile uint64_t *doneSeqsCpu; // [nRanks] CPU-mapped pointer to doneSeqsDev
70-
void *doneSeqsGdrHandle; // GDR handle for cleanup
71-
72-
// GPU-visible ready sequence counters for stream→proxy handshake.
73-
// Written by GPU stream (streamWriteValue64 or host-func callback) to signal
74-
// that source buffer data is committed. Proxy polls before issuing RDMA.
75-
uint64_t *readySeqsDev; // [nRanks] device pointer (GPU-visible)
76-
volatile uint64_t *readySeqsCpu; // [nRanks] CPU-mapped pointer
77-
void *readySeqsGdrHandle; // GDR handle for cleanup
65+
// GPU-visible done sequence counters for STREAM_OPS mode.
66+
// GPU stream waits on doneSeqsDev via streamWaitValue64.
67+
uint64_t *doneSeqsDev; // [nRanks] device pointer (GPU-visible)
68+
// CPU-side done sequence counters for HOST_FUNC mode.
69+
// Written by proxy thread, polled by host-func callback.
70+
volatile uint64_t *doneSeqsCpu; // [nRanks] host memory
71+
72+
// GPU-visible ready sequence counters for STREAM_OPS mode.
73+
// GPU stream writes readySeqsDev via streamWriteValue64 to signal data ready.
74+
uint64_t *readySeqsDev; // [nRanks] device pointer (GPU-visible)
75+
// CPU-side ready sequence counters for HOST_FUNC mode.
76+
// Written by host-func callback, polled by proxy thread.
77+
volatile uint64_t *readySeqsCpu; // [nRanks] host memory
7878

7979
// Synchronization method: HOST_FUNC (default) or STREAM_OPS (opt-in via env)
8080
int useStreamOps; // 0 = HOST_FUNC (default), 1 = STREAM_OPS

‎flagcx/flagcx.cc‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -695,11 +695,16 @@ flagcxResult_t flagcxOneSideSignalRegister(const flagcxComm_t comm, void *buff,
695695
}
696696
}
697697

698-
// NOTE: Signal buffer IPC registration for D2D bypass is NOT done here.
699-
// buildIpcPeerPointers() requires a collective allGather which may deadlock
700-
// if the ptrType/localRanks guard doesn't hold uniformly across all ranks.
701-
// The D2D signal path uses VMM (symWin) when available; otherwise the proxy
702-
// + RDMA path handles signalling correctly.
698+
// Register signal buffer in IPC table for intra-node D2D bypass.
699+
// Guard on ptrType only (global env var, uniform across all ranks) to ensure
700+
// all ranks enter the collective allGather inside buildIpcPeerPointers.
701+
if (ptrType == FLAGCX_PTR_CUDA) {
702+
int idx = buildIpcPeerPointers(comm, buff, size);
703+
if (idx >= 0) {
704+
INFO(FLAGCX_REG, "Signal buffer IPC registered (slot %d) for D2D bypass",
705+
idx);
706+
}
707+
}
703708

704709
return flagcxSuccess;
705710

‎test/script/rma_test.sh‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,23 +28,23 @@ echo "--- test_put ---"
2828
mpirun -np "$NP" --allow-run-as-root \
2929
-x FLAGCX_USE_HETERO_COMM=1 \
3030
-x FLAGCX_MEM_ENABLE=1 \
31-
-x FLAGCX_VMM_ENABLE=1 \
31+
-x FLAGCX_VMM_ENABLE=0 \
3232
"$PERF_BIN/perf_put" -b 8 -e 1M -f 2 -R 2
3333

3434
echo ""
3535
echo "--- test_get ---"
3636
mpirun -np "$NP" --allow-run-as-root \
3737
-x FLAGCX_USE_HETERO_COMM=1 \
3838
-x FLAGCX_MEM_ENABLE=1 \
39-
-x FLAGCX_VMM_ENABLE=1 \
39+
-x FLAGCX_VMM_ENABLE=0 \
4040
"$PERF_BIN/perf_get" -b 8 -e 1M -f 2 -R 2
4141

4242
echo ""
4343
echo "--- test_one_side_register ---"
4444
mpirun -np "$NP" --allow-run-as-root \
4545
-x FLAGCX_USE_HETERO_COMM=1 \
4646
-x FLAGCX_MEM_ENABLE=1 \
47-
-x FLAGCX_VMM_ENABLE=1 \
47+
-x FLAGCX_VMM_ENABLE=0 \
4848
"$PERF_BIN/perf_one_side_register" -b 8 -e 1M -f 2 -R 2
4949

5050
echo ""

0 commit comments

Comments
 (0)