[UIL] Add scalar IR API, on-device net construction, and per-peer staging fix

bindings/ir/public_symbols.txt

@@ -1,68 +1,124 @@
-flagcxCoopAnyInitBlock
-flagcxCoopAnyInitLanes
-flagcxCoopAnyInitThread
-flagcxCoopAnyInitTileSpan
-flagcxCoopAnyInitWarp
+flagcxCoopAnyInitBlockC
+flagcxCoopAnyInitLanesC
+flagcxCoopAnyInitThreadC
+flagcxCoopAnyInitTileSpanC
+flagcxCoopAnyInitWarpC
 flagcxCoopSizeC
+flagcxCoopSizeExS
+flagcxCoopSizeS
 flagcxCoopSyncC
+flagcxCoopSyncExS
+flagcxCoopSyncS
 flagcxCoopThreadRankC
+flagcxCoopThreadRankExS
+flagcxCoopThreadRankS
 flagcxDataTypeSizeDevice
 flagcxDevCommGetIntraRank
 flagcxDevCommGetIntraSize
 flagcxDevCommGetRank
 flagcxDevCommGetSize
-flagcxDevNetFlush
-flagcxDevNetGet
+flagcxDevNetFlushC
+flagcxDevNetFlushS
+flagcxDevNetGetC
+flagcxDevNetGetFromCommS
+flagcxDevNetGetS
 flagcxDevNetIncreaseSignalShadow
 flagcxDevNetInitC
-flagcxDevNetPut
-flagcxDevNetPutValue
-flagcxDevNetPutValue_RCtrInc
-flagcxDevNetPutValue_RSigAdd
-flagcxDevNetPutValue_RSigInc
-flagcxDevNetPut_LCtrInc
-flagcxDevNetPut_LSigAdd
-flagcxDevNetPut_LSigInc
-flagcxDevNetPut_RCtrInc
-flagcxDevNetPut_RCtrInc_LCtrInc
-flagcxDevNetPut_RCtrInc_LSigAdd
-flagcxDevNetPut_RCtrInc_LSigInc
-flagcxDevNetPut_RSigAdd
-flagcxDevNetPut_RSigAdd_LCtrInc
-flagcxDevNetPut_RSigAdd_LSigAdd
-flagcxDevNetPut_RSigAdd_LSigInc
-flagcxDevNetPut_RSigInc
-flagcxDevNetPut_RSigInc_LCtrInc
-flagcxDevNetPut_RSigInc_LSigAdd
-flagcxDevNetPut_RSigInc_LSigInc
-flagcxDevNetReadCounter
-flagcxDevNetReadSignal
-flagcxDevNetRecv
+flagcxDevNetPutC
+flagcxDevNetPutS
+flagcxDevNetPutS_LCtrInc
+flagcxDevNetPutS_LSigAdd
+flagcxDevNetPutS_LSigInc
+flagcxDevNetPutS_RCtrInc
+flagcxDevNetPutS_RCtrInc_LCtrInc
+flagcxDevNetPutS_RCtrInc_LSigAdd
+flagcxDevNetPutS_RCtrInc_LSigInc
+flagcxDevNetPutS_RSigAdd
+flagcxDevNetPutS_RSigAdd_LCtrInc
+flagcxDevNetPutS_RSigAdd_LSigAdd
+flagcxDevNetPutS_RSigAdd_LSigInc
+flagcxDevNetPutS_RSigInc
+flagcxDevNetPutS_RSigInc_LCtrInc
+flagcxDevNetPutS_RSigInc_LSigAdd
+flagcxDevNetPutS_RSigInc_LSigInc
+flagcxDevNetPutValueC
+flagcxDevNetPutValueS
+flagcxDevNetPutValueS_RCtrInc
+flagcxDevNetPutValueS_RSigAdd
+flagcxDevNetPutValueS_RSigInc
+flagcxDevNetPutValue_RCtrIncC
+flagcxDevNetPutValue_RSigAddC
+flagcxDevNetPutValue_RSigIncC
+flagcxDevNetPut_LCtrIncC
+flagcxDevNetPut_LSigAddC
+flagcxDevNetPut_LSigIncC
+flagcxDevNetPut_RCtrIncC
+flagcxDevNetPut_RCtrInc_LCtrIncC
+flagcxDevNetPut_RCtrInc_LSigAddC
+flagcxDevNetPut_RCtrInc_LSigIncC
+flagcxDevNetPut_RSigAddC
+flagcxDevNetPut_RSigAdd_LCtrIncC
+flagcxDevNetPut_RSigAdd_LSigAddC
+flagcxDevNetPut_RSigAdd_LSigIncC
+flagcxDevNetPut_RSigIncC
+flagcxDevNetPut_RSigInc_LCtrIncC
+flagcxDevNetPut_RSigInc_LSigAddC
+flagcxDevNetPut_RSigInc_LSigIncC
+flagcxDevNetReadCounterC
+flagcxDevNetReadCounterS
+flagcxDevNetReadSignalC
+flagcxDevNetReadSignalS
+flagcxDevNetRecvC
+flagcxDevNetRecvS
 flagcxDevNetResetCounter
 flagcxDevNetResetSignal
