Also free the tasks in space_regrid before partitioning.

Lots of refactoring to tidyup code reuse

src/engine.c

@@ -894,28 +894,10 @@ void engine_repartition(struct engine *e) {
    * tasks as these will be regenerated at the next rebuild. */
 
   /* Sorting indices. */
-  if (e->s->cells_top != NULL) {
-      threadpool_map(&e->threadpool, space_rebuild_recycle_mapper,
-                     e->s->cells_top, e->s->nr_cells, sizeof(struct cell),
-                     0, e->s);
-      e->s->maxdepth = 0;
-  }
+  if (e->s->cells_top != NULL) space_free_cells(e->s);
 
-  /* Tasks. */
-  struct scheduler *s = &e->sched;
-  if (s->tasks != NULL) {
-      free(s->tasks);
-      s->tasks = NULL;
-  }
-  if (s->tasks_ind != NULL) {
-      free(s->tasks_ind);
-      s->tasks_ind = NULL;
-  }
-  if (s->tid_active != NULL) {
-      free(s->tid_active);
-      s->tid_active = NULL;
-  }
-  s->size = 0;
+  /* Task arrays. */
+  scheduler_free_tasks(&e->sched);
 
   /* Now comes the tricky part: Exchange particles between all nodes.
      This is done in two steps, first allreducing a matrix of

src/scheduler.c

@@ -983,9 +983,7 @@ void scheduler_reset(struct scheduler *s, int size) {
   if (size > s->size) {
 
     /* Free existing task lists if necessary. */
-    if (s->tasks != NULL) free(s->tasks);
-    if (s->tasks_ind != NULL) free(s->tasks_ind);
-    if (s->tid_active != NULL) free(s->tid_active);
+    scheduler_free_tasks(s);
 
     /* Allocate the new lists. */
     if (posix_memalign((void *)&s->tasks, task_align,
@@ -1661,11 +1659,29 @@ void scheduler_print_tasks(const struct scheduler *s, const char *fileName) {
  */
 void scheduler_clean(struct scheduler *s) {
 
-  free(s->tasks);
-  free(s->tasks_ind);
+  scheduler_free_tasks(s);
   free(s->unlocks);
   free(s->unlock_ind);
-  free(s->tid_active);
   for (int i = 0; i < s->nr_queues; ++i) queue_clean(&s->queues[i]);
   free(s->queues);
 }
+
+/**
+ * @brief Free the task arrays allocated by this #scheduler.
+ */
+void scheduler_free_tasks(struct scheduler *s) {
+
+  if (s->tasks != NULL) {
+    free(s->tasks);
+    s->tasks = NULL;
+  }
+  if (s->tasks_ind != NULL) {
+    free(s->tasks_ind);
+    s->tasks_ind = NULL;
+  }
+  if (s->tid_active != NULL) {
+    free(s->tid_active);
+    s->tid_active = NULL;
+  }
+  s->size = 0;
+}

src/scheduler.h

@@ -143,5 +143,6 @@ void scheduler_set_unlocks(struct scheduler *s);
 void scheduler_dump_queue(struct scheduler *s);
 void scheduler_print_tasks(const struct scheduler *s, const char *fileName);
 void scheduler_clean(struct scheduler *s);
+void scheduler_free_tasks(struct scheduler *s);
 
 #endif /* SWIFT_SCHEDULER_H */

src/space.c

@@ -252,6 +252,15 @@ void space_rebuild_recycle_mapper(void *map_data, int num_elements,
   }
 }
 
+/**
+ * @brief Free up any allocated cells.
+ */
+void space_free_cells(struct space *s) {
+  threadpool_map(&s->e->threadpool, space_rebuild_recycle_mapper,
+                 s->cells_top, s->nr_cells, sizeof(struct cell), 0, s);
+  s->maxdepth = 0;
+}
+
 /**
  * @brief Re-build the top-level cell grid.
  *
@@ -378,13 +387,15 @@ void space_regrid(struct space *s, int verbose) {
 
     /* Free the old cells, if they were allocated. */
     if (s->cells_top != NULL) {
-      threadpool_map(&s->e->threadpool, space_rebuild_recycle_mapper,
-                     s->cells_top, s->nr_cells, sizeof(struct cell), 0, s);
+      space_free_cells(s);
       free(s->cells_top);
       free(s->multipoles_top);
-      s->maxdepth = 0;
     }
 
+    /* Also free the task arrays, these will be regenerated and we can use the
+     * memory while copying the particle arrays. */
+    if (s->e != NULL) scheduler_free_tasks(&s->e->sched);
+
     /* Set the new cell dimensions only if smaller. */
     for (int k = 0; k < 3; k++) {
       s->cdim[k] = cdim[k];
@@ -491,9 +502,7 @@ void space_regrid(struct space *s, int verbose) {
   else { /* Otherwise, just clean up the cells. */
 
     /* Free the old cells, if they were allocated. */
-    threadpool_map(&s->e->threadpool, space_rebuild_recycle_mapper,
-                   s->cells_top, s->nr_cells, sizeof(struct cell), 0, s);
-    s->maxdepth = 0;
+    space_free_cells(s);
   }
 
   if (verbose)

src/space.h

@@ -199,8 +199,6 @@ void space_recycle_list(struct space *s, struct cell *cell_list_begin,
                         struct cell *cell_list_end,
                         struct gravity_tensors *multipole_list_begin,
                         struct gravity_tensors *multipole_list_end);
-void space_rebuild_recycle_mapper(void *map_data, int num_elements,
-                                  void *extra_data);
 void space_split(struct space *s, struct cell *cells, int nr_cells,
                  int verbose);
 void space_split_mapper(void *map_data, int num_elements, void *extra_data);
@@ -226,5 +224,6 @@ void space_check_timesteps(struct space *s);
 void space_replicate(struct space *s, int replicate, int verbose);
 void space_reset_task_counters(struct space *s);
 void space_clean(struct space *s);
+void space_free_cells(struct space *s);
 
 #endif /* SWIFT_SPACE_H */