@@ -73,6 +73,41 @@ flagcxMakeTeamFromKind(const flagcxDevComm &comm, flagcxTeamKind_t kind) {
7373 }
7474}
7575
76+ /* ================================================================
77+ * Internal helper: grid-wide barrier (sense-reversing)
78+ *
79+ * Synchronizes all blocks on the same GPU. Uses a monotonic
80+ * sense-flip protocol so it's safely reusable across iterations.
81+ * gridSyncState[0] = arrive counter, gridSyncState[1] = sense.
82+ * ================================================================ */
83+
84+ static FLAGCX_DEVICE_INLINE_DECORATOR void
85+ flagcxGridSync (unsigned int * gridSyncState ) {
86+ __syncthreads ();
87+ if (threadIdx .x == 0 ) {
88+ unsigned int curSense = * (volatile unsigned int * )& gridSyncState [1 ];
89+ unsigned int arrived = atomicAdd (& gridSyncState [0 ], 1u ) + 1 ;
90+ if (arrived == gridDim .x ) {
91+ // Last block: reset counter, flip sense to release others
92+ atomicExch (& gridSyncState [0 ], 0u );
93+ __threadfence ();
94+ atomicExch (& gridSyncState [1 ], 1u - curSense );
95+ } else {
96+ // Spin until sense flips (all blocks arrived)
97+ #ifndef NDEBUG
98+ unsigned int __spins = 0 ;
99+ #endif
100+ while (* (volatile unsigned int * )& gridSyncState [1 ] == curSense ) {
101+ #ifndef NDEBUG
102+ if (++ __spins >= 100000000u )
103+ __trap (); // grid barrier timeout — likely a hang
104+ #endif
105+ }
106+ }
107+ }
108+ __syncthreads ();
109+ }
110+
76111/* ================================================================
77112 * Category 2: Scalar Cooperative Group (6)
78113 * ================================================================ */
@@ -231,6 +266,9 @@ flagcxIntraBarrierSyncS(const void *commOpaque, flagcxCoopKind_t coopKind,
231266 flagcxDevBarrier < flagcxTeamTagIntra , flagcxCoopAny > bar (coop , * comm , team ,
232267 index , multimem );
233268 bar .sync (order );
269+ if (comm -> _gridBarrierState ) {
270+ flagcxGridSync (comm -> _gridBarrierState );
271+ }
234272}
235273
236274/* ================================================================
@@ -274,6 +312,9 @@ flagcxInterBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
274312 flagcxDevBarrier < flagcxTeamTagInter , flagcxCoopAny > bar (coop , * net , team ,
275313 index );
276314 bar .sync (order , fence );
315+ if (net -> _gridBarrierState ) {
316+ flagcxGridSync (net -> _gridBarrierState );
317+ }
277318}
278319
279320/* ================================================================
@@ -317,6 +358,9 @@ flagcxWorldBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
317358 flagcxDevBarrier < flagcxTeamTagWorld , flagcxCoopAny > bar (
318359 coop , flagcxTeamTagWorld {}, * net , index , multimem );
319360 bar .sync (order , fence );
361+ if (net -> _gridBarrierState ) {
362+ flagcxGridSync (net -> _gridBarrierState );
363+ }
320364}
321365
322366/* ================================================================
0 commit comments