Memory footprint improvements

examples/parameter_example.yml

@@ -138,8 +138,12 @@ Scheduler:
   engine_max_parts_per_ghost:    1000  # (Optional) Maximum number of parts per ghost.
   engine_max_sparts_per_ghost:   1000  # (Optional) Maximum number of sparts per ghost.
   engine_max_parts_per_cooling: 10000  # (Optional) Maximum number of parts per cooling task.
+  engine_redist_alloc_margin:     1.2  # (Optional) Multiplier factor for the number of local particles to allocate on a given rank.
+  engine_foreign_alloc_margin:    1.05 # (Optional) Multiplier factor for the number of foreign particles to allocate on a given rank.
   dependency_graph_frequency:       0  # (Optional) Dumping frequency of the dependency graph. By default, writes only at the first step.
   task_level_output_frequency:      0  # (Optional) Dumping frequency of the task level data. By default, writes only at the first step.
+  free_foreign_during_restart:      0  # (Optional) Should the code free the foreign data when dumping restart files in order to get breathing space?
+  free_foreign_during_rebuild:      0  # (Optional) Should the code free the foreign data when calling a rebuld in order to get breathing space?
 
 # Parameters governing the time integration (Set dt_min and dt_max to the same value for a fixed time-step run.)
 TimeIntegration:

src/engine.c

