Fix S20 multi-context deadlock: enable all blocks to participate in signaling

Author: MC952-arch

bindings/ir/flagcx_device_scalar_ir_impl.h

@@ -266,9 +266,6 @@ flagcxIntraBarrierSyncS(const void *commOpaque, flagcxCoopKind_t coopKind,
   flagcxDevBarrier<flagcxTeamTagIntra, flagcxCoopAny> bar(coop, *comm, team,
                                                           index, multimem);
   bar.sync(order);
-  if (comm->_gridBarrierState) {
-    flagcxGridSync(comm->_gridBarrierState);
-  }
 }
 
 /* ================================================================
@@ -312,9 +309,6 @@ flagcxInterBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
   flagcxDevBarrier<flagcxTeamTagInter, flagcxCoopAny> bar(coop, *net, team,
                                                           index);
   bar.sync(order, fence);
-  if (net->_gridBarrierState) {
-    flagcxGridSync(net->_gridBarrierState);
-  }
 }
 
 /* ================================================================
@@ -358,9 +352,6 @@ flagcxWorldBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
   flagcxDevBarrier<flagcxTeamTagWorld, flagcxCoopAny> bar(
       coop, flagcxTeamTagWorld{}, *net, index, multimem);
   bar.sync(order, fence);
-  if (net->_gridBarrierState) {
-    flagcxGridSync(net->_gridBarrierState);
-  }
 }
 
 /* ================================================================

flagcx/adaptor/device_api/default_dev_api_backend.cc

@@ -897,10 +897,23 @@ static flagcxResult_t defaultDevApiCommGetDevicePtr(flagcxDevComm_t devComm,
   flagcxDevComm hostCopy(*devComm);
   hostCopy._netContexts = nullptr;
 
-  // Step 1: Copy flagcxDevComm to device
+  // Step 1: Allocate grid sync state (2 x unsigned int, zero-initialized)
   void *dPtr = nullptr;
   void *netDevPtr = nullptr;
+  void *gridSyncPtr = nullptr;
   flagcxResult_t res = flagcxSuccess;
+  {
+    size_t gsSize = 2 * sizeof(unsigned int);
+    FLAGCXCHECKGOTO(deviceAdaptor->deviceMalloc(&gridSyncPtr, gsSize,
+                                                flagcxMemDevice, NULL),
+                    res, fail);
+    FLAGCXCHECKGOTO(deviceAdaptor->deviceMemset(gridSyncPtr, 0, gsSize,
+                                                flagcxMemDevice, NULL),
+                    res, fail);
+  }
+  hostCopy._gridBarrierState = (unsigned int *)gridSyncPtr;
+
+  // Step 2: Copy flagcxDevComm to device
   FLAGCXCHECKGOTO(deviceAdaptor->deviceMalloc(&dPtr, sizeof(flagcxDevComm),
                                               flagcxMemDevice, NULL),
                   res, fail);
@@ -909,7 +922,7 @@ static flagcxResult_t defaultDevApiCommGetDevicePtr(flagcxDevComm_t devComm,
                                   flagcxMemcpyHostToDevice, NULL, NULL),
       res, fail);
 
-  // Step 2: Allocate + construct net array on device
+  // Step 3: Allocate + construct net array on device
   if (hostCopy._contextCount > 0 && flagcxDevNetSizeOf() > 0) {
     size_t netArraySize = hostCopy._contextCount * flagcxDevNetSizeOf();
     FLAGCXCHECKGOTO(deviceAdaptor->deviceMalloc(&netDevPtr, netArraySize,
@@ -927,12 +940,16 @@ static flagcxResult_t defaultDevApiCommGetDevicePtr(flagcxDevComm_t devComm,
 
   devComm->cachedDevicePtr = dPtr;
   devComm->cachedNetContextsPtr = netDevPtr;
+  devComm->cachedGridBarrierPtr = gridSyncPtr;
   *devPtr = dPtr;
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
   return flagcxSuccess;
 
 fail:
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
+  if (gridSyncPtr) {
+    deviceAdaptor->deviceFree(gridSyncPtr, flagcxMemDevice, NULL);
+  }
   if (netDevPtr) {
     deviceAdaptor->deviceFree(netDevPtr, flagcxMemDevice, NULL);
   }
@@ -949,10 +966,15 @@ static flagcxResult_t defaultDevApiCommFreeDevicePtr(flagcxDevComm_t devComm) {
   pthread_mutex_lock(&devComm->cachedPtrMutex);
   void *ptr = devComm->cachedDevicePtr;
   void *netPtr = devComm->cachedNetContextsPtr;
+  void *gridPtr = devComm->cachedGridBarrierPtr;
   devComm->cachedDevicePtr = nullptr;
   devComm->cachedNetContextsPtr = nullptr;
+  devComm->cachedGridBarrierPtr = nullptr;
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
 
+  if (gridPtr) {
+    FLAGCXCHECK(deviceAdaptor->deviceFree(gridPtr, flagcxMemDevice, NULL));
+  }
   if (netPtr) {
     FLAGCXCHECK(deviceAdaptor->deviceFree(netPtr, flagcxMemDevice, NULL));
   }