Closer and closer to global work stealing done

use Queue;
use Random;
use Time;
use BlockDist;
use BigInteger;
use LocalAtomicObject;

/*
  A highly parallel segmented queue. Each node gets its own queue, and in each queue
  its segmented into 'here.maxTaskPar' segments. Segments allow for actual parallelism
  while operating on the queue in that it enables us to manage 'best-case', 'average-case',
  and 'worst-case' scenarios by making multiple passes over each segment. Examples
  of 'best-case' scenarios would be when a segment is unlocked. 'average-case' would be
  when any unlocked/locked segment, and so on.

  This queue employs work stealing, where after multiple failed passes over the
  segments in the queue we attempt to steal work from other queues. Work stealing
  employs a nice helper algorithm where any task performing dequeue operations,
  after also making multiple failed consecutive passes over the queue,
  will enlist their help in terms of finding elements from other nodes.

  TODO/BUG: Global work stealing is non-deterministic in terms of whether or not
  the queue is empty. If we process a node and they steal from another node ahead
  of us, then we skip over it as if the queue is empty when it really isn't. Its
  possible a rather large amount has been stolen as well, which makes this extremely
  problematic. Requires a better global synchronization scheme, definitely need
  the employment for GlobalAtomicObject.
*/

/*
  Below are segment statuses, which is a way to make visible to outsiders the
  current ongoing operation.
*/
// Free to contest
const UNLOCKED : uint = 0;
// These are normal operations which are for more-verbose debugging purposes.
const ENQUEUE : uint = 1;
const DEQUEUE : uint = 2;
// The work stealer is initializing fields for work stealing, and will be ready soon.
const STEALING_INIT : uint = 3;
// The work stealer is currently stealing work, and may require assistance.
const STEALING_WORK : uint = 4;
// The work stealer is currently merging work that is stolen, and no longer requires assistance.
const STEALING_MERGE : uint = 5;

/*
  Below are statuses specific work queue segment slots.
*/
// The work stealer has reinitialized the slot for use.
const INITIALIZED : int = -1;
// A helper task has finished but has not found any work...
const FINISHED_WITH_NO_WORK : int = 0;
// A helper task has finished with work.
const FINISHED_WITH_WORK : int = 1;

// The amount of elements per unroll block; defines the tuple size.
config param ELEMS_PER_BLOCK = 1024;
// The amount we attempt to steal from each segment of other nodes. Helps to
// further ensure a proper load balance by not leaving segments empty by stealing
// on average 1/4 of their work, leaving them with 3/4. As we steal from *all* other
// segments, it is imperative that this stays low to prevent a ping-pong effect on
// work propagating from excessive work stealing. We steal at least one element, but
// no more than ELEMS_PER_BLOCK per segment.
config param WORK_STEALING_RATIO = 0.25;
config param logMPMCQueue = false;

class WorkQueueSegmentBlock {
  type eltType;
  var nElems : int;
  var elems : ELEMS_PER_BLOCK * eltType;
  var next : WorkQueueSegmentBlock(eltType);

  iter drain(n) {
    while n > 0 {
      if nElems == 0 {
        return;
      }

      var elem = elems[nElems];
      nElems = headBlock.nElems - 1;
      nElems.sub(1);

      yield elem;
      n = n - 1;
    }
  }
}

record WorkQueueSegmentSlot {
  type eltType;

  // The locale we are assigned to steal from...
  var locIdx : int;
  var status : atomic int;
  var block : WorkQueueSegmentBlock(eltType);
}

record WorkQueueSegment {
  type eltType;

  var status : atomic uint;

  var headBlock : WorkQueueSegmentBlock(eltType);
  var tailBlock : WorkQueueSegmentBlock(eltType);

  /*
    Fields specific to work stealing for this segment (given 'ws' prefix).
    A helper must first 'register' themselves in our status table, which keeps
    track of which cells are currently taken and which are not. In work stealing,
    we attempt to steal ELEMS_PER_BLOCK * here.maxTaskPar elements, but never assign
    more than one task to a given node.
  */

  // Number of potential helpers. We don't want more workers than there are nodes.
  const wsNumSlots = numLocales - 1;

  // Used to atomically obtain a slot. A task knows they have a slot if it in
  // [0, wsNumSlots). If it is outside of this range, then all slots are filled
  // or this work stealing session is over before the next is initialized.
  var wsSlotsTaken : atomic int;

