Re-introduce the 'end_force' function as a task in order for the cooling...

Re-introduce the 'end_force' function as a task in order for the cooling function to be called at the right place in the time-step.

examples/analyse_tasks.py

@@ -53,7 +53,7 @@ infile = args.input
 #  Tasks and subtypes. Indexed as in tasks.h.
 TASKTYPES = ["none", "sort", "self", "pair", "sub_self", "sub_pair",
              "init_grav", "ghost_in", "ghost", "ghost_out", "extra_ghost", "drift_part", "drift_gpart",
-             "kick1", "kick2", "timestep", "send", "recv", "grav_top_level",
+             "end_force", "kick1", "kick2", "timestep", "send", "recv", "grav_top_level",
              "grav_long_range", "grav_ghost_in", "grav_ghost_out", "grav_mm", "grav_down", "cooling",
              "sourceterms", "count"]
 SUBTYPES = ["none", "density", "gradient", "force", "grav", "external_grav",

examples/plot_tasks.py

@@ -111,7 +111,7 @@ pl.rcParams.update(PLOT_PARAMS)
 #  Tasks and subtypes. Indexed as in tasks.h.
 TASKTYPES = ["none", "sort", "self", "pair", "sub_self", "sub_pair",
              "init_grav", "ghost_in", "ghost", "ghost_out", "extra_ghost", "drift_part", "drift_gpart",
-             "kick1", "kick2", "timestep", "send", "recv", "grav_top_level",
+             "end_force", "kick1", "kick2", "timestep", "send", "recv", "grav_top_level",
              "grav_long_range", "grav_ghost_in", "grav_ghost_out", "grav_mm", "grav_down", "cooling",
              "sourceterms", "count"]
 

src/cell.c

@@ -2102,6 +2102,7 @@ int cell_unskip_hydro_tasks(struct cell *c, struct scheduler *s) {
     if (c->kick1 != NULL) scheduler_activate(s, c->kick1);
     if (c->kick2 != NULL) scheduler_activate(s, c->kick2);
     if (c->timestep != NULL) scheduler_activate(s, c->timestep);
+    if (c->end_force != NULL) scheduler_activate(s, c->end_force);
     if (c->cooling != NULL) scheduler_activate(s, c->cooling);
     if (c->sourceterms != NULL) scheduler_activate(s, c->sourceterms);
   }
@@ -2213,6 +2214,7 @@ int cell_unskip_gravity_tasks(struct cell *c, struct scheduler *s) {
     if (c->kick1 != NULL) scheduler_activate(s, c->kick1);
     if (c->kick2 != NULL) scheduler_activate(s, c->kick2);
     if (c->timestep != NULL) scheduler_activate(s, c->timestep);
+    if (c->end_force != NULL) scheduler_activate(s, c->end_force);
     if (c->grav_down != NULL) scheduler_activate(s, c->grav_down);
     if (c->grav_long_range != NULL) scheduler_activate(s, c->grav_long_range);
   }

src/cell.h

@@ -249,6 +249,9 @@ struct cell {
   /*! The second kick task */
   struct task *kick2;
 
+  /*! The task to end the force calculation */
+  struct task *end_force;
+
   /*! The task to compute time-steps */
   struct task *timestep;
 

src/engine.c

@@ -170,6 +170,7 @@ void engine_add_ghosts(struct engine *e, struct cell *c, struct task *ghost_in,
 void engine_make_hierarchical_tasks_common(struct engine *e, struct cell *c) {
 
   struct scheduler *s = &e->sched;
+  const int is_with_cooling = (e->policy & engine_policy_cooling);
 
   /* Are we in a super-cell ? */
   if (c->super == c) {
@@ -188,6 +189,11 @@ void engine_make_hierarchical_tasks_common(struct engine *e, struct cell *c) {
       c->timestep = scheduler_addtask(s, task_type_timestep, task_subtype_none,
                                       0, 0, c, NULL);
 
+      /* Add the task finishing the force calculation */
+      c->end_force = scheduler_addtask(s, task_type_end_force,
+                                       task_subtype_none, 0, 0, c, NULL);
+
+      if (!is_with_cooling) scheduler_addunlock(s, c->end_force, c->kick2);
       scheduler_addunlock(s, c->kick2, c->timestep);
       scheduler_addunlock(s, c->timestep, c->kick1);
     }
@@ -254,6 +260,7 @@ void engine_make_hierarchical_tasks_hydro(struct engine *e, struct cell *c) {
         c->cooling = scheduler_addtask(s, task_type_cooling, task_subtype_none,
                                        0, 0, c, NULL);
 
+        scheduler_addunlock(s, c->super->end_force, c->cooling);
         scheduler_addunlock(s, c->cooling, c->super->kick2);
       }
 
@@ -319,7 +326,7 @@ void engine_make_hierarchical_tasks_gravity(struct engine *e, struct cell *c) {
         if (periodic) scheduler_addunlock(s, c->grav_ghost_out, c->grav_down);
         scheduler_addunlock(s, c->init_grav, c->grav_long_range);
         scheduler_addunlock(s, c->grav_long_range, c->grav_down);
-        scheduler_addunlock(s, c->grav_down, c->super->kick2);
+        scheduler_addunlock(s, c->grav_down, c->super->end_force);
       }
     }
 
@@ -1169,7 +1176,7 @@ void engine_addtasks_send_hydro(struct engine *e, struct cell *ci,
 
 #else
       /* The send_rho task should unlock the super_hydro-cell's kick task. */
-      scheduler_addunlock(s, t_rho, ci->super->kick2);
+      scheduler_addunlock(s, t_rho, ci->super->end_force);
 
       /* The send_rho task depends on the cell's ghost task. */
       scheduler_addunlock(s, ci->super_hydro->ghost_out, t_rho);
@@ -2564,7 +2571,7 @@ static inline void engine_make_external_gravity_dependencies(
 
   /* init --> external gravity --> kick */
   scheduler_addunlock(sched, c->super_gravity->drift_gpart, gravity);
-  scheduler_addunlock(sched, gravity, c->super->kick2);
+  scheduler_addunlock(sched, gravity, c->super->end_force);
 }
 
 /**
@@ -2749,7 +2756,7 @@ void engine_make_extra_hydroloop_tasks_mapper(void *map_data, int num_elements,
 
       /* Now, build all the dependencies for the hydro */
       engine_make_hydro_loops_dependencies(sched, t, t2, t->ci, with_cooling);
-      scheduler_addunlock(sched, t2, t->ci->super->kick2);
+      scheduler_addunlock(sched, t2, t->ci->super->end_force);
 #endif
     }
 
@@ -2808,14 +2815,14 @@ void engine_make_extra_hydroloop_tasks_mapper(void *map_data, int num_elements,
       /* that are local and are not descendant of the same super_hydro-cells */
       if (t->ci->nodeID == nodeID) {
         engine_make_hydro_loops_dependencies(sched, t, t2, t->ci, with_cooling);
-        scheduler_addunlock(sched, t2, t->ci->super->kick2);
+        scheduler_addunlock(sched, t2, t->ci->super->end_force);
       }
       if (t->cj->nodeID == nodeID) {
         if (t->ci->super_hydro != t->cj->super_hydro)
           engine_make_hydro_loops_dependencies(sched, t, t2, t->cj,
                                                with_cooling);
         if (t->ci->super != t->cj->super)
-          scheduler_addunlock(sched, t2, t->cj->super->kick2);
+          scheduler_addunlock(sched, t2, t->cj->super->end_force);
       }
 
 #endif
@@ -2865,7 +2872,7 @@ void engine_make_extra_hydroloop_tasks_mapper(void *map_data, int num_elements,
       /* that are local and are not descendant of the same super_hydro-cells */
       if (t->ci->nodeID == nodeID) {
         engine_make_hydro_loops_dependencies(sched, t, t2, t->ci, with_cooling);
-        scheduler_addunlock(sched, t2, t->ci->super->kick2);
+        scheduler_addunlock(sched, t2, t->ci->super->end_force);
       } else
         error("oo");
 #endif
@@ -2930,14 +2937,14 @@ void engine_make_extra_hydroloop_tasks_mapper(void *map_data, int num_elements,
       /* that are local and are not descendant of the same super_hydro-cells */
       if (t->ci->nodeID == nodeID) {
         engine_make_hydro_loops_dependencies(sched, t, t2, t->ci, with_cooling);
-        scheduler_addunlock(sched, t2, t->ci->super->kick2);
+        scheduler_addunlock(sched, t2, t->ci->super->end_force);
       }
       if (t->cj->nodeID == nodeID) {
         if (t->ci->super_hydro != t->cj->super_hydro)
           engine_make_hydro_loops_dependencies(sched, t, t2, t->cj,
                                                with_cooling);
         if (t->ci->super != t->cj->super)
-          scheduler_addunlock(sched, t2, t->cj->super->kick2);
+          scheduler_addunlock(sched, t2, t->cj->super->end_force);
       }
 #endif
     }
@@ -3417,7 +3424,14 @@ void engine_marktasks_mapper(void *map_data, int num_elements,
       }
     }
 
-    /* Kick/init ? */
+    /* End force ? */
+    else if (t->type == task_type_end_force) {
+
+      if (cell_is_active_hydro(t->ci, e) || cell_is_active_gravity(t->ci, e))
+        scheduler_activate(s, t);
+    }
+
+    /* Kick ? */
     else if (t->type == task_type_kick1 || t->type == task_type_kick2) {
 
       if (cell_is_active_hydro(t->ci, e) || cell_is_active_gravity(t->ci, e))
@@ -3554,7 +3568,7 @@ int engine_estimate_nr_tasks(struct engine *e) {
   int n1 = 0;
   int n2 = 0;
   if (e->policy & engine_policy_hydro) {
-    n1 += 36;
+    n1 += 37;
     n2 += 2;
 #ifdef WITH_MPI
     n1 += 6;
@@ -4046,7 +4060,7 @@ void engine_skip_force_and_kick(struct engine *e) {
     if (t->type == task_type_drift_part || t->type == task_type_drift_gpart ||
         t->type == task_type_kick1 || t->type == task_type_kick2 ||
         t->type == task_type_timestep || t->subtype == task_subtype_force ||
-        t->subtype == task_subtype_grav ||
+        t->subtype == task_subtype_grav || t->type == task_type_end_force ||
         t->type == task_type_grav_long_range ||
         t->type == task_type_grav_ghost_in ||
         t->type == task_type_grav_ghost_out ||

src/partition.c

@@ -580,6 +580,7 @@ static void repart_edge_metis(int partweights, int bothweights, int timebins,
     if (t->type == task_type_drift_part || t->type == task_type_drift_gpart ||
         t->type == task_type_ghost || t->type == task_type_extra_ghost ||
         t->type == task_type_kick1 || t->type == task_type_kick2 ||
+	t->type == task_type_end_force || t->type == task_type_cooling ||
         t->type == task_type_timestep || t->type == task_type_init_grav ||
         t->type == task_type_grav_down ||
         t->type == task_type_grav_long_range) {

src/runner.c

@@ -1979,10 +1979,11 @@ void *runner_main(void *data) {
           runner_do_kick1(r, ci, 1);
           break;
         case task_type_kick2:
-          if (!(e->policy & engine_policy_cooling))
-            runner_do_end_force(r, ci, 1);
           runner_do_kick2(r, ci, 1);
           break;
+        case task_type_end_force:
+          runner_do_end_force(r, ci, 1);
+          break;
         case task_type_timestep:
           runner_do_timestep(r, ci, 1);
           break;
@@ -2024,7 +2025,6 @@ void *runner_main(void *data) {
           runner_do_grav_long_range(r, t->ci, 1);
           break;
         case task_type_cooling:
-          if (e->policy & engine_policy_cooling) runner_do_end_force(r, ci, 1);
           runner_do_cooling(r, t->ci, 1);
           break;
         case task_type_sourceterms:

src/scheduler.c

@@ -225,6 +225,14 @@ void scheduler_write_dependencies(struct scheduler *s, int verbose) {
         if (tb->implicit)
           fprintf(f, "\t %s [style = filled];\n\t %s [color = lightgrey];\n",
                   tb_name, tb_name);
+
+        /* Change shape of MPI communications */
+        if (ta->type == task_type_send || ta->type == task_type_recv)
+          fprintf(f, "\t \"%s %s\" [shape = diamond];\n",
+                  taskID_names[ta->type], subtaskID_names[ta->subtype]);
+        if (tb->type == task_type_send || tb->type == task_type_recv)
+          fprintf(f, "\t \"%s %s\" [shape = diamond];\n",
+                  taskID_names[tb->type], subtaskID_names[tb->subtype]);
       }
     }
   }
@@ -234,7 +242,7 @@ void scheduler_write_dependencies(struct scheduler *s, int verbose) {
   int force_cluster[4] = {0};
   int gravity_cluster[4] = {0};
 
-  /* Modify the style of some tasks on the plot */
+  /* Check whether we need to construct a group of tasks */
   for (int type = 0; type < task_type_count; ++type) {
 
     for (int subtype = 0; subtype < task_subtype_count; ++subtype) {
@@ -244,11 +252,6 @@ void scheduler_write_dependencies(struct scheduler *s, int verbose) {
       /* Does this task/sub-task exist? */
       if (table[ind] != -1) {
 
-        /* Make MPI tasks a different shape */
-        if (type == task_type_send || type == task_type_recv)
-          fprintf(f, "\t \"%s %s\" [shape = diamond];\n", taskID_names[type],
-                  subtaskID_names[subtype]);
-
         for (int k = 0; k < 4; ++k) {
           if (type == task_type_self + k && subtype == task_subtype_density)
             density_cluster[k] = 1;
@@ -1258,6 +1261,9 @@ void scheduler_reweight(struct scheduler *s, int verbose) {
       case task_type_grav_long_range:
         cost = wscale * t->ci->gcount;
         break;
+      case task_type_end_force:
+        cost = wscale * t->ci->count + wscale * t->ci->gcount;
+        break;
       case task_type_kick1:
         cost = wscale * t->ci->count + wscale * t->ci->gcount;
         break;

src/space.c

@@ -224,6 +224,7 @@ void space_rebuild_recycle_mapper(void *map_data, int num_elements,
     c->kick1 = NULL;
     c->kick2 = NULL;
     c->timestep = NULL;
+    c->end_force = NULL;
     c->drift_part = NULL;
     c->drift_gpart = NULL;
     c->cooling = NULL;

src/task.c

@@ -48,13 +48,15 @@
 
 /* Task type names. */
 const char *taskID_names[task_type_count] = {
-    "none",          "sort",           "self",           "pair",
-    "sub_self",      "sub_pair",       "init_grav",      "ghost_in",
-    "ghost",         "ghost_out",      "extra_ghost",    "drift_part",
-    "drift_gpart",   "kick1",          "kick2",          "timestep",
-    "send",          "recv",           "grav_top_level", "grav_long_range",
-    "grav_ghost_in", "grav_ghost_out", "grav_mm",        "grav_down",
-    "cooling",       "sourceterms"};
+    "none",          "sort",           "self",
+    "pair",          "sub_self",       "sub_pair",
+    "init_grav",     "ghost_in",       "ghost",
+    "ghost_out",     "extra_ghost",    "drift_part",
+    "drift_gpart",   "end_force",      "kick1",
+    "kick2",         "timestep",       "send",
+    "recv",          "grav_top_level", "grav_long_range",
+    "grav_ghost_in", "grav_ghost_out", "grav_mm",
+    "grav_down",     "cooling",        "sourceterms"};
 
 /* Sub-task type names. */
 const char *subtaskID_names[task_subtype_count] = {
@@ -152,6 +154,7 @@ __attribute__((always_inline)) INLINE static enum task_actions task_acts_on(
       }
       break;
 
+    case task_type_end_force:
     case task_type_kick1:
     case task_type_kick2:
     case task_type_timestep:
@@ -272,6 +275,7 @@ void task_unlock(struct task *t) {
   /* Act based on task type. */
   switch (type) {
 
+    case task_type_end_force:
     case task_type_kick1:
     case task_type_kick2:
     case task_type_timestep:
@@ -361,6 +365,7 @@ int task_lock(struct task *t) {
 #endif
       break;
 
+    case task_type_end_force:
     case task_type_kick1:
     case task_type_kick2:
     case task_type_timestep:

src/task.h

@@ -52,6 +52,7 @@ enum task_types {
   task_type_extra_ghost,
   task_type_drift_part,
   task_type_drift_gpart,
+  task_type_end_force,
   task_type_kick1,
   task_type_kick2,
   task_type_timestep,