Avoid calling destroy event inside host func

.github/workflows/test.yml

@@ -53,12 +53,9 @@ jobs:
           if [ -n "$GITHUB_HEAD_REF" ] && [ -n "$GITHUB_BASE_REF" ]; then
             from_ref="origin/$GITHUB_HEAD_REF"
             to_ref="origin/$GITHUB_BASE_REF"
-          else
-            from_ref="HEAD^"
-            to_ref="HEAD"
+            echo "From reference: $from_ref; To reference: $to_ref"
+            pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
           fi
-          echo "From reference: $from_ref; To reference: $to_ref"
-          pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
         continue-on-error: false
 
       - name: Check the current working directory

.github/workflows/torch-api-test.yml

@@ -50,12 +50,9 @@ jobs:
           if [ -n "$GITHUB_HEAD_REF" ] && [ -n "$GITHUB_BASE_REF" ]; then
             from_ref="origin/$GITHUB_HEAD_REF"
             to_ref="origin/$GITHUB_BASE_REF"
-          else
-            from_ref="HEAD^"
-            to_ref="HEAD"
+            echo "From reference: $from_ref; To reference: $to_ref"
+            pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
           fi
-          echo "From reference: $from_ref; To reference: $to_ref"
-          pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
         continue-on-error: false
 
       - name: Run `make` to build the project

.github/workflows/unit-test.yml

@@ -53,12 +53,9 @@ jobs:
           if [ -n "$GITHUB_HEAD_REF" ] && [ -n "$GITHUB_BASE_REF" ]; then
             from_ref="origin/$GITHUB_HEAD_REF"
             to_ref="origin/$GITHUB_BASE_REF"
-          else
-            from_ref="HEAD^"
-            to_ref="HEAD"
+            echo "From reference: $from_ref; To reference: $to_ref"
+            pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
           fi
-          echo "From reference: $from_ref; To reference: $to_ref"
-          pre-commit run --from-ref "$from_ref" --to-ref "$to_ref"
         continue-on-error: false
 
       - name: Build Google Test

flagcx/core/group.cc

@@ -109,10 +109,8 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
 
   // Each groupLaunch we create a semaphore to track the p2p ops
   // and a stream to launch host or device func
-  struct flagcxHostSemaphore *semaphore;
-  FLAGCXCHECK(flagcxCalloc(&semaphore, 1));
-  semaphore->flag = 0;
-  semaphore->counter = 0;
+  std::shared_ptr<flagcxHostSemaphore> semaphore =
+      std::make_shared<flagcxHostSemaphore>();
   flagcxStream_t launchStream = nullptr;
 
   if (groupCommPreconnectHeadMain != nullptr) {
@@ -312,7 +310,7 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
     }
   } else {
     FLAGCXCHECK(deviceAdaptor->launchHostFunc(launchStream, cpuAsyncKernel,
-                                              (void *)semaphore));
+                                              (void *)semaphore.get()));
   }
   // deprecated code path for host func, since the previous
   // hang issue may be walked around by using zero copy

flagcx/core/launch_kernel.cc

