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_
@@ -323,20 +321,24 @@ flagcxWorldBarrierSyncS(const void *netOpaque, flagcxCoopKind_t coopKind,
 /* ================================================================
  * Category 9: Net — Obtain Pre-Allocated Transport (1)
  *
- * Constructs a flagcxDevNet on the stack and returns a pointer.
- * NOTE: This returns a stack pointer — the caller must use it
- * immediately within the same kernel invocation scope. The returned
- * pointer is valid for the duration of the calling function.
+ * Returns a pointer into the comm-owned pre-allocated flagcxDevNet[]
+ * array (device-resident, built by flagcxDevCommGetDevicePtr).
+ * Lifetime: valid as long as the DevComm device pointer is alive
+ * (freed on flagcxDevCommDestroy or flagcxDevCommFreeDevicePtr).
+ * The returned pointer is read-only and safe to use from any thread.
  * ================================================================ */
 
 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;
+  if (!nets || comm->_contextCount <= 0)
+    return (const void *)0;
+  int safeIdx = (int)((unsigned)idx % (unsigned)comm->_contextCount);
+  return &nets[safeIdx];
 }
 
 /* ================================================================

flagcx/adaptor/flagcx_device.cc

@@ -25,6 +25,12 @@
 #include <new>
 #include <sched.h> // sched_yield
 
+// Host-visible helpers implemented in device_api_host_helpers.cu (compiled by
+// nvcc)
+extern "C" size_t flagcxDevNetSizeOf();
+extern "C" void flagcxDevNetConstructArray(void *dst, const void *comm,
+                                           int count);
+
 // ==========================================================================
 // Shared: IPC peer pointer exchange (used by both tiers)
 // ==========================================================================
@@ -1056,6 +1062,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 +1277,32 @@ 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 * flagcxDevNetSizeOf();
+    // Build array on host heap
+    void *hostNets = malloc(netArraySize);
+    if (hostNets) {
+      flagcxDevNetConstructArray(hostNets, &hostCopy, hostCopy._contextCount);
+      // 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 +1315,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 +1337,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,8 +99,10 @@ struct flagcxDevCommInternal {
   flagcxInnerDevComm_t devComm; // Typed vendor handle (per-adaptor struct)
 
   // ---- Device pointer cache (for Triton integration) ----
-  void *cachedDevicePtr; // Lazily allocated by flagcxDevCommGetDevicePtr
-  pthread_mutex_t cachedPtrMutex; // Protects lazy init of cachedDevicePtr
+  void *cachedDevicePtr;      // Lazily allocated by flagcxDevCommGetDevicePtr
+  void *cachedNetContextsPtr; // Device memory for pre-allocated flagcxDevNet[]
+  pthread_mutex_t cachedPtrMutex; // Protects lazy init of cachedDevicePtr and
+                                  // cachedNetContextsPtr
 };
 
 // ============================================================

flagcx/kernels/device_api_host_helpers.cu

@@ -0,0 +1,21 @@
+/*************************************************************************
+ * Copyright (c) 2026 BAAI. All rights reserved.
+ *
+ * Host-callable helpers that require device-compiler visibility of
+ * device-only types (e.g. flagcxDevNet). Compiled by nvcc/hipcc,
+ * linked into libflagcx.so, called from flagcx_device.cc via extern "C".
+ ************************************************************************/
+
+#include "device_api/flagcx_device_core.h"
+#include <new>
+
+extern "C" size_t flagcxDevNetSizeOf() { return sizeof(flagcxDevNet); }
+
+extern "C" void flagcxDevNetConstructArray(void *dst, const void *comm,
+                                           int count) {
+  const flagcxDevComm *devComm = (const flagcxDevComm *)comm;
+  flagcxDevNet *nets = (flagcxDevNet *)dst;
+  for (int i = 0; i < count; i++) {
+    ::new (&nets[i]) flagcxDevNet(*devComm, i);
+  }
+}