mpi_gatherv_scatterv.cpp (2003B)
1 #include <algorithm> 2 #include <iostream> 3 #include <mpi.h> 4 #include <sstream> 5 #include <vector> 6 7 int main(int argc, char* argv[]) 8 { 9 int rank, size; 10 MPI_Init(&argc, &argv); 11 MPI_Comm_rank(MPI_COMM_WORLD, &rank); 12 MPI_Comm_size(MPI_COMM_WORLD, &size); 13 const bool isroot = (0 == rank); 14 15 std::vector<int> sendbuf(rank + 1, rank); 16 std::vector<int> recvbuf, recvcounts, displacement; 17 if (isroot) { 18 recvbuf.resize(size * (size + 1) / 2); 19 recvcounts.resize(size); 20 displacement.resize(size); 21 for (int i = 0; i < size; ++i) { 22 recvcounts[i] = i + 1; 23 displacement[i] = i * (i + 1) / 2; 24 } 25 } 26 27 MPI_Gatherv(sendbuf.data(), 28 sendbuf.size(), 29 MPI_INT, 30 recvbuf.data(), 31 recvcounts.data(), 32 displacement.data(), 33 MPI_INT, 34 0, 35 MPI_COMM_WORLD); 36 37 if (isroot) { 38 std::cout << "Result of MPI_Gatherv (on root rank)\n"; 39 for (int r = 0; r < size; ++r) { 40 const int offset = displacement[r]; 41 for (int i = 0; i < recvcounts[r]; ++i) { 42 std::cout << ' ' << recvbuf[offset + i]; 43 } 44 std::cout << '\n'; 45 } 46 } 47 48 // inverse operation 49 std::vector<int> sendcounts; 50 sendbuf.swap(recvbuf); 51 if (isroot) { 52 sendcounts.swap(recvcounts); 53 } 54 55 MPI_Scatterv(sendbuf.data(), 56 sendcounts.data(), 57 displacement.data(), 58 MPI_INT, 59 recvbuf.data(), 60 recvbuf.size(), 61 MPI_INT, 62 0, 63 MPI_COMM_WORLD); 64 65 if (isroot) { 66 std::cout << "\nResult of MPI_Scatterv\n"; 67 } 68 std::ostringstream out; 69 out << "Rank " << rank << ":"; 70 for (const auto v : recvbuf) { 71 out << ' ' << v; 72 } 73 out << '\n'; 74 std::cout << out.str(); 75 76 MPI_Finalize(); 77 return 0; 78 }