@@ -50,8 +50,7 @@ void cpuAsyncLoadWithMaxSpinCount(void *args) {
 
 void cpuAsyncKernel(void *args) {
   flagcxHostSemaphore *semaphore = (flagcxHostSemaphore *)args;
-  semaphore->signalFlag();
+  semaphore->signalStart();
   semaphore->wait();
-  free(semaphore);
-  semaphore = nullptr;
+  semaphore->signalEnd();
 }

flagcx/core/launch_kernel.h

@@ -36,28 +36,33 @@ void cpuAsyncLoad(void *args);
 void cpuAsyncLoadWithMaxSpinCount(void *args);
 
 struct flagcxHostSemaphore {
-  int flag;    // if ready to be triggered
-  int counter; // total operations to wait for inside the group
+  int start = 0;   // started or not
+  int end = 0;     // ended or not
+  int counter = 0; // total operations to wait for inside the group
   std::vector<flagcxEvent_t> events;
 
+  ~flagcxHostSemaphore() {
+    for (auto event : events) {
+      deviceAdaptor->eventDestroy(event);
+    }
+  }
   flagcxEvent_t getEvent() {
     events.push_back(nullptr);
     auto &event = events.back();
     deviceAdaptor->eventCreate(&event);
     return event;
   }
-  void signalFlag() { __atomic_store_n(&flag, 1, __ATOMIC_RELEASE); }
+  void signalStart() { __atomic_store_n(&start, 1, __ATOMIC_RELEASE); }
+  void signalEnd() { __atomic_store_n(&end, 1, __ATOMIC_RELEASE); }
   void signalCounter(int value) {
     __atomic_fetch_sub(&counter, value, __ATOMIC_RELEASE);
   }
-  int poll() { return __atomic_load_n(&flag, __ATOMIC_ACQUIRE); }
+  int pollStart() { return __atomic_load_n(&start, __ATOMIC_ACQUIRE); }
+  int pollEnd() { return __atomic_load_n(&end, __ATOMIC_ACQUIRE); }
   void wait() {
     while (__atomic_load_n(&counter, __ATOMIC_ACQUIRE) > 0) {
       sched_yield();
     }
-    for (auto event : events) {
-      deviceAdaptor->eventDestroy(event);
-    }
   }
 };
 

flagcx/core/net.cc

@@ -136,7 +136,7 @@ flagcxResult_t flagcxNetInit(struct flagcxHeteroComm *comm) {
 
 flagcxResult_t flagcxProxySend(sendNetResources *resources, void *data,
                                size_t size, flagcxProxyArgs *args) {
-  if (!args->semaphore->poll()) {
+  if (!args->semaphore->pollStart()) {
     return flagcxSuccess;
   }
   if (args->transmitted < args->chunkSteps) {
@@ -220,7 +220,7 @@ flagcxResult_t flagcxProxySend(sendNetResources *resources, void *data,
 
 flagcxResult_t flagcxProxyRecv(recvNetResources *resources, void *data,
                                size_t size, flagcxProxyArgs *args) {
-  if (!args->semaphore->poll()) {
+  if (!args->semaphore->pollStart()) {
     return flagcxSuccess;
   }
   if (args->copied < args->chunkSteps) {

flagcx/core/proxy.cc

@@ -196,7 +196,8 @@ static flagcxResult_t progressOps(struct flagcxProxyState *proxyState,
                 }
               }
             } else {
-              if (op->args.done == 1) {
+              if (op->args.done == 1 && op->args.semaphore->pollEnd()) {
+                op->args.semaphore.reset();
                 flagcxIntruQueueDelete(queue, op);
                 free(op);
               }
@@ -220,7 +221,9 @@ static flagcxResult_t progressOps(struct flagcxProxyState *proxyState,
                 }
               }
             } else {
-              if (op->args.done == 1) {
+              if (op->args.done == 1 && op->args.semaphore->pollEnd()) {
+                // update refcount and delete semaphore when refcount = 0
+                op->args.semaphore.reset();
                 flagcxIntruQueueDelete(queue, op);
                 free(op);
               }

flagcx/core/proxy.h

@@ -15,6 +15,7 @@
 #include "net.h"
 #include "reg_pool.h"
 #include "socket.h"
+#include <memory>
 #include <pthread.h>
 
 enum flagcxProxyOpState {
@@ -117,7 +118,7 @@ struct flagcxProxyArgs {
 
   /*for launch*/
   int deviceFuncRelaxedOrdering = 0;
-  struct flagcxHostSemaphore *semaphore = nullptr;
+  std::shared_ptr<flagcxHostSemaphore> semaphore;
   // only for device func, to be deprecated
   volatile bool eventRecorded = false;
   volatile bool hlArgs = false;