scheduler.c

/*******************************************************************************
 * This file is part of SWIFT.
 * Copyright (c) 2012 Pedro Gonnet (pedro.gonnet@durham.ac.uk)
 *                    Matthieu Schaller (matthieu.schaller@durham.ac.uk)
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 ******************************************************************************/

/* Config parameters. */
#include "../config.h"

/* Some standard headers. */
#include <limits.h>
#include <math.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* MPI headers. */
#ifdef WITH_MPI
#include <mpi.h>
#endif

/* This object's header. */
#include "scheduler.h"

/* Local headers. */
#include "atomic.h"
#include "const.h"
#include "cycle.h"
#include "error.h"
#include "intrinsics.h"
#include "kernel_hydro.h"
#include "timers.h"

/**
 * @brief Add an unlock_task to the given task.
 *
 * @param s The #scheduler.
 * @param ta The unlocking #task.
 * @param tb The #task that will be unlocked.
 */

void scheduler_addunlock(struct scheduler *s, struct task *ta,
                         struct task *tb) {
  /* Get an index at which to store this unlock. */
  const int ind = atomic_inc(&s->nr_unlocks);

  /* Does the buffer need to be grown? */
  if (ind == s->size_unlocks) {
    /* Allocate the new buffer. */
    struct task **unlocks_new;
    int *unlock_ind_new;
    const int size_unlocks_new = s->size_unlocks * 2;
    if ((unlocks_new = (struct task **)malloc(
             sizeof(struct task *) *size_unlocks_new)) == NULL ||
        (unlock_ind_new = (int *)malloc(sizeof(int) * size_unlocks_new)) ==
            NULL)
      error("Failed to re-allocate unlocks.");

    /* Wait for all writes to the old buffer to complete. */
    while (s->completed_unlock_writes < ind)
      ;

    /* Copy the buffers. */
    memcpy(unlocks_new, s->unlocks, sizeof(struct task *) * ind);
    memcpy(unlock_ind_new, s->unlock_ind, sizeof(int) * ind);
    free(s->unlocks);
    free(s->unlock_ind);
    s->unlocks = unlocks_new;
    s->unlock_ind = unlock_ind_new;

    /* Publish the new buffer size. */
    s->size_unlocks = size_unlocks_new;
  }

  /* Wait for there to actually be space at my index. */
  while (ind > s->size_unlocks)
    ;

  /* Write the unlock to the scheduler. */
  s->unlocks[ind] = tb;
  s->unlock_ind[ind] = ta - s->tasks;
  atomic_inc(&s->completed_unlock_writes);
}

/**
 * @brief Split tasks that may be too large.
 *
 * @param s The #scheduler we are working in.
 */

