92 lines
3.0 KiB
C++
92 lines
3.0 KiB
C++
|
|
#include <iostream>
|
|
#include <initializer_list>
|
|
#include <mpi.h>
|
|
|
|
/* Quick and dirty primality test */
|
|
bool is_prime(int num) {
|
|
if (num == 2 || num == 3) { return true; }
|
|
if ((num % 2) == 0 || (num % 3) == 0) { return false; }
|
|
|
|
for (int div = 5; div * div <= num; div += 6) {
|
|
if ((num % div) == 0 || ((num + 2) % div) == 0) { return false; }
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
int main(int argn, char* argv[]) {
|
|
|
|
MPI_Init(nullptr, nullptr);
|
|
|
|
int mpi_size;
|
|
MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
|
|
|
|
// Check if a 2D grid is (more or less) possible
|
|
if ((mpi_size > 2) && is_prime(mpi_size)) {
|
|
// Print warning only once (by rank 0)
|
|
int mpi_rank;
|
|
MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
|
|
if (mpi_rank == 0) {
|
|
std::cout << "Warning: 2D grid degenerates to 1D." << std::endl;
|
|
}
|
|
}
|
|
|
|
// Group processes into a cartesian communication topology
|
|
int mpi_dims[2] = {0, 0};
|
|
MPI_Dims_create(mpi_size, 2, mpi_dims);
|
|
|
|
// Setup a cartesian topology communicator (NON-cyclic)
|
|
const int mpi_periods[2] = {false, false};
|
|
MPI_Comm mpi_comm_grid;
|
|
MPI_Cart_create(
|
|
MPI_COMM_WORLD, // Old Communicator
|
|
2, // number of dimensions
|
|
mpi_dims, // grid dimensions
|
|
mpi_periods, // grid periodicity
|
|
true, // allow process reordering
|
|
&mpi_comm_grid
|
|
);
|
|
|
|
// Get rank with respect to the grid communicator
|
|
int mpi_rank;
|
|
MPI_Comm_rank(mpi_comm_grid, &mpi_rank);
|
|
|
|
// Get coordinates with respect to the grid communicator
|
|
int mpi_coords[2];
|
|
MPI_Cart_coords(mpi_comm_grid, mpi_rank, 2, mpi_coords);
|
|
|
|
// Get direct neightbours in the communication grid
|
|
struct { int north; int east; int south; int west; } mpi_neighbours;
|
|
// Get X-direction (dim 0) neightbours
|
|
MPI_Cart_shift(
|
|
mpi_comm_grid, // grid communicator
|
|
0, // axis index (0 <-> X)
|
|
1, // offset
|
|
&(mpi_neighbours.west), // negated offset neightbour
|
|
&(mpi_neighbours.east) // offset neightbour
|
|
);
|
|
// Get Y-direction (dim 1) neightbours
|
|
MPI_Cart_shift(
|
|
mpi_comm_grid,
|
|
1, // axis index (1 <-> X)
|
|
1,
|
|
&(mpi_neighbours.south),
|
|
&(mpi_neighbours.north)
|
|
);
|
|
|
|
// Print grid communicator location and with direct neighbors (offset +-1)
|
|
std::cout << "Rank " << mpi_rank << ": Coords ("
|
|
<< mpi_coords[0] << ", "
|
|
<< mpi_coords[1] << ") - Neighbours: ";
|
|
if (mpi_neighbours.west != MPI_PROC_NULL) { std::cout << mpi_neighbours.west << ' '; }
|
|
if (mpi_neighbours.east != MPI_PROC_NULL) { std::cout << mpi_neighbours.east << ' '; }
|
|
if (mpi_neighbours.south != MPI_PROC_NULL) { std::cout << mpi_neighbours.south << ' '; }
|
|
if (mpi_neighbours.north != MPI_PROC_NULL) { std::cout << mpi_neighbours.north << ' '; }
|
|
std::cout << std::endl;
|
|
|
|
MPI_Finalize();
|
|
|
|
return 0;
|
|
}
|