#include "paddle/phi/core/distributed/comm_context_manager.h"

#include "paddle/phi/core/distributed/collective/process_group.h"

auto map = phi::distributed::ProcessGroupMapFromGid::getInstance();

if (map->has(ring_id)) {

// Use ProcessGroup

phi::distributed::ProcessGroup* pg = map->get(ring_id);

std::vector<phi::DenseTensor> in_tensor;

std::vector<phi::DenseTensor> out_tensor;

in_tensor.push_back(num_classes_per_device);

out_tensor.push_back(num_classes_per_device);

phi::distributed::AllreduceOptions opts;

opts.reduce_op = phi::distributed::ReduceOp::SUM;

auto task = pg->AllReduce(in_tensor, out_tensor, opts);

task->Wait();

} else {

phi::distributed::NCCLCommContext* comm_ctx = nullptr;

// use global calculate stream

auto stream = dev_ctx.stream();

const auto& comm_context_manager =

phi::distributed::CommContextManager::GetInstance();

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();

comm_ctx->AllReduce(

&num_classes_per_device, num_classes_per_device, ncclSum, stream);

phi::backends::gpu::GpuStreamSync(stream);

}