void scheduler_splittasks_mapper(void *map_data, void *extra_data) {

  /* Static constants. */
  const static int pts[7][8] = {{-1, 12, 10, 9, 4, 3, 1, 0},
                                {-1, -1, 11, 10, 5, 4, 2, 1},
                                {-1, -1, -1, 12, 7, 6, 4, 3},
                                {-1, -1, -1, -1, 8, 7, 5, 4},
                                {-1, -1, -1, -1, -1, 12, 10, 9},
                                {-1, -1, -1, -1, -1, -1, 11, 10},
                                {-1, -1, -1, -1, -1, -1, -1, 12}};
  const static float sid_scale[13] = {0.1897, 0.4025, 0.1897, 0.4025, 0.5788,
                                      0.4025, 0.1897, 0.4025, 0.1897, 0.4025,
                                      0.5788, 0.4025, 0.5788};

  /* Extract the parameters. */
  struct scheduler *s = (struct scheduler *)extra_data;
  struct task *t = (struct task *)map_data;

  /* Iterate on this task until we're done with it. */
  int redo = 1;
  while (redo) {

    /* Reset the redo flag. */
    redo = 0;

    /* Non-splittable task? */
    if ((t->ci == NULL || (t->type == task_type_pair && t->cj == NULL)) ||
        ((t->type == task_type_kick) && t->ci->nodeID != s->nodeID) ||
        ((t->type == task_type_drift) && t->ci->nodeID != s->nodeID) ||
        ((t->type == task_type_init) && t->ci->nodeID != s->nodeID)) {
      t->type = task_type_none;
      t->skip = 1;
      return;
    }

    /* Self-interaction? */
    if (t->type == task_type_self) {

      /* Get a handle on the cell involved. */
      struct cell *ci = t->ci;

      /* Foreign task? */
      if (ci->nodeID != s->nodeID) {
        t->skip = 1;
        return;
      }

      /* Is this cell even split? */
      if (ci->split) {

        /* Make a sub? */
        if (scheduler_dosub && ci->count < space_subsize / ci->count) {

          /* convert to a self-subtask. */
          t->type = task_type_sub;

          /* Otherwise, make tasks explicitly. */
        } else {

          /* Take a step back (we're going to recycle the current task)... */
          redo = 1;

          /* Add the self tasks. */
          int first_child = 0;
          while (ci->progeny[first_child] == NULL) first_child++;
          t->ci = ci->progeny[first_child];
          for (int k = first_child + 1; k < 8; k++)
            if (ci->progeny[k] != NULL)
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_self, t->subtype, 0, 0,
                                    ci->progeny[k], NULL, 0),
                  s);

          /* Make a task for each pair of progeny. */
          for (int j = 0; j < 8; j++)
            if (ci->progeny[j] != NULL)
              for (int k = j + 1; k < 8; k++)
                if (ci->progeny[k] != NULL)
                  scheduler_splittasks_mapper(
                      scheduler_addtask(s, task_type_pair, t->subtype,
                                        pts[j][k], 0, ci->progeny[j],
                                        ci->progeny[k], 0),
                      s);
        }
      }

      /* Pair interaction? */
    } else if (t->type == task_type_pair) {

      /* Get a handle on the cells involved. */
      struct cell *ci = t->ci;
      struct cell *cj = t->cj;
      const double hi = ci->dmin;
      const double hj = cj->dmin;

      /* Foreign task? */
      if (ci->nodeID != s->nodeID && cj->nodeID != s->nodeID) {
        t->skip = 1;
        return;
      }

      /* Get the sort ID, use space_getsid and not t->flags
         to make sure we get ci and cj swapped if needed. */
      double shift[3];
      int sid = space_getsid(s->space, &ci, &cj, shift);

      /* Should this task be split-up? */
      if (ci->split && cj->split &&
          ci->h_max * kernel_gamma * space_stretch < hi / 2 &&
          cj->h_max * kernel_gamma * space_stretch < hj / 2) {

        /* Replace by a single sub-task? */
        if (scheduler_dosub &&
            ci->count * sid_scale[sid] < space_subsize / cj->count &&
            sid != 0 && sid != 2 && sid != 6 && sid != 8) {

          /* Make this task a sub task. */
          t->type = task_type_sub;

          /* Otherwise, split it. */
        } else {

          /* Take a step back (we're going to recycle the current task)... */
          redo = 1;

          /* For each different sorting type... */
          switch (sid) {

            case 0: /* (  1 ,  1 ,  1 ) */
              t->ci = ci->progeny[7];
              t->cj = cj->progeny[0];
              t->flags = 0;
              break;

            case 1: /* (  1 ,  1 ,  0 ) */
              t->ci = ci->progeny[6];
              t->cj = cj->progeny[0];
              t->flags = 1;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 1, 0,
                                    ci->progeny[7], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[6], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              break;

            case 2: /* (  1 ,  1 , -1 ) */
              t->ci = ci->progeny[6];
              t->cj = cj->progeny[1];
              t->flags = 2;
              t->tight = 1;
              break;

            case 3: /* (  1 ,  0 ,  1 ) */
              t->ci = ci->progeny[5];
              t->cj = cj->progeny[0];
              t->flags = 3;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 3, 0,
                                    ci->progeny[7], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[5], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 6, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              break;

            case 4: /* (  1 ,  0 ,  0 ) */
              t->ci = ci->progeny[4];
              t->cj = cj->progeny[0];
              t->flags = 4;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 5, 0,
                                    ci->progeny[5], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[6], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 3, 0,
                                    ci->progeny[4], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 4, 0,
                                    ci->progeny[5], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 6, 0,
                                    ci->progeny[6], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[7], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 1, 0,
                                    ci->progeny[4], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[5], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 4, 0,
                                    ci->progeny[6], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 5, 0,
                                    ci->progeny[7], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[4], cj->progeny[3], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 1, 0,
                                    ci->progeny[5], cj->progeny[3], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 3, 0,
                                    ci->progeny[6], cj->progeny[3], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 4, 0,
                                    ci->progeny[7], cj->progeny[3], 1),
                  s);
              break;

            case 5: /* (  1 ,  0 , -1 ) */
              t->ci = ci->progeny[4];
              t->cj = cj->progeny[1];
              t->flags = 5;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 5, 0,
                                    ci->progeny[6], cj->progeny[3], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[4], cj->progeny[3], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[6], cj->progeny[1], 1),
                  s);
              break;

            case 6: /* (  1 , -1 ,  1 ) */
              t->ci = ci->progeny[5];
              t->cj = cj->progeny[2];
              t->flags = 6;
              t->tight = 1;
              break;

            case 7: /* (  1 , -1 ,  0 ) */
              t->ci = ci->progeny[4];
              t->cj = cj->progeny[3];
              t->flags = 6;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[5], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[4], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[5], cj->progeny[3], 1),
                  s);
              break;

            case 8: /* (  1 , -1 , -1 ) */
              t->ci = ci->progeny[4];
              t->cj = cj->progeny[3];
              t->flags = 8;
              t->tight = 1;
              break;

            case 9: /* (  0 ,  1 ,  1 ) */
              t->ci = ci->progeny[3];
              t->cj = cj->progeny[0];
              t->flags = 9;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 9, 0,
                                    ci->progeny[7], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[3], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              break;

            case 10: /* (  0 ,  1 ,  0 ) */
              t->ci = ci->progeny[2];
              t->cj = cj->progeny[0];
              t->flags = 10;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 11, 0,
                                    ci->progeny[3], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[6], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 6, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 9, 0,
                                    ci->progeny[2], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 10, 0,
                                    ci->progeny[3], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[6], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 7, 0,
                                    ci->progeny[7], cj->progeny[1], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 1, 0,
                                    ci->progeny[2], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[3], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 10, 0,
                                    ci->progeny[6], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 11, 0,
                                    ci->progeny[7], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[2], cj->progeny[5], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 1, 0,
                                    ci->progeny[3], cj->progeny[5], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 9, 0,
                                    ci->progeny[6], cj->progeny[5], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 10, 0,
                                    ci->progeny[7], cj->progeny[5], 1),
                  s);
              break;

            case 11: /* (  0 ,  1 , -1 ) */
              t->ci = ci->progeny[2];
              t->cj = cj->progeny[1];
              t->flags = 11;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 11, 0,
                                    ci->progeny[6], cj->progeny[5], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[2], cj->progeny[5], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 6, 0,
                                    ci->progeny[6], cj->progeny[1], 1),
                  s);
              break;

            case 12: /* (  0 ,  0 ,  1 ) */
              t->ci = ci->progeny[1];
              t->cj = cj->progeny[0];
              t->flags = 12;
              t->tight = 1;
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 11, 0,
                                    ci->progeny[3], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 5, 0,
                                    ci->progeny[5], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 2, 0,
                                    ci->progeny[7], cj->progeny[0], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 9, 0,
                                    ci->progeny[1], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 12, 0,
                                    ci->progeny[3], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 8, 0,
                                    ci->progeny[5], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 5, 0,
                                    ci->progeny[7], cj->progeny[2], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 3, 0,
                                    ci->progeny[1], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 6, 0,
                                    ci->progeny[3], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 12, 0,
                                    ci->progeny[5], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 11, 0,
                                    ci->progeny[7], cj->progeny[4], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                    ci->progeny[1], cj->progeny[6], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 3, 0,
                                    ci->progeny[3], cj->progeny[6], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 9, 0,
                                    ci->progeny[5], cj->progeny[6], 1),
                  s);
              scheduler_splittasks_mapper(
                  scheduler_addtask(s, task_type_pair, t->subtype, 12, 0,
                                    ci->progeny[7], cj->progeny[6], 1),
                  s);
              break;
          } /* switch(sid) */
        }

        /* Otherwise, break it up if it is too large? */
      } else if (scheduler_doforcesplit && ci->split && cj->split &&
                 (ci->count > space_maxsize / cj->count)) {

        // message( "force splitting pair with %i and %i parts." , ci->count ,
        // cj->count );

        /* Replace the current task. */
        t->type = task_type_none;

        for (int j = 0; j < 8; j++)
          if (ci->progeny[j] != NULL)
            for (int k = 0; k < 8; k++)
              if (cj->progeny[k] != NULL) {
                scheduler_splittasks_mapper(
                    scheduler_addtask(s, task_type_pair, t->subtype, 0, 0,
                                      ci->progeny[j], cj->progeny[k], 0),
                    s);
                t->flags = space_getsid(s->space, &t->ci, &t->cj, shift);
              }

        /* Otherwise, if not spilt, stitch-up the sorting. */
      } else {

        /* Create the sort for ci. */
        lock_lock( &ci->lock );
        if (ci->sorts == NULL)
          ci->sorts =
              scheduler_addtask(s, task_type_sort, 0, 1 << sid, 0, ci, NULL, 0);
        else
          ci->sorts->flags |= (1 << sid);
        lock_unlock_blind( &ci->lock );
        scheduler_addunlock(s, ci->sorts, t);

        /* Create the sort for cj. */
        lock_lock( &cj->lock );
        if (cj->sorts == NULL)
          cj->sorts =
              scheduler_addtask(s, task_type_sort, 0, 1 << sid, 0, cj, NULL, 0);
        else
          cj->sorts->flags |= (1 << sid);
        lock_unlock_blind( &cj->lock );
        scheduler_addunlock(s, cj->sorts, t);
      }

    } /* pair interaction? */

  } /* iterate over the current task. */
}

