fix: add EventFlag enum, fix mutex usage in ProcessGroupFactory::Instance()

Author: Chamberlain0w0

infini_train/include/core/runtime/device_guard.h

@@ -94,7 +94,7 @@ class DeviceGuardImpl {
 
     virtual void EventCreate(Event **event) const;
 
-    virtual void EventCreateWithFlags(Event **event, uint32_t flags) const;
+    virtual void EventCreateWithFlags(Event **event, EventFlag flags) const;
 
     virtual void EventDestroy(Event *event) const;
 

infini_train/include/core/runtime/runtime_common.h

@@ -22,6 +22,13 @@ class Stream {
     virtual ~Stream() = default;
 };
 
+enum class EventFlag : uint32_t {
+    kDefault = 0x0,
+    kBlockingSync = 0x1,
+    kDisableTiming = 0x2,
+    kInterprocess = 0x4,
+};
+
 // Generic runtime status for backend-agnostic control flow.
 #define INFINI_TRAIN_RUNTIME_STATUS_LIST(X)                                                                            \
     X(kSuccess, 0)                                                                                                     \

infini_train/src/core/runtime/cpu/cpu_guard_impl.cc

@@ -65,7 +65,7 @@ void CpuGuardImpl::EventCreate(Event **event) const {
     *event = nullptr;
 }
 
-void CpuGuardImpl::EventCreateWithFlags(Event **event, uint32_t flags) const {
+void CpuGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const {
     CHECK_NOTNULL(event);
     LOG(WARNING) << "CpuGuardImpl::EventCreateWithFlags is not supported. Returning nullptr event.";
     *event = nullptr;

infini_train/src/core/runtime/cpu/cpu_guard_impl.h

@@ -33,7 +33,7 @@ class CpuGuardImpl final : public DeviceGuardImpl {
     // Event management (explicitly unsupported for now)
     void EventCreate(Event **event) const override;
 
-    void EventCreateWithFlags(Event **event, uint32_t flags) const override;
+    void EventCreateWithFlags(Event **event, EventFlag flags) const override;
 
     void EventDestroy(Event *event) const override;
 

infini_train/src/core/runtime/cuda/cuda_guard_impl.cc

@@ -139,7 +139,7 @@ void CudaGuardImpl::GetStreamPriorityRange(int *low, int *high) const {
 // event
 void CudaGuardImpl::EventCreate(Event **event) const { *event = new CudaEvent(); }
 
-void CudaGuardImpl::EventCreateWithFlags(Event **event, uint32_t flags) const { *event = new CudaEvent(flags); }
+void CudaGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const { *event = new CudaEvent(flags); }
 
 void CudaGuardImpl::EventDestroy(Event *event) const {
     if (event == nullptr) {

infini_train/src/core/runtime/cuda/cuda_guard_impl.h

@@ -41,7 +41,7 @@ class CudaGuardImpl final : public DeviceGuardImpl {
     // event
     void EventCreate(Event **event) const override;
 
-    void EventCreateWithFlags(Event **event, uint32_t flags) const override;
+    void EventCreateWithFlags(Event **event, EventFlag flags) const override;
 
     void EventDestroy(Event *event) const override;
 

infini_train/src/core/runtime/cuda/cuda_runtime_common.cc

@@ -3,8 +3,26 @@
 #include "infini_train/include/common/cuda/common_cuda.h"
 
 namespace infini_train::core::cuda {
+namespace {
+uint32_t ToCudaEventFlags(EventFlag flags) {
+    switch (flags) {
+    case EventFlag::kDefault:
+        return cudaEventDefault;
+    case EventFlag::kBlockingSync:
+        return cudaEventBlockingSync;
+    case EventFlag::kDisableTiming:
+        return cudaEventDisableTiming;
+    case EventFlag::kInterprocess:
+        // CUDA requires cudaEventDisableTiming with interprocess events.
+        return cudaEventInterprocess | cudaEventDisableTiming;
+    default:
+        LOG(FATAL) << "Unsupported EventFlag value: " << static_cast<uint32_t>(flags);
+    }
+    return cudaEventDefault;
+}
+} // namespace
 
-CudaEvent::CudaEvent(uint32_t flags) { CUDA_CHECK(cudaEventCreateWithFlags(&event_, flags)); }
+CudaEvent::CudaEvent(EventFlag flags) { CUDA_CHECK(cudaEventCreateWithFlags(&event_, ToCudaEventFlags(flags))); }
 
 CudaEvent::~CudaEvent() {
     if (event_ != nullptr) {

infini_train/src/core/runtime/cuda/cuda_runtime_common.h

@@ -15,7 +15,7 @@ namespace infini_train::core::cuda {
 
 class CudaEvent final : public Event {
 public:
-    explicit CudaEvent(uint32_t flags = cudaEventDefault);
+    explicit CudaEvent(EventFlag flags = EventFlag::kDefault);
     ~CudaEvent() override;
 
     cudaEvent_t cuda_event() const;

infini_train/src/core/runtime/device_guard.cc

@@ -43,7 +43,7 @@ void DeviceGuardImpl::GetStreamPriorityRange(int *, int *) const {
 
 void DeviceGuardImpl::EventCreate(Event **) const { LOG(FATAL) << "DeviceGuardImpl::EventCreate is not implemented."; }
 
-void DeviceGuardImpl::EventCreateWithFlags(Event **, uint32_t) const {
+void DeviceGuardImpl::EventCreateWithFlags(Event **, EventFlag) const {
     LOG(FATAL) << "DeviceGuardImpl::EventCreateWithFlags is not implemented.";
 }
 

infini_train/src/nn/parallel/process_group.cc

@@ -425,20 +425,27 @@ std::shared_ptr<Tensor> ProcessGroup::Gather(const std::vector<std::shared_ptr<T
 
 ProcessGroupFactory *ProcessGroupFactory::Instance() {
     // NOTE(zbl): Instance() with no arguments only gets initialized instance with a certain backend
-    std::lock_guard<std::mutex> lock(g_process_group_factory_mutex);
     auto &instance = g_process_group_factory_instance;
     if (instance == nullptr) {
-        LOG(FATAL) << "ProcessGroupFactory is not initialized with backend. "
-                   << "Call ProcessGroupFactory::Instance(backend) first.";
+        std::lock_guard<std::mutex> lock(g_process_group_factory_mutex);
+        if (instance == nullptr) {
+            LOG(FATAL) << "ProcessGroupFactory is not initialized with backend. "
+                       << "Call ProcessGroupFactory::Instance(backend) first.";
+        }
     }
     return instance.get();
 }
 
 ProcessGroupFactory *ProcessGroupFactory::Instance(Device::DeviceType backend) {
-    std::lock_guard<std::mutex> lock(g_process_group_factory_mutex);
     auto &instance = g_process_group_factory_instance;
     if (instance == nullptr) {
-        instance.reset(new ProcessGroupFactory(backend));
+        std::lock_guard<std::mutex> lock(g_process_group_factory_mutex);
+        if (instance == nullptr) {
+            instance.reset(new ProcessGroupFactory(backend));
+        } else if (instance->backend_ != backend) {
+            LOG(FATAL) << "ProcessGroupFactory backend mismatch. initialized=" << static_cast<int>(instance->backend_)
+                       << ", requested=" << static_cast<int>(backend);
+        }
     } else if (instance->backend_ != backend) {
         LOG(FATAL) << "ProcessGroupFactory backend mismatch. initialized=" << static_cast<int>(instance->backend_)
                    << ", requested=" << static_cast<int>(backend);