-flagcxDevNetSend
-flagcxDevNetSignalCtrInc
-flagcxDevNetSignalSigAdd
-flagcxDevNetSignalSigInc
-flagcxDevNetTerm
-flagcxDevNetWait
-flagcxDevNetWaitCounter
-flagcxDevNetWaitSignal
-flagcxDevNetWaitSignalMeetShadow
+flagcxDevNetSendC
+flagcxDevNetSendS
+flagcxDevNetSignalCtrIncC
+flagcxDevNetSignalCtrIncS
+flagcxDevNetSignalSigAddC
+flagcxDevNetSignalSigAddS
+flagcxDevNetSignalSigIncC
+flagcxDevNetSignalSigIncS
+flagcxDevNetTermC
+flagcxDevNetTermS
+flagcxDevNetWaitC
+flagcxDevNetWaitCounterC
+flagcxDevNetWaitCounterS
+flagcxDevNetWaitS
+flagcxDevNetWaitSignalC
+flagcxDevNetWaitSignalMeetShadowC
+flagcxDevNetWaitSignalMeetShadowS
+flagcxDevNetWaitSignalS
 flagcxGetIntraPointerC
+flagcxGetIntraPointerS
 flagcxGetLocalPointerC
+flagcxGetLocalPointerS
 flagcxGetMulticastPointerC
+flagcxGetMulticastPointerS
 flagcxGetPeerPointerC
-flagcxGetTeamInter
-flagcxGetTeamIntra
-flagcxGetTeamWorld
-flagcxInterBarrierSessionInit
-flagcxInterBarrierSessionSync
-flagcxIntraBarrierSessionArrive
-flagcxIntraBarrierSessionInit
-flagcxIntraBarrierSessionSync
-flagcxIntraBarrierSessionWait
+flagcxGetPeerPointerS
+flagcxGetTeamInterC
+flagcxGetTeamIntraC
+flagcxGetTeamWorldC
+flagcxInterBarrierArriveS
+flagcxInterBarrierSessionInitC
+flagcxInterBarrierSessionSyncC
+flagcxInterBarrierSyncS
+flagcxInterBarrierWaitS
+flagcxIntraBarrierArriveS
+flagcxIntraBarrierSessionArriveC
+flagcxIntraBarrierSessionInitC
+flagcxIntraBarrierSessionSyncC
+flagcxIntraBarrierSessionWaitC
+flagcxIntraBarrierSyncS
+flagcxIntraBarrierWaitS
 flagcxTeamRankToIntraC
+flagcxTeamRankToIntraS
 flagcxTeamRankToWorldC
-flagcxWorldBarrierSessionInit
-flagcxWorldBarrierSessionSync
+flagcxTeamRankToWorldS
+flagcxWorldBarrierArriveS
+flagcxWorldBarrierSessionInitC
+flagcxWorldBarrierSessionSyncC
+flagcxWorldBarrierSyncS
+flagcxWorldBarrierWaitS

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

@@ -19,6 +19,9 @@
 #include "device_utils.h"
 #include "flagcx.h"
 
+// Scalar enum types for IR/Triton integration (shared with native path).
+#include "flagcx_device_enums.h"
+
 // Device traits — provides DeviceAPI with all type/function dispatch.
 #include "comm_traits.h"
 
@@ -68,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_enums.h

@@ -0,0 +1,44 @@
+/*************************************************************************
+ * Copyright (c) 2026 BAAI. All rights reserved.
+ *
+ * FlagCX Device API Enums — Scalar enum types for IR/Triton integration.
+ *
+ * These enums encode cooperative group kind and team kind as plain
+ * integers, enabling LLVM IR callers (e.g. Triton) to express these
+ * concepts without instantiating C++ structs.
+ *
+ * Safe to include from: CUDA device code, host code, LLVM bitcode builds.
+ ************************************************************************/
+
+#ifndef FLAGCX_DEVICE_ENUMS_H_
+#define FLAGCX_DEVICE_ENUMS_H_
+
+#include <stdint.h>
+
+/* ================================================================
+ * Cooperative Group Kind
+ *
+ * Identifies the cooperation scope for collective device operations.
+ * Used by the scalar IR API in place of struct-based cooperative groups.
+ * ================================================================ */
+typedef enum {
+  FLAGCX_COOP_THREAD = 0,    /* Single thread (no cooperation) */
+  FLAGCX_COOP_WARP = 1,      /* Full warp (FLAGCX_SIMT_WIDTH threads) */
+  FLAGCX_COOP_BLOCK = 2,     /* Entire CTA */
+  FLAGCX_COOP_TILE_SPAN = 3, /* Consecutive tile span (needs t0, nTiles, id) */
+  FLAGCX_COOP_LANES = 4,     /* Arbitrary lane bitmask */
+} flagcxCoopKind_t;
+
+/* ================================================================
+ * Team Kind
+ *
+ * Identifies the team scope within a communicator.
+ * Used by the scalar IR API in place of struct-based teams.
+ * ================================================================ */
+typedef enum {
+  FLAGCX_TEAM_INTRA = 0, /* Intra-node ranks */
+  FLAGCX_TEAM_INTER = 1, /* Inter-node representatives */
+  FLAGCX_TEAM_WORLD = 2, /* All ranks */
+} flagcxTeamKind_t;
+
+#endif /* FLAGCX_DEVICE_ENUMS_H_ */

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);
+  }
+}