void scheduler_splittasks(struct scheduler *s) {

  /* Call the mapper on each current task. */
  threadpool_map(s->threadpool, scheduler_splittasks_mapper, s->tasks,
                 s->nr_tasks, sizeof(struct task), s);
}

/**
 * @brief Add a #task to the #scheduler.
 *
 * @param s The #scheduler we are working in.
 * @param type The type of the task.
 * @param subtype The sub-type of the task.
 * @param flags The flags of the task.
 * @param wait
 * @param ci The first cell to interact.
 * @param cj The second cell to interact.
 * @param tight
 */

struct task *scheduler_addtask(struct scheduler *s, int type, int subtype,
                               int flags, int wait, struct cell *ci,
                               struct cell *cj, int tight) {

  /* Get the next free task. */
  const int ind = atomic_inc(&s->tasks_next);

  /* Overflow? */
  if (ind >= s->size) error("Task list overflow.");

  /* Get a pointer to the new task. */
  struct task *t = &s->tasks[ind];

  /* Copy the data. */
  t->type = type;
  t->subtype = subtype;
  t->flags = flags;
  t->wait = wait;
  t->ci = ci;
  t->cj = cj;
  t->skip = 0;
  t->tight = tight;
  t->implicit = 0;
  t->weight = 0;
  t->rank = 0;
  t->tic = 0;
  t->toc = 0;
  t->nr_unlock_tasks = 0;
  t->rid = -1;
  t->last_rid = -1;

  /* Init the lock. */
  lock_init(&t->lock);

  /* Add an index for it. */
  // lock_lock( &s->lock );
  s->tasks_ind[atomic_inc(&s->nr_tasks)] = ind;
  // lock_unlock_blind( &s->lock );

  /* Return a pointer to the new task. */
  return t;
}

