Refactor METIS cell graph initialisation code into one function.

src/engine.c

@@ -34,6 +34,10 @@
 /* MPI headers. */
 #ifdef WITH_MPI
 #include <mpi.h>
+/* METIS headers only used when MPI is also available. */
+#ifdef HAVE_METIS
+#include <metis.h>
+#endif
 #endif
 
 #ifdef HAVE_LIBNUMA
@@ -311,49 +315,18 @@ void engine_repartition(struct engine *e) {
      * drop the vertices for _EDGE and _VERTEX_EDGE. */
     int nr_cells = s->nr_cells;
     struct cell *cells = s->cells;
-    int ind[3], *cdim = s->cdim;
     struct task *t, *tasks = e->sched.tasks;
     float wscale = 1e-3, vscale = 1e-3, wscale_buff;
     int wtot = 0;
     int wmax = 1e9 / e->nr_nodes;
     int wmin;
 
-    /* Allocate and fill the adjcny indexing array. */
+    /* Allocate and fill the adjncy indexing array defining the graph of
+     * cells.*/
     int *inds;
-    if ((inds = (int *)malloc(sizeof(int) * 26 * nr_cells)) == NULL)
+    if ((inds = (int *)malloc(sizeof(idx_t) * 26 * nr_cells)) == NULL)
       error("Failed to allocate the inds array");
-
-    for (int cid = 0; cid < nr_cells; cid++) {
-      ind[0] = cells[cid].loc[0] / cells[cid].h[0] + 0.5;
-      ind[1] = cells[cid].loc[1] / cells[cid].h[1] + 0.5;
-      ind[2] = cells[cid].loc[2] / cells[cid].h[2] + 0.5;
-      int l = 0;
-      for (int i = -1; i <= 1; i++) {
-        int ii = ind[0] + i;
-        if (ii < 0)
-          ii += cdim[0];
-        else if (ii >= cdim[0])
-          ii -= cdim[0];
-        for (int j = -1; j <= 1; j++) {
-          int jj = ind[1] + j;
-          if (jj < 0)
-            jj += cdim[1];
-          else if (jj >= cdim[1])
-            jj -= cdim[1];
-          for (int k = -1; k <= 1; k++) {
-            int kk = ind[2] + k;
-            if (kk < 0)
-              kk += cdim[2];
-            else if (kk >= cdim[2])
-              kk -= cdim[2];
-            if (i || j || k) {
-              inds[cid * 26 + l] = cell_getid(cdim, ii, jj, kk);
-              l += 1;
-            }
-          }
-        }
-      }
-    }
+    part_graph_init_metis(s, inds, NULL);
 
     /* Allocate and init weights. */
     int *weights_v = NULL;

src/partition.c

@@ -154,6 +154,74 @@ void part_split_vector(struct space *s, int nregions, int *samplecells) {
  * weighted by the number of particles scheme. Note METIS is optional.
  */
 
+/** 
+ * @brief Fill the METIS xadj and adjncy arrays defining the graph of cells
+ *        in a space.
+ *
+ * See the METIS manual if you want to understand this format. The cell graph
+ * consists of all nodes as vertices with edges as the connections to all
+ * neighbours, so we have 26 per vertex.
+ *
+ * @param s the space of cells.
+ * @param adjncy the METIS adjncy array to fill, must be of size
+ *               26 * the number of cells in the space.
+ * @param xadj the METIS xadj array to fill, must be of size
+ *             number of cells in space + 1. NULL for not used.
+ */
+#ifdef HAVE_METIS
+void part_graph_init_metis(struct space *s, idx_t *adjncy, idx_t *xadj) {
+
+  /* Loop over all cells in the space. */
+  int cid = 0;
+  for (int l = 0; l < s->cdim[0]; l++) {
+    for (int m = 0; m < s->cdim[1]; m++) {
+      for (int n = 0; n < s->cdim[2]; n++) {
+
+        /* Visit all neighbours of this cell, wrapping space at edges. */
+        int p = 0;
+        for (int i = -1; i <= 1; i++) {
+          int ii = l + i;
+          if (ii < 0)
+            ii += s->cdim[0];
+          else if (ii >= s->cdim[0])
+            ii -= s->cdim[0];
+          for (int j = -1; j <= 1; j++) {
+            int jj = m + j;
+            if (jj < 0)
+              jj += s->cdim[1];
+            else if (jj >= s->cdim[1])
+              jj -= s->cdim[1];
+            for (int k = -1; k <= 1; k++) {
+              int kk = n + k;
+              if (kk < 0)
+                kk += s->cdim[2];
+              else if (kk >= s->cdim[2])
+                kk -= s->cdim[2];
+
+              /* If not self, record id of neighbour. */
+              if (i || j || k) {
+                adjncy[cid * 26 + p] = cell_getid(s->cdim, ii, jj, kk);
+                p++;
+              }
+            }
+          }
+        }
+
+        /* Next cell. */
+        cid++;
+      }
+    }
+  }
+
+  /* If given set xadj. */
+  if (xadj != NULL) {
+    xadj[0] = 0;
+    for (int k = 0; k < s->nr_cells; k++) xadj[k + 1] = xadj[k] + 26;
+
+  }
+}
+#endif
+
 /**
  * @brief Partition the given space into a number of connected regions.
  *
@@ -166,13 +234,13 @@ void part_split_vector(struct space *s, int nregions, int *samplecells) {
  *        NULL for unit weights.
  * @param edgew weights for the graph edges between all cells, sizeof number
  *        of cells * 26 if used, NULL for unit weights. Need to be packed
- *        in CSR format, so same as adjcny array.
+ *        in CSR format, so same as adjncy array.
  * @param celllist on exit this contains the ids of the selected regions,
  *        sizeof number of cells.
  */
 void part_pick_metis(struct space *s, int nregions, int *vertexw, int *edgew,
                      int *celllist) {
-#if defined(HAVE_METIS)
+#ifdef HAVE_METIS
 
   /* Total number of cells. */
   int ncells = s->cdim[0] * s->cdim[1] * s->cdim[2];
@@ -203,49 +271,8 @@ void part_pick_metis(struct space *s, int nregions, int *vertexw, int *edgew,
   if ((regionid = (idx_t *)malloc(sizeof(idx_t) * ncells)) == NULL)
     error("Failed to allocate regionid array");
 
-  /* Fill the xadj and adjncy array to define the graph of cells. */
-  /* Loop over all cells. */
-  int cid = 0;
-  for (int l = 0; l < s->cdim[0]; l++) {
-    for (int m = 0; m < s->cdim[1]; m++) {
-      for (int n = 0; n < s->cdim[2]; n++) {
-        /* Visit all neighbours of this cell, wrapping space at edges. */
-        int p = 0;
-        for (int i = -1; i <= 1; i++) {
-          int ii = l + i;
-          if (ii < 0)
-            ii += s->cdim[0];
-          else if (ii >= s->cdim[0])
-            ii -= s->cdim[0];
-          for (int j = -1; j <= 1; j++) {
-            int jj = m + j;
-            if (jj < 0)
-              jj += s->cdim[1];
-            else if (jj >= s->cdim[1])
-              jj -= s->cdim[1];
-            for (int k = -1; k <= 1; k++) {
-              int kk = n + k;
-              if (kk < 0)
-                kk += s->cdim[2];
-              else if (kk >= s->cdim[2])
-                kk -= s->cdim[2];
-
-              /* If not self, record id of neighbour. */
-              if (i || j || k) {
-                adjncy[cid * 26 + p] = cell_getid(s->cdim, ii, jj, kk);
-                p++;
-              }
-            }
-          }
-        }
-
-        /* Next cell. */
-        cid++;
-      }
-    }
-  }
-  xadj[0] = 0;
-  for (int k = 0; k < ncells; k++) xadj[k + 1] = xadj[k] + 26;
+  /* Define the cell graph. */
+  part_graph_init_metis(s, adjncy, xadj);
 
   /* Init the vertex weights array. */
   if (vertexw != NULL) {

src/partition.h

@@ -21,6 +21,9 @@
 
 #include "space.h"
 #include "cell.h"
+#ifdef HAVE_METIS
+#include <metis.h>
+#endif
 
 /* Initial partitioning types. */
 enum initpart_type {
@@ -58,6 +61,10 @@ void part_pick_metis(struct space *s, int nregions, int *vertexw, int *edgew,
                      int *celllist);
 void part_split_metis(struct space *s, int nregions, int *celllist);
 
+#ifdef HAVE_METIS
+void part_graph_init_metis(struct space *s, idx_t *adjncy, idx_t *xadj);
+#endif
+
 int part_check_complete(struct space *s, int verbose, int nregions);
 
 #endif /* SWIFT_PARTITION_H */