initial commit containing changes that allow for inputting a custom MPI...

initial commit containing changes that allow for inputting a custom MPI topology and rank-grid mapping for usage with amrex

Distribution.hpp

@@ -75,28 +75,57 @@ class Distribution {
 public:
 
 
-
+  //
   // constructors
+  //
+  
+  // standard setup
 
   Distribution(MPI_Comm comm, 
 	       int const n[], 
 	       bool debug = false)
-    : m_comm(comm), m_debug(debug)
+    : m_comm(comm), m_rmap(NULL), m_debug(debug)
   {
-    initialize(comm, n);
+    int Ndims[3] = { 0, 0, 0 };
+    initialize(comm, n, Ndims);
   }
 
   Distribution(MPI_Comm comm, 
 	       int ng, 
 	       bool debug = false)
-    : m_comm(comm), m_debug(debug)
+    : m_comm(comm), m_rmap(NULL), m_debug(debug)
   {
     int n[3] = { ng, ng, ng };
-    initialize(comm, n);
+    int Ndims[3] = { 0, 0, 0 };
+    initialize(comm, n, Ndims);
+  }
+
+  // custom setup with 3d decomposition, grid to rank map, and column-major flag 
+
+  Distribution(MPI_Comm comm,
+           int const n[],
+           int const Ndims[],
+           int *rmap,
+           bool debug = false)
+    : m_comm(comm), m_rmap(rmap), m_debug(debug)
+  {
+    initialize(comm, n, Ndims);
   }
 
+  Distribution(MPI_Comm comm,
+           int ng,
+           int const Ndims[],
+           int *rmap,
+           bool debug = false)
+    : m_comm(comm), m_rmap(rmap), m_debug(debug)
+  {
+    int n[3] = { ng, ng, ng };
+    initialize(comm, n, Ndims);
+  }
 
+  //
   // destructor
+  //
 
   virtual ~Distribution() 
   {
@@ -106,13 +135,14 @@ public:
 
   // initialization
   
-  void initialize(MPI_Comm comm, int const n[]) {
+  void initialize(MPI_Comm comm, int const n[], int const Ndims[]) {
     int flag;
     MPI_Initialized(&flag);
     if (flag == 0) {
       MPI_Init(0, 0);
     }
-    distribution_init(comm, n, &m_d, m_debug);
+
+    distribution_init(comm, n, Ndims, &m_d, m_rmap, m_debug);
   }
 
 
@@ -248,6 +278,7 @@ public:
 
 protected:
   MPI_Comm m_comm;
+  int* m_rmap;
   bool m_debug;
 };
 

distribution.c

@@ -245,7 +245,9 @@ void Rank_z_pencils(int * myrank,
 // create 1-, 2- and 3-d cartesian data distributions comm MPI Communicator
 void distribution_init(MPI_Comm comm, 
 		       const int n[], 
+		       const int Ndims[], 
 		       distribution_t *d, 
+               const int* rmap,
 		       bool debug)
 {
 /* 
@@ -287,6 +289,23 @@ void distribution_init(MPI_Comm comm,
   MPI_Comm_rank(comm, &self);
   MPI_Comm_size(comm, &nproc);
 
+  // Construct the rankmap[grid] --> rank map (converts grids to ranks) from the input 
+  // (if none is provided then we assume the trivial map {0, 1, ..., nproc}).
+  d->rankmap  = (int *) malloc(sizeof(int)*nproc);
+  if(rmap) for(int i=0; i<nproc; i++) d->rankmap[i] = rmap[i];
+  else for(int i=0; i<nproc; i++) d->rankmap[i] = i;
+
+  // Construct the inverse gridmap[rank] --> grid map (converts ranks to grids)
+  d->gridmap = (int *) malloc(sizeof(int)*nproc);
+  for(int i=0; i<nproc; i++) {
+    for(int j=0; j<nproc; j++) {
+      if(i == d->rankmap[j]) { d->gridmap[i] = j; break; }
+    }
+  }
+
+  // Map this rank to the correct grid
+  self = d->gridmap[self]; 
+
   if (!self) 
     printf("Initializing redistribution using a %s layout on %d ranks.\n",
 #ifdef PENCIL
@@ -336,7 +355,7 @@ void distribution_init(MPI_Comm comm,
   d->process_topology_1.n[0] = n[0] / d->process_topology_1.nproc[0];
   d->process_topology_1.n[1] = n[1] / d->process_topology_1.nproc[1];
   d->process_topology_1.n[2] = n[2] / d->process_topology_1.nproc[2];
-  
+
 
 
   // set up process grid with 3d decomposition (CUBE)
@@ -344,7 +363,7 @@ void distribution_init(MPI_Comm comm,
   d->process_topology_3.nproc[1] = 0;
   d->process_topology_3.nproc[2] = 0;
   period[0] = period[1] = period[2] = 1;
-  MPI_Dims_create(nproc, ndim, d->process_topology_3.nproc);
+  Custom3D_Dims_create(Ndims, nproc, ndim, d->process_topology_3.nproc);
 
   if(self == 0) {
     printf("distribution 3D: [%d:%d:%d]\n",
@@ -964,7 +983,22 @@ void distribution_init(MPI_Comm comm,
   d->d3_chunk=(complex_t *) malloc(sizeof(complex_t)*buff_size);
 }
 
+// Use MPI_Dims_create to create a 3D decomposition, or use a user-provided decomposition
+// if it is appropriate (it has the required number of processors)
+void Custom3D_Dims_create(const int Ndims[], int nproc, int ndims, int dims[])
+{
+
+  int check = 1;
+  for(int i=0; i<ndims; i++) check *= Ndims[i];
+
+  if(check == nproc) {
+    for(int i=0; i<ndims; i++) dims[i] = Ndims[i];
+  }
+  else {
+    MPI_Dims_create(nproc, ndims, dims);
+  } 
 
+}
 
 // create 1-, 2- and 3-d cartesian data distributions with explicitly
 // provided dimension lists
@@ -1192,6 +1226,8 @@ void distribution_fini(distribution_t *d)
   MPI_Comm_free(&d->process_topology_3.cart);
   free(d->d2_chunk);
   free(d->d3_chunk);
+  free(d->gridmap);
+  free(d->rankmap);
 }
 
 
@@ -1829,13 +1865,13 @@ static void redistribute_2_and_3(const complex_t *a,
     }
     if((self == me) && print_me)fprintf(stderr, "%d, %d, %d Made it before comm!...\n", self,p, npeers);
     
-    // record the communication to be done in a schedule
+    // record the communication to be done in a schedule. Make sure to map each grid to the correct rank
     if (direction == REDISTRIBUTE_3_TO_2) {
-      recv_peer = d3_peer;
-      send_peer = d2_peer;
+      recv_peer = d->rankmap[d3_peer];
+      send_peer = d->rankmap[d2_peer];
     } else if (direction == REDISTRIBUTE_2_TO_3) {
-      recv_peer = d2_peer;
-      send_peer = d3_peer;
+      recv_peer = d->rankmap[d2_peer];
+      send_peer = d->rankmap[d3_peer];
     } else {
       abort();
     }

distribution.h

@@ -100,6 +100,8 @@ typedef struct {
   process_topology_t process_topology_3;
   complex_t *d2_chunk;
   complex_t *d3_chunk;
+  int *gridmap;
+  int *rankmap;
   MPI_Comm parent;
 } distribution_t;
 
@@ -109,11 +111,15 @@ typedef struct {
 //   comm   MPI Communicator
 //   d      distribution descriptor
 //   n      (global) grid dimensions (3 element array)
+//   Ndims  3d process grid (3 element array: x, y, z)
+//   rmap   pointer to grid->rank map
 //   debug  debugging output
 ///
 void distribution_init(MPI_Comm comm, 
 		       const int n[], 
+               const int Ndims[],
 		       distribution_t *d, 
+               const int* rmap,
 		       bool debug);
 
 
@@ -138,6 +144,12 @@ void distribution_init_explicit(MPI_Comm comm,
                                 distribution_t *d,
                                 bool debug);
 
+///
+// creates a custom 3D decomposition or uses MPI_Dims_create to do so
+///
+void Custom3D_Dims_create(const int Ndims[], int nproc, int ndims, int dims[]);
+
+
 ///
 // clean up the data distribution
 //   d    distribution descriptor