Push the cooling task to a lower level to gain more parallelism

doc/RTD/source/ParameterFiles/parameter_description.rst

@@ -946,15 +946,18 @@ which stops these from being done at the scale of the leaf cells, of which
 there can be a large number. In this case cells with gravity tasks must be at
 least 4 levels above the leaf cells (when possible).
 
-To control the depth at which the ghost tasks are placed, there are
-two parameters (one for the gas, one for the stars). These specify the
-maximum number of particles allowed in such a task before splitting
-into finer ones. These parameters are:
+To control the depth at which the ghost tasks are placed, there are two
+parameters (one for the gas, one for the stars). These specify the maximum
+number of particles allowed in such a task before splitting into finer ones. A
+similar parameter exists for the cooling tasks, which can be useful to tweak for
+models in which the cooling operations are expensive. These three parameters
+are:
 
 .. code:: YAML
 
-  engine_max_parts_per_ghost:   1000
-  engine_max_sparts_per_ghost:  1000
+  engine_max_parts_per_ghost:    1000
+  engine_max_sparts_per_ghost:   1000
+  engine_max_parts_per_cooling: 10000
 
 
 Extra space is required when particles are created in the system (to the time

doc/RTD/source/Planetary/index.rst

@@ -32,3 +32,4 @@ chosen from the several available equations of state.
    
    Hydro Scheme <hydro_scheme>
    Equations of State <equations_of_state>
+   Initial Conditions <initial_conditions>

examples/parameter_example.yml

@@ -117,8 +117,9 @@ Scheduler:
   tasks_per_cell:            0.0       # (Optional) The average number of tasks per cell. If not large enough the simulation will fail (means guess...).
   links_per_tasks:           25        # (Optional) The average number of links per tasks (before adding the communication tasks). If not large enough the simulation will fail (means guess...). Defaults to 10.
   mpi_message_limit:         4096      # (Optional) Maximum MPI task message size to send non-buffered, KB.
-  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_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.
 
 # Parameters governing the time integration (Set dt_min and dt_max to the same value for a fixed time-step run.)
 TimeIntegration:

src/cell.c

@@ -2613,6 +2613,50 @@ void cell_activate_hydro_ghosts(struct cell *c, struct scheduler *s,
   cell_recursively_activate_hydro_ghosts(c, s, e);
 }
 
+/**
+ * @brief Recursively activate the cooling (and implicit links) in a cell
+ * hierarchy.
+ *
+ * @param c The #cell.
+ * @param s The #scheduler.
+ * @param e The #engine.
+ */
+void cell_recursively_activate_cooling(struct cell *c, struct scheduler *s,
+                                       const struct engine *e) {
+  /* Early abort? */
+  if ((c->hydro.count == 0) || !cell_is_active_hydro(c, e)) return;
+
+  /* Is the ghost at this level? */
+  if (c->hydro.cooling != NULL) {
+    scheduler_activate(s, c->hydro.cooling);
+  } else {
+
+#ifdef SWIFT_DEBUG_CHECKS
+    if (!c->split)
+      error("Reached the leaf level without finding a cooling task!");
+#endif
+
+    /* Keep recursing */
+    for (int k = 0; k < 8; k++)
+      if (c->progeny[k] != NULL)
+        cell_recursively_activate_cooling(c->progeny[k], s, e);
+  }
+}
+
+/**
+ * @brief Activate the cooling tasks (and implicit links) in a cell hierarchy.
+ *
+ * @param c The #cell.
+ * @param s The #scheduler.
+ * @param e The #engine.
+ */
+void cell_activate_cooling(struct cell *c, struct scheduler *s,
+                           const struct engine *e) {
+  scheduler_activate(s, c->hydro.cooling_in);
+  scheduler_activate(s, c->hydro.cooling_out);
+  cell_recursively_activate_cooling(c, s, e);
+}
+
 /**
  * @brief Recurse down in a cell hierarchy until the hydro.super level is
  * reached and activate the spart drift at that level.
@@ -3722,7 +3766,7 @@ int cell_unskip_hydro_tasks(struct cell *c, struct scheduler *s) {
     if (c->kick2 != NULL) scheduler_activate(s, c->kick2);
     if (c->timestep != NULL) scheduler_activate(s, c->timestep);
     if (c->hydro.end_force != NULL) scheduler_activate(s, c->hydro.end_force);
-    if (c->hydro.cooling != NULL) scheduler_activate(s, c->hydro.cooling);
+    if (c->hydro.cooling_in != NULL) cell_activate_cooling(c, s, e);
 #ifdef WITH_LOGGER
     if (c->logger != NULL) scheduler_activate(s, c->logger);
 #endif

src/cell.h

@@ -390,6 +390,12 @@ struct cell {
     /*! The task to end the force calculation */
     struct task *end_force;
 
+    /*! Dependency implicit task for cooling (in->cooling->out) */
+    struct task *cooling_in;
+
+    /*! Dependency implicit task for cooling (in->cooling->out) */
+    struct task *cooling_out;
+
     /*! Task for cooling */
     struct task *cooling;
 

src/engine.h

@@ -114,6 +114,7 @@ enum engine_step_properties {
 #define engine_default_timesteps_file_name "timesteps"
 #define engine_max_parts_per_ghost_default 1000
 #define engine_max_sparts_per_ghost_default 1000
+#define engine_max_parts_per_cooling_default 10000
 #define engine_star_resort_task_depth_default 2
 #define engine_tasks_per_cell_margin 1.2
 #define engine_default_stf_subdir_per_output "."

src/engine_maketasks.c

@@ -56,6 +56,7 @@
 extern int engine_max_parts_per_ghost;
 extern int engine_max_sparts_per_ghost;
 extern int engine_star_resort_task_depth;
+extern int engine_max_parts_per_cooling;
 
 /**
  * @brief Add send tasks for the gravity pairs to a hierarchy of cells.
@@ -1055,6 +1056,33 @@ void engine_add_ghosts(struct engine *e, struct cell *c, struct task *ghost_in,
   }
 }
 
+/**
+ * @brief Recursively add non-implicit cooling tasks to a cell hierarchy.
+ */
+void engine_add_cooling(struct engine *e, struct cell *c,
+                        struct task *cooling_in, struct task *cooling_out) {
+
+  /* Abort as there are no hydro particles here? */
+  if (c->hydro.count_total == 0) return;
+
+  /* If we have reached the leaf OR have to few particles to play with*/
+  if (!c->split || c->hydro.count_total < engine_max_parts_per_cooling) {
+
+    /* Add the cooling task and its dependencies */
+    struct scheduler *s = &e->sched;
+    c->hydro.cooling = scheduler_addtask(s, task_type_cooling,
+                                         task_subtype_none, 0, 0, c, NULL);
+    scheduler_addunlock(s, cooling_in, c->hydro.cooling);
+    scheduler_addunlock(s, c->hydro.cooling, cooling_out);
+
+  } else {
+    /* Keep recursing */
+    for (int k = 0; k < 8; k++)
+      if (c->progeny[k] != NULL)
+        engine_add_cooling(e, c->progeny[k], cooling_in, cooling_out);
+  }
+}
+
 /**
  * @brief Generate the hydro hierarchical tasks for a hierarchy of cells -
  * i.e. all the O(Npart) tasks -- hydro version
@@ -1163,11 +1191,17 @@ void engine_make_hierarchical_tasks_hydro(struct engine *e, struct cell *c,
       /* Subgrid tasks: cooling */
       if (with_cooling) {
 
-        c->hydro.cooling = scheduler_addtask(s, task_type_cooling,
-                                             task_subtype_none, 0, 0, c, NULL);
+        c->hydro.cooling_in =
+            scheduler_addtask(s, task_type_cooling_in, task_subtype_none, 0,
+                              /*implicit=*/1, c, NULL);
+        c->hydro.cooling_out =
+            scheduler_addtask(s, task_type_cooling_out, task_subtype_none, 0,
+                              /*implicit=*/1, c, NULL);
+
+        engine_add_cooling(e, c, c->hydro.cooling_in, c->hydro.cooling_out);
 
-        scheduler_addunlock(s, c->hydro.end_force, c->hydro.cooling);
-        scheduler_addunlock(s, c->hydro.cooling, c->super->kick2);
+        scheduler_addunlock(s, c->hydro.end_force, c->hydro.cooling_in);
+        scheduler_addunlock(s, c->hydro.cooling_out, c->super->kick2);
 
       } else {
         scheduler_addunlock(s, c->hydro.end_force, c->super->kick2);

src/engine_marktasks.c

@@ -978,7 +978,8 @@ void engine_marktasks_mapper(void *map_data, int num_elements,
     }
 
     /* Subgrid tasks: cooling */
-    else if (t_type == task_type_cooling) {
+    else if (t_type == task_type_cooling || t_type == task_type_cooling_in ||
+             t_type == task_type_cooling_out) {
       if (cell_is_active_hydro(t->ci, e)) scheduler_activate(s, t);
     }
 

src/space.c

@@ -94,6 +94,7 @@ int space_extra_gparts = space_extra_gparts_default;
 /*! Maximum number of particles per ghost */
 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;
 
 /*! Maximal depth at which the stars resort task can be pushed */
 int engine_star_resort_task_depth = engine_star_resort_task_depth_default;
@@ -252,6 +253,8 @@ void space_rebuild_recycle_mapper(void *map_data, int num_elements,
     c->black_holes.black_holes_out = NULL;
     c->grav.drift = NULL;
     c->grav.drift_out = NULL;
+    c->hydro.cooling_in = NULL;
+    c->hydro.cooling_out = NULL;
     c->hydro.cooling = NULL;
     c->grav.long_range = NULL;
     c->grav.down_in = NULL;
@@ -4986,6 +4989,10 @@ void space_init(struct space *s, struct swift_params *params,
       parser_get_opt_param_int(params, "Scheduler:engine_max_sparts_per_ghost",
                                engine_max_sparts_per_ghost_default);
 
+  engine_max_parts_per_cooling =
+      parser_get_opt_param_int(params, "Scheduler:engine_max_parts_per_cooling",
+                               engine_max_parts_per_cooling_default);
+
   if (verbose) {
     message("max_size set to %d split_size set to %d", space_maxsize,
             space_splitsize);
@@ -5853,6 +5860,9 @@ void space_struct_dump(struct space *s, FILE *stream) {
   restart_write_blocks(&engine_max_sparts_per_ghost, sizeof(int), 1, stream,
                        "engine_max_sparts_per_ghost",
                        "engine_max_sparts_per_ghost");
+  restart_write_blocks(&engine_max_parts_per_cooling, sizeof(int), 1, stream,
+                       "engine_max_parts_per_cooling",
+                       "engine_max_parts_per_cooling");
   restart_write_blocks(&engine_star_resort_task_depth, sizeof(int), 1, stream,
                        "engine_star_resort_task_depth",
                        "engine_star_resort_task_depth");
@@ -5920,6 +5930,8 @@ void space_struct_restore(struct space *s, FILE *stream) {
                       "engine_max_parts_per_ghost");
   restart_read_blocks(&engine_max_sparts_per_ghost, sizeof(int), 1, stream,
                       NULL, "engine_max_sparts_per_ghost");
+  restart_read_blocks(&engine_max_parts_per_cooling, sizeof(int), 1, 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");
 

src/task.c

@@ -81,6 +81,8 @@ const char *taskID_names[task_type_count] = {"none",
                                              "grav_mesh",
                                              "grav_end_force",
                                              "cooling",
+                                             "cooling_in",
+                                             "cooling_out",
                                              "star_formation",
                                              "star_formation_in",
                                              "star_formation_out",

src/task.h

@@ -75,6 +75,8 @@ enum task_types {
   task_type_grav_mesh,
   task_type_end_grav_force,
   task_type_cooling,
+  task_type_cooling_in,  /* Implicit */
+  task_type_cooling_out, /* Implicit */
   task_type_star_formation,
   task_type_star_formation_in,  /* Implicit */
   task_type_star_formation_out, /* Implicit */

tools/task_plots/analyse_tasks.py

@@ -103,6 +103,8 @@ TASKTYPES = [
     "grav_mesh",
     "grav_end_force",
     "cooling",
+    "cooling_in",
+    "cooling_out",
     "star_formation",
     "star_formation_in",
     "star_formation_out",

tools/task_plots/iplot_tasks.py

@@ -157,6 +157,8 @@ TASKTYPES = [
     "grav_mesh",
     "grav_end_force",
     "cooling",
+    "cooling_in",
+    "cooling_out",
     "star_formation",
     "star_formation_in",
     "star_formation_out",

tools/task_plots/plot_tasks.py

@@ -180,6 +180,8 @@ TASKTYPES = [
     "grav_mesh",
     "grav_end_force",
     "cooling",
+    "cooling_in",
+    "cooling_out",
     "star_formation",
     "star_formation_in",
     "star_formation_out",