| 187 | |
| 188 | template<typename Dtype> |
| 189 | void NCCL<Dtype>::Broadcast() { |
| 190 | if (barrier_) { // NULL in multi process case |
| 191 | barrier_->wait(); |
| 192 | } |
| 193 | NCCL_CHECK(ncclBcast(data_, static_cast<int>(size_), |
| 194 | nccl::dataType<Dtype>::type, 0, |
| 195 | comm_, cudaStreamDefault)); |
| 196 | if (barrier_) { |
| 197 | barrier_->wait(); |
| 198 | } |
| 199 | } |
| 200 | |
| 201 | template<typename Dtype> |
| 202 | void NCCL<Dtype>::run(int layer) { |