Refactor group management logic (#284)

Author: MC952-arch

flagcx/core/group.cc

@@ -370,3 +370,70 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
 fail:
   goto exit;
 }
+
+static flagcxResult_t groupCleanup(struct flagcxAsyncJob *job_) {
+  struct flagcxGroupJob *gjob = (struct flagcxGroupJob *)job_;
+  struct flagcxHeteroComm *groupCommHeadMain = *gjob->groupCommHeadPtr;
+  struct flagcxHeteroComm *groupCommPreconnectHeadMain =
+      *gjob->groupCommPreconnectHeadPtr;
+  struct flagcxIntruQueue<struct flagcxAsyncJob, &flagcxAsyncJob::next>
+      *asyncJobsMain = gjob->asyncJobsPtr;
+
+  // clean up preconnect comms
+  while (groupCommPreconnectHeadMain != nullptr) {
+    struct flagcxHeteroComm *comm = groupCommPreconnectHeadMain;
+    struct flagcxHeteroComm *next = comm->preconnectNext;
+    comm->preconnectNext = reinterpret_cast<struct flagcxHeteroComm *>(0x1);
+    groupCommPreconnectHeadMain = next;
+  }
+
+  // clean up async jobs
+  while (!flagcxIntruQueueEmpty(asyncJobsMain)) {
+    struct flagcxAsyncJob *job = flagcxIntruQueueDequeue(asyncJobsMain);
+    free(job);
+  }
+
+  // clean up comms
+  while (groupCommHeadMain != nullptr) {
+    struct flagcxHeteroComm *comm = groupCommHeadMain;
+    struct flagcxHeteroComm *next = comm->groupNext;
+    (void)flagcxGroupCommLeave(comm);
+    groupCommHeadMain = next;
+  }
+
+  return flagcxSuccess;
+}
+
+static inline void groupResetJobState() {
+  flagcxGroupBlocking = 0;
+  flagcxGroupJobMainPtr = NULL;
+  flagcxGroupCommPreconnectHead = nullptr;
+  flagcxGroupCommHead = nullptr;
+  memset(&flagcxGroupJobMain, 0, sizeof(struct flagcxGroupJob));
+}
+
+flagcxResult_t flagcxGroupEndInternal() {
+  flagcxResult_t ret = flagcxSuccess;
+  flagcxGroupDepth--;
+  if (flagcxGroupDepth < 0)
+    return flagcxSystemError;
+  if (flagcxGroupDepth == 0) {
+    if (flagcxGroupCommPreconnectHead || flagcxGroupCommHead) {
+      flagcxGroupJobMain.groupCommHeadPtr = &flagcxGroupCommHead;
+      flagcxGroupJobMain.groupCommPreconnectHeadPtr =
+          &flagcxGroupCommPreconnectHead;
+      flagcxGroupJobMain.asyncJobsPtr = &flagcxAsyncJobs;
+      flagcxGroupJobMain.initialized = true;
+      flagcxGroupJobMainPtr = &flagcxGroupJobMain;
+      FLAGCXCHECKGOTO(groupLaunch(&flagcxGroupJobMainPtr->base), ret, fail);
+      groupResetJobState();
+    }
+  }
+
+exit:
+  return ret;
+fail:
+  groupCleanup(&flagcxGroupJobMainPtr->base);
+  groupResetJobState();
+  goto exit;
+}

flagcx/core/group.h

@@ -10,13 +10,6 @@
 #include "assert.h"
 #include "comm.h"
 
-flagcxResult_t flagcxGroupErrCheck(flagcxResult_t ret);
-void flagcxGroupCommJoin(struct flagcxHeteroComm *comm);
-void flagcxGroupCommPreconnect(struct flagcxHeteroComm *comm);
-flagcxResult_t flagcxGroupCommLeave(struct flagcxHeteroComm *comm);
-flagcxResult_t flagcxGroupJobAbort(struct flagcxGroupJob *groupJob);
-flagcxResult_t flagcxGroupJobComplete(struct flagcxGroupJob *groupJob);
-
 typedef flagcxResult_t (*flagcxInitFunc_t)(flagcxHeteroComm_t *newcomm,
                                            int ndev, flagcxUniqueId commId,
                                            int myrank, int cudaDev);
