fix(device-ir): pre-allocate net contexts in DevComm, expand scalar IR tests

Author: MC952-arch

bindings/ir/flagcx_device_scalar_ir_impl.h

@@ -1,13 +1,11 @@
 /*************************************************************************
  * Copyright (c) 2026 BAAI. All rights reserved.
  *
- * FlagCX Device Scalar IR API — Implementation (Step 2: Coop + Team).
+ * FlagCX Device Scalar IR API — Implementation.
  *
- * This file implements the scalar cooperative group and team functions.
+ * This file implements the scalar (S-suffixed) IR functions.
  * It is included by the bitcode compilation unit alongside the existing
  * flagcx_device_wrapper_impl.h.
- *
- * Remaining categories (barrier, net) will be added in subsequent steps.
  ************************************************************************/
 #ifndef FLAGCX_DEVICE_SCALAR_IR_IMPL_H_
 #define FLAGCX_DEVICE_SCALAR_IR_IMPL_H_
@@ -332,11 +330,12 @@ flagcxWorldBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
 FLAGCX_IR_EXTERN_C FLAGCX_DEVICE_INLINE_DECORATOR const void *
 flagcxDevNetGetFromCommS(const void *commOpaque, int idx) {
   const flagcxDevComm *comm = (const flagcxDevComm *)commOpaque;
-  // When always-inlined, the static local is promoted into the caller's
-  // scope by LLVM. This provides stable storage for the returned pointer.
-  static flagcxDevNet storage;
-  ::new (&storage) flagcxDevNet(*comm, idx);
-  return &storage;
+  // Return pointer into pre-allocated device array (built by
+  // flagcxDevCommGetDevicePtr). Each entry is a fully-constructed
+  // flagcxDevNet for that context index — no per-call construction.
+  const flagcxDevNet *nets = (const flagcxDevNet *)comm->_netContexts;
+  int safeIdx = (comm->_contextCount > 0) ? (idx % comm->_contextCount) : 0;
+  return &nets[safeIdx];
 }
 
 /* ================================================================

flagcx/adaptor/flagcx_device.cc

@@ -1056,6 +1056,9 @@ extern "C" flagcxResult_t flagcxDevCommDestroy(flagcxComm_t comm,
   }
 
   // Device pointer cache cleanup
+  if (devComm->cachedNetContextsPtr) {
+    flagcxCommDeferFree(comm, devComm->cachedNetContextsPtr, flagcxMemDevice);
+  }
   if (devComm->cachedDevicePtr) {
     flagcxCommDeferFree(comm, devComm->cachedDevicePtr, flagcxMemDevice);
   }
@@ -1268,6 +1271,34 @@ extern "C" flagcxResult_t flagcxDevCommGetDevicePtr(flagcxDevComm_t devComm,
   // Construct value struct on host stack
   flagcxDevComm hostCopy(*devComm);
 
+  // Pre-allocate net contexts (one flagcxDevNet per context) in device memory
+  void *netDevPtr = nullptr;
+  if (hostCopy._contextCount > 0) {
+    size_t netArraySize = hostCopy._contextCount * sizeof(flagcxDevNet);
+    // Build array on host heap
+    flagcxDevNet *hostNets = (flagcxDevNet *)malloc(netArraySize);
+    if (hostNets) {
+      for (int i = 0; i < hostCopy._contextCount; i++) {
+        ::new (&hostNets[i]) flagcxDevNet(hostCopy, i);
+      }
+      // Copy to device
+      flagcxResult_t netRes = deviceAdaptor->deviceMalloc(
+          &netDevPtr, netArraySize, flagcxMemDevice, NULL);
+      if (netRes == flagcxSuccess) {
+        netRes =
+            deviceAdaptor->deviceMemcpy(netDevPtr, hostNets, netArraySize,
+                                        flagcxMemcpyHostToDevice, NULL, NULL);
+      }
+      if (netRes != flagcxSuccess) {
+        if (netDevPtr)
+          deviceAdaptor->deviceFree(netDevPtr, flagcxMemDevice, NULL);
+        netDevPtr = nullptr;
+      }
+      free(hostNets);
+    }
+    hostCopy._netContexts = netDevPtr;
+  }
+
   // Allocate device memory and copy
   void *dPtr = nullptr;
   flagcxResult_t res = flagcxSuccess;
@@ -1280,12 +1311,16 @@ extern "C" flagcxResult_t flagcxDevCommGetDevicePtr(flagcxDevComm_t devComm,
       res, fail);
 
   devComm->cachedDevicePtr = dPtr;
+  devComm->cachedNetContextsPtr = netDevPtr;
   *devPtr = dPtr;
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
   return flagcxSuccess;
 
 fail:
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
+  if (netDevPtr) {
+    deviceAdaptor->deviceFree(netDevPtr, flagcxMemDevice, NULL);
+  }
   if (dPtr) {
     deviceAdaptor->deviceFree(dPtr, flagcxMemDevice, NULL);
   }
@@ -1298,9 +1333,14 @@ extern "C" flagcxResult_t flagcxDevCommFreeDevicePtr(flagcxDevComm_t devComm) {
 
   pthread_mutex_lock(&devComm->cachedPtrMutex);
   void *ptr = devComm->cachedDevicePtr;
+  void *netPtr = devComm->cachedNetContextsPtr;
   devComm->cachedDevicePtr = nullptr;
+  devComm->cachedNetContextsPtr = nullptr;
   pthread_mutex_unlock(&devComm->cachedPtrMutex);
 
+  if (netPtr) {
+    FLAGCXCHECK(deviceAdaptor->deviceFree(netPtr, flagcxMemDevice, NULL));
+  }
   if (ptr) {
     FLAGCXCHECK(deviceAdaptor->deviceFree(ptr, flagcxMemDevice, NULL));
   }

flagcx/adaptor/include/device_api/flagcx_device_core.h

@@ -71,14 +71,20 @@ struct flagcxDevComm {
   int _contextCount;
   int _nInterPeers;
 
+  // Pre-allocated net contexts (one per _contextCount, device memory).
+  // Set by flagcxDevCommGetDevicePtr; nullptr when contextCount == 0.
+  // Actually flagcxDevNet[] but kept as void* for C/opaque linkage.
+  void *_netContexts;
+
   FLAGCX_HOST_DEVICE_INLINE flagcxDevComm()
       : _commBase(), _signalCount(0), _counterCount(0), _contextCount(0),
-        _nInterPeers(0) {}
+        _nInterPeers(0), _netContexts(nullptr) {}
 
 #ifndef __clang_llvm_bitcode_lib__
   FLAGCX_HOST_DEVICE_INLINE flagcxDevComm(const flagcxDevCommInternal &di)
       : _signalCount(di.signalCount), _counterCount(di.counterCount),
-        _contextCount(di.contextCount), _nInterPeers(di.nInterPeers) {
+        _contextCount(di.contextCount), _nInterPeers(di.nInterPeers),
+        _netContexts(nullptr) {
     if (di.devComm) {
       _commBase = *(typename DeviceAPI::Comm *)di.devComm;
     } else {

flagcx/adaptor/include/device_api/flagcx_device_internal.h

@@ -99,7 +99,8 @@ struct flagcxDevCommInternal {
   flagcxInnerDevComm_t devComm; // Typed vendor handle (per-adaptor struct)
 
   // ---- Device pointer cache (for Triton integration) ----
-  void *cachedDevicePtr; // Lazily allocated by flagcxDevCommGetDevicePtr
+  void *cachedDevicePtr;      // Lazily allocated by flagcxDevCommGetDevicePtr
+  void *cachedNetContextsPtr; // Device memory for pre-allocated flagcxDevNet[]
   pthread_mutex_t cachedPtrMutex; // Protects lazy init of cachedDevicePtr
 };