| 66 | } |
| 67 | |
| 68 | void Network::Allreduce(char* input, comm_size_t input_size, int type_size, char* output, const ReduceFunction& reducer) { |
| 69 | if (num_machines_ <= 1) { |
| 70 | Log::Fatal("Please initilize the network interface first"); |
| 71 | } |
| 72 | comm_size_t count = input_size / type_size; |
| 73 | // if small package or small count , do it by all gather.(reduce the communication times.) |
| 74 | if (count < num_machines_ || input_size < 4096) { |
| 75 | AllreduceByAllGather(input, input_size, type_size, output, reducer); |
| 76 | return; |
| 77 | } |
| 78 | // assign the blocks to every rank. |
| 79 | comm_size_t step = (count + num_machines_ - 1) / num_machines_; |
| 80 | if (step < 1) { |
| 81 | step = 1; |
| 82 | } |
| 83 | block_start_[0] = 0; |
| 84 | for (int i = 0; i < num_machines_ - 1; ++i) { |
| 85 | block_len_[i] = std::min<comm_size_t>(step * type_size, input_size - block_start_[i]); |
| 86 | block_start_[i + 1] = block_start_[i] + block_len_[i]; |
| 87 | } |
| 88 | block_len_[num_machines_ - 1] = input_size - block_start_[num_machines_ - 1]; |
| 89 | // do reduce scatter |
| 90 | ReduceScatter(input, input_size, type_size, block_start_.data(), block_len_.data(), output, input_size, reducer); |
| 91 | // do all gather |
| 92 | Allgather(output, block_start_.data(), block_len_.data(), output, input_size); |
| 93 | } |
| 94 | |
| 95 | void Network::AllreduceByAllGather(char* input, comm_size_t input_size, int type_size, char* output, const ReduceFunction& reducer) { |
| 96 | if (num_machines_ <= 1) { |