Applied the same trick to the foreign gparts.

src/cell.c

@@ -69,11 +69,6 @@
 /* Global variables. */
 int cell_next_tag = 0;
 
-int depth_0_counter = 0;
-int depth_1_counter = 0;
-int depth_2_counter = 0;
-int depth_3_counter = 0;
-
 /**
  * @brief Get the size of the cell subtree.
  *
@@ -104,6 +99,9 @@ int cell_getsize(struct cell *c) {
 int cell_link_parts(struct cell *c, struct part *parts) {
 
 #ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Linking foreign particles in a local cell!");
+
   if(c->hydro.parts != NULL)
     error("Linking parts into a cell that was already linked");
 #endif
@@ -123,57 +121,23 @@ int cell_link_parts(struct cell *c, struct part *parts) {
   return c->hydro.count;
 }
 
-
 /**
- * @brief Link the cells recursively to the given #part array.
+ * @brief Link the cells recursively to the given #gpart array.
  *
  * @param c The #cell.
- * @param parts The #part array.
+ * @param gparts The #gpart array.
  *
  * @return The number of particles linked.
  */
-int cell_link_foreign_parts(struct cell *c, struct part *parts) {
+int cell_link_gparts(struct cell *c, struct gpart *gparts) {
 
 #ifdef SWIFT_DEBUG_CHECKS
   if (c->nodeID == engine_rank)
     error("Linking foreign particles in a local cell!");
-#endif
-
-  /* Do we have a hydro task at this level? */
-  if (c->hydro.density != NULL) {
 
-    /* Recursively attach the parts */
-    const int counts = cell_link_parts(c, parts);
-#ifdef SWIFT_DEBUG_CHECKS
-    if (counts != c->hydro.count)
-      error("Something is wrong with the foreign counts");
+  if(c->grav.parts != NULL)
+    error("Linking gparts into a cell that was already linked");
 #endif
-    return counts;
-  }
-
-  /* Go deeper to find the level where the tasks are */
-  if (c->split) {
-    int count = 0;
-    for (int k = 0; k < 8; k++) {
-      if (c->progeny[k] != NULL) {
-        count += cell_link_foreign_parts(c->progeny[k], &parts[count]);
-      }
-    }
-    return count;
-  } else {
-    return 0;
-  }
-}
-
-/**
- * @brief Link the cells recursively to the given #gpart array.
- *
- * @param c The #cell.
- * @param gparts The #gpart array.
- *
- * @return The number of particles linked.
- */
-int cell_link_gparts(struct cell *c, struct gpart *gparts) {
 
   c->grav.parts = gparts;
 
@@ -200,6 +164,14 @@ int cell_link_gparts(struct cell *c, struct gpart *gparts) {
  */
 int cell_link_sparts(struct cell *c, struct spart *sparts) {
 
+#ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Linking foreign particles in a local cell!");
+
+  if(c->stars.parts != NULL)
+    error("Linking sparts into a cell that was already linked");
+#endif
+
   c->stars.parts = sparts;
 
   /* Fill the progeny recursively, depth-first. */
@@ -215,6 +187,72 @@ int cell_link_sparts(struct cell *c, struct spart *sparts) {
   return c->stars.count;
 }
 
+int cell_link_foreign_parts(struct cell *c, struct part *parts) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Linking foreign particles in a local cell!");
+#endif
+
+  /* Do we have a hydro task at this level? */
+  if (c->hydro.density != NULL) {
+
+    /* Recursively attach the parts */
+    const int counts = cell_link_parts(c, parts);
+#ifdef SWIFT_DEBUG_CHECKS
+    if (counts != c->hydro.count)
+      error("Something is wrong with the foreign counts");
+#endif
+    return counts;
+  }
+
+  /* Go deeper to find the level where the tasks are */
+  if (c->split) {
+    int count = 0;
+    for (int k = 0; k < 8; k++) {
+      if (c->progeny[k] != NULL) {
+        count += cell_link_foreign_parts(c->progeny[k], &parts[count]);
+      }
+    }
+    return count;
+  } else {
+    return 0;
+  }
+}
+
+int cell_link_foreign_gparts(struct cell *c, struct gpart *gparts) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Linking foreign particles in a local cell!");
+#endif
+
+  /* Do we have a hydro task at this level? */
+  if (c->grav.grav != NULL) {
+
+    /* Recursively attach the parts */
+    const int counts = cell_link_gparts(c, gparts);
+#ifdef SWIFT_DEBUG_CHECKS
+    if (counts != c->grav.count)
+      error("Something is wrong with the foreign counts");
+#endif
+    return counts;
+  }
+
+  /* Go deeper to find the level where the tasks are */
+  if (c->split) {
+    int count = 0;
+    for (int k = 0; k < 8; k++) {
+      if (c->progeny[k] != NULL) {
+        count += cell_link_foreign_gparts(c->progeny[k], &gparts[count]);
+      }
+    }
+    return count;
+  } else {
+    return 0;
+  }
+}
+
 int cell_count_parts_for_tasks(const struct cell *c) {
 
 #ifdef SWIFT_DEBUG_CHECKS
@@ -224,10 +262,6 @@ int cell_count_parts_for_tasks(const struct cell *c) {
 
   /* Do we have a hydro task at this level? */
   if (c->hydro.density != NULL) {
-    if (c->depth == 0) ++depth_0_counter;
-    if (c->depth == 1) ++depth_1_counter;
-    if (c->depth == 2) ++depth_2_counter;
-    if (c->depth == 3) ++depth_3_counter;
     return c->hydro.count;
   }
 
@@ -244,6 +278,32 @@ int cell_count_parts_for_tasks(const struct cell *c) {
   }
 }
 
+int cell_count_gparts_for_tasks(const struct cell *c) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Counting foreign particles in a local cell!");
+#endif
+
+  /* Do we have a hydro task at this level? */
+  if (c->grav.grav != NULL) {
+    return c->grav.count;
+  }
+
+  if (c->split) {
+    int count = 0;
+    for (int k = 0; k < 8; ++k) {
+      if (c->progeny[k] != NULL) {
+        count += cell_count_gparts_for_tasks(c->progeny[k]);
+      }
+    }
+    return count;
+  } else {
+    return 0;
+  }
+}
+
+
 /**
  * @brief Pack the data of the given cell and all it's sub-cells.
  *

src/cell.h

@@ -689,10 +689,12 @@ int cell_pack_multipoles(struct cell *c, struct gravity_tensors *m);
 int cell_unpack_multipoles(struct cell *c, struct gravity_tensors *m);
 int cell_getsize(struct cell *c);
 int cell_link_parts(struct cell *c, struct part *parts);
-int cell_link_foreign_parts(struct cell *c, struct part *parts);
 int cell_link_gparts(struct cell *c, struct gpart *gparts);
 int cell_link_sparts(struct cell *c, struct spart *sparts);
+int cell_link_foreign_parts(struct cell *c, struct part *parts);
+int cell_link_foreign_gparts(struct cell *c, struct gpart *gparts);
 int cell_count_parts_for_tasks(const struct cell *c);
+int cell_count_gparts_for_tasks(const struct cell *c);
 void cell_clean_links(struct cell *c, void *data);
 void cell_make_multipoles(struct cell *c, integertime_t ti_current);
 void cell_check_multipole(struct cell *c);

src/engine.c

@@ -1750,11 +1750,6 @@ void engine_exchange_proxy_multipoles(struct engine *e) {
 #endif
 }
 
-extern int depth_0_counter;
-extern int depth_1_counter;
-extern int depth_2_counter;
-extern int depth_3_counter;
-
 void engine_allocate_foreign_particles(struct engine *e) {
 
 #ifdef WITH_MPI
@@ -1765,23 +1760,26 @@ void engine_allocate_foreign_particles(struct engine *e) {
 
   /* Count the number of particles we need to import and re-allocate
      the buffer if needed. */
-  size_t count_parts_in_old = 0;
-  size_t count_parts_in = 0;
+  size_t count_parts_in_old = 0, count_gparts_in_old = 0;
+  size_t count_parts_in = 0, count_gparts_in = 0;
   for (int k = 0; k < nr_proxies; k++) {
     for (int j = 0; j < e->proxies[k].nr_cells_in; j++) {
 
       if (e->proxies[k].cells_in_type[j] & proxy_cell_type_hydro) {
         count_parts_in_old += e->proxies[k].cells_in[j]->hydro.count;
-
         count_parts_in += cell_count_parts_for_tasks(e->proxies[k].cells_in[j]);
       }
+
+      if (e->proxies[k].cells_in_type[j] & proxy_cell_type_gravity) {
+        count_gparts_in_old += e->proxies[k].cells_in[j]->grav.count;
+        count_gparts_in += cell_count_gparts_for_tasks(e->proxies[k].cells_in[j]);
+      }
+
     }
   }
 
-  message("New count: %zd old count: %zd", count_parts_in, count_parts_in_old);
-
-  message("counters: %d %d %d %d", depth_0_counter, depth_1_counter,
-          depth_2_counter, depth_3_counter);
+  message("New parts count: %zd old parts count: %zd", count_parts_in, count_parts_in_old);
+  message("New gparts count: %zd old gparts count: %zd", count_gparts_in, count_gparts_in_old);
 
   /* Allocate space for the foreign particles we will receive */
   if (count_parts_in > s->size_parts_foreign) {
@@ -1791,9 +1789,19 @@ void engine_allocate_foreign_particles(struct engine *e) {
                        sizeof(struct part) * s->size_parts_foreign) != 0)
       error("Failed to allocate foreign part data.");
   }
+  /* Allocate space for the foreign particles we will receive */
+  if (count_gparts_in > s->size_gparts_foreign) {
+    if (s->gparts_foreign != NULL) free(s->gparts_foreign);
+    s->size_gparts_foreign = 1.1 * count_gparts_in;
+    if (posix_memalign((void **)&s->gparts_foreign, part_align,
+                       sizeof(struct gpart) * s->size_gparts_foreign) != 0)
+      error("Failed to allocate foreign gpart data.");
+  }
 
   struct part *parts = s->parts_foreign;
+  struct gpart *gparts = s->gparts_foreign;
   size_t total_count_parts = 0;
+  size_t total_count_gparts = 0;
   for (int k = 0; k < nr_proxies; k++) {
     for (int j = 0; j < e->proxies[k].nr_cells_in; j++) {
 
@@ -1803,11 +1811,19 @@ void engine_allocate_foreign_particles(struct engine *e) {
         parts = &parts[count_parts];
 	total_count_parts += count_parts;
       }
+
+      if (e->proxies[k].cells_in_type[j] & proxy_cell_type_gravity) {
+
+        const size_t count_gparts = cell_link_foreign_gparts(e->proxies[k].cells_in[j], gparts);
+        gparts = &gparts[count_gparts];
+	total_count_gparts += count_gparts;
+      }
     }
   }
 
   /* Update the counters */
   s->nr_parts_foreign = parts - s->parts_foreign;
+  s->nr_gparts_foreign = gparts - s->gparts_foreign;
 
   
   message("count_parts: %zd %zd", count_parts_in, total_count_parts);