Skip to content

Commit e8d1b8c

Browse files
MC952-archMC952-arch
committed
fix: harden barrier, proxy, and NVSHMEM grid sync
1 parent cd35dc0 commit e8d1b8c

16 files changed

Lines changed: 156 additions & 218 deletions

‎bindings/ir/flagcx_device_scalar_ir_impl.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,14 @@ flagcxGridSync(unsigned int *gridSyncState) {
9494
atomicExch(&gridSyncState[1], 1u - curSense);
9595
} else {
9696
// Spin until sense flips (all blocks arrived)
97+
#ifndef NDEBUG
98+
unsigned int __spins = 0;
99+
#endif
97100
while (*(volatile unsigned int *)&gridSyncState[1] == curSense) {
101+
#ifndef NDEBUG
102+
if (++__spins >= 100000000u)
103+
__trap(); // grid barrier timeout — likely a hang
104+
#endif
98105
}
99106
}
100107
}

‎flagcx/adaptor/device_api/default_dev_api_backend.cc‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,12 @@ defaultDevApiCommCreate(flagcxComm_t comm,
421421
// barrier needs nTeamRanks slots per barrier instance (one per CTA)
422422
int userSignals = reqs->interSignalCount;
423423
int barrierSlots = devComm->nTeamRanks * reqs->interBarrierCount;
424+
if (reqs->interBarrierCount > 0 &&
425+
barrierSlots / reqs->interBarrierCount != devComm->nTeamRanks) {
426+
WARN("barrierSignalBase overflow: nTeamRanks=%d interBarrierCount=%d",
427+
devComm->nTeamRanks, reqs->interBarrierCount);
428+
return flagcxInternalError;
429+
}
424430
devComm->signalCount = userSignals + barrierSlots;
425431
devComm->barrierSignalBase = userSignals;
426432
size_t sigSize =

‎flagcx/adaptor/device_api/nvshmem_dev_api_backend.cc‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,11 @@ static flagcxResult_t nvshmemDevApiCommGetDevicePtr(flagcxDevComm_t devComm,
126126
return flagcxSuccess;
127127
}
128128

129-
// Construct value struct on host stack, then copy to device
129+
// Construct value struct on host stack, then copy to device.
130+
// Note: _gridBarrierState is intentionally left nullptr for NVSHMEM.
131+
// NVSHMEM barriers use nvshmemx_barrier_block() + per-block arrive/release
132+
// flags internally, so the IR-level flagcxGridSync (sense-reversing) is not
133+
// needed and the null check in the IR functions will skip it.
130134
flagcxDevComm hostCopy(*devComm);
131135
hostCopy._netContexts = nullptr;
132136

‎flagcx/adaptor/include/device_api/default_comm_traits.h‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#define FLAGCX_FALLBACK_DEVICE_TRAITS_H_
1919

2020
#include "flagcx_kernel_core.h"
21+
#include <cassert>
2122
#ifndef __CUDACC__
2223
#include "sym_heap.h"
2324
#endif
@@ -831,7 +832,9 @@ struct Barrier<DefaultBackend<P>, flagcxTeamTagIntra, Coop> {
831832
_myRank(team.rank), _nBarriers(dc.nBarriers), _ctaIndex(index),
832833
_epochBuffer(dc.epochBuffer),
833834
_epoch(Atomic::load(&dc.epochBuffer[index],
834-
flagcxDeviceMemoryOrderAcquire)) {}
835+
flagcxDeviceMemoryOrderAcquire)) {
836+
assert(index < FLAGCX_DEVICE_CTA_COUNT);
837+
}
835838

