Support batch_isend_irecv operation in torch plugin by adding group s…

…emantics (#28)
FlagOpen · Feb 7, 2025 · 0869dcf · 0869dcf
1 parent 5ae71db
commit 0869dcf
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/plugin/torch/src/backend_flagcx.cpp b/plugin/torch/src/backend_flagcx.cpp
@@ -1,5 +1,6 @@
 #include "backend_flagcx.hpp"
 #include <iostream>
+#include <c10/core/DeviceGuard.h>
 
 namespace c10d
 {
@@ -254,14 +255,19 @@ namespace c10d
 
     void BackendFlagcx::groupStart()
     {
-        // flagcxGroupStart();
-        // ++flagcxActiveGroupCounter_;
+#if defined(USE_NVIDIA_ADAPTOR) || defined(USE_ILUVATAR_COREX_ADAPTOR)
+        initComm(c10::impl::getDeviceGuardImpl(at::DeviceType::CUDA)->getDevice());
+#elif defined(USE_CAMBRICON_ADAPTOR)
+        initComm(c10::impl::getDeviceGuardImpl(at::DeviceType::PrivateUse1)->getDevice());
+#endif
+        flagcxGroupStart();
+        ++flagcxActiveGroupCounter_;
     }
 
     void BackendFlagcx::groupEnd()
     {
-        // flagcxGroupEnd();
-        // --flagcxActiveGroupCounter_;
+        flagcxGroupEnd();
+        --flagcxActiveGroupCounter_;
     }
 
     void BackendFlagcx::startCoalescing()
@@ -816,4 +822,4 @@ namespace c10d
         m.def("createBackendFlagcx", &BackendFlagcx::createBackendFlagcx);
     }
 
-} // namespace c10d
+} // namespace c10d