@@ -79,69 +72,17 @@ extern __thread struct flagcxIntruQueue<struct flagcxAsyncJob,
                                         &flagcxAsyncJob::next>
     flagcxAsyncJobs;
 
-static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_);
+flagcxResult_t flagcxGroupErrCheck(flagcxResult_t ret);
+void flagcxGroupCommJoin(struct flagcxHeteroComm *comm);
+void flagcxGroupCommPreconnect(struct flagcxHeteroComm *comm);
+flagcxResult_t flagcxGroupCommLeave(struct flagcxHeteroComm *comm);
+// Not implemented
+flagcxResult_t flagcxGroupJobAbort(struct flagcxGroupJob *groupJob);
+// Not implemented
+flagcxResult_t flagcxGroupJobComplete(struct flagcxGroupJob *groupJob);
 flagcxResult_t flagcxHeteroGroupStart();
 flagcxResult_t flagcxHeteroGroupEnd();
 
-static inline void groupResetJobState() {
-  flagcxGroupBlocking = 0;
-  flagcxGroupJobMainPtr = NULL;
-  flagcxGroupCommPreconnectHead = nullptr;
-  flagcxGroupCommHead = nullptr;
-  memset(&flagcxGroupJobMain, 0, sizeof(struct flagcxGroupJob));
-  return;
-}
-
-static inline flagcxResult_t groupJobComplete(struct flagcxGroupJob *job) {
-  flagcxResult_t ret = flagcxSuccess;
-  if (job) {
-    ret = flagcxAsyncJobComplete(&job->base);
-    groupResetJobState();
-  }
-  return ret;
-}
-
-inline flagcxResult_t flagcxGroupStartInternal() {
-  flagcxGroupDepth++;
-  return flagcxSuccess;
-}
-
-inline flagcxResult_t flagcxGroupEndInternal() {
-  flagcxResult_t ret = flagcxSuccess;
-  flagcxGroupDepth--;
-  if (flagcxGroupDepth < 0)
-    return flagcxSystemError;
-  if (flagcxGroupDepth == 0) {
-
-    /**
-     * TODO: do all jobs at Groups
-     **/
-    if (flagcxGroupCommPreconnectHead || flagcxGroupCommHead) {
-
-      flagcxGroupJobMain.groupCommHeadPtr = &flagcxGroupCommHead;
-      flagcxGroupJobMain.groupCommPreconnectHeadPtr =
-          &flagcxGroupCommPreconnectHead;
-      flagcxGroupJobMain.asyncJobsPtr = &flagcxAsyncJobs;
-      flagcxGroupJobMain.initialized = true;
-      flagcxGroupJobMainPtr = &flagcxGroupJobMain;
-
-      FLAGCXCHECKGOTO(groupLaunch(&flagcxGroupJobMainPtr->base), ret, fail);
-      groupResetJobState();
-    }
-  }
-
-exit:
-  return ret;
-fail:
-  /**
-   * TODO: add groupCleanup()
-   **/
-  // groupCleanup(&flagcxGroupCommHead, &flagcxGroupCommPreconnectHead,
-  // &flagcxAsyncJobs, &flagcxGroupError, &flagcxGroupBlocking,
-  // &flagcxGroupJobAbortFlag, ret);
-  goto exit;
-}
-
 inline flagcxResult_t flagcxGroupErrCheck(flagcxResult_t ret) {
   if (flagcxGroupDepth > 0) {
     if (ret != flagcxSuccess && ret != flagcxInProgress)
@@ -173,4 +114,9 @@ inline flagcxResult_t flagcxGroupCommLeave(struct flagcxHeteroComm *comm) {
   return flagcxSuccess;
 }
 
+inline flagcxResult_t flagcxGroupStartInternal() {
+  flagcxGroupDepth++;
+  return flagcxSuccess;
+}
+
 #endif