/**
 * @brief Set the unlock pointers in each task.
 *
 * @param s The #scheduler.
 */

void scheduler_set_unlocks(struct scheduler *s) {

  /* Store the counts for each task. */
  int *counts;
  if ((counts = (int *)malloc(sizeof(int) * s->nr_tasks)) == NULL)
    error("Failed to allocate temporary counts array.");
  bzero(counts, sizeof(int) * s->nr_tasks);
  for (int k = 0; k < s->nr_unlocks; k++) counts[s->unlock_ind[k]] += 1;

  /* Compute the offset for each unlock block. */
  int *offsets;
  if ((offsets = (int *)malloc(sizeof(int) * (s->nr_tasks + 1))) == NULL)
    error("Failed to allocate temporary offsets array.");
  offsets[0] = 0;
  for (int k = 0; k < s->nr_tasks; k++) offsets[k + 1] = offsets[k] + counts[k];

  /* Create and fill a temporary array with the sorted unlocks. */
  struct task **unlocks;
  if ((unlocks = (struct task **)malloc(sizeof(struct task *) *
                                        s->size_unlocks)) == NULL)
    error("Failed to allocate temporary unlocks array.");
  for (int k = 0; k < s->nr_unlocks; k++) {
    const int ind = s->unlock_ind[k];
    unlocks[offsets[ind]] = s->unlocks[k];
    offsets[ind] += 1;
  }

  /* Swap the unlocks. */
  free(s->unlocks);
  s->unlocks = unlocks;

  /* Re-set the offsets. */
  offsets[0] = 0;
  for (int k = 1; k < s->nr_tasks; k++)
    offsets[k] = offsets[k - 1] + counts[k - 1];

  /* Set the unlocks in the tasks. */
  for (int k = 0; k < s->nr_tasks; k++) {
    struct task *t = &s->tasks[k];
    t->nr_unlock_tasks = counts[k];
    t->unlock_tasks = &s->unlocks[offsets[k]];
  }

  /* Verify that there are no duplicate unlocks. */
  /* for (int k = 0; k < s->nr_tasks; k++) {
    struct task *t = &s->tasks[k];
    for (int i = 0; i < t->nr_unlock_tasks; i++) {
      for (int j = i + 1; j < t->nr_unlock_tasks; j++) {
        if (t->unlock_tasks[i] == t->unlock_tasks[j])
          error("duplicate unlock!");
      }
    }
  } */

  /* Clean up. */
  free(counts);
  free(offsets);
}