836839
// arrive: thread-striped store epoch+1 to each peer's inbox slot for me
837840
FLAGCX_DEVICE_INLINE_DECORATOR void
@@ -916,7 +919,9 @@ struct Barrier<DefaultBackend<P>, flagcxTeamTagInter, Coop> {
916919
_signalCount(net.signalCount), _contextId(net.contextId),
917920
_teamRank(net.teamRank), _nTeamRanks(net.nTeamRanks),
918921
_barrierSignal0(net.barrierSignalBase + (int)index * net.nTeamRanks),
919-
_stride(dc.intraSize), _localRank(dc.intraRank) {}
922+
_stride(dc.intraSize), _localRank(dc.intraRank) {
923+
assert(index < FLAGCX_DEVICE_CTA_COUNT);
924+
}
920925

921926
// arrive: signal all remote peers "I have arrived"
922927
FLAGCX_DEVICE_INLINE_DECORATOR void

‎flagcx/adaptor/include/device_api/nvshmem_comm_traits.h‎

Lines changed: 79 additions & 158 deletions
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,61 @@ struct CommTraits<NvshmemBackend> {
485485

486486
// ============================================================
487487
// Barrier specializations for NvshmemBackend
488+
// ============================================================
489+
// Grid-wide synchronization helpers for NVSHMEM barriers.
490+
// Block 0 collects arrive flags from all blocks, performs the
491+
// nvshmemx_barrier_block PE-level barrier, then releases others.
492+
// ============================================================
493+
494+
template <typename Coop>
495+
FLAGCX_DEVICE_INLINE_DECORATOR void
496+
nvshmemGridArrive(Coop &coop, nvshmem_team_t team, volatile uint64_t *arrive,
497+
volatile uint64_t *release) {
498+
#ifdef __CUDACC__
499+
int numBlocks = FLAGCX_GRID_DIM_X;
500+
coop.sync();
501+
if (coop.threadRank() == 0) {
502+
arrive[FLAGCX_BLOCK_IDX_X]++;
503+
}
504+
if (FLAGCX_BLOCK_IDX_X == 0) {
505+
coop.sync();
506+
uint64_t expected = arrive[0];
507+
for (int i = coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
508+
if (i == 0)
509+
continue;
510+
while (arrive[i] < expected) {
511+
}
512+
}
513+
coop.sync();
514+
}
515+
#endif
516+
}
517+
518+
template <typename Coop>
519+
FLAGCX_DEVICE_INLINE_DECORATOR void
520+
nvshmemGridWait(Coop &coop, nvshmem_team_t team, volatile uint64_t *arrive,
521+
volatile uint64_t *release) {
522+
#ifdef __CUDACC__
523+
int numBlocks = FLAGCX_GRID_DIM_X;
524+
if (FLAGCX_BLOCK_IDX_X == 0) {
525+
coop.sync();
526+
nvshmemx_barrier_block(team);
527+
coop.sync();
528+
for (int i = coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
529+
release[i]++;
530+
}
531+
} else {
532+
if (coop.threadRank() == 0) {
533+
uint64_t cur = release[FLAGCX_BLOCK_IDX_X];
534+
while (release[FLAGCX_BLOCK_IDX_X] == cur) {
535+
}
536+
}
537+
}
538+
coop.sync();
539+
#endif
540+
}
541+
542+
// ============================================================
488543
// Signal-based split-phase barriers using nvshmemx_signal_op.
489544
// ============================================================
490545

@@ -513,70 +568,24 @@ struct Barrier<NvshmemBackend, flagcxTeamTagIntra, Coop> {
513568
_teamRank(dc.intraRank),
514569
_gridSyncState((volatile uint64_t *)dc.gridSyncState) {}
515570

516-
// Grid arrive: each block writes own flag, block 0 checks all in parallel
517-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridArrive(nvshmem_team_t team,
518-
volatile uint64_t *arrive,
519-
volatile uint64_t *release) {
520-
#ifdef __CUDACC__
521-
int numBlocks = FLAGCX_GRID_DIM_X;
522-
_coop.sync();
523-
if (_coop.threadRank() == 0) {
524-
arrive[FLAGCX_BLOCK_IDX_X]++;
525-
}
526-
if (FLAGCX_BLOCK_IDX_X == 0) {
527-
_coop.sync();
528-
uint64_t expected = arrive[0];
529-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
530-
if (i == 0)
531-
continue;
532-
while (arrive[i] < expected) {
533-
}
534-
}
535-
_coop.sync();
536-
}
537-
#endif
538-
}
539-
540-
// Grid wait: block 0 does PE barrier then releases per-block flags
541-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridWait(nvshmem_team_t team,
542-
volatile uint64_t *arrive,
543-
volatile uint64_t *release) {
544-
#ifdef __CUDACC__
545-
int numBlocks = FLAGCX_GRID_DIM_X;
546-
if (FLAGCX_BLOCK_IDX_X == 0) {
547-
_coop.sync();
548-
nvshmemx_barrier_block(team);
549-
_coop.sync();
550-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
551-
release[i]++;
552-
}
553-
} else {
554-
if (_coop.threadRank() == 0) {
555-
uint64_t cur = release[FLAGCX_BLOCK_IDX_X];
556-
while (release[FLAGCX_BLOCK_IDX_X] == cur) {
557-
}
558-
}
559-
}
560-
_coop.sync();
561-
#endif
562-
}
563-
564571
FLAGCX_DEVICE_INLINE_DECORATOR void
565572
arrive(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel) {
566-
_gridArrive(_team, _gridSyncState,
567-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
573+
nvshmemGridArrive(_coop, _team, _gridSyncState,
574+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
568575
}
569576

570577
FLAGCX_DEVICE_INLINE_DECORATOR void
571578
wait(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel) {
572-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
579+
nvshmemGridWait(_coop, _team, _gridSyncState,
580+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
573581
}
574582

575583
FLAGCX_DEVICE_INLINE_DECORATOR void
576584
sync(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel) {
577-
_gridArrive(_team, _gridSyncState,
578-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
579-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
585+
nvshmemGridArrive(_coop, _team, _gridSyncState,
586+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
587+
nvshmemGridWait(_coop, _team, _gridSyncState,
588+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
580589
}
581590
};
582591

@@ -607,71 +616,27 @@ struct Barrier<NvshmemBackend, flagcxTeamTagInter, Coop> {
607616
_gridSyncState((volatile uint64_t *)(dc.gridSyncState +
608617
2 * FLAGCX_DEVICE_CTA_COUNT)) {}
609618

610-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridArrive(nvshmem_team_t team,
611-
volatile uint64_t *arrive,
612-
volatile uint64_t *release) {
613-
#ifdef __CUDACC__
614-
int numBlocks = FLAGCX_GRID_DIM_X;
615-
_coop.sync();
616-
if (_coop.threadRank() == 0) {
617-
arrive[FLAGCX_BLOCK_IDX_X]++;
618-
}
619-
if (FLAGCX_BLOCK_IDX_X == 0) {
620-
_coop.sync();
621-
uint64_t expected = arrive[0];
622-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
623-
if (i == 0)
624-
continue;
625-
while (arrive[i] < expected) {
626-
}
627-
}
628-
_coop.sync();
629-
}
630-
#endif
631-
}
632-
633-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridWait(nvshmem_team_t team,
634-
volatile uint64_t *arrive,
635-
volatile uint64_t *release) {
636-
#ifdef __CUDACC__
637-
int numBlocks = FLAGCX_GRID_DIM_X;
638-
if (FLAGCX_BLOCK_IDX_X == 0) {
639-
_coop.sync();
640-
nvshmemx_barrier_block(team);
641-
_coop.sync();
642-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
643-
release[i]++;
644-
}
645-
} else {
646-
if (_coop.threadRank() == 0) {
647-
uint64_t cur = release[FLAGCX_BLOCK_IDX_X];
648-
while (release[FLAGCX_BLOCK_IDX_X] == cur) {
649-
}
650-
}
651-
}
652-
_coop.sync();
653-
#endif
654-
}
655-
656619
FLAGCX_DEVICE_INLINE_DECORATOR void
657620
arrive(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
658621
flagcxDevNetFenceLevel = flagcxDevNetFenceLevel::Relaxed) {
659-
_gridArrive(_team, _gridSyncState,
660-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
622+
nvshmemGridArrive(_coop, _team, _gridSyncState,
623+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
661624
}
662625

663626
FLAGCX_DEVICE_INLINE_DECORATOR void
664627
wait(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
665628
flagcxDevNetFenceLevel = flagcxDevNetFenceLevel::Relaxed) {
666-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
629+
nvshmemGridWait(_coop, _team, _gridSyncState,
630+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
667631
}
668632

669633
FLAGCX_DEVICE_INLINE_DECORATOR void
670634
sync(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
671635
flagcxDevNetFenceLevel fence = flagcxDevNetFenceLevel::Relaxed) {
672-
_gridArrive(_team, _gridSyncState,
673-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
674-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
636+
nvshmemGridArrive(_coop, _team, _gridSyncState,
637+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
638+
nvshmemGridWait(_coop, _team, _gridSyncState,
639+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
675640
}
676641
};
677642

@@ -715,71 +680,27 @@ struct Barrier<NvshmemBackend, flagcxTeamTagWorld, Coop> {
715680
_gridSyncState((volatile uint64_t *)(dc.gridSyncState +
716681
2 * FLAGCX_DEVICE_CTA_COUNT)) {}
717682

718-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridArrive(nvshmem_team_t team,
719-
volatile uint64_t *arrive,
720-
volatile uint64_t *release) {
721-
#ifdef __CUDACC__
722-
int numBlocks = FLAGCX_GRID_DIM_X;
723-
_coop.sync();
724-
if (_coop.threadRank() == 0) {
725-
arrive[FLAGCX_BLOCK_IDX_X]++;
726-
}
727-
if (FLAGCX_BLOCK_IDX_X == 0) {
728-
_coop.sync();
729-
uint64_t expected = arrive[0];
730-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
731-
if (i == 0)
732-
continue;
733-
while (arrive[i] < expected) {
734-
}
735-
}
736-
_coop.sync();
737-
}
738-
#endif
739-
}
740-
741-
FLAGCX_DEVICE_INLINE_DECORATOR void _gridWait(nvshmem_team_t team,
742-
volatile uint64_t *arrive,
743-
volatile uint64_t *release) {
744-
#ifdef __CUDACC__
745-
int numBlocks = FLAGCX_GRID_DIM_X;
746-
if (FLAGCX_BLOCK_IDX_X == 0) {
747-
_coop.sync();
748-
nvshmemx_barrier_block(team);
749-
_coop.sync();
750-
for (int i = _coop.threadRank(); i < numBlocks; i += FLAGCX_BLOCK_DIM_X) {
751-
release[i]++;
752-
}
753-
} else {
754-
if (_coop.threadRank() == 0) {
755-
uint64_t cur = release[FLAGCX_BLOCK_IDX_X];
756-
while (release[FLAGCX_BLOCK_IDX_X] == cur) {
757-
}
758-
}
759-
}
760-
_coop.sync();
761-
#endif
762-
}
763-
764683
FLAGCX_DEVICE_INLINE_DECORATOR void
765684
arrive(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
766685
flagcxDevNetFenceLevel fence = flagcxDevNetFenceLevel::Relaxed) {
767-
_gridArrive(_team, _gridSyncState,
768-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
686+
nvshmemGridArrive(_coop, _team, _gridSyncState,
687+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
769688
}
770689

771690
FLAGCX_DEVICE_INLINE_DECORATOR void
772691
wait(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
773692
flagcxDevNetFenceLevel fence = flagcxDevNetFenceLevel::Relaxed) {
774-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
693+
nvshmemGridWait(_coop, _team, _gridSyncState,
694+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
775695
}
776696

777697
FLAGCX_DEVICE_INLINE_DECORATOR void
778698
sync(flagcxDeviceMemoryOrder_t order = flagcxDeviceMemoryOrderAcqRel,
779699
flagcxDevNetFenceLevel fence = flagcxDevNetFenceLevel::Relaxed) {
780-
_gridArrive(_team, _gridSyncState,
781-
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
782-
_gridWait(_team, _gridSyncState, _gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
700+
nvshmemGridArrive(_coop, _team, _gridSyncState,
701+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
702+
nvshmemGridWait(_coop, _team, _gridSyncState,
703+
_gridSyncState + FLAGCX_DEVICE_CTA_COUNT);
783704
}
784705
};
785706

0 commit comments

Comments
 (0)