Reduce the SWIFT memory usage during a VR call by freeing the foreign particle buffers

src/cell.c

@@ -409,6 +409,37 @@ int cell_link_foreign_gparts(struct cell *c, struct gpart *gparts) {
 #endif
 }
 
+/**
+ * @brief Recursively nullify all the particle pointers in a cell hierarchy.
+ *
+ * Should only be used on foreign cells!
+ *
+ * This will make any task or action running on these cells likely crash.
+ * Recreating the foreign links will be necessary.
+ *
+ * @param c The #cell to act on.
+ */
+void cell_unlink_foreign_particles(struct cell *c) {
+
+#ifdef SWIFT_DEBUG_CHECKS
+  if (c->nodeID == engine_rank)
+    error("Unlinking foreign particles in a local cell!");
+#endif
+
+  c->grav.parts = NULL;
+  c->hydro.parts = NULL;
+  c->stars.parts = NULL;
+  c->black_holes.parts = NULL;
+
+  if (c->split) {
+    for (int k = 0; k < 8; k++) {
+      if (c->progeny[k] != NULL) {
+        cell_unlink_foreign_particles(c->progeny[k]);
+      }
+    }
+  }
+}
+
 /**
  * @brief Recursively count the number of #part in foreign cells that
  * are in cells with hydro-related tasks.

src/cell.h

@@ -860,6 +860,7 @@ int cell_link_sparts(struct cell *c, struct spart *sparts);
 int cell_link_bparts(struct cell *c, struct bpart *bparts);
 int cell_link_foreign_parts(struct cell *c, struct part *parts);
 int cell_link_foreign_gparts(struct cell *c, struct gpart *gparts);
+void cell_unlink_foreign_particles(struct cell *c);
 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);

src/engine.c

@@ -211,8 +211,9 @@ void engine_repartition(struct engine *e) {
   /* Task arrays. */
   scheduler_free_tasks(&e->sched);
 
-  /* Foreign parts. */
-  space_free_foreign_parts(e->s);
+  /* Foreign parts. (no need to nullify the cell pointers as the cells
+   * will be regenerated) */
+  space_free_foreign_parts(e->s, /*clear_cell_pointers=*/0);
 
   /* Now comes the tricky part: Exchange particles between all nodes.
      This is done in two steps, first allreducing a matrix of

src/space.c

@@ -57,6 +57,7 @@
 #include "minmax.h"
 #include "multipole.h"
 #include "pressure_floor.h"
+#include "proxy.h"
 #include "restart.h"
 #include "sort_part.h"
 #include "star_formation.h"
@@ -313,8 +314,10 @@ void space_free_cells(struct space *s) {
  * @brief Free any memory in use for foreign particles.
  *
  * @param s The #space.
+ * @param clear_cell_pointers Are we also setting all the foreign cell particle
+ * pointers to NULL?
  */
-void space_free_foreign_parts(struct space *s) {
+void space_free_foreign_parts(struct space *s, const int clear_cell_pointers) {
 
 #ifdef WITH_MPI
   if (s->parts_foreign != NULL) {
@@ -337,6 +340,13 @@ void space_free_foreign_parts(struct space *s) {
     s->size_bparts_foreign = 0;
     s->bparts_foreign = NULL;
   }
+  if (clear_cell_pointers) {
+    for (int k = 0; k < s->e->nr_proxies; k++) {
+      for (int j = 0; j < s->e->proxies[k].nr_cells_in; j++) {
+        cell_unlink_foreign_particles(s->e->proxies[k].cells_in[j]);
+      }
+    }
+  }
 #endif
 }
 

src/space.h

@@ -379,7 +379,7 @@ void space_reset_task_counters(struct space *s);
 void space_clean(struct space *s);
 void space_free_cells(struct space *s);
 
-void space_free_foreign_parts(struct space *s);
+void space_free_foreign_parts(struct space *s, const int clear_cell_pointers);
 
 void space_struct_dump(struct space *s, FILE *stream);
 void space_struct_restore(struct space *s, FILE *stream);

src/velociraptor_interface.c

@@ -408,6 +408,12 @@ void velociraptor_invoke(struct engine *e, const int linked_with_snap) {
   const int nr_cells = s->nr_cells;
   const struct cell *cells_top = s->cells_top;
 
+  /* Start by freeing some of the unnecessary memory to give VR some breathing
+     space */
+#ifdef WITH_MPI
+  space_free_foreign_parts(e->s, /*clear_cell_pointers=*/1);
+#endif
+
   /* Allow thread to run on any core for the duration of the call to
    * VELOCIraptor so that  when OpenMP threads are spawned
    * they can run on any core on the processor. */
@@ -697,6 +703,12 @@ void velociraptor_invoke(struct engine *e, const int linked_with_snap) {
   /* Record we have ran stf this timestep */
   e->stf_this_timestep = 1;
 
+  /* Reallocate the memory that was freed earlier */
+#ifdef WITH_MPI
+
+  engine_allocate_foreign_particles(e);
+#endif
+
 #else
   error("SWIFT not configure to run with VELOCIraptor.");
 #endif /* HAVE_VELOCIRAPTOR */