/**
 * @brief #threadpool_map function which runs through the task
 *        graph and re-computes the task wait counters.
 */

void scheduler_simple_rewait_mapper(void *map_data, void *extra_data) {

  struct task *t = (struct task *)map_data;

  /* Increment the waits of the dependances */
  for (int k = 0; k < t->nr_unlock_tasks; k++) {
    struct task *u = t->unlock_tasks[k];
    atomic_inc(&u->wait);
  }
}
/**
 * @brief Sort the tasks in topological order over all queues.
 *
 * @param s The #scheduler.
 */

void scheduler_ranktasks(struct scheduler *s) {

  struct task *tasks = s->tasks;
  int *tid = s->tasks_ind;
  const int nr_tasks = s->nr_tasks;

  /* Run through the tasks and get all the waits right. */
  threadpool_map(s->threadpool, scheduler_simple_rewait_mapper, tasks, nr_tasks,
                 sizeof(struct task), NULL);

  /* Load the tids of tasks with no waits. */
  int left = 0;
  for (int k = 0; k < nr_tasks; k++)
    if (tasks[k].wait == 0) {
      tid[left] = k;
      left += 1;
    }

  /* Main loop. */
  for (int j = 0, rank = 0; left < nr_tasks; rank++) {

    /* Did we get anything? */
    if (j == left) error("Unsatisfiable task dependencies detected.");
    const int left_old = left;

    /* Unlock the next layer of tasks. */
    for (; j < left_old; j++) {
      struct task *t = &tasks[tid[j]];
      t->rank = rank;
      /* message( "task %i of type %s has rank %i." , i ,
          (t->type == task_type_self) ? "self" : (t->type == task_type_pair) ?
         "pair" : "sort" , rank ); */
      for (int k = 0; k < t->nr_unlock_tasks; k++) {
        struct task *u = t->unlock_tasks[k];
        if (--u->wait == 0) {
          tid[left] = u - tasks;
          left += 1;
        }
      }
    }

    /* Move back to the old left (like Sanders). */
    j = left_old;
  }

  /* Verify that the tasks were ranked correctly. */
  /* for (int k = 1; k < s->nr_tasks; k++)
    if (tasks[tid[k - 1]].rank > tasks[tid[k - 1]].rank)
      error("Task ranking failed."); */
}

/**
 * @brief (Re)allocate the task arrays.
 *
 * @param s The #scheduler.
 * @param size The maximum number of tasks in the #scheduler.
 */

void scheduler_reset(struct scheduler *s, int size) {

  /* Do we need to re-allocate? */
  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);

    /* Allocate the new lists. */
    if ((s->tasks = (struct task *)malloc(sizeof(struct task) *size)) == NULL ||
        (s->tasks_ind = (int *)malloc(sizeof(int) * size)) == NULL)
      error("Failed to allocate task lists.");
  }

  /* Reset the counters. */
  s->size = size;
  s->nr_tasks = 0;
  s->tasks_next = 0;
  s->waiting = 0;
  s->mask = 0;
  s->submask = 0;
  s->nr_unlocks = 0;
  s->completed_unlock_writes = 0;

  /* Set the task pointers in the queues. */
  for (int k = 0; k < s->nr_queues; k++) s->queues[k].tasks = s->tasks;
}

/**
 * @brief Compute the task weights
 *
 * @param s The #scheduler.
 */

