| 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) { |
| 97 | Log::Fatal("Please initilize the network interface first"); |
| 98 | } |
| 99 | // assign blocks |
| 100 | comm_size_t all_size = input_size * num_machines_; |
| 101 | block_start_[0] = 0; |
| 102 | block_len_[0] = input_size; |
| 103 | for (int i = 1; i < num_machines_; ++i) { |
| 104 | block_start_[i] = block_start_[i - 1] + block_len_[i - 1]; |
| 105 | block_len_[i] = input_size; |
| 106 | } |
| 107 | // need use buffer here, since size of "output" is smaller than size after all gather |
| 108 | if (input_size*num_machines_ > buffer_size_) { |
| 109 | buffer_size_ = input_size*num_machines_; |
| 110 | buffer_.resize(buffer_size_); |
| 111 | } |
| 112 | |
| 113 | Allgather(input, block_start_.data(), block_len_.data(), buffer_.data(), all_size); |
| 114 | for (int i = 1; i < num_machines_; ++i) { |
| 115 | reducer(buffer_.data() + block_start_[i], buffer_.data() + block_start_[0], type_size, input_size); |
| 116 | } |
| 117 | // copy back |
| 118 | std::memcpy(output, buffer_.data(), input_size); |
| 119 | } |
| 120 | |
| 121 | void Network::Allgather(char* input, comm_size_t send_size, char* output) { |
| 122 | if (num_machines_ <= 1) { |