| 247 | } |
| 248 | |
| 249 | void Network::ReduceScatterRecursiveHalving(char* input, comm_size_t input_size, int type_size, |
| 250 | const comm_size_t* block_start, const comm_size_t* block_len, char* output, |
| 251 | comm_size_t, const ReduceFunction& reducer) { |
| 252 | if (!recursive_halving_map_.is_power_of_2) { |
| 253 | if (recursive_halving_map_.type == RecursiveHalvingNodeType::Other) { |
| 254 | // send local data to neighbor first |
| 255 | linkers_->Send(recursive_halving_map_.neighbor, input, input_size); |
| 256 | } else if (recursive_halving_map_.type == RecursiveHalvingNodeType::GroupLeader) { |
| 257 | // receive neighbor data first |
| 258 | int need_recv_cnt = input_size; |
| 259 | linkers_->Recv(recursive_halving_map_.neighbor, output, need_recv_cnt); |
| 260 | // reduce |
| 261 | reducer(output, input, type_size, input_size); |
| 262 | } |
| 263 | } |
| 264 | if (recursive_halving_map_.type != RecursiveHalvingNodeType::Other) { |
| 265 | for (int i = 0; i < recursive_halving_map_.k; ++i) { |
| 266 | // get target |
| 267 | int target = recursive_halving_map_.ranks[i]; |
| 268 | comm_size_t send_block_start = recursive_halving_map_.send_block_start[i]; |
| 269 | comm_size_t recv_block_start = recursive_halving_map_.recv_block_start[i]; |
| 270 | // get send information |
| 271 | comm_size_t send_size = 0; |
| 272 | for (int j = 0; j < recursive_halving_map_.send_block_len[i]; ++j) { |
| 273 | send_size += block_len[send_block_start + j]; |
| 274 | } |
| 275 | // get recv information |
| 276 | comm_size_t need_recv_cnt = 0; |
| 277 | for (int j = 0; j < recursive_halving_map_.recv_block_len[i]; ++j) { |
| 278 | need_recv_cnt += block_len[recv_block_start + j]; |
| 279 | } |
| 280 | // send and recv at same time |
| 281 | linkers_->SendRecv(target, input + block_start[send_block_start], send_size, target, output, need_recv_cnt); |
| 282 | // reduce |
| 283 | reducer(output, input + block_start[recv_block_start], type_size, need_recv_cnt); |
| 284 | } |
| 285 | } |
| 286 | if (!recursive_halving_map_.is_power_of_2) { |
| 287 | if (recursive_halving_map_.type == RecursiveHalvingNodeType::GroupLeader) { |
| 288 | // send result to neighbor |
| 289 | linkers_->Send(recursive_halving_map_.neighbor, |
| 290 | input + block_start[recursive_halving_map_.neighbor], |
| 291 | block_len[recursive_halving_map_.neighbor]); |
| 292 | } else if (recursive_halving_map_.type == RecursiveHalvingNodeType::Other) { |
| 293 | // receive result from neighbor |
| 294 | int need_recv_cnt = block_len[rank_]; |
| 295 | linkers_->Recv(recursive_halving_map_.neighbor, output, need_recv_cnt); |
| 296 | return; |
| 297 | } |
| 298 | } |
| 299 | // copy result |
| 300 | std::memcpy(output, input + block_start[rank_], block_len[rank_]); |
| 301 | } |
| 302 | |
| 303 | void Network::ReduceScatterRing(char* input, comm_size_t, int type_size, |
| 304 | const comm_size_t* block_start, const comm_size_t* block_len, char* output, |