void scheduler_reweight(struct scheduler *s) {

  const int nr_tasks = s->nr_tasks;
  int *tid = s->tasks_ind;
  struct task *tasks = s->tasks;
  const int nodeID = s->nodeID;
  const float sid_scale[13] = {0.1897, 0.4025, 0.1897, 0.4025, 0.5788,
                               0.4025, 0.1897, 0.4025, 0.1897, 0.4025,
                               0.5788, 0.4025, 0.5788};
  const float wscale = 0.001;
  // ticks tic;

  /* Run through the tasks backwards and set their waits and
     weights. */
  // tic = getticks();
  for (int k = nr_tasks - 1; k >= 0; k--) {
    struct task *t = &tasks[tid[k]];
    t->weight = 0;
    for (int j = 0; j < t->nr_unlock_tasks; j++)
      if (t->unlock_tasks[j]->weight > t->weight)
        t->weight = t->unlock_tasks[j]->weight;
    if (!t->implicit && t->tic > 0)
      t->weight += wscale * (t->toc - t->tic);
    else
      switch (t->type) {
        case task_type_sort:
          t->weight += wscale * intrinsics_popcount(t->flags) * t->ci->count *
                       (sizeof(int) * 8 - intrinsics_clz(t->ci->count));
          break;
        case task_type_self:
          t->weight += 1 * t->ci->count * t->ci->count;
          break;
        case task_type_pair:
          if (t->ci->nodeID != nodeID || t->cj->nodeID != nodeID)
            t->weight +=
                3 * wscale * t->ci->count * t->cj->count * sid_scale[t->flags];
          else
            t->weight +=
                2 * wscale * t->ci->count * t->cj->count * sid_scale[t->flags];
          break;
        case task_type_sub:
          if (t->cj != NULL) {
            if (t->ci->nodeID != nodeID || t->cj->nodeID != nodeID) {
              if (t->flags < 0)
                t->weight += 3 * wscale * t->ci->count * t->cj->count;
              else
                t->weight += 3 * wscale * t->ci->count * t->cj->count *
                             sid_scale[t->flags];
            } else {
              if (t->flags < 0)
                t->weight += 2 * wscale * t->ci->count * t->cj->count;
              else
                t->weight += 2 * wscale * t->ci->count * t->cj->count *
                             sid_scale[t->flags];
            }
          } else
            t->weight += 1 * wscale * t->ci->count * t->ci->count;
          break;
        case task_type_ghost:
          if (t->ci == t->ci->super) t->weight += wscale * t->ci->count;
          break;
        case task_type_kick:
          t->weight += wscale * t->ci->count;
          break;
        case task_type_drift:
          t->weight += wscale * t->ci->count;
          break;
        case task_type_init:
          t->weight += wscale * t->ci->count;
          break;
        default:
          break;
      }
    if (t->type == task_type_send) t->weight = INT_MAX / 8;
    if (t->type == task_type_recv) t->weight *= 1.41;
  }
  // message( "weighting tasks took %.3f %s." ,
  // clocks_from_ticks( getticks() - tic ), clocks_getunit());

  /* int min = tasks[0].weight, max = tasks[0].weight;
  for ( k = 1 ; k < nr_tasks ; k++ )
      if ( tasks[k].weight < min )
          min = tasks[k].weight;
      else if ( tasks[k].weight > max )
          max = tasks[k].weight;
  message( "task weights are in [ %i , %i ]." , min , max ); */
}

/**
 * @brief #threadpool_map function which runs through the task
 *        graph and re-computes the task wait counters.
 */

void scheduler_rewait_mapper(void *map_data, void *extra_data) {

  struct scheduler *s = (struct scheduler *)extra_data;
  struct task *t = (struct task *)map_data;

  if (t->skip) return;

  /* Skip tasks not in the mask */
  if (!((1 << t->type) & s->mask) || !((1 << t->subtype) & s->submask)) return;

  /* Skip sort tasks that have already been performed */
  if (t->type == task_type_sort && t->flags == 0) return;

  /* Sets the waits of the dependances */
  for (int k = 0; k < t->nr_unlock_tasks; k++) {
    struct task *u = t->unlock_tasks[k];
    atomic_inc(&u->wait);
  }
}

/**
 * @brief Start the scheduler, i.e. fill the queues with ready tasks.
 *
 * @param s The #scheduler.
 * @param mask The task types to enqueue.
 * @param submask The sub-task types to enqueue.
 */

void scheduler_start(struct scheduler *s, unsigned int mask,
                     unsigned int submask) {

  const int nr_tasks = s->nr_tasks;
  int *tid = s->tasks_ind;
  struct task *tasks = s->tasks;
  // ticks tic;

  /* Store the masks */
  s->mask = mask | (1 << task_type_comm_root);