Minor improvement on device syncer

include/mscclpp/concurrency_device.hpp

@@ -4,6 +4,7 @@
 #ifndef MSCCLPP_CONCURRENCY_DEVICE_HPP_
 #define MSCCLPP_CONCURRENCY_DEVICE_HPP_
 
+#include "atomic_device.hpp"
 #include "poll_device.hpp"
 
 namespace mscclpp {
@@ -29,17 +30,11 @@ struct DeviceSyncer {
     if (threadIdx.x == 0) {
       // Need a `__threadfence()` before to flip `flag`.
       __threadfence();
-      int tmp = isIncFlag_ ^ 1;
-      if (tmp) {
-        if (atomicInc(&count_, maxOldCnt) == maxOldCnt) {
-          flag_ = 1;
-        }
-        POLL_MAYBE_JAILBREAK(!flag_, maxSpinCount);
+      unsigned int tmp = isIncFlag_ ^ 1;
+      if (atomicInc(&count_, maxOldCnt) == maxOldCnt) {
+        atomicStore(&flag_, tmp, memoryOrderRelaxed);
       } else {
-        if (atomicInc(&count_, maxOldCnt) == maxOldCnt) {
-          flag_ = 0;
-        }
-        POLL_MAYBE_JAILBREAK(flag_, maxSpinCount);
+        POLL_MAYBE_JAILBREAK((atomicLoad(&flag_, memoryOrderRelaxed) != tmp), maxSpinCount);
       }
       isIncFlag_ = tmp;
     }
@@ -51,11 +46,11 @@ struct DeviceSyncer {
 
  private:
   /// The flag to indicate whether the barrier is reached by the latest thread.
-  volatile int flag_;
+  unsigned int flag_;
   /// The counter of synchronized blocks.
   unsigned int count_;
   /// The flag to indicate whether to increase or decrease @ref flag_.
-  int isIncFlag_;
+  unsigned int isIncFlag_;
 };
 
 }  // namespace mscclpp