DeepLearningで使われるallreduceのアルゴリズム 試して理解

先日書いた記事で書いたallreduceを簡単に実装してみて，理解できているか試した． kuroko1t.hatenablog.com

コードはここ． （2D-Torus all-reduceとdouble binary treesはまだ．．）

hierachical

最初のグループ内のReduce用にMPI_Comm_splitを利用してcomm分割．
Reduceして集めた各グループのSUM値をallreduceできるように再度MPI_Comm_splitを利用してcommを分けallreduce.

 if (local_rank == 0) {
      MPI_Allreduce(&out.front(), &out1.front(), size, MPI_INT, MPI_SUM, group_comm);
 }

ring-allreduce

MPI_Send, MPI_Recvを利用して実装してみた． 同期処理の1対１通信を利用したのでrankの偶奇でsend,recv入れ替えた実装を行った． Isend, Irecvを利用すればこのような処理は必要なさそう．各プロセスが１つのSUM値を持っているので最後にMPI_Bcastして完了とした．

// ring reduce for P-1 times
if (rank %2 == 0) {
  MPI_CHECK(MPI_Send(&out1.front() + int64_t(data_send),
                     1, MPI_INT, send_rank, rank, MPI_COMM_WORLD));
  MPI_CHECK(MPI_Recv(&out_tmp.front() + int64_t(data_recv),
                     1, MPI_INT, recv_rank, recv_rank, MPI_COMM_WORLD, NULL));
  out1[data_recv] += out_tmp[data_recv];
} else {
  MPI_CHECK(MPI_Recv(&out_tmp.front()+ int64_t(data_recv),
                     1, MPI_INT, recv_rank, recv_rank, MPI_COMM_WORLD, NULL));
  MPI_CHECK(MPI_Send(&out1.front() + int64_t(data_send),
                     1, MPI_INT, send_rank, rank, MPI_COMM_WORLD));
  out1[data_recv] += out_tmp[data_recv];
}