[NewComm] No.1 compatiable upgrade for class_center_sample op

paddle/fluid/operators/class_center_sample_op.cu

@@ -30,13 +30,17 @@ namespace cub = hipcub;
 #include <random>
 
 #include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/distributed/comm_context_manager.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/phi/core/tensor_utils.h"
 
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
 #include "paddle/fluid/distributed/collective/process_group.h"
 #include "paddle/fluid/platform/collective_helper.h"
 #include "paddle/fluid/platform/device/gpu/nccl_helper.h"
+#include "paddle/phi/core/distributed/nccl_comm_context.h"
+#include "paddle/phi/core/flags.h"
+PHI_DECLARE_bool(dynamic_static_unified_comm);
 #endif
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -351,6 +355,8 @@ void ClassCenterSampleKernel(const Context& dev_ctx,
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
   if (nranks > 1) {
     auto map = paddle::distributed::ProcessGroupMapFromGid::getInstance();
+    gpuStream_t stream = nullptr;
+
     if (map->has(ring_id)) {
       // Use ProcessGroup
       paddle::distributed::ProcessGroup* pg = map->get(ring_id);
@@ -364,21 +370,46 @@ void ClassCenterSampleKernel(const Context& dev_ctx,
       auto task = pg->AllReduce(in_tensor, out_tensor, opts);
       task->Wait();
     } else {
-      const auto& comm = paddle::platform::NCCLCommContext::Instance().Get(
-          ring_id, dev_ctx.GetPlace());
-      // use global calculate stream
-      const auto calcu_stream =
-          static_cast<GPUContext*>(
-              phi::DeviceContextPool::Instance().Get(dev_ctx.GetPlace()))
-              ->stream();
-      PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclAllReduce(
-          num_classes_per_device_ptr,
-          num_classes_per_device_ptr,
-          num_classes_per_device.numel(),
-          phi::ToNCCLDataType(num_classes_per_device.dtype()),
-          ncclSum,
-          comm->comm(),
-          calcu_stream));
+      paddle::platform::NCCLComm* comm = nullptr;
+      phi::distributed::NCCLCommContext* comm_ctx = nullptr;
+      const auto& comm_context_manager =
+          phi::distributed::CommContextManager::GetInstance();
+      if (FLAGS_dynamic_static_unified_comm) {
+        PADDLE_ENFORCE_EQ(comm_context_manager.Has(std::to_string(ring_id)),
+                          true,
+                          errors::InvalidArgument(
+                              "You choose to use new communication library by "
+                              "setting environment "
+                              "variable FLAGS_dynamic_static_unified_comm "
+                              "True. But ring_id(%d) is "
+                              "not found in comm_context_manager.",
+                              std::to_string(ring_id)));
+        comm_ctx = static_cast<phi::distributed::NCCLCommContext*>(
+            comm_context_manager.Get(std::to_string(ring_id)));
+        stream = comm_ctx->GetStream();
+      } else {
+        comm = paddle::platform::NCCLCommContext::Instance().Get(
+            ring_id, dev_ctx.GetPlace());
+        // use global calculate stream
+        stream = static_cast<GPUContext*>(
+                     phi::DeviceContextPool::Instance().Get(dev_ctx.GetPlace()))
+                     ->stream();
+      }
+
+      if (comm_ctx) {
+        comm_ctx->AllReduce(
+            &num_classes_per_device, num_classes_per_device, ncclSum, stream);
+        paddle::platform::GpuStreamSync(stream);
+      } else {
+        PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::ncclAllReduce(
+            num_classes_per_device_ptr,
+            num_classes_per_device_ptr,
+            num_classes_per_device.numel(),
+            phi::ToNCCLDataType(num_classes_per_device.dtype()),
+            ncclSum,
+            comm->comm(),
+            stream));
+      }
     }
   }
 #endif

test/legacy_test/test_class_center_sample_op.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import unittest
 
 import numpy as np
@@ -87,7 +88,11 @@ def init_dtype(self):
     def init_fix_seed(self):
         self.fix_seed = True
 
+    def with_new_comm(self):
+        os.environ["FLAGS_dynamic_static_unified_comm"] = "0"
+
     def setUp(self):
+        self.with_new_comm()
         self.initParams()
         self.init_dtype()
         self.init_fix_seed()
@@ -126,6 +131,11 @@ def init_fix_seed(self):
         self.fix_seed = True
 
 
+class TestClassCenterSampleOpWithNewComm(TestClassCenterSampleOp):
+    def with_new_comm(self):
+        os.environ["FLAGS_dynamic_static_unified_comm"] = "1"
+
+
 class TestClassCenterSampleV2(unittest.TestCase):
     def setUp(self):
         self.initParams()