| 164 | |
| 165 | template<typename Dtype> |
| 166 | void NCCL<Dtype>::InitSingleProcess(vector<NCCL<Dtype>*>* nccls) { |
| 167 | ncclComm_t* comms = new ncclComm_t[nccls->size()]; |
| 168 | int* gpu_list = new int[nccls->size()]; |
| 169 | for (int i = 0; i < nccls->size(); ++i) { |
| 170 | gpu_list[i] = (*nccls)[i]->solver_->param().device_id(); |
| 171 | } |
| 172 | NCCL_CHECK(ncclCommInitAll(comms, static_cast<int>(nccls->size()), gpu_list)); |
| 173 | for (int i = 0; i < nccls->size(); ++i) { |
| 174 | (*nccls)[i]->comm_ = comms[i]; |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | template<typename Dtype> |
| 179 | string NCCL<Dtype>::new_uid() { |