@@ -1199,6 +1199,12 @@ void engine_rebuild(struct engine *e, const int repartitioned,
   /* Give some breathing space */
   scheduler_free_tasks(&e->sched);
 
+  /* Free the foreign particles to get more breathing space. */
+#ifdef WITH_MPI
+  if (e->free_foreign_when_rebuilding)
+    space_free_foreign_parts(e->s, /*clear_cell_pointers=*/1);
+#endif
+
   /* Re-build the space. */
   space_rebuild(e->s, repartitioned, e->verbose);
 
@@ -1312,6 +1318,12 @@ void engine_rebuild(struct engine *e, const int repartitioned,
   /* Re-build the tasks. */
   engine_maketasks(e);
 
+  /* Reallocate freed memory */
+#ifdef WITH_MPI
+  if (e->free_foreign_when_rebuilding)
+    engine_allocate_foreign_particles(e, /*fof=*/0);
+#endif
+
   /* Make the list of top-level cells that have tasks */
   space_list_useful_top_level_cells(e->s);
 
@@ -2879,6 +2891,10 @@ void engine_init(
       parser_get_opt_param_int(params, "FOF:dump_catalogue_when_seeding", 0);
   e->snapshot_units = (struct unit_system *)malloc(sizeof(struct unit_system));
   units_init_default(e->snapshot_units, params, "Snapshots", internal_units);
+  e->free_foreign_when_dumping_restart = parser_get_opt_param_int(
+      params, "Scheduler:free_foreign_during_restart", 0);
+  e->free_foreign_when_rebuilding = parser_get_opt_param_int(
+      params, "Scheduler:free_foreign_during_rebuild", 0);
   e->snapshot_output_count = 0;
   e->stf_output_count = 0;
   e->los_output_count = 0;

src/engine.h

@@ -108,9 +108,9 @@ enum engine_step_properties {
 #define engine_maxproxies 64
 #define engine_tasksreweight 1
 #define engine_parts_size_grow 1.05
-#define engine_redistribute_alloc_margin 1.2
+#define engine_redistribute_alloc_margin_default 1.2
 #define engine_rebuild_link_alloc_margin 1.2
-#define engine_foreign_alloc_margin 1.05
+#define engine_foreign_alloc_margin_default 1.05
 #define engine_default_energy_file_name "statistics"
 #define engine_default_timesteps_file_name "timesteps"
 #define engine_max_parts_per_ghost_default 1000
@@ -524,6 +524,12 @@ struct engine {
   /* Number of Lustre OSTs on the system to use as rank-based striping offset */
   int restart_lustre_OST_count;
 
+  /* Do we free the foreign data before writing restart files? */
+  int free_foreign_when_dumping_restart;
+
+  /* Do we free the foreign data before rebuilding the tree? */
+  int free_foreign_when_rebuilding;
+
   /* Name of the restart file. */
   const char *restart_file;
 

src/engine_io.c

@@ -78,6 +78,13 @@ void engine_dump_restarts(struct engine *e, int drifted_all, int force) {
 
       /* Drift all particles first (may have just been done). */
       if (!drifted_all) engine_drift_all(e, /*drift_mpole=*/1);
+
+        /* Free the foreign particles to get more breathing space. */
+#ifdef WITH_MPI
+      if (e->free_foreign_when_dumping_restart)
+        space_free_foreign_parts(e->s, /*clear_cell_pointers=*/1);
+#endif
+
       restart_write(e, e->restart_file);
 
 #ifdef WITH_MPI
@@ -85,6 +92,10 @@ void engine_dump_restarts(struct engine *e, int drifted_all, int force) {
        * sets of restart files should the code crash before all the ranks
        * are done */
       MPI_Barrier(MPI_COMM_WORLD);
+
+      /* Reallocate freed memory */
+      if (e->free_foreign_when_dumping_restart)
+        engine_allocate_foreign_particles(e, /*fof=*/0);
 #endif
 
       if (e->verbose)

src/space.c

@@ -90,6 +90,11 @@ int engine_max_parts_per_ghost = engine_max_parts_per_ghost_default;
 int engine_max_sparts_per_ghost = engine_max_sparts_per_ghost_default;
 int engine_max_parts_per_cooling = engine_max_parts_per_cooling_default;
 
+/*! Allocation margins */
+double engine_redistribute_alloc_margin =
+    engine_redistribute_alloc_margin_default;
+double engine_foreign_alloc_margin = engine_foreign_alloc_margin_default;
+
 /*! Maximal depth at which the stars resort task can be pushed */
 int engine_star_resort_task_depth = engine_star_resort_task_depth_default;
 
@@ -1267,6 +1272,14 @@ void space_init(struct space *s, struct swift_params *params,
       parser_get_opt_param_int(params, "Scheduler:engine_max_parts_per_cooling",
                                engine_max_parts_per_cooling_default);
 
+  engine_redistribute_alloc_margin = parser_get_opt_param_double(
+      params, "Scheduler:engine_redist_alloc_margin",
+      engine_redistribute_alloc_margin_default);
+
+  engine_foreign_alloc_margin = parser_get_opt_param_double(
+      params, "Scheduler:engine_foreign_alloc_margin",
+      engine_foreign_alloc_margin_default);
+
   if (verbose) {
     message("max_size set to %d split_size set to %d", space_maxsize,
             space_splitsize);
@@ -2414,6 +2427,12 @@ void space_struct_dump(struct space *s, FILE *stream) {
   restart_write_blocks(&engine_star_resort_task_depth, sizeof(int), 1, stream,
                        "engine_star_resort_task_depth",
                        "engine_star_resort_task_depth");
+  restart_write_blocks(&engine_redistribute_alloc_margin, sizeof(double), 1,
+                       stream, "engine_redistribute_alloc_margin",
+                       "engine_redistribute_alloc_margin");
+  restart_write_blocks(&engine_foreign_alloc_margin, sizeof(double), 1, stream,
+                       "engine_foreign_alloc_margin",
+                       "engine_foreign_alloc_margin");
 
   /* More things to write. */
   if (s->nr_parts > 0) {
@@ -2488,6 +2507,10 @@ void space_struct_restore(struct space *s, FILE *stream) {
                       NULL, "engine_max_parts_per_cooling");
   restart_read_blocks(&engine_star_resort_task_depth, sizeof(int), 1, stream,
                       NULL, "engine_star_resort_task_depth");
+  restart_read_blocks(&engine_redistribute_alloc_margin, sizeof(double), 1,
+                      stream, NULL, "engine_redistribute_alloc_margin");
+  restart_read_blocks(&engine_foreign_alloc_margin, sizeof(double), 1, stream,
+                      NULL, "engine_foreign_alloc_margin");
 
   /* Things that should be reconstructed in a rebuild. */
   s->cells_top = NULL;

src/space.h

@@ -84,6 +84,8 @@ extern int space_extra_gparts;
 extern int space_extra_sparts;
 extern int space_extra_bparts;
 extern int space_extra_sinks;
+extern double engine_redistribute_alloc_margin;
+extern double engine_foreign_alloc_margin;
 
 /**
  * @brief The space in which the cells and particles reside.