[CRL] Deprecate flagcxHandlerGroup, separate device handle/uniqueId/comm lifecycle (#480)

Author: MC952-arch

flagcx/adaptor/include/device_api/sunrise_comm_traits.h

@@ -7,13 +7,6 @@
 #ifndef FLAGCX_SUNRISE_COMM_TRAITS_H_
 #define FLAGCX_SUNRISE_COMM_TRAITS_H_
 
-// Sunrise/PCCL inner-type stubs: never populated at runtime on the
-// sunrise path; only here to satisfy public type signatures.
-struct flagcxInnerWindow {
-  int winFlags;
-};
-struct flagcxInnerDevComm {};
-
 // Sunrise default backend: reuse Default<DefaultPlatform> (IPC barriers +
 // FIFO one-sided). PCCL/PTPU already provide collectives, so FlagCX needs
 // no SIMT kernels of its own.

flagcx/adaptor/include/sunrise_adaptor.h

@@ -7,6 +7,12 @@
 #include "tang.h"
 #include "tang_runtime.h"
 #include <map>
+
+struct flagcxInnerWindow {
+  int winFlags;
+};
+struct flagcxInnerDevComm {};
+
 struct flagcxInnerComm {
   pcclComm_t base;
 };

flagcx/flagcx.cc

@@ -156,44 +156,61 @@ bool useHeteroComm() {
   return false;
 }
 
-flagcxResult_t flagcxHandleInit(flagcxHandlerGroup_t *handler) {
+flagcxResult_t flagcxDeviceHandleInit(flagcxDeviceHandle_t *devHandle) {
+  if (devHandle == NULL) {
+    WARN("flagcxDeviceHandleInit: devHandle is NULL");
+    return flagcxInvalidArgument;
+  }
   flagcxResult_t res = flagcxSuccess;
   flagcxDeviceAdaptorPluginInit();
   flagcxCCLAdaptorPluginInit();
+  (*devHandle) = NULL;
+  FLAGCXCHECKGOTO(flagcxCalloc(devHandle, 1), res, fail);
+  **devHandle = globalDeviceHandle;
+  return flagcxSuccess;
+
+fail:
+  if (*devHandle) {
+    free(*devHandle);
+    *devHandle = NULL;
+  }
+  flagcxCCLAdaptorPluginFinalize();
+  flagcxDeviceAdaptorPluginFinalize();
+  return res;
+}
+
+flagcxResult_t flagcxDeviceHandleFree(flagcxDeviceHandle_t devHandle) {
+  if (devHandle == NULL)
+    return flagcxSuccess;
+  free(devHandle);
+  flagcxCCLAdaptorPluginFinalize();
+  flagcxDeviceAdaptorPluginFinalize();
+  return flagcxSuccess;
+}
+
+flagcxResult_t flagcxHandleInit(flagcxHandlerGroup_t *handler) {
+  flagcxResult_t res = flagcxSuccess;
   (*handler) = NULL;
   FLAGCXCHECKGOTO(flagcxCalloc(handler, 1), res, fail);
-  FLAGCXCHECKGOTO(flagcxCalloc(&(*handler)->uniqueId, 1), res, fail);
-  // comm left NULL — allocated by flagcxCommInitRank
-  FLAGCXCHECKGOTO(flagcxCalloc(&(*handler)->devHandle, 1), res, fail);
-  *(*handler)->devHandle = globalDeviceHandle;
+  FLAGCXCHECKGOTO(flagcxDeviceHandleInit(&(*handler)->devHandle), res, fail);
   return flagcxSuccess;
 
 fail:
   if (*handler) {
-    free((*handler)->uniqueId);
-    free((*handler)->devHandle);
     free(*handler);
     *handler = NULL;
   }
-  flagcxCCLAdaptorPluginFinalize();
-  flagcxDeviceAdaptorPluginFinalize();
   return res;
 }
 
 flagcxResult_t flagcxHandleFree(flagcxHandlerGroup_t handler) {
   if (handler != NULL) {
-    if (handler->uniqueId)
-      free(handler->uniqueId);
-    if (handler->devHandle)
-      free(handler->devHandle);
+    flagcxDeviceHandleFree(handler->devHandle);
+    handler->devHandle = NULL;
     handler->uniqueId = NULL;
     handler->comm = NULL;
-    handler->devHandle = NULL;
     free(handler);
-    handler = NULL;
   }
-  flagcxCCLAdaptorPluginFinalize();
-  flagcxDeviceAdaptorPluginFinalize();
   return flagcxSuccess;
 }
 
@@ -1180,10 +1197,10 @@ flagcxResult_t flagcxGetVersion(int *version) {
   return flagcxHeteroGetVersion(version);
 }
 
-flagcxResult_t flagcxGetUniqueId(flagcxUniqueId_t *uniqueId) {
-  // Allocate if caller passed a NULL pointer
-  if (*uniqueId == nullptr) {
-    FLAGCXCHECK(flagcxCalloc(uniqueId, 1));
+flagcxResult_t flagcxGetUniqueId(flagcxUniqueId_t uniqueId) {
+  if (uniqueId == NULL) {
+    WARN("flagcxGetUniqueId: uniqueId is NULL");
+    return flagcxInvalidArgument;
   }
 
   // Init bootstrap net
@@ -1194,9 +1211,9 @@ flagcxResult_t flagcxGetUniqueId(flagcxUniqueId_t *uniqueId) {
   FLAGCXCHECK(bootstrapGetUniqueId(&handle));
   // flagcxUniqueId and bootstrapHandle don't have the same size and alignment
   // reset to 0 to avoid undefined data
-  memset((void *)*uniqueId, 0, sizeof(**uniqueId));
+  memset((void *)uniqueId, 0, sizeof(*uniqueId));
   // copy to avoid alignment mismatch
-  memcpy((void *)*uniqueId, &handle, sizeof(handle));
+  memcpy((void *)uniqueId, &handle, sizeof(handle));
   return flagcxSuccess;
 }
 
@@ -1523,6 +1540,14 @@ flagcxResult_t flagcxCommInitRank(flagcxComm_t *comm, int nranks,
     WARN("Invalid rank requested : %d/%d", rank, nranks);
     return flagcxInvalidArgument;
   }
+  if (commId == NULL || comm == NULL) {
+    WARN("flagcxCommInitRank: commId or comm is NULL");
+    return flagcxInvalidArgument;
+  }
+
+  // Ensure device/CCL plugins are loaded (idempotent, ref-counted)
+  flagcxDeviceAdaptorPluginInit();
+  flagcxCCLAdaptorPluginInit();
 
   (*comm) = NULL;
   flagcxCalloc(comm, 1);
@@ -1729,7 +1754,8 @@ flagcxResult_t flagcxCommInitRank(flagcxComm_t *comm, int nranks,
   INFO(FLAGCX_INIT, "Flagcx USE_TUNER flag set to %d", useTuner);
   if (useTuner) {
     (*comm)->tuner = &internalTuner;
-    (*comm)->commId = commId;
+    FLAGCXCHECK(flagcxCalloc(&(*comm)->commId, 1));
+    memcpy((*comm)->commId, commId, sizeof(flagcxUniqueId));
     (*comm)->uniqueIdData = uniqueIdData;
     (*comm)->tunerInnerComm = NULL;
     (*comm)->isTunningComm = false;
@@ -2013,13 +2039,18 @@ flagcxResult_t flagcxCommDestroy(flagcxComm_t comm) {
   // Destroy tuner
   if (comm->tuner) {
     comm->tuner->destroy(comm->tunerContext);
-    // Free uniqueIdData
+    // Free uniqueIdData and commId
     free(comm->uniqueIdData);
+    free(comm->commId);
   }
 
   // Finalize net adaptor plugin (dlclose)
   FLAGCXCHECK(flagcxNetAdaptorPluginFinalize());
 
+  // Finalize device/CCL adaptor plugins (ref-counted)
+  flagcxCCLAdaptorPluginFinalize();
+  flagcxDeviceAdaptorPluginFinalize();
+
   free(comm);
   return flagcxSuccess;
 }

flagcx/include/flagcx.h

@@ -174,16 +174,28 @@ struct flagcxDeviceHandle {
 };
 typedef struct flagcxDeviceHandle *flagcxDeviceHandle_t;
 
+/* Initialize device handle — loads device/CCL adaptor plugins, returns a
+ * populated device handle. Call before flagcxCommInitRank if you need device
+ * operations (setDevice, getDeviceCount, streamCreate, etc.) early.
+ * If not called explicitly, flagcxCommInitRank will load plugins internally. */
+flagcxResult_t flagcxDeviceHandleInit(flagcxDeviceHandle_t *devHandle);
+
+/* Free device handle — finalizes device/CCL adaptor plugins. */
+flagcxResult_t flagcxDeviceHandleFree(flagcxDeviceHandle_t devHandle);
+
+/* Deprecated: use flagcxDeviceHandleInit/Free and manage comm/uniqueId
+ * separately */
 struct flagcxHandlerGroup {
   flagcxUniqueId_t uniqueId;
   flagcxComm_t comm;
   flagcxDeviceHandle_t devHandle;
 };
 typedef struct flagcxHandlerGroup *flagcxHandlerGroup_t;
 
-/* Init and free FlagCX handls including flagcxComm_t, flagcxStream_t */
+/* Deprecated: use flagcxDeviceHandleInit instead */
 flagcxResult_t flagcxHandleInit(flagcxHandlerGroup_t *handler);
 
+/* Deprecated: use flagcxDeviceHandleFree instead */
 flagcxResult_t flagcxHandleFree(flagcxHandlerGroup_t handler);
 
 /* User buffer registration functions. The actual allocated size might
@@ -224,8 +236,9 @@ flagcxResult_t flagcxGetVersion(int *version);
 
 /* Generates an Id to be used in flagcxCommInitRank. flagcxGetUniqueId should be
  * called once and the Id should be distributed to all ranks in the
- * communicator before calling flagcxCommInitRank. */
-flagcxResult_t flagcxGetUniqueId(flagcxUniqueId_t *uniqueId);
+ * communicator before calling flagcxCommInitRank.
+ * The caller must provide a valid pointer to a flagcxUniqueId struct. */
+flagcxResult_t flagcxGetUniqueId(flagcxUniqueId_t uniqueId);
 
 /* Creates a new communicator (multi thread/process version).
  * rank must be between 0 and nranks-1 and unique within a communicator clique.

plugin/interservice/flagcx_wrapper.py

@@ -21,7 +21,6 @@
 flagcxEventType_t = ctypes.c_int
 flagcxIpcMemHandle_t = ctypes.c_void_p
 
-flagcxHandlerGroup_t = ctypes.c_void_p
 flagcxComm_t = ctypes.c_void_p
 flagcxEvent_t = ctypes.c_void_p
 flagcxStream_t = ctypes.c_void_p
@@ -212,15 +211,15 @@ class Function:
 
 class FLAGCXLibrary:
     exported_functions = [
-        Function("flagcxHandleInit", flagcxResult_t,
-                [ctypes.POINTER(flagcxHandlerGroup_t)]),
-        Function("flagcxHandleFree", flagcxResult_t,
-                [flagcxHandlerGroup_t]),
+        Function("flagcxDeviceHandleInit", flagcxResult_t,
+                [ctypes.POINTER(flagcxDeviceHandle_t)]),
+        Function("flagcxDeviceHandleFree", flagcxResult_t,
+                [flagcxDeviceHandle_t]),
         Function("flagcxGetErrorString", ctypes.c_char_p, [flagcxResult_t]),
         Function("flagcxGetVersion", flagcxResult_t,
                  [ctypes.POINTER(ctypes.c_int)]),
         Function("flagcxGetUniqueId", flagcxResult_t,
-                [ctypes.POINTER(ctypes.POINTER(flagcxUniqueId))]),
+                [ctypes.POINTER(flagcxUniqueId)]),
         # Note that flagcxComm_t is a pointer type, so the first argument
         # is a pointer to a pointer
         Function("flagcxCommInitRank", flagcxResult_t, [
@@ -432,13 +431,13 @@ def __init__(self, so_file: Optional[str] = None):
             FLAGCXLibrary.path_to_dict_mapping[so_file] = _funcs
         self._funcs = FLAGCXLibrary.path_to_dict_mapping[so_file]
 
-        # init flagcx handler to call device-related apis
-        self.handler = flagcxHandlerGroup_t()
-        self.FLAGCX_CHECK(self._funcs["flagcxHandleInit"](ctypes.byref(self.handler)))
-    
+        # init flagcx device handle to call device-related apis
+        self.devHandle = flagcxDeviceHandle_t()
+        self.FLAGCX_CHECK(self._funcs["flagcxDeviceHandleInit"](ctypes.byref(self.devHandle)))
+
     def __del__(self):
-        # free flagcx handler
-        self.FLAGCX_CHECK(self._funcs["flagcxHandleFree"](self.handler))
+        # free flagcx device handle
+        self.FLAGCX_CHECK(self._funcs["flagcxDeviceHandleFree"](self.devHandle))
 
     def flagcxGetErrorString(self, result: flagcxResult_t) -> str:
         return self._funcs["flagcxGetErrorString"](result).decode("utf-8")
@@ -459,7 +458,7 @@ def flagcxGetVersion(self) -> str:
         return f"{major}.{minor}.{patch}"
 
     def flagcxGetUniqueId(self) -> flagcxUniqueId:
-        unique_id = ctypes.POINTER(flagcxUniqueId)()
+        unique_id = flagcxUniqueId()
         self.FLAGCX_CHECK(self._funcs["flagcxGetUniqueId"](
             ctypes.byref(unique_id)))
         return unique_id
@@ -486,7 +485,7 @@ def flagcxCommInitRank(self, world_size: int, unique_id: flagcxUniqueId,
                          rank: int) -> flagcxComm_t:
         comm = flagcxComm_t()
         self.FLAGCX_CHECK(self._funcs["flagcxCommInitRank"](ctypes.byref(comm),
-                                                        world_size, unique_id,
+                                                        world_size, ctypes.byref(unique_id),
                                                         rank))
         return comm
 
@@ -715,26 +714,26 @@ def flagcxDevMemFreeDevicePtr(self, dev_mem: flagcxDevMem_t) -> None:
 
     def adaptor_stream_create(self):
         new_stream = flagcxStream_t()
-        self.FLAGCX_CHECK(self.handler.contents.devHandle.contents.streamCreate(ctypes.byref(new_stream)))
+        self.FLAGCX_CHECK(self.devHandle.contents.streamCreate(ctypes.byref(new_stream)))
         return new_stream
 
     def adaptor_stream_copy(self, old_stream):
         new_stream = flagcxStream_t()
         raw_stream = getattr(old_stream, 'musa_stream', old_stream.cuda_stream)
-        self.FLAGCX_CHECK(self.handler.contents.devHandle.contents.streamCopy(ctypes.byref(new_stream), ctypes.c_void_p(raw_stream)))
+        self.FLAGCX_CHECK(self.devHandle.contents.streamCopy(ctypes.byref(new_stream), ctypes.c_void_p(raw_stream)))
         return new_stream
 
     def adaptor_stream_free(self, stream):
-        self.FLAGCX_CHECK(self.handler.contents.devHandle.contents.streamFree(stream))
+        self.FLAGCX_CHECK(self.devHandle.contents.streamFree(stream))
 
     def adaptor_stream_destroy(self, stream):
-        self.FLAGCX_CHECK(self.handler.contents.devHandle.contents.streamDestroy(stream))
+        self.FLAGCX_CHECK(self.devHandle.contents.streamDestroy(stream))
 
     def sync_stream(self, stream):
-        self.FLAGCX_CHECK(self.handler.contents.devHandle.contents.streamSynchronize(stream))
+        self.FLAGCX_CHECK(self.devHandle.contents.streamSynchronize(stream))
 
 
 __all__ = [
     "FLAGCXLibrary", "flagcxDataTypeEnum", "flagcxRedOpTypeEnum", "flagcxUniqueId",
-    "flagcxHandlerGroup_t", "flagcxComm_t", "flagcxStream_t", "flagcxEvent_t", "buffer_type"
+    "flagcxDeviceHandle_t", "flagcxComm_t", "flagcxStream_t", "flagcxEvent_t", "buffer_type"
 ]