Method PrepareNCCLCommunicator

paddle/fluid/framework/compiled_program.cc:801–847

}

void CompiledProgram::PrepareNCCLCommunicator(Scope *global_scope) {
  if (member_->build_strategy_.reduce_ ==
      BuildStrategy::ReduceStrategy::kNoReduce) {
    return;
  }

  if (member_->IsUseCUDA(member_->use_device_) && member_->nranks_ > 1) {
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
    member_->InitOrGetNCCLCommunicator(global_scope, &member_->build_strategy_);

    // Initialize device context's nccl comm, will be used by normal
    // Operators like sync_batch_norm, and collective ops.
    // NOTE: more than one CompiledProgram with same place, the nccl comm will
    // be rewrite and there will be some problem.
    // NOTE: NCCL group-calls and non-group-calls can not use the same
    // NCCL communicator, so for ParallelGraph and Multi-Process mode, re-use
    // same communicators.
    auto *nccl_ctxs = member_->nccl_ctxs_->GetSyncBatchNormCtx(
        global_scope, member_->places_);
    auto &pool = phi::DeviceContextPool::Instance();
    for (auto &place : member_->places_) {
      auto *dev_ctx = static_cast<phi::GPUContext *>(pool.Get(place));
      auto &nccl_ctx = nccl_ctxs->at(place);
      dev_ctx->set_nccl_comm(nccl_ctx.comm());
    }
#else
    PADDLE_THROW(common::errors::PreconditionNotMet("Not compiled with CUDA."));
#endif
  }
  if (member_->use_device_ == p::kXPU && member_->nranks_ > 1) {
#if defined(PADDLE_WITH_XPU_BKCL)
    member_->InitOrGetBKCLCommunicator(global_scope, member_->build_strategy_);

    auto *bkcl_ctxs = member_->bkcl_ctxs_->GetSyncBatchNormCtx(
        global_scope, member_->places_);
    auto &pool = phi::DeviceContextPool::Instance();
    for (size_t dev_id = 0; dev_id < member_->places_.size(); ++dev_id) {
      auto *dev_ctx =
          static_cast<phi::XPUContext *>(pool.Get(member_->places_[dev_id]));
      auto &bkcl_ctx = bkcl_ctxs->at(member_->places_[dev_id]);
      dev_ctx->SetBkclContext(bkcl_ctx.comm());
    }
#else
    PADDLE_THROW(common::errors::PreconditionNotMet("Not compiled with XPU."));
#endif
  }
}

std::vector<ir::Graph *> CompiledProgram::CompileGraphWithBuildStrategy(
    ir::Graph *graph,