  // Slots which keep track of data stolen as well as the current status of the helper task.
  var wsSlots : [{0..#wsNumSlots}] WorkQueueSegmentSlot(eltType);

  // Number of elements currently in the list to allow faster determinism of whether
  // or not an operation should attempt this bucket.
  var nElems : atomic uint;

  inline proc isEmpty {
    return nElems.read() == 0;
  }

  inline proc acquireWithStatus(newStatus) {
    return status.compareExchangeStrong(UNLOCKED, newStatus);
  }

  inline proc isUnlocked {
    return status.read() == UNLOCKED;
  }

  inline proc currentStatus {
    return status.read();
  }

  inline proc releaseStatus() {
    status.write(UNLOCKED);
  }

  // TODO: Increase efficiency of bulk taking
  iter proc takeElements(nElems) {
    while true {
      if isEmpty {
        return;
      }

      var elem = headBlock.elems[headBlock.nElems];
      headBlock.nElems = headBlock.nElems - 1;
      nElems.sub(1);

      // Fix list if we consumed last one...
      if headBlock.nElems == 0 {
        var block = headBlock;
        headBlock = headBlock.next;
        delete block;
      }

      yield elem;
    }
  }

  // TODO: Increase efficiency of bulk adding
  proc addElements(elts) {
    var block = tailBlock;

    // Empty?
    if block == nil then {
      tailBlock = new WorkQueueSegmentBlock(eltType);
      headBlock = tailBlock;
      block = tailBlock;
    }

    for elt in elts {
      // Full?
      if block.nElems == ELEMS_PER_BLOCK {
        block.next = new WorkQueueSegmentBlock(eltType);
        tailBlock = block.next;
        block = block.next;
      }

      block.nElems = block.nElems + 1;
      block.elems[block.nElems] = elt;

      // We update the count in each iteration to increase the visibility of this
      // operation (as dequeue operations would skip over us during 2nd pass)
      nElems.add(1);
    }
  }
}

class WorkQueue : Queue {
  // A handle to our parent 'distributed' work queue, which is needed for work stealing.
  var parentHandle : DistributedWorkQueue(eltType);

  /*
    Helps evenly distribute and balance placement of elements in a best-effort
    round-robin approach. In the case where we have parallel enqueues or dequeues,
    they are less likely overlap with each other. Furthermore, it increases our
    chance to find our 'ideal' segment.
  */
  var startIdxEnq : atomic uint;
  var startIdxDeq : atomic uint;

  /*
    We only allow one segment to be nominated as the 'work stealer' for our queue.
    This reduces the complexity of having to manage multiple work stealing segments
    while also providing more concurrency through work stealing.
  */
  var workStealer : atomic int;

  var maxParallelSegmentSpace = {0 .. #here.maxTaskPar};
  var segments : [maxParallelSegmentSpace] WorkQueueSegment(eltType);

  const maxIterationsPerPhase = here.maxTaskPar;

  inline nextStartIdxEnq {
    return (startIdxEnq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  inline nextStartIdxDeq {
    return (startIdxDeq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  proc WorkQueue(type eltType) {
    workStealer.write(-1);
  }

  /*
    Enqueue...
  */

  proc enqueue(elt : eltType) {
    var startIdx = nextStartIdxEnq;
    var idx = startIdx : int;
    var iterations = 0;

    /*
      Pass 1: Best Case

      Find a segment that is unlocked and attempt to acquire it. As we are adding
      elements, we don't care how many elements there are, just that we find
      some place to add ours.
    */
    while iterations < maxIterationsPerPhase {
      ref segment = segments[idx];

      // Attempt to acquire...
      if segment.acquireWithStatus(ENQUEUE) {
        segment.addElements(elt);
        segment.releaseStatus();
        return;
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    iterations = 0;
    idx = startIdx : int;

    /*
      Pass 2: Average Case

      Find any segment (locked or unlocked) and make an attempt to acquire it.
      If the segment is actively 'work stealing' then we skip over it to minimize
      the excessive wait-time; given that there can be only *one* work-stealer
      for a node, we know we'll find a better one nearby.
    */
    while true {
      ref segment = segments[idx];

      while true {
        select segment.currentStatus {
          // Quick acquire...
          when UNLOCKED do {
            if segment.acquireWithStatus(ENQUEUE) {
              segment.addElements(elt);
              segment.releaseStatus();
              return;
            }
          }
          // If someone is stealing, then that means that we'd be waiting for a while
          // anyway, so we may as well attempt to find a node with a shorter wait time.
          when STEALING_INIT do break;
          when STEALING_WORK do break;
          when STEALING_MERGE do break;
        }

        chpl_task_yield();
      }

      idx = (idx + 1) % here.maxTaskPar;
    }

    halt("Almost exited enqueueAcquire before getting a segment...");
  }

  /*
    Dequeue...
  */

  proc dequeue() : (bool, eltType) {
    var startIdx = nextStartIdxDeq;
    var idx = startIdx;
    var iterations = 0;

    /*
      Pass 1: Best Case

      Find the first bucket that is both unlocked and contains elements. This is
      extremely helpful in the case where we have a good distribution of elements
      in each segment.
    */
    while iterations < maxIterationsPerPhase {
      ref segment = segments[idx];

      // Attempt to acquire...
      if !segment.isEmpty && segment.acquireWithStatus(DEQUEUE) {
        var (hasElem, elem) : (bool, eltType);
        for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
        segment.releaseStatus();

        if hasElem {
          return (hasElem, elem);
        }
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    idx = startIdx;
    iterations = 0;

    /*
      Pass 2: Average Case

      Find the first bucket containing elements. We don't care if it is locked
      or unlocked this time, just that it contains elements; this handles majority
      of cases where we have elements anywhere in any segment.
    */
    while iterations < maxIterations {
      ref segment = segments[idx];

      // Attempt to acquire...
      while !segment.isEmpty {
        if segment.isUnlocked && segment.acquireWithStatus(DEQUEUE) {
          var (hasElem, elem) : (bool, eltType);
          for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
          segment.releaseStatus();

          if hasElem {
            return (hasElem, elem);
          }
        }

        // Backoff
        chpl_task_yield();
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    idx = startIdx;

    /*
      Pass 3: Worst Case

      After two full iterations, we're sure the queue is full at this point, so we
      can attempt to steal work from other nodes. In this pass, we find *any* segment
      and if it is empty, we attempt to become the work-stealer; if someone else is the
      current work stealer we assist them instead and lift an element for ourselves.

      Furthermore, in this phase we loop indefinitely until we are 100% certain it is
      empty or we get an item, so introduce some backoff here.
    */
    var backoff = 0;
    while true {
      ref segment = segments[idx];

      select segment.currentStatus {
        // Quick acquire
        when UNLOCKED {
          backoff = 0;
          if segment.acquireWithStatus(DEQUEUE) {
            // Steal if empty... because we help if we are not the one set, we only
            // know whether the queue is empty if we were the work stealer.
            if seg.isEmpty {
              var (_idx, hasElem, elem) = stealWork(idx);

              // If we have lifted an element, we've retrieved an item.
              if hasElem {
                return (hasElem, elem);
              }

              // If the index returned is equivalent to ours, then we were the 'work stealer'
              // for our queue, and since only we personally know how many elements there are
              // (as we were in charge of merging all work stolen), we know the queue is empty.
              // Otherwise, if the index differs then it means we helped assist another
              // segment in work stealing and we should move to it as if there is any work
              // available, it would be there.
              if _idx == idx && !hasElem {
                return (false, _defaultOf(eltType));
              } else {
                idx = _idx;
              }

              continue;
            }

            // We're lucky; another element has been added to the current segment,
            // take it and leave like normal...
            var (hasElem, elem) : (bool, eltType);
            for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
            segment.releaseStatus();
            return (hasElem, elem);
          }
        }
        // If someone else is stealing work, we add ourself as a potential helper
        // to help speed this along, and hopefully we get an element for ourself.
        when STEALING_WORK {
          var (hasElem, elem) = helpStealWork(startIdx);

          // We have our element...
          if hasElem {
            return (hasElem, elem);
          }
          backoff = 0;
        }
      }

      // Backoff to maximum...
      if backoff == 0 then chpl_task_yield();
      else sleep(max(1000, 2 ** backoff), TimeUnits.microseconds);
      backoff = backoff + 1;
    }

    halt("Somehow exited without acquiring an item for dequeue...");
  }

  /*
    Work Stealing
  */

  /*
    Attempt to steal work as the segment referenced by 'idx'. As we may only have
    one work stealing segment, we must first contest for that right. If we obtain
    that right, we become responsible with coordinating and distributing work to other
    helper tasks, and then merging all of the work together. If another segment has
    that right, we move to that segment and become one of their helpers.
  */
  proc stealWork(idx) : (int, bool, eltType) {
    ref segment = segments[idx];

    // We can't steal if we're the only ones, so return as empty
    if numLocales == 1 {
        return (idx, false, _defaultOf(eltType));
    }

    if logMPMCQueue then writeln("Attempting to steal work for segment #", idx);

    // Attempt to become the sole work stealing segment. If we fail, then we enlist
    // our aid to the current work stealing segment as a helper.
    while true {
      // We are the new work stealing segment...
      if workStealer.compareExchangeStrong(-1, idx) {
        break;
      }

      // Another segment is the work stealing segment; help them out...
      // If we help steal work, it is possible we also lift an element while helping
      // Furthermore, we also return the index of the work stealing segment, as it is
      // likely to have more work.
      var currentStealer = workStealer.read();
      if currentStealer != -1 {
        segment.releaseStatus();
        var (hasElem, elem) = helpStealWork(currentStealer);
        return (currentStealer, hasElem, elem);
      }
    }

    // Initialize each slot for potential helpers. It is imperative that we initialize
    // the slots *before* resetting the wsSlotsTaken to avoid race conditions, as
    // helpers only help if they obtain a valid index.
    var locIdx = 0;
    for slot in segment.wsSlots {
      if locIdx == here.id then locIdx = locIdx + 1;
      slot.locIdx = locIdx;
      slot.status.write(INITIALIZED);
      locIdx = locIdx + 1;
    }
    segment.wsSlotsTaken.write(0);
    segment.status.write(STEALING_WORK);

    // Help yourself...
    var (hasElem, elem) = helpStealWork(idx);
    // Lets others know that we no longer require help.
    segment.status.write(STEALING_MERGE);
    if logMPMCQueue then writeln("Merging work stolen on segment #", idx);

    // At this point, it is possible that other helpers are still working, and so
    // we wait for them to finish. Merge what has been stolen.
    for slot in segment.wsSlots {
      // Wait for the current helper to finish, and when they do skip any without work...
      while slot.status.read() == INITIALIZED do chpl_task_yield();
      if slot.status.read() == FINISHED_WITH_NO_WORK then continue;

      ref stolenBlock = slot.block;
      var tailBlock = seg.tailBlock;
      var blockElems = stolenBlock.nElems;
      segment.nElems.add(blockElems : uint);

      // Drain elements we've stolen and merge it into our segment.
      for elt in stolenBlock.drain(ELEMS_PER_BLOCK) {
        // If we have not lifted an element for ourself by now, do so.
        if !hasElem {
          elem = elt;
          hasElem = true;
          segment.nElems.sub(1);
          if logMPMCQueue then writeln("Lifted element ", elem, " from block stolen from segment #", i, " as segment #", idx);
          continue;
        }

        segment.addElements(elt);
      }

      delete stolenBlock;
    }

    // At this point, we have stolen from all other nodes, and we may relinquish
    // our role as primary work-stealer.
    workStealer.write(-1);
    segment.releaseStatus();
    return (idx, hasElem, elem);
  }

  /*
    Enlist our aid as a helper thread for the current work stealer. To do so, we
    must claim a slot (atomically),
  */
  proc helpStealWork(idx) : (bool, eltType) {
    ref segment = segments[idx];

    while true {
      select segment.currentStatus {
        // They are still initializing...
        when STEALING_INIT {
          // Tight spin, we'll be able to help soon...
        }
        // They are initialized, we can now attempt to help...
        when STEALING_WORK {
          break;
        }
        // We're not stealing work...
        otherwise do return (false, _defaultOf(eltType));
      }
    }

    if logMPMCQueue then writeln("Helping steal work for segment #", idx);
    var (hasElem, elem) : (bool, eltType);

    // Attempt to help...
    while true {
      var ourIdx = segment.wsSlotsTaken.fetchAdd(1);
      // If the index we get is not in a valid range, we're done helping.
      if ourIdx >= segment.wsNumSlots {
        return (hasElem, elem);
      }

      ref slot = segment.wsSlots[ourIdx];
      if logMPMCQueue then writeln("Assigned to steal work from segment #", ourIdx, " for segment #", idx);

      // We know which locale we are going to be working on, but we must setup some
      // temporaries that can make it easier for back-and-forth communication.
      // We create an array of tuples, one for each segment on the other node.
      var stolenWork : [{0..#Locales[slot.locIdx].maxTaskPar}] (int, ELEMS_PER_BLOCK * eltType);

      on Locales[slot.locIdx] {
        // We keep track of a simple bitmap to keep track of which segments we have
        // already processed, since on the first pass we skip over any locked segments.
        var segmentBitmap : bigint;
        var queue = parentHandle.localQueue;
        var iterations = 0;
        var otherStartIdx = queue.nextStartIdxDeq;
        var otherSegmentIdx = otherStartIdx;

        /*
          Pass 1: Best Case

          Find any unlocked segment that contains elements...
        */
        while iterations < here.maxTaskPar {
          ref otherSegment = segments[otherSegmentIdx];

          // Attempt to acquire...
          if !otherSegment.isEmpty && otherSegment.acquireWithStatus(DEQUEUE) {
            var toSteal = max(1, min(ELEMS_PER_BLOCK, otherSegment.nElems.read() * WORK_STEALING_RATIO));
            var stolen : ELEMS_PER_BLOCK * eltType;
            var stolenIdx = 1;
            for elt in otherSegment.takeElements(toSteal) {
                stolen[stolenIdx] = elt;
                stolenIdx = stolenIdx + 1;
            }
            otherSegment.releaseStatus();

            // Mark as processed and move stolen work...
            segmentBitmap.setbit(otherSegmentIdx);
            stolenWork[otherSegmentIdx] = (stolenIdx, stolen);
          }

          iterations = iterations + 1;
          otherSegmentIdx = (otherSegmentIdx + 1) % here.maxTaskPar;
        }

        otherSegmentIdx = otherStartIdx;
        iterations = 0;

        /*
          Pass 2: Average Case

          Find any segment not already processed that contains elements...
        */
        while iterations < maxIterations {
          if segmentBitmap.tstbit(otherSegmentIdx) {
            iterations = iterations + 1;
            otherSegmentIdx = (otherSegmentIdx + 1) % here.maxTaskPar;
            continue;
          }
          ref otherSegment = segments[otherSegmentIdx];

          // Attempt to acquire...
          while !otherSegment.isEmpty {
            // If they are stealing work as well, we back out.
            if otherSegment.status.read() == STEALING_WORK then break;

            if otherSegment.isUnlocked && otherSegment.acquireWithStatus(DEQUEUE) {
              var toSteal = max(1, min(ELEMS_PER_BLOCK, otherSegment.nElems.read() * WORK_STEALING_RATIO));
              var stolen : ELEMS_PER_BLOCK * eltType;
              var stolenIdx = 1;
              for elt in otherSegment.takeElements(toSteal) {
                  stolen[stolenIdx] = elt;
                  stolenIdx = stolenIdx + 1;
              }
              otherSegment.releaseStatus();

              // Mark as processed and move stolen work...
              segmentBitmap.setbit(otherSegmentIdx);
              stolenWork[otherSegmentIdx] = (stolenIdx, stolen);
            }

            // Backoff
            chpl_task_yield();
          }

          iterations = iterations + 1;
          otherSegmentIdx = (otherSegmentIdx + 1) % here.maxTaskPar;
        }

        // At this point we are done, we've stolen from all segments that aren't
        // currently stealing work...
      }

      // Link all of the stolen work together (only if they contain any work)

      if logMPMCQueue then writeln("Finished stealing ", nStolen, " amounts of work for segment #", idx);
    }

    halt("Exited helpStealWork loop...");
  }
}

class DistributedWorkQueue : Queue {
  var targetLocales : [] locale = Locales;
  var perLocaleDomain = domain(1) dmapped Block(boundingBox=targetLocales.domain, targetLocales=targetLocales);
  var workQueues : [perLocaleDomain] WorkQueue(eltType);

  proc DistributedWorkQueue(type eltType) {
    forall loc in Locales {
      on loc {
        workQueues[workQueues.domain.subDomain().first] = new WorkQueue(eltType, parentHandle=this);
      }
    }
  }

  inline proc localQueue {
    return workQueues[workQueues.domain.subDomain().first];
  }
}

proc proof_of_correctness() {
  var nElementsForMPMC = 1000000;
  writeln("Starting MPMCQueue Proof of Correctness (Work Depletion) Test ~ nElementsForMPMC: ", nElementsForMPMC);

  var queue = new WorkQueue(int);
  var counter : atomic int;
  var expected = (nElementsForMPMC * (nElementsForMPMC + 1)) / 2;

  forall i in 1 .. nElementsForMPMC {
    queue.enqueue(i);
  }

  coforall loc in Locales do on loc {
    var localeLocalCounter : atomic int;
    coforall tid in 0 .. #here.maxTaskPar {
      var taskLocalCounter : int;
      while true {
        var (present, elem) = queue.dequeue();
        if !present {
          break;
        }

        taskLocalCounter = taskLocalCounter + elem;
      }
      localeLocalCounter.add(taskLocalCounter);
    }
    writeln(here, ": Counter=", localeLocalCounter.read());
    counter.add(localeLocalCounter.read());
  }

  if counter.read() == expected {
    writeln("PASSED TEST!");
  } else {
    writeln("FAILED TEST! Expected: ", expected, ", but received: ", counter.read());
    writeln(queue);
  }
}

proc benchmark() {
  var nJitter = 0;
  var nComputations = 0;
  var nElements = 1000000;
  var nTrials = 8;
  var enqueueTrialTime : [1 .. nTrials] real;
  var dequeueTrialTime : [1 .. nTrials] real;

  // Obtain average time for enqueue followed by dequeued...
  for i in 1 .. nTrials {
    var queue = new WorkQueue(int);
    // We only use one or the other, but we must declare both.
    // TODO: Have both implement some parent 'Queue' class/interface?
    var timer = new Timer();
    timer.start();

    coforall loc in Locales do on loc {
      var iterations = nElements;

      coforall tid in 0 .. #here.maxTaskPar {
        var x : atomic int;
        var randStr = makeRandomStream(int);
        for j in 1 .. iterations / here.maxTaskPar {
          queue.enqueue(j);
          var nComps = nComputations + (if nJitter then (randStr.getNext() % nJitter) else 0);
          for i in 1 .. nComps {
            // Hopefully compiler doesn't throw away?
            x.write(sin(i) : int);
          }
        }
      }
    }

    timer.stop();
    enqueueTrialTime[i] = nElements / timer.elapsed();
    writeln(i, "/", nTrials, ": ", (+ reduce enqueueTrialTime) / i);
    timer.clear();
    timer.start();

    coforall loc in Locales do on loc {
      var iterations = nElements;

      coforall tid in 0 .. #here.maxTaskPar {
        var x : atomic int;
        var randStr = makeRandomStream(int);
        for j in 1 .. iterations / here.maxTaskPar {
          var retval = queue.dequeue();
          var nComps = nComputations + (if nJitter then (randStr.getNext() % nJitter) else 0);
          for i in 1 .. nComps {
            // Hopefully compiler doesn't throw away?
            x.write(sin(i) : int);
          }
        }
      }
    }

    timer.stop();
    dequeueTrialTime[i] = nElements / timer.elapsed();
    writeln(i, "/", nTrials, ": ", (+ reduce dequeueTrialTime) / i);

    delete queue;
  }
}

proc main() {
  benchmark();
}

use Queue;
use Random;
use Time;
use BlockDist;
use BigInteger;
use LocalAtomicObject;

/*
  A highly parallel segmented queue. Each node gets its own queue, and in each queue
  its segmented into 'here.maxTaskPar' segments. Segments allow for actual parallelism
  while operating on the queue in that it enables us to manage 'best-case', 'average-case',
  and 'worst-case' scenarios by making multiple passes over each segment. Examples
  of 'best-case' scenarios would be when a segment is unlocked. 'average-case' would be
  when any unlocked/locked segment, and so on.

  This queue employs work stealing, where after multiple failed passes over the
  segments in the queue we attempt to steal work from other queues. Work stealing
  employs a nice helper algorithm where any task performing dequeue operations,
  after also making multiple failed consecutive passes over the queue,
  will enlist their help in terms of finding elements from other nodes.

  TODO/BUG: Global work stealing is non-deterministic in terms of whether or not
  the queue is empty. If we process a node and they steal from another node ahead
  of us, then we skip over it as if the queue is empty when it really isn't. Its
  possible a rather large amount has been stolen as well, which makes this extremely
  problematic. Requires a better global synchronization scheme, definitely need
  the employment for GlobalAtomicObject.
*/

/*
  Below are segment statuses, which is a way to make visible to outsiders the
  current ongoing operation.
*/
// Free to contest
const UNLOCKED : uint = 0;
// These are normal operations which are for more-verbose debugging purposes.
const ENQUEUE : uint = 1;
const DEQUEUE : uint = 2;
// The work stealer is initializing fields for work stealing, and will be ready soon.
const STEALING_INIT : uint = 3;
// The work stealer is currently stealing work, and may require assistance.
const STEALING_WORK : uint = 4;
// The work stealer is currently merging work that is stolen, and no longer requires assistance.
const STEALING_MERGE : uint = 5;

/*
  Below are statuses specific work queue segment slots.
*/
// The work stealer has reinitialized the slot for use.
const INITIALIZED : int = -1;
// A helper task has finished but has not found any work...
const FINISHED_WITH_NO_WORK : int = 0;
// A helper task has finished with work.
const FINISHED_WITH_WORK : int = 1;

// The amount of elements per unroll block; defines the tuple size.
config param ELEMS_PER_BLOCK = 1024;
// The amount we attempt to steal from each segment of other nodes. Helps to
// further ensure a proper load balance by not leaving segments empty by stealing
// on average 1/4 of their work, leaving them with 3/4. As we steal from *all* other
// segments, it is imperative that this stays low to prevent a ping-pong effect on
// work propagating from excessive work stealing. We steal at least one element, but
// no more than ELEMS_PER_BLOCK per segment.
config param WORK_STEALING_RATIO = 0.25;
config param logMPMCQueue = false;

class WorkQueueSegmentBlock {
  type eltType;
  var nElems : int;
  var elems : ELEMS_PER_BLOCK * eltType;
  var next : WorkQueueSegmentBlock(eltType);

  iter drain(n) {
    while n > 0 {
      if nElems == 0 {
        return;
      }

      var elem = elems[nElems];
      nElems = headBlock.nElems - 1;
      nElems.sub(1);

      yield elem;
      n = n - 1;
    }
  }
}

record WorkQueueSegmentSlot {
  type eltType;

  var status : atomic int;
  var block : WorkQueueSegmentBlock(eltType);
}

record WorkQueueSegment {
  type eltType;

  var status : atomic uint;

  var headBlock : WorkQueueSegmentBlock(eltType);
  var tailBlock : WorkQueueSegmentBlock(eltType);

  /*
    Fields specific to work stealing for this segment (given 'ws' prefix).
    A helper must first 'register' themselves in our status table, which keeps
    track of which cells are currently taken and which are not. In work stealing,
    we attempt to steal ELEMS_PER_BLOCK * here.maxTaskPar elements, but never assign
    more than one task to a given node.
  */

  // Number of potential helpers. We don't want more workers than there are nodes.
  const wsNumSlots = min(numLocales - 1, here.maxTaskPar);

  // A randomly chosen offset we begin stealing at. Needed to prevent starvation
  // of nodes inherent in choosing some arbitrary fixed startpoint.
  var wsStartOffset : int;

  // Used to atomically obtain a slot. A task knows they have a slot if it in
  // [0, wsNumSlots). If it is outside of this range, then all slots are filled
  // or this work stealing session is over before the next is initialized.
  var wsSlotsTaken : atomic int;

  // Slots which keep track of data stolen as well as the current status of the helper task.
  var wsSlots : [{0..#wsNumSlots}] WorkQueueSegmentSlot(eltType);

  // Number of elements currently in the list to allow faster determinism of whether
  // or not an operation should attempt this bucket.
  var nElems : atomic uint;

  inline proc isEmpty {
    return nElems.read() == 0;
  }

  inline proc acquireWithStatus(newStatus) {
    return status.compareExchangeStrong(UNLOCKED, newStatus);
  }

  inline proc isUnlocked {
    return status.read() == UNLOCKED;
  }

  inline proc currentStatus {
    return status.read();
  }

  inline proc releaseStatus() {
    status.write(UNLOCKED);
  }

  // TODO: Increase efficiency of bulk taking
  iter proc takeElements(nElems) {
    while true {
      if isEmpty {
        return;
      }

      var elem = headBlock.elems[headBlock.nElems];
      headBlock.nElems = headBlock.nElems - 1;
      nElems.sub(1);

      // Fix list if we consumed last one...
      if headBlock.nElems == 0 {
        var block = headBlock;
        headBlock = headBlock.next;
        delete block;
      }

      yield elem;
    }
  }

  // TODO: Increase efficiency of bulk adding
  proc addElements(elts) {
    var block = tailBlock;

    // Empty?
    if block == nil then {
      tailBlock = new WorkQueueSegmentBlock(eltType);
      headBlock = tailBlock;
      block = tailBlock;
    }

    for elt in elts {
      // Full?
      if block.nElems == ELEMS_PER_BLOCK {
        block.next = new WorkQueueSegmentBlock(eltType);
        tailBlock = block.next;
        block = block.next;
      }

      block.nElems = block.nElems + 1;
      block.elems[block.nElems] = elt;

      // We update the count in each iteration to increase the visibility of this
      // operation (as dequeue operations would skip over us during 2nd pass)
      nElems.add(1);
    }
  }
}

class WorkQueue : Queue {
  /*
    Helps evenly distribute and balance placement of elements in a best-effort
    round-robin approach. In the case where we have parallel enqueues or dequeues,
    they are less likely overlap with each other. Furthermore, it increases our
    chance to find our 'ideal' segment.
  */
  var startIdxEnq : atomic uint;
  var startIdxDeq : atomic uint;

  /*
    We only allow one segment to be nominated as the 'work stealer' for our queue.
    This reduces the complexity of having to manage multiple work stealing segments
    while also providing more concurrency through work stealing.
  */
  var workStealer : atomic int;

  var maxParallelSegmentSpace = {0 .. #here.maxTaskPar};
  var segments : [maxParallelSegmentSpace] WorkQueueSegment(eltType);

  const maxIterationsPerPhase = here.maxTaskPar;

  inline nextStartIdxEnq {
    return (startIdxEnq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  inline nextStartIdxDeq {
    return (startIdxDeq.fetchAdd(1) % here.maxTaskPar : uint) : int;
  }

  proc WorkQueue(type eltType) {
    workStealer.write(-1);
  }

  /*
    Enqueue...
  */

  proc enqueue(elt : eltType) {
    var startIdx = nextStartIdxEnq;
    var idx = startIdx : int;
    var iterations = 0;

    /*
      Pass 1: Best Case

      Find a segment that is unlocked and attempt to acquire it. As we are adding
      elements, we don't care how many elements there are, just that we find
      some place to add ours.
    */
    while iterations < maxIterationsPerPhase {
      ref segment = segments[idx];

      // Attempt to acquire...
      if segment.acquireWithStatus(ENQUEUE) {
        segment.addElements(elt);
        segment.releaseStatus();
        return;
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    iterations = 0;
    idx = startIdx : int;

    /*
      Pass 2: Average Case

      Find any segment (locked or unlocked) and make an attempt to acquire it.
      If the segment is actively 'work stealing' then we skip over it to minimize
      the excessive wait-time; given that there can be only *one* work-stealer
      for a node, we know we'll find a better one nearby.
    */
    while true {
      ref segment = segments[idx];

      while true {
        select segment.currentStatus {
          // Quick acquire...
          when UNLOCKED do {
            if segment.acquireWithStatus(ENQUEUE) {
              segment.addElements(elt);
              segment.releaseStatus();
              return;
            }
          }
          // If someone is stealing, then that means that we'd be waiting for a while
          // anyway, so we may as well attempt to find a node with a shorter wait time.
          when STEALING_INIT do break;
          when STEALING_WORK do break;
          when STEALING_MERGE do break;
        }

        chpl_task_yield();
      }

      idx = (idx + 1) % here.maxTaskPar;
    }

    halt("Almost exited enqueueAcquire before getting a segment...");
  }

  /*
    Dequeue...
  */

  proc dequeue() : (bool, eltType) {
    var startIdx = nextStartIdxDeq;
    var idx = startIdx;
    var iterations = 0;

    /*
      Pass 1: Best Case

      Find the first bucket that is both unlocked and contains elements. This is
      extremely helpful in the case where we have a good distribution of elements
      in each segment.
    */
    while iterations < maxIterationsPerPhase {
      ref segment = segments[idx];

      // Attempt to acquire...
      if !segment.isEmpty && segment.acquireWithStatus(DEQUEUE) {
        var (hasElem, elem) : (bool, eltType);
        for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
        segment.releaseStatus();

        if hasElem {
          return (hasElem, elem);
        }
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    idx = startIdx;
    iterations = 0;

    /*
      Pass 2: Average Case

      Find the first bucket containing elements. We don't care if it is locked
      or unlocked this time, just that it contains elements; this handles majority
      of cases where we have elements anywhere in any segment.
    */
    while iterations < maxIterations {
      ref segment = segments[idx];

      // Attempt to acquire...
      while !segment.isEmpty {
        if segment.isUnlocked && segment.acquireWithStatus(DEQUEUE) {
          var (hasElem, elem) : (bool, eltType);
          for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
          segment.releaseStatus();

          if hasElem {
            return (hasElem, elem);
          }
        }

        // Backoff
        chpl_task_yield();
      }

      iterations = iterations + 1;
      idx = (idx + 1) % here.maxTaskPar;
    }

    idx = startIdx;

    /*
      Pass 3: Worst Case

      After two full iterations, we're sure the queue is full at this point, so we
      can attempt to steal work from other nodes. In this pass, we find *any* segment
      and if it is empty, we attempt to become the work-stealer; if someone else is the
      current work stealer we assist them instead and lift an element for ourselves.

      Furthermore, in this phase we loop indefinitely until we are 100% certain it is
      empty or we get an item, so introduce some backoff here.
    */
    var backoff = 0;
    while true {
      ref segment = segments[idx];

      select segment.currentStatus {
        // Quick acquire
        when UNLOCKED {
          backoff = 0;
          if segment.acquireWithStatus(DEQUEUE) {
            // Steal if empty... because we help if we are not the one set, we only
            // know whether the queue is empty if we were the work stealer.
            if seg.isEmpty {
              var (_idx, hasElem, elem) = stealWork(idx);

              // If we have lifted an element, we've retrieved an item.
              if hasElem {
                return (hasElem, elem);
              }

              // If the index returned is equivalent to ours, then we were the 'work stealer'
              // for our queue, and since only we personally know how many elements there are
              // (as we were in charge of merging all work stolen), we know the queue is empty.
              // Otherwise, if the index differs then it means we helped assist another
              // segment in work stealing and we should move to it as if there is any work
              // available, it would be there.
              if _idx == idx && !hasElem {
                return (false, _defaultOf(eltType));
              } else {
                idx = _idx;
              }

              continue;
            }

            // We're lucky; another element has been added to the current segment,
            // take it and leave like normal...
            var (hasElem, elem) : (bool, eltType);
            for elt in segment.takeElements(1) do (hasElem, elem) = (true, elt);
            segment.releaseStatus();
            return (hasElem, elem);
          }
        }
        // If someone else is stealing work, we add ourself as a potential helper
        // to help speed this along, and hopefully we get an element for ourself.
        when STEALING_WORK {
          var (hasElem, elem) = helpStealWork(startIdx);

          // We have our element...
          if hasElem {
            return (hasElem, elem);
          }
          backoff = 0;
        }
      }

      // Backoff to maximum...
      if backoff == 0 then chpl_task_yield();
      else sleep(max(1000, 2 ** backoff), TimeUnits.microseconds);
      backoff = backoff + 1;
    }

    halt("Somehow exited without acquiring an item for dequeue...");
  }

  /*
    Work Stealing
  */

  /*
    Attempt to steal work as the segment referenced by 'idx'. As we may only have
    one work stealing segment, we must first contest for that right. If we obtain
    that right, we become responsible with coordinating and distributing work to other
    helper tasks, and then merging all of the work together. If another segment has
    that right, we move to that segment and become one of their helpers.
  */
  proc stealWork(idx) : (int, bool, eltType) {
    ref segment = segments[idx];

    // We can't steal if we're the only ones, so return as empty
    if numLocales == 1 {
        return (idx, false, _defaultOf(eltType));
    }

    if logMPMCQueue then writeln("Attempting to steal work for segment #", idx);

    // Attempt to become the sole work stealing segment. If we fail, then we enlist
    // our aid to the current work stealing segment as a helper.
    while true {
      // We are the new work stealing segment...
      if workStealer.compareExchangeStrong(-1, idx) {
        break;
      }

      // Another segment is the work stealing segment; help them out...
      // If we help steal work, it is possible we also lift an element while helping
      // Furthermore, we also return the index of the work stealing segment, as it is
      // likely to have more work.
      var currentStealer = workStealer.read();
      if currentStealer != -1 {
        segment.releaseStatus();
        var (hasElem, elem) = helpStealWork(currentStealer);
        return (currentStealer, hasElem, elem);
      }
    }

    // Initialize each slot for potential helpers. It is imperative that we initialize
    // the slots *before* resetting the wsSlotsTaken to avoid race conditions, as
    // helpers only help if they obtain a valid index.
    for slot in segment.wsSlots {
      slot.status.write(INITIALIZED);
    }
    segment.wsStartOffset = makeRandomStream(int).getNext() % numLocales;
    segment.wsSlotsTaken.write(0);
    segment.status.write(STEALING_WORK);

    // Help yourself...
    var (hasElem, elem) = helpStealWork(idx);
    // Lets others know that we no longer require help.
    segment.status.write(STEALING_MERGE);
    if logMPMCQueue then writeln("Merging work stolen on segment #", idx);

    // At this point, it is possible that other helpers are still working, and so
    // we wait for them to finish. Merge what has been stolen.
    for slot in segment.wsSlots {
      // Wait for the current helper to finish, and when they do skip any without work...
      while slot.status.read() == INITIALIZED do chpl_task_yield();
      if slot.status.read() == FINISHED_WITH_NO_WORK then continue;

      ref stolenBlock = slot.block;
      var tailBlock = seg.tailBlock;
      var blockElems = stolenBlock.nElems;
      segment.nElems.add(blockElems : uint);

      // Drain elements we've stolen and merge it into our segment.
      for elt in stolenBlock.drain(ELEMS_PER_BLOCK) {
        // If we have not lifted an element for ourself by now, do so.
        if !hasElem {
          elem = elt;
          hasElem = true;
          segment.nElems.sub(1);
          if logMPMCQueue then writeln("Lifted element ", elem, " from block stolen from segment #", i, " as segment #", idx);
          continue;
        }

        segment.addElements(elt);
      }

      delete stolenBlock;
    }

    // At this point, we have stolen from all other nodes, and we may relinquish
    // our role as primary work-stealer.
    workStealer.write(-1);
    segment.releaseStatus();
    return (idx, hasElem, elem);
  }

  proc helpStealGlobalWork(idx) : (bool, eltType) {
    ref seg = segments[idx : int];

    while true {
      var currStatus = seg.status.read();

      select currStatus {
        // They are still initializing...
        when STEALING_GLOBAL_INIT {
          // Do nothing...
        }
        // They are initialized, we can now attempt to help...
        when STEALING_GLOBAL_WORK {
          break;
        }
        // We're not stealing work (globally)...
        otherwise do return (false, _defaultOf(eltType));
      }
    }

    if logMPMCQueue then writeln("Helping steal work (globally) for Locale-segment #", here.id, "-", idx);
    var (hasElem, elem) : (bool, eltType);
    var stolenBlock = new WorkQueueSegmentBlock(eltType);

    // Obtain our working index. If it is above the maximum allowed for concurrent
    // helpers, we just return early in hope to find work elsewhere. It should be
    // noted that the work stealer only resets the counter after reinitializing the
    // segment's work stealing state. If we get preempted long enough to miss the
    // current work stealing session, then it would have to be after the previous ended
    // and the current has been initialized.
    while true {
      var ourIdx = seg.workStealingIdx.fetchAdd(1);
      if ourIdx == idx then continue;
      if ourIdx >= here.maxTaskPar {
        return (hasElem, elem);
      }
      ref wsStatus = seg.workStealingStatus[ourIdx];
      wsStatus.write(CLAIMED);

      // When we steal globally, we take enough to fill one block. We begin at a
      // randomized locale that is not our own (in order to reduce the chances of too
      // many of our helper tasks hitting the same locale over and over again, depleting it).
      var rand = makeRandomStream(int);
      var localeIdx = rand.getNext() % (targetLocales.domain.high + 1);
      var maxIterations = targetLocales.domain.high;
      var iterations = 0;
      var nElemsToSteal = ELEMS_PER_BLOCK;

      while iterations < maxIterations {
        var startSegmentIdx = rand.getNext() % targetLocales[localeIdx].maxTaskPar;

        on targetLocales[localeIdx] {
          // Localize some of the data at the beginning to avoid excessive communication
          var segmentIdx = startSegmentIdx;
          var (_hasElem, _elem) : (bool, eltType);
          _hasElem = hasElem;
          var stolenElems : ELEMS_PER_BLOCK * eltType;
          var theirQueue = workQueues[workQueues.domain.subDomain().first];

          var maxSegmentIterations = here.maxTaskPar;
          var segmentIterations = 0;
          // Search all segments...
          while segmentIterations < maxSegmentIterations {
            ref _seg = theirQueue.segments[segmentIdx];

            // Attempt to acquire...
            while !_seg.isEmpty {
              if _seg.status.read() == UNLOCKED && _seg.status.compareExchangeStrong(UNLOCKED, DEQUEUE) {
                var (__hasElem, __elem) = _seg.take;
                if !_hasElem {
                  _hasElem = __hasElem;
                  _elem = __elem;
                } else {
                  stolenElem
                }

                _seg.status.write(UNLOCKED);
                if hasElem {
                  return (hasElem, elem);
                }
              }

              // Backoff
              chpl_task_yield();
            }

            segmentIterations = segmentIterations + 1;
            segmentIdx = segmentIdx % here.maxTaskPar;
          }
        }

        localeIdx = (localeIdx + 1) % (targetLocales.domain.high + 1);
        iterations = iterations + 1;
      }

      // Do our part by stealing work...
      ref otherSeg = segments[ourIdx];
      while !otherSeg.isEmpty {
        if otherSeg.status.read() != UNLOCKED {
          chpl_task_yield();
        } else if otherSeg.status.compareExchangeStrong(UNLOCKED, VICTIM) {
          break;
        }
      }

      if otherSeg.isEmpty {
        wsStatus.write(FINISHED_WITH_NO_WORK);
        otherSeg.status.write(UNLOCKED);
        continue;
      }

      if !hasElem {
        // Steal an item...
        (hasElem, elem) = (true, otherSeg.headBlock.elems[otherSeg.headBlock.nElems]);
        otherSeg.headBlock.nElems = otherSeg.headBlock.nElems - 1;
        otherSeg.nElems.sub(1);
        if logMPMCQueue then writeln("Lifted element ", elem, " from block stolen from segment #", ourIdx, " as helper for segment #", idx);

        // Fix list if we consumed last one...
        if otherSeg.headBlock.nElems == 0 {
          var block = otherSeg.headBlock;
          otherSeg.headBlock = otherSeg.headBlock.next;
          delete block;
        }

        if otherSeg.isEmpty {
          wsStatus.write(FINISHED_WITH_NO_WORK);
          otherSeg.status.write(UNLOCKED);
          continue;
        }
      }

      // Report what has been stolen...
      var nStolen = otherSeg.headBlock.nElems : uint;
      seg.workStolen[ourIdx] = otherSeg.headBlock;
      otherSeg.nElems.sub(nStolen);
      otherSeg.headBlock = otherSeg.headBlock.next;
      wsStatus.write(FINISHED_WITH_WORK);
      otherSeg.status.write(UNLOCKED);

      if logMPMCQueue then writeln("Finished stealing ", nStolen, " amounts of work for segment #", idx);
    }

    halt("Exited helpStealWork loop...");
  }

  proc helpStealWork(idx) : (bool, eltType) {
    ref seg = segments[idx : int];

    while true {
      var currStatus = seg.status.read();

      select currStatus {
        // They are still initializing...
        when STEALING_INIT {
          // Do nothing...
        }
        // They are initialized, we can now attempt to help...
        when STEALING_WORK {
          break;
        }
        // We're not stealing work...
        otherwise do return (false, _defaultOf(eltType));
      }
    }

    if logMPMCQueue then writeln("Helping steal work for segment #", idx);
    var (hasElem, elem) : (bool, eltType);

    while true {
      // Obtain our working index. If it is above the maximum allowed for concurrent
      // helpers, we just return early in hope to find work elsewhere. It should be
      // noted that the work stealer only resets the counter after reinitializing the
      // segment's work stealing state. If we get preempted long enough to miss the
      // current work stealing session, then it would have to be after the previous ended
      // and the current has been initialized.
      var ourIdx = seg.workStealingIdx.fetchAdd(1);
      if ourIdx == idx then continue;
      if ourIdx >= here.maxTaskPar {
        return (hasElem, elem);
      }
      ref wsStatus = seg.workStealingStatus[ourIdx];
      wsStatus.write(CLAIMED);

      if logMPMCQueue then writeln("Assigned to steal work from segment #", ourIdx, " for segment #", idx);

      // Do our part by stealing work...
      ref otherSeg = segments[ourIdx];
      while !otherSeg.isEmpty {
        if otherSeg.status.read() != UNLOCKED {
          chpl_task_yield();
        } else if otherSeg.status.compareExchangeStrong(UNLOCKED, VICTIM) {
          break;
        }
      }

      if otherSeg.isEmpty {
        wsStatus.write(FINISHED_WITH_NO_WORK);
        otherSeg.status.write(UNLOCKED);
        continue;
      }

      if !hasElem {
        // Steal an item...
        (hasElem, elem) = (true, otherSeg.headBlock.elems[otherSeg.headBlock.nElems]);
        otherSeg.headBlock.nElems = otherSeg.headBlock.nElems - 1;
        otherSeg.nElems.sub(1);
        if logMPMCQueue then writeln("Lifted element ", elem, " from block stolen from segment #", ourIdx, " as helper for segment #", idx);

        // Fix list if we consumed last one...
        if otherSeg.headBlock.nElems == 0 {
          var block = otherSeg.headBlock;
          otherSeg.headBlock = otherSeg.headBlock.next;
          delete block;
        }

        if otherSeg.isEmpty {
          wsStatus.write(FINISHED_WITH_NO_WORK);
          otherSeg.status.write(UNLOCKED);
          continue;
        }
      }

      // Report what has been stolen...
      var nStolen = otherSeg.headBlock.nElems : uint;
      seg.workStolen[ourIdx] = otherSeg.headBlock;
      otherSeg.nElems.sub(nStolen);
      otherSeg.headBlock = otherSeg.headBlock.next;
      wsStatus.write(FINISHED_WITH_WORK);
      otherSeg.status.write(UNLOCKED);

      if logMPMCQueue then writeln("Finished stealing ", nStolen, " amounts of work for segment #", idx);
    }

    halt("Exited helpStealWork loop...");
  }
}

class DistributedWorkQueue : Queue {
  var targetLocales : [] locale = Locales;
  var perLocaleDomain = domain(1) dmapped Block(boundingBox=targetLocales.domain, targetLocales=targetLocales);
  var workQueues : [perLocaleDomain] WorkQueue(eltType);

  proc DistributedWorkQueue(type eltType) {
    forall loc in Locales {
      on loc {
        workQueues[workQueues.domain.subDomain().first] = new WorkQueue(eltType);
      }
    }
  }
}

proc proof_of_correctness() {
  var nElementsForMPMC = 1000000;
  writeln("Starting MPMCQueue Proof of Correctness (Work Depletion) Test ~ nElementsForMPMC: ", nElementsForMPMC);

  var queue = new WorkQueue(int);
  var counter : atomic int;
  var expected = (nElementsForMPMC * (nElementsForMPMC + 1)) / 2;

  forall i in 1 .. nElementsForMPMC {
    queue.enqueue(i);
  }

  coforall loc in Locales do on loc {
    var localeLocalCounter : atomic int;
    coforall tid in 0 .. #here.maxTaskPar {
      var taskLocalCounter : int;
      while true {
        var (present, elem) = queue.dequeue();
        if !present {
          break;
        }

        taskLocalCounter = taskLocalCounter + elem;
      }
      localeLocalCounter.add(taskLocalCounter);
    }
    writeln(here, ": Counter=", localeLocalCounter.read());
    counter.add(localeLocalCounter.read());
  }

  if counter.read() == expected {
    writeln("PASSED TEST!");
  } else {
    writeln("FAILED TEST! Expected: ", expected, ", but received: ", counter.read());
    writeln(queue);
  }
}

proc benchmark() {
  var nJitter = 0;
  var nComputations = 0;
  var nElements = 1000000;
  var nTrials = 8;
  var enqueueTrialTime : [1 .. nTrials] real;
  var dequeueTrialTime : [1 .. nTrials] real;

  // Obtain average time for enqueue followed by dequeued...
  for i in 1 .. nTrials {
    var queue = new WorkQueue(int);
    // We only use one or the other, but we must declare both.
    // TODO: Have both implement some parent 'Queue' class/interface?
    var timer = new Timer();
    timer.start();

    coforall loc in Locales do on loc {
      var iterations = nElements;

      coforall tid in 0 .. #here.maxTaskPar {
        var x : atomic int;
        var randStr = makeRandomStream(int);
        for j in 1 .. iterations / here.maxTaskPar {
          queue.enqueue(j);
          var nComps = nComputations + (if nJitter then (randStr.getNext() % nJitter) else 0);
          for i in 1 .. nComps {
            // Hopefully compiler doesn't throw away?
            x.write(sin(i) : int);
          }
        }
      }
    }

    timer.stop();
    enqueueTrialTime[i] = nElements / timer.elapsed();
    writeln(i, "/", nTrials, ": ", (+ reduce enqueueTrialTime) / i);
    timer.clear();
    timer.start();

    coforall loc in Locales do on loc {
      var iterations = nElements;

      coforall tid in 0 .. #here.maxTaskPar {
        var x : atomic int;
        var randStr = makeRandomStream(int);
        for j in 1 .. iterations / here.maxTaskPar {
          var retval = queue.dequeue();
          var nComps = nComputations + (if nJitter then (randStr.getNext() % nJitter) else 0);
          for i in 1 .. nComps {
            // Hopefully compiler doesn't throw away?
            x.write(sin(i) : int);
          }
        }
      }
    }

    timer.stop();
    dequeueTrialTime[i] = nElements / timer.elapsed();
    writeln(i, "/", nTrials, ": ", (+ reduce dequeueTrialTime) / i);

    delete queue;
  }
}

proc main() {
  benchmark();
}


Haven't compiled it yet, but this is the proof-of-concept distributed queue!

use CyclicDist;
use local_lock_free_queue;

class Distributed_FIFO {
  type eltType;

  // Our position counter: Upper 32 bits = Head, Lower 32 bits = Tail
  var position : atomic uint(64);

  // Our domain: We use a cyclic distribution
  var domainMapping = {1 .. numLocales};
  var cyclicDomain : domainMapping dmapped Cyclic(startIdx=1);
  var localQueues : [cyclicDomain] queue(eltType);

  proc Distributed_FIFO(type eltType) {

  }

  proc enqueue(elem : eltType) {
    // The index we are going to be working on...
    var idx : uint(32);

    // Position is queried on the locale that owns the queue
    on position {
      var pos : uint(64);
      var head, tail : uint(32);

      // Obtain our desired position
      while (true) {
        pos = position.read();
        head = pos & 0xFFFFFFFF00000000
        tail = pos & 0x00000000FFFFFFFF

        // Some platforms yield higher performance due to reduced memory contention
        // under concurrent writes to the same cache line. We update the position
        // atomically.
        if (pos.compareExchangeWeak(pos, (head << 32) | (tail + 1))) {
          break;
        }
      }

      // Translates to a single network call to write to the orginal locale's index
      idx = tail % numLocales;

      writeln("Head: ", head, "; Tail: ", tail);
    }

    writeln(here, ": Assigned index: ", idx);

    // Now we add our item
    ref queue = localQueues[idx];
    on queue queue.enqueue(elem);
  }

  proc dequeue() : (bool, eltType) {
    // The index we are going to be working on...
    var idx : uint(32);
    var hasElem : bool = true;

    // Position is queried on the locale that owns the queue
    on position {
      var pos : uint(64);
      var head, tail : uint(32);

      // Obtain our desired position
      while (true) {
        pos = position.read();
        head = pos & 0xFFFFFFFF00000000
        tail = pos & 0x00000000FFFFFFFF

        if (head == tail) {
          writeln("Queue is empty... Pos: ", pos);
          hasElem = false;
          break;
        }

        // Some platforms yield higher performance due to reduced memory contention
        // under concurrent writes to the same cache line. We update the position
        // atomically.
        if (pos.compareExchangeWeak(pos, (head << 32) | (tail + 1))) {
          break;
        }
      }

      if hasElem {
        // Translates to a single network call to write to the orginal locale's index
        idx = tail % numLocales;

        writeln("Head: ", head, "; Tail: ", tail);
      }
    }

    if !hasElem {
      var dummy : eltType;
      return (false, dummy);
    }

    writeln(here, ": Assigned index: ", idx);

    // Now we get our item from the queue
    // Note that at the index given, its possible that an enqueueing task has not
    // finished yet, but we know there *should* be at least something for us, so we can
    // spin until it has what we want.
    ref queue = localQueues[idx];
    var elem : eltType;
    on queue {
      var retval : (bool, eltType);
      while retval[1] {
        retval = queue.enqueue(elem);
      }

      // We have our value... this also translates into a network call (PUT)
      elem = retval[2];
    }

    return (true, elem);
  }
}

error: pathspec 'src/Distributed_FIFO.chpl' did not match any file(s) known to git


Chapel version by Peter Kjellström

Chapel version by Peter Kjellström

https://twitter.com/nsccap/status/1233424385471664128

/*
   Example provided by Peter Kjellström (@nsccap on Twitter):

   https://twitter.com/nsccap/status/1233424385471664128
*/

config const N = 1000000;
var pi = sqrt(6.0 * (+ reduce [ i in 1..N by - 1 ] 1.0/(i*i) ));
writeln(pi);

error: pathspec 'nsccap_pi.chpl' did not match any file(s) known to git


11 - largest product in a grid

11 - largest product in a grid

error: pathspec '011.chpl' did not match any file(s) known to git


2 - even fibonacci numbers

2 - even fibonacci numbers

/*
 * Even Fibonacci number
 */

config const max = 4000000;
var sum = 0;

// Print the Fibonacci number less than and equal to 4 million.
iter fibonacci() {
  var current = 1,
    next = 1;
  while current <= max {
    yield current;
    current += next;
    current <=> next;
  }
}

for i in fibonacci() {
  if i % 2 == 0 {
    sum += i;
  }
}

writeln(sum);

error: pathspec '002.chpl' did not match any file(s) known to git


The module AdvancedIters.chpl contains the dynamic, guided and adaptive
work-stealing iterators, supporting non-strided and strided ranges as well as
parallel zippered iterations.

- dynamic(range, chunkSize, nTasks) represents the equivalent to OpenMP
dynamic scheduling approach: given an input range, each task (of nTasks)
splits chunks of fixed size (chunkSize) when it finishes its previous chunk.

- guided(range, nTasks) represents the OpenMP guided policy: given an input
range, each task splits chunks of variable size when it finishes its previous
chunk. The size of each chunk is the number of unassigned iterations divided
by the number of tasks (nTasks). These sizes decrease approximately
exponentially to 1. The splitting strategy is therefore adaptive.

- adaptive(range, nTasks) represents a naive adative work-stealing leader
iterator: initially it distributes the range to split among the tasks
(nTasks). Then it performs adaptive splitting locally on each task. When a
task exhausts its local range, it steals and splits from the range of another
task (victim). The Splitting method on the local range and on the victim range
is Binary (the size of each chunk is the number of unassigned
iterations divided by 2). There are three stealing strategies that can be
selected at compile time with config param methodStealing:

* methodStealing=0: Each task without work tries to steal from its neighbor
     range until it exhausts that range. Then it continues with the next
     neighbor range, and so on until there is not more work. It is the default
     policy.

* methodStealing=1: Each task without work tries to steal once from its
       neighbor range, next from the following neighbor range and so on in a
       round-robin way until there is not more work.

* methodStealing=2: As method 0, but now the splitting in the victi range is
       performed from its tail.

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@19030 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

config param verbose:bool=false;

//************************* Dynamic iterator

// Serial iterator

iter dynamic(c:range(?), chunkSize:int, numTasks:int) {

  compilerWarning("Working with serial dynamic Iterator");
  if verbose then
    writeln("Serial dynamic Iterator. Working with range ", c);

  for i in c do yield i;
}

// Parallel iterator

// Leader iterator

iter dynamic(param tag:iterator, c:range(?), chunkSize:int, numTasks:int= if dataParTasksPerLocale==0 then here.numCores else dataParTasksPerLocale)
where tag == iterator.leader
{
  type rType=c.type;
  // Check the size and do it serial if not enough work
  if c.length == 0 then halt("The range is empty");
  if c.length/chunkSize < numTasks then {
    writeln("Dynamic Iterator: serial execution because there is not enough work");
    const totalRange:rType= densify(c,c);
    //yield tuple(c.translate(-c.low)) ;
    yield tuple(totalRange);

  }
  else {
    var remain:rType=c;
    var moreWork=true;
    var splitLock$:sync bool=true;   // sync variables to control the splitting

    coforall tid in 0..#numTasks do {
      while moreWork do {
	  // There is local work in remain
	  splitLock$; // Perform the splitting in a critical section
	  const current:rType=splitChunk(remain, chunkSize, moreWork);
	  splitLock$=true;
	  if current.length !=0 then {
	    const zeroBasedIters:rType= densify(current,c);
	    //const zeroBasedIters= current.translate(-c.low);
	    if verbose then
	      writeln("Parallel dynamic Iterator. Working at tid ", tid, " with range ", current, " yielded as ", zeroBasedIters);

	    yield tuple(zeroBasedIters);}
	}
    }
  }
}

// Follower
iter dynamic(param tag:iterator, c:range(?), chunkSize:int, numTasks:int, follower)
where tag == iterator.follower
{
  type rType=c.type;
  const current:rType=unDensify(follower(1),c);
  //follower(1).translate(c.low);
  if (verbose) then
    writeln("Follower received range ", follower, " ; shifting to ", current);
  for i in current do {
    yield i;
  }
}

//************************* Guided iterator
//Serial iterator
iter guided(c:range(?), numTasks:int) {

  compilerWarning("Working with serial guided Iterator");
  if verbose then
    writeln("Serial guided Iterator. Working with range ", c);

  for i in c do yield i;
}

// Parallel iterator

// Leader iterator
iter guided(param tag:iterator, c:range(?), numTasks:int = if dataParTasksPerLocale==0 then here.numCores else dataParTasksPerLocale)
where tag == iterator.leader
{
  type rType=c.type;
  // Check the size and do it serial if not enough work
  if c.length == 0 then halt("The range is empty");
  if c.length/numTasks < numTasks then {
    writeln("Guided Iterator: serial execution because there is not enoguh work");
    const totalRange:rType= densify(c,c);
    yield tuple(totalRange);
  }

  else {
    var remain:rType=c;
    var undone=true;
    const factor=numTasks;
    var splitLock$:sync bool=true;  // synnc variables to control the splitting

    coforall tid in 0..#numTasks do {
      while undone do {
	  // There is local work in remain(tid)
	  splitLock$; // Perform the splitting in a critical section
	  const current:rType=adaptSplit(remain, factor, undone);
	  splitLock$=true;
	  if current.length !=0 then {
	    const zeroBasedIters:rType = densify(current,c);
	    if verbose then
	      writeln("Parallel guided Iterator. Working at tid ", tid, " with range ", current, " yielded as ", zeroBasedIters);

	    yield tuple(zeroBasedIters);}
	}
    }
  }
}

// Follower
iter guided(param tag:iterator, c:range(?), numTasks:int, follower)
where tag == iterator.follower

{
  type rType=c.type;
  const current:rType=unDensify(follower(1),c);
  if verbose then
    writeln("Follower received range ", follower, " ; shifting to ", current);
  for i in current do {
    yield i;
  }
}

//************************* Adaptive work-stealing iterator
//Serial iterator
iter adaptive(c:range(?), numTasks:int) {

  compilerWarning("Working with serial adaptive Iterator");
  if verbose then
    writeln("Serial adaptive work-stealing Iterator. Working with range ", c);

  for i in c do yield i;

}

// Parallel iterator

config param methodStealing:int=0;

// Leader iterator
iter adaptive(param tag:iterator, c:range(?), numTasks:int = if dataParTasksPerLocale==0 then here.numCores else dataParTasksPerLocale)
where tag == iterator.leader

{
  if (methodStealing > 2 || methodStealing < 0) then
    halt("methodStealing value must be between 0 and 2");

  type rType=c.type;
  // Check the size and do it serial if not enough work
  if c.length == 0 then halt("The range is empty");
  if c.length < numTasks then {
    writeln("Adaptive work-stealing Iterator: serial execution because there is not enough work");
    const totalRange:rType = densify(c,c);
    //    compilerWarning("First yield type:", typeToString(totalRange.type));
    yield tuple(totalRange);

  }
  else {
    const r:rType=densify(c,c);
    const SpaceThreads:domain(1)=0..#numTasks;
    var localWork:[SpaceThreads] rType; // The remaining range to split on each Thread
    var moreLocalWork:[SpaceThreads] bool=true; // bool var to signal there is still work on each local range
    var splitLock$:[SpaceThreads] sync bool=true; // sync var to control the splitting on each Thread
    const factor:int=2;

    const factorSteal:int=2;
    var moreWork:bool=true; // A global var to control the termination
    var nVisitedVictims: int=0;

    // Variables to put a barrier to ensure the initial range is computed on each Thread
    var barrierCount$:sync int=0;
    var wait$:single bool;

    // Start the parallel work
    coforall tid in 0..#numTasks do {

      // Step 1: Initial range per Thread/Task

      localWork[tid]=splitRange(r,tid);  // Initial Local range in localWork[tid]
      if verbose then
	writeln("Parallel adaptive work-stealing Iterator. Working at tid ", tid, " with initial range ", localWork[tid]);

      // A barrier waiting for each thread to finish the initial assignment
      barrier;

      // Step 2: While there is work at tid, do splitting

      while moreLocalWork[tid] do {
	  // There is local work
	  const zeroBasedIters:rType=splitWork(tid); // The current range we get after splitting locally
	  if zeroBasedIters.length !=0 then {
	    if verbose then
	      writeln("Parallel adaptive Iterator. Working locally at tid ", tid, " with range yielded as ", zeroBasedIters);
	    yield tuple(zeroBasedIters);}
	}

      // Step3: Task tid finished its work, so it will try to steal from a neighbor

      var victim=(tid+1) % numTasks;
      var stealFailed:bool=false;

      while moreWork do {
	  if verbose then
	    writeln("Entering at Stealing phase in tid ", tid," with victim ", victim, " using method of Stealing ", methodStealing);

	  // Perform the spliting at the victim remaining range

	  select methodStealing {
	    when 0 do { // Steals from the victim until it exhausts its range

	      while moreLocalWork[victim] do {
		  // There is work in victim
		  const zeroBasedIters2:rType=splitWork(victim); // The current range we get
		                                        //after splitting from a victim range
		  if zeroBasedIters2.length !=0 then {
		    if verbose then
		      writeln("Range stealed at victim ", victim," yielded as ", zeroBasedIters2," by tid ", tid);
		    yield tuple(zeroBasedIters2);}
		}
	    }

	    when 1 do { // Steals once from victim. If it has sucess on the stealing,
	                // it chooses another victim and tries to steal again.

	      if moreLocalWork[victim] then {
		// There is work in victim
		const zeroBasedIters2:rType=splitWork(victim); // The current range we get
		                                   //after splitting from a victim range
		if zeroBasedIters2.length !=0 then {
		  if verbose then
		    writeln("Range stealed at victim ", victim," yielded as ", zeroBasedIters2," by tid ", tid);
		  yield tuple(zeroBasedIters2);}
	      }
	      else {
		stealFailed=true;
	      }
	    }
	    otherwise { // Tt steals from the victim, but now splitting
	                //from the tail of the victim's range

	      while moreLocalWork[victim] do {
		  // There is work in victim
		  const zeroBasedIters2:rType=splitWork(victim, true); // The current range we get
		                        // after splitting from a victim range:
		                       // now we set splitTatil to true (second parameter in the call)
		  if zeroBasedIters2.length !=0 then {
		    if verbose then
		      writeln("Range stealed at victim ", victim," yielded as ", zeroBasedIters2," by tid ", tid);
		    yield tuple(zeroBasedIters2);}
		}
	    }
	    }
	  // If here, then it can have failed the stealing intent at the victim (method 1),
	  // or we have exhausted the victim range (methods 0, 2)

	  if (methodStealing==0 || methodStealing==2 || methodStealing==1 && stealFailed) then {
	    nVisitedVictims += 1; // Signal that there is no more work in victim
	    if verbose then
	      writeln("Failed Stealing intent at tid ", tid," with victim ", victim, " and total no. of visited victims ", nVisitedVictims);
	    stealFailed=false;
	  }
	  // Check if there is no more work
	  if nVisitedVictims >= numTasks-1 then
	    moreWork=false; // Signal that there is no more work in any victim
	  else {  // There can be still work in other victim
	    victim=(victim+1) % numTasks; // New victim to steal
	    if methodStealing==1 then
	      if victim==tid then victim=(victim+1) % numTasks;
	  }
	}
    }

    proc splitRange(c:range(?),tid:int)
  {
    const size= c.length/numTasks;
    var rLocal:rType=(c+tid*size)#size;
    if tid==numTasks-1 then
      rLocal=c#(size*(numTasks-1)-c.length);
    return rLocal;
  }

  proc barrier
  {
    var tmp=barrierCount$;
    barrierCount$=tmp+1;
    if tmp+1==numTasks then wait$=true;
    wait$;
  }

    proc splitWork(taskid:int, splitTail:bool=false)
  {
    //var current:rType;
    if splitTail then { // Splitting from the tail of the remaining range (method 2)
      splitLock$[taskid]; // Perform the splitting in a critical section
      const current:rType=adaptSplit(localWork[taskid], factorSteal, moreLocalWork[taskid], splitTail);
      splitLock$[taskid]=true;
      const zeroBasedCurrent3:rType=current;
      return zeroBasedCurrent3;
    }
    else { // Splitting from the head of the remaining range (methods 0 and 1)
      splitLock$[taskid]; // Perform the splitting in a critical section
      const current:rType=adaptSplit(localWork[taskid], factorSteal, moreLocalWork[taskid]);
      splitLock$[taskid]=true;
      //const zeroBasedCurrent = densify(current,c);
      const zeroBasedCurrent:rType=current;
      return zeroBasedCurrent;
      }

  }
  }
}

// Follower
iter adaptive(param tag:iterator, c:range(?), numTasks:int, follower)
where tag == iterator.follower
{
  type rType=c.type;
  var current:rType=unDensify(follower(1),c);
  //follower(1).translate(c.low);
  if (verbose) then
    writeln("Follower received range ", follower, " ; shifting to ", current);
  for i in current do {
    yield i;
  }
}

//************************* Helper functions
proc splitChunk(inout rangeToSplit:range(?), chunkSize:int, inout itLeft:bool)
{
  type rType=rangeToSplit.type;
  type lenType=rangeToSplit.length.type;
  var totLen, size:lenType;

  totLen=rangeToSplit.length;
  size=chunkSize;
  if totLen <= size then
    {
    size=totLen;
    itLeft=false;
    }
  const firstRange:rType=rangeToSplit#size;
  rangeToSplit=rangeToSplit#(size-totLen);
  return firstRange;
}

proc adaptSplit(inout rangeToSplit:range(?), splitFactor:int, inout itLeft:bool, splitTail:bool=false)
{
  type rType=rangeToSplit.type;
  type lenType=rangeToSplit.length.type;
  var totLen, size:lenType;
  var direction: int;
  const profThreshold=1;

  totLen=rangeToSplit.length;
  if totLen > profThreshold then
    size=max(totLen/splitFactor, profThreshold);
  else {
    size=totLen;
    itLeft=false;}
  if splitTail then
    direction=-1;
  else direction=1;
  const firstRange:rType=rangeToSplit#(direction*size);
  rangeToSplit=rangeToSplit#(direction*(size-totLen));
  return firstRange;
}

error: pathspec 'modules/standard/AdvancedIters.chpl' did not match any file(s) known to git


add localdisk experiment

add localdisk experiment

use Logging, IO, Partitions;

proc main() {
  initPartitions();

  var t: Timer;
  t.start();

  for partition in Partitions {
    on partition {
      var infile = open("words" + here.id + ".txt", iomode.cwr);
      var writer = infile.writer();
      writer.writeln("I'm here at ", here.id);
    }
  }

  t.stop();
  timing("complete in ",t.elapsed(TimeUnits.microseconds), " microseconds");
}

error: pathspec 'chapel/localdisk.chpl' did not match any file(s) known to git


Added my 2nd go at Jacobi.

Added my 2nd go at Jacobi.

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@6260 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

config const n = 4;
config const epsilon = 1.0e-5;

const stencil = (/(-1, -1), (-1, 1), (1, -1), (1, 1)/);

const D = [1..n, 1..n] distributed(block);
const BigD = D.expand(1);
const SouthD = D.exterior(1, 0);

var A, Temp : [BigD] float;

Temp(SouthD) = 1.0;

do {
  [i in D] Temp(i) = A(i + stencil) / + / 4.0;
  if abs(A - Temp) / max > epsilon {
    A = Temp;
    break;
  }
  [i in D] A(i) = Temp(i + stencil) / + / 4.0;
} while abs(A - Temp) / max > epsilon;

writeln(A);

error: pathspec 'test/studies/jacobi/deitz/jacobi2.chpl' did not match any file(s) known to git