Refactored the 'stealWork' portion, only helpStealWork portion left

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();
}

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;
config param logMPMCQueue = false;

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

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 MAX_ELEMS * 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
  */

  // We have the lock but it is empty, try to set ourself as work stealer...
  // As our index can change based on whether or not we become the work stealer
  // or help someone else become the work stealer, we return our acquired index.
  proc stealWork(idx) : (int, bool, eltType) {
    ref seg = segments[idx];

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

    // If we set ourself as work stealer, we may proceed, otherwise if someone else
    // is set, we help them out instead...
    while true {
      if workStealer.compareExchangeStrong(-1, idx) then break;

      // Someone else is set...
      var currStealer = workStealer.read();
      if currStealer != -1 {
        seg.status.write(UNLOCKED);
        var (hasElem, elem) = helpStealWork(currStealer);
        return (currStealer, hasElem, elem);
      }
    }

    // We are the current work stealer, so our responsibility is to initialize
    // all work statuses and signal that we may begin our work stealing
    for status in seg.workStealingStatus do status.write(INITIALIZED);
    seg.workStealingIdx.write(0);
    seg.status.write(STEALING_LOCAL_WORK);

    // Help steal work with everyone else
    var (hasElem, elem) = helpStealWork(idx);
    seg.status.write(STEALING_MERGE);
    if logMPMCQueue then writeln("Merging work stolen on segment #", idx);

    // We are finished, all work stolen (if any) are now available...
    // As work stealer, our next job is merge it all together...
    for i in 0..#here.maxTaskPar {
      // Spin until they are finished... if they finish but have no work, skip them
      while seg.workStealingStatus[i].read() <= CLAIMED do chpl_task_yield();
      if i == idx || seg.workStealingStatus[i].read() != FINISHED_WITH_WORK then continue;

      // TODO: Find a way copy data from one tuple to another that is faster
      ref stolenBlock = seg.workStolen[i];
      var tailBlock = seg.tailBlock;

      var blockElems = stolenBlock.nElems;
      seg.nElems.add(blockElems : uint);
      while stolenBlock.nElems > 0 {
        // If we haven't stolen an element by now, do so
        if !hasElem {
          elem = stolenBlock.elems[stolenBlock.nElems];
          hasElem = true;
          stolenBlock.nElems = stolenBlock.nElems - 1;
          seg.nElems.sub(1);
          if logMPMCQueue then writeln("Lifted element ", elem, " from block stolen from segment #", i, " as segment #", idx);
          continue;
        }

        if tailBlock.nElems == ELEMS_PER_BLOCK {
          tailBlock.next = new WorkQueueSegmentBlock(eltType);
          seg.tailBlock = tailBlock.next;
          tailBlock = tailBlock.next;
        }

        tailBlock.nElems = tailBlock.nElems + 1;
        tailBlock.elems[tailBlock.nElems] = stolenBlock.elems[stolenBlock.nElems];
        stolenBlock.nElems = stolenBlock.nElems - 1;
      }

      delete stolenBlock;
    }

    // Check to see if we have actually stolen any work. (I.E: We haven't lifted any)
    if !hasElem {
      // Since we can verify that the queue is empty, we need to look elsewhere.

    }

    // No longer work stealer
    workStealer.write(-1);
    seg.status.write(UNLOCKED);

    // At this point, we have successfully stolen work. Furthermore, we have
    // stolen at least one element for ourself.
    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();
}


Fixed a typo.

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

/*-------------------------------------------------
Code 5 Task pool of integral blocks - Chapel
-------------------------------------------------*/
class taskpool {
  const poolSize;
  var taskarr : [0..poolSize-1] sync blockIndices; // describe blockIndices?
  var head, tail : sync int = 0;
  def add(blk: blockIndices) {
    const pos = tail;
    tail = (pos+1)%poolSize;
    taskarr(pos) = blk;
  }
  def remove() {
    const pos = head;
    head = (pos+1)%poolSize;
    return taskarr(pos);
  }
}

/*-------------------------------------------------
Code 6 Top-level driver for task pool - Chapel
-------------------------------------------------*/
config const numConsumers = numLocales * Locales(0).numCores,
             poolSize = numConsumers;
const t = taskpool(poolSize);
cobegin {
  coforall loc in 1..numConsumers do
    consumer();
  producer();
}

/*-------------------------------------------------
Code 7 Producer of integral blocks - Chapel
-------------------------------------------------*/
def producer() {
  forall blk in genBlocks() do
    t.add(blk);
}

/*-------------------------------------------------
Code 8 Fock index space iterator - Chapel
-------------------------------------------------*/
def genBlocks() {
  forall (iat, jat, kat) in [1..natom, 1..iat, 1..iat] {
    const lattop = if (kat==iat) then jat
                                 else kat;
    forall lat in 1..lattop do
      yield blockIndices(...);  // what is "..."?
  }
  forall loc in 1..numConsumers do
    yield blockIndices(0,0,0,0,0,0,0,0); // should/could this be four ranges?
}

/*-------------------------------------------------
Code 9 Consumer of integral blocks - Chapel
-------------------------------------------------*/
def consumer() {
  var blk = t.remove();
  while (blk.ilo != 0) {
    const copyofblk = blk;
    cobegin {
      buildjk_atom4(copyofblk);
      blk = t.remove();
    }
  }
}

/*-------------------------------------------------
Code 10 Symmetrization of J and K - Chapel
-------------------------------------------------*/

cobegin {
  [(i,j) in D] jmat2T(i,j) = jmat2(j,i);
  [(i,j) in D] kmat2T(i,j) = kmat2(j,i);
}
jmat2 = (jmat2+jmat2T)*2;
kmat2 += kmat2T;

/*-------------------------------------------------
Code 5 Task pool of integral blocks - Chapel
-------------------------------------------------*/
class taskpool {
  const poolSize;
  var taskarr : [0..poolSize-1] sync blockIndices; // describe blockIndices?
  var head, tail : sync int = 0;
  def add(blk: blockIndices) {
    const pos = tail;
    tail = (pos+1)%poolSize;
    taskarr(pos) = blk;
  }
  def remove() {
    const pos = head;
    head = (pos+1)%poolSize;
    return taskarr(pos);
  }
}

/*-------------------------------------------------
Code 6 Top-level driver for task pool - Chapel
-------------------------------------------------*/
config const numConsumers = numLocales * Locale(0).numCores,
             poolSize = numConsumers;
const t = taskpool(poolSize);
cobegin {
  coforall loc in 1..numConsumers do
    consumer();
  producer();
}

/*-------------------------------------------------
Code 7 Producer of integral blocks - Chapel
-------------------------------------------------*/
def producer() {
  forall blk in genBlocks() do
    t.add(blk);
}

/*-------------------------------------------------
Code 8 Fock index space iterator - Chapel
-------------------------------------------------*/
def genBlocks() {
  forall (iat, jat, kat) in [1..natom, 1..iat, 1..iat] {
    const lattop = if (kat==iat) then jat
                                 else kat;
    forall lat in 1..lattop do
      yield blockIndices(...);  // what is "..."?
  }
  forall loc in 1..numConsumers do
    yield blockIndices(0,0,0,0,0,0,0,0); // should/could this be four ranges?
}

/*-------------------------------------------------
Code 9 Consumer of integral blocks - Chapel
-------------------------------------------------*/
def consumer() {
  var blk = t.remove();
  while (blk.ilo != 0) {
    const copyofblk = blk;
    cobegin {
      buildjk_atom4(copyofblk);
      blk = t.remove();
    }
  }
}

/*-------------------------------------------------
Code 10 Symmetrization of J and K - Chapel
-------------------------------------------------*/

cobegin {
  [(i,j) in D] jmat2T(i,j) = jmat2(j,i);
  [(i,j) in D] kmat2T(i,j) = kmat2(j,i);
}
jmat2 = (jmat2+jmat2T)*2;
kmat2 += kmat2T;


Fixed dfs_count function so that it returns a value on all paths.

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

/*
**  UTS -- The Unbalanced Tree Search
**
**  James Dinan <dinan@cray.com>
**  July, 2007
*/

use Math;
use Time;
use sha1_rng;

/**** UTS TYPES ****/
enum NodeDistrib { Binomial, Geometric, Hybrid, Constant };
enum GeoDistrib  { GeoFixed, GeoLinear, GeoPoly, GeoCyclic };

/**** COMPILE PARAMS ****/
config param debug:    bool = false;
config param parallel: bool = true;
param uts_version:   string = "2.1";

/**** UTS CONFIG VARS ****/
config const SEED:        uint = 16;    // Seed for the hash of the root node
config const MAX_CHILDREN: int = 10;    // Max. Children a node may have
config const B_0:          int = 4;     // Branching Factor
config const MAX_DEPTH:    int = 6;     // Depth limit
config const nonLeafProb: real = 0.10;  // Probability of a non-leaf (binomial)
config const nonLeafBF:    int = 4;     // Non-leaf branching factor (binomial)
config const shiftDepth:  real = 0.5;   // Depth at which hybrid trees go from Geo=>Bin

config const distrib: NodeDistrib = Geometric;
config const geoDist: GeoDistrib  = GeoFixed;
config const testMode             = false;

config const MAX_THREADS:  int = 20;
config const MIN_THREADS:  int = MAX_THREADS/4;

// Global thread counter
var thread_cnt: sync int = 0;
var threads_spawned: int = 0; // "Locked" by thread_cnt

/**** UTS TreeNode Class ****/
class TreeNode {
  var depth:   int;
  var hash:    [1..20] uint(8);

  // By default, children will be empty since it has range [1..0]
  var nChildren: int = 0;
  var childDom = [1..nChildren];
  var children:  [childDom] TreeNode;

  // Generate this node's children
  def genChildren(): int {
    select distrib {
      when Geometric do
        nChildren = numGeoChildren(geoDist);
      when Binomial do
        nChildren = numBinChildren();
      when Hybrid do
        nChildren = numHybridChildren();
      when Constant do
        nChildren = numConstChildren();
    }

    if debug then writeln("Constructing ", nChildren, " children: ", childDom);

    childDom  = [1..nChildren];

    forall i in childDom {
      if debug then writeln("  + (", depth, ", ", i, ")");
      children[i]    = TreeNode(depth+1);
      rng_spawn(hash[1], children[i].hash[1], i-1);
    }

    return nChildren;
  }

  // Constant Distribution: Find the number of children
  //  Note: This is a _balanced_ tree.
  def numConstChildren(): int {
    return if depth < MAX_DEPTH then B_0 else 0;
  }

  // Hybrid trees are geometric at the top of the tree and binomial
  // below a certain depth.
  def numHybridChildren(): int {
    if (depth < shiftDepth * MAX_DEPTH) then
      return numGeoChildren(geoDist);

    else
      return numBinChildren();
  }

  // Binomial Distribution: Find the number of children
  //  Note: distribution is identical everywhere below root
  def numBinChildren(): int {
    if (depth == 0) then
      return B_0;

    else
      return if to_prob(rng_rand(hash[1])) < nonLeafProb then nonLeafBF else 0;
  }

  // Geometric Distribution: Find the number of children
  def numGeoChildren(shape: GeoDistrib): int {
    var b_i: real = B_0;

    // use shape function to compute target b_i
    if (depth > 0) {
      select shape {

        // Expected size is polynomial in depth
        when GeoPoly do
          b_i = B_0 * (depth ** (-log(B_0:real)/log(MAX_DEPTH:real)));

        // Expected size is cyclic
        when GeoCyclic {
          if (depth > 5 * MAX_DEPTH) then
            b_i = 0.0;
          else
            b_i = B_0 ** sin(2.0*3.141592653589793*depth/MAX_DEPTH);
        }

        // Expected size is the same at all nodes up to max depth
        when GeoFixed do
          b_i = if (depth < MAX_DEPTH) then B_0 else 0;

        // Expected size decreases linearly in b_i
        when GeoLinear do
          b_i =  B_0 * (1.0 - depth:real / MAX_DEPTH:real);

      }
    }

    // given target b_i, find prob p so expected value of
    // geometric distribution is b_i.
    var p: real = 1.0 / (1.0 + b_i);

    // get uniform random number on [0,1)
    var u: real = to_prob(rng_rand(hash[1]));

    // max number of children at this cumulative probability
    // (from inverse geometric cumulative density function)
    return floor(log(1 - u) / log(1 - p)):int;
  }
}

/*
** Print out search parameters
*/
def uts_showSearchParams() {
  writeln("UTS v", uts_version, " - Unbalanced Tree Search (",
      if parallel then "Parallel Chapel" else "Sequential Chapel", ")");

  writeln("Tree type: ", distrib);
  writeln("Tree shape parameters:");
  writeln("  root branching factor b_0 = ", B_0, ", root seed = ", SEED);

  if (distrib == Geometric || distrib == Hybrid) then
    writeln("  Geo. params: Max_Depth = ", MAX_DEPTH, ", shape function = ", geoDist);

  if (distrib == Binomial || distrib == Hybrid) then
    writeln("  Bin. params: nonLeafProb=", nonLeafProb, " nonLeafBF=", nonLeafBF);

  if (distrib == Hybrid) then
    writeln("  Hybrid params: shiftDepth=", shiftDepth);

  writeln("Random number generator: ", rng_getName());

  if (parallel) {
    writeln("Parallel execution parameters:");
    writeln("  MAX_THREADS = ", MAX_THREADS, ", MIN_THREADS = ", MIN_THREADS);
  }

  writeln();
}

/*
** Request n threads.  Get back the number of threads granted.
*/
def requestThreads(n:int): int {
  if (parallel) {
    // Trade some imbalance here for blocking overhead
    if (readXX(thread_cnt) < MIN_THREADS) {
      var thread_cnt_l = thread_cnt;

      // Try to get a ticket to run in parallel
      if (thread_cnt_l < MAX_THREADS) {
        var threads_avail   = MAX_THREADS-thread_cnt_l;
        var threads_granted = if (n <= threads_avail) then n else threads_avail;

        thread_cnt = thread_cnt_l + threads_granted;
        threads_spawned += threads_granted;
        return threads_granted;
      } else {
        thread_cnt = thread_cnt_l;
      }
    }
  }

  return 0;
}

/*
**  Parallel DFS
*/
def dfs_count(n: TreeNode, wasParallel: bool = false):int {
  if n != nil {
    if (parallel) {
      var threads_granted = requestThreads(n.nChildren);
      var count: sync int = 0;

      coforall i in [1..threads_granted] {
        count += dfs_count(n.children[i], true);
      }

      for i in [threads_granted+1..n.nChildren] {
        count += dfs_count(n.children[i], false);
      }

      if (wasParallel) then thread_cnt -= 1;
      return count+1;

    } else {
      var count: int;

      for i in n.childDom {
        count += dfs_count(n.children[i], false);
      }
      return count+1;
    }
  } else {
    return 0;
  }
}

/*
**  Parallel Tree Creation
*/
def create_tree(parent: TreeNode, wasParallel: bool = false): int {
  if parent == nil {
    writeln("create_tree(): Warning, called with nil");
    return 0;
  }

  if (parallel) {
    var count: sync int = parent.genChildren();
    var threads_granted = requestThreads(parent.nChildren);

    // Spawn threads
    coforall i in [1..threads_granted] {
      count += create_tree(parent.children[i], true);
    }

    // Run the rest sequentially
    for i in [threads_granted+1..parent.nChildren] {
      count += create_tree(parent.children[i], false);
    }

    if (wasParallel) then thread_cnt -= 1;
    return count;

  } else {
    var count: int = parent.genChildren();

    for i in parent.childDom do
      count += create_tree(parent.children[i], false);
    return count;
  }
}

def main {
  var t_create: Timer();
  var t_dfs   : Timer();
  var counted, created: int;
  var root: TreeNode;

  root = TreeNode(0);
  rng_init(root.hash[1], SEED:int);

  uts_showSearchParams();

  writeln("Performing tree creation..");
  t_create.start();
  created = create_tree(root);
  t_create.stop();

  if parallel && !testMode then writeln("  threads_spawned= ", threads_spawned);
  threads_spawned = 0;

  writeln("Performing tree traversal..");
  t_dfs.start();
  counted = dfs_count(root);
  t_dfs.stop();

  if parallel && !testMode then writeln("  threads_spawned= ", threads_spawned);

  writeln();
  writeln("Total nodes: Created = ", created, ", ", "Counted = ", counted);
  if !testMode then writeln("Time: t_create= ", t_create.elapsed(),
          " (", created/t_create.elapsed(), " nodes/sec)", ", t_dfs = ", t_dfs.elapsed(),
          " (", counted/t_dfs.elapsed(), " nodes/sec)");
}

/*
**  UTS -- The Unbalanced Tree Search
**
**  James Dinan <dinan@cray.com>
**  July, 2007
*/

use Math;
use Time;
use sha1_rng;

/**** UTS TYPES ****/
enum NodeDistrib { Binomial, Geometric, Hybrid, Constant };
enum GeoDistrib  { GeoFixed, GeoLinear, GeoPoly, GeoCyclic };

/**** COMPILE PARAMS ****/
config param debug:    bool = false;
config param parallel: bool = true;
param uts_version:   string = "2.1";

/**** UTS CONFIG VARS ****/
config const SEED:        uint = 16;    // Seed for the hash of the root node
config const MAX_CHILDREN: int = 10;    // Max. Children a node may have
config const B_0:          int = 4;     // Branching Factor
config const MAX_DEPTH:    int = 6;     // Depth limit
config const nonLeafProb: real = 0.10;  // Probability of a non-leaf (binomial)
config const nonLeafBF:    int = 4;     // Non-leaf branching factor (binomial)
config const shiftDepth:  real = 0.5;   // Depth at which hybrid trees go from Geo=>Bin

config const distrib: NodeDistrib = Geometric;
config const geoDist: GeoDistrib  = GeoFixed;
config const testMode             = false;

config const MAX_THREADS:  int = 20;
config const MIN_THREADS:  int = MAX_THREADS/4;

// Global thread counter
var thread_cnt: sync int = 0;
var threads_spawned: int = 0; // "Locked" by thread_cnt

/**** UTS TreeNode Class ****/
class TreeNode {
  var depth:   int;
  var hash:    [1..20] uint(8);

  // By default, children will be empty since it has range [1..0]
  var nChildren: int = 0;
  var childDom = [1..nChildren];
  var children:  [childDom] TreeNode;

  // Generate this node's children
  def genChildren(): int {
    select distrib {
      when Geometric do
        nChildren = numGeoChildren(geoDist);
      when Binomial do
        nChildren = numBinChildren();
      when Hybrid do
        nChildren = numHybridChildren();
      when Constant do
        nChildren = numConstChildren();
    }

    if debug then writeln("Constructing ", nChildren, " children: ", childDom);

    childDom  = [1..nChildren];

    forall i in childDom {
      if debug then writeln("  + (", depth, ", ", i, ")");
      children[i]    = TreeNode(depth+1);
      rng_spawn(hash[1], children[i].hash[1], i-1);
    }

    return nChildren;
  }

  // Constant Distribution: Find the number of children
  //  Note: This is a _balanced_ tree.
  def numConstChildren(): int {
    return if depth < MAX_DEPTH then B_0 else 0;
  }

  // Hybrid trees are geometric at the top of the tree and binomial
  // below a certain depth.
  def numHybridChildren(): int {
    if (depth < shiftDepth * MAX_DEPTH) then
      return numGeoChildren(geoDist);

    else
      return numBinChildren();
  }

  // Binomial Distribution: Find the number of children
  //  Note: distribution is identical everywhere below root
  def numBinChildren(): int {
    if (depth == 0) then
      return B_0;

    else
      return if to_prob(rng_rand(hash[1])) < nonLeafProb then nonLeafBF else 0;
  }

  // Geometric Distribution: Find the number of children
  def numGeoChildren(shape: GeoDistrib): int {
    var b_i: real = B_0;

    // use shape function to compute target b_i
    if (depth > 0) {
      select shape {

        // Expected size is polynomial in depth
        when GeoPoly do
          b_i = B_0 * (depth ** (-log(B_0:real)/log(MAX_DEPTH:real)));

        // Expected size is cyclic
        when GeoCyclic {
          if (depth > 5 * MAX_DEPTH) then
            b_i = 0.0;
          else
            b_i = B_0 ** sin(2.0*3.141592653589793*depth/MAX_DEPTH);
        }

        // Expected size is the same at all nodes up to max depth
        when GeoFixed do
          b_i = if (depth < MAX_DEPTH) then B_0 else 0;

        // Expected size decreases linearly in b_i
        when GeoLinear do
          b_i =  B_0 * (1.0 - depth:real / MAX_DEPTH:real);

      }
    }

    // given target b_i, find prob p so expected value of
    // geometric distribution is b_i.
    var p: real = 1.0 / (1.0 + b_i);

    // get uniform random number on [0,1)
    var u: real = to_prob(rng_rand(hash[1]));

    // max number of children at this cumulative probability
    // (from inverse geometric cumulative density function)
    return floor(log(1 - u) / log(1 - p)):int;
  }
}

/*
** Print out search parameters
*/
def uts_showSearchParams() {
  writeln("UTS v", uts_version, " - Unbalanced Tree Search (",
      if parallel then "Parallel Chapel" else "Sequential Chapel", ")");

  writeln("Tree type: ", distrib);
  writeln("Tree shape parameters:");
  writeln("  root branching factor b_0 = ", B_0, ", root seed = ", SEED);

  if (distrib == Geometric || distrib == Hybrid) then
    writeln("  Geo. params: Max_Depth = ", MAX_DEPTH, ", shape function = ", geoDist);

  if (distrib == Binomial || distrib == Hybrid) then
    writeln("  Bin. params: nonLeafProb=", nonLeafProb, " nonLeafBF=", nonLeafBF);

  if (distrib == Hybrid) then
    writeln("  Hybrid params: shiftDepth=", shiftDepth);

  writeln("Random number generator: ", rng_getName());

  if (parallel) {
    writeln("Parallel execution parameters:");
    writeln("  MAX_THREADS = ", MAX_THREADS, ", MIN_THREADS = ", MIN_THREADS);
  }

  writeln();
}

/*
** Request n threads.  Get back the number of threads granted.
*/
def requestThreads(n:int): int {
  if (parallel) {
    // Trade some imbalance here for blocking overhead
    if (readXX(thread_cnt) < MIN_THREADS) {
      var thread_cnt_l = thread_cnt;

      // Try to get a ticket to run in parallel
      if (thread_cnt_l < MAX_THREADS) {
        var threads_avail   = MAX_THREADS-thread_cnt_l;
        var threads_granted = if (n <= threads_avail) then n else threads_avail;

        thread_cnt = thread_cnt_l + threads_granted;
        threads_spawned += threads_granted;
        return threads_granted;
      } else {
        thread_cnt = thread_cnt_l;
      }
    }
  }

  return 0;
}

/*
**  Parallel DFS
*/
def dfs_count(n: TreeNode, wasParallel: bool = false):int {
  if n != nil {
    if (parallel) {
      var threads_granted = requestThreads(n.nChildren);
      var count: sync int = 0;

      coforall i in [1..threads_granted] {
        count += dfs_count(n.children[i], true);
      }

      for i in [threads_granted+1..n.nChildren] {
        count += dfs_count(n.children[i], false);
      }

      if (wasParallel) then thread_cnt -= 1;
      return count+1;

    } else {
      var count: int;

      for i in n.childDom {
        count += dfs_count(n.children[i], false);
      }
      return count+1;
    }
  }
}

/*
**  Parallel Tree Creation
*/
def create_tree(parent: TreeNode, wasParallel: bool = false): int {
  if parent == nil {
    writeln("create_tree(): Warning, called with nil");
    return 0;
  }

  if (parallel) {
    var count: sync int = parent.genChildren();
    var threads_granted = requestThreads(parent.nChildren);

    // Spawn threads
    coforall i in [1..threads_granted] {
      count += create_tree(parent.children[i], true);
    }

    // Run the rest sequentially
    for i in [threads_granted+1..parent.nChildren] {
      count += create_tree(parent.children[i], false);
    }

    if (wasParallel) then thread_cnt -= 1;
    return count;

  } else {
    var count: int = parent.genChildren();

    for i in parent.childDom do
      count += create_tree(parent.children[i], false);
    return count;
  }
}

def main {
  var t_create: Timer();
  var t_dfs   : Timer();
  var counted, created: int;
  var root: TreeNode;

  root = TreeNode(0);
  rng_init(root.hash[1], SEED:int);

  uts_showSearchParams();

  writeln("Performing tree creation..");
  t_create.start();
  created = create_tree(root);
  t_create.stop();

  if parallel && !testMode then writeln("  threads_spawned= ", threads_spawned);
  threads_spawned = 0;

  writeln("Performing tree traversal..");
  t_dfs.start();
  counted = dfs_count(root);
  t_dfs.stop();

  if parallel && !testMode then writeln("  threads_spawned= ", threads_spawned);

  writeln();
  writeln("Total nodes: Created = ", created, ", ", "Counted = ", counted);
  if !testMode then writeln("Time: t_create= ", t_create.elapsed(),
          " (", created/t_create.elapsed(), " nodes/sec)", ", t_dfs = ", t_dfs.elapsed(),
          " (", counted/t_dfs.elapsed(), " nodes/sec)");
}


Replace uses of sleep with sync variables

As has been pointed out in several different tests, using sleeps
to try and order tasks is an unsafe practice.  This test was
failing on my Mac using array views which seemed surprising
given that it doesn't even use arrays.  Turns out that it's
just a race condition.  Replacing the sleeps with sync variables
to avoid the race.

The comment states that this practice was done for debugging
which is all well and good when debugging, but once tests
are being committed for long-term testing, let's try and get
the potential for races out of them please.

// Verifies updates to 'this' within tasks/on blocks within
// a constructor or a method.

// Style note: I use chpl_task_sleep to avoid 'sync' statements.
// I avoid 'sync' statements and do not import the Time module
// - in order to reduce the size of the AST and the generated code.
// That is useful for ease of debugging.

extern proc chpl_task_sleep(t:uint): void;
const loc = Locales[numLocales-1];

record RR {
  var xx, yy: int;
  // the default constructor
  proc RR() {
    var done$: sync bool;
    on loc {
      this.xx = 555;
    }
    begin {
      this.yy = 666;
      done$ = true;
    }
    done$;
  }
  // method with args
  proc modify(ee: int, ff: int) {
    var done$: sync bool;
    on loc {
      this.xx = ee;
    }
    begin {
      this.yy = ff;
      done$ = true;
    }
    done$;
  }
} // record RR

record QQ {
  var aa, bb: int;
  // non-default constructor
  proc QQ(cc: int, dd: int) {
    var done$: sync bool;
    on loc {
      this.aa = cc;
    }
    begin {
      this.bb = dd;
      done$ = true;
    }
    done$;
  }
  // method with no args
  proc modify() {
    var done$: sync bool;
    on loc {
      this.aa = 171717;
    }
    begin {
      this.bb = 181818;
      done$ = true;
    }
    done$;
  }
} // record QQ

writeln("start");

var rr = new RR();
var qq = new QQ(777, 888);
writeln(rr, qq);

rr.modify(151515, 161616);
qq.modify();
writeln(rr, qq);

// Verifies updates to 'this' within tasks/on blocks within
// a constructor or a method.

// Style note: I use chpl_task_sleep to avoid 'sync' statements.
// I avoid 'sync' statements and do not import the Time module
// - in order to reduce the size of the AST and the generated code.
// That is useful for ease of debugging.

extern proc chpl_task_sleep(t:uint): void;
const loc = Locales[numLocales-1];

record RR {
  var xx, yy: int;
  // the default constructor
  proc RR() {
    on loc {
      this.xx = 555;
    }
    begin {
      this.yy = 666;
    }
    chpl_task_sleep(1);
  }
  // method with args
  proc modify(ee: int, ff: int) {
    on loc {
      this.xx = ee;
    }
    begin {
      this.yy = ff;
    }
    chpl_task_sleep(1);
  }
} // record RR

record QQ {
  var aa, bb: int;
  // non-default constructor
  proc QQ(cc: int, dd: int) {
    on loc {
      this.aa = cc;
    }
    begin {
      this.bb = dd;
    }
    chpl_task_sleep(1);
  }
  // method with no args
  proc modify() {
    on loc {
      this.aa = 171717;
    }
    begin {
      this.bb = 181818;
    }
    chpl_task_sleep(1);
  }
} // record QQ

writeln("start");

var rr = new RR();
var qq = new QQ(777, 888);
writeln(rr, qq);

rr.modify(151515, 161616);
qq.modify();
writeln(rr, qq);


Removed a superflous "if".  :)

Removed a superflous "if".  :)

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

config param eps            = 1e-16;
config param truncateAtEps  = false;

// Truncate values to 0 if they are below machine precision
def truncate(x) {
    if truncateAtEps && x <= eps then return 0.0;
    else return x;
}

// Copy matrix B's transpose into matrix A
def transposeCopy(A: [] real, B: [] real) where A.rank == 2 && B.rank == 2 {
    forall (i, j) in A.domain do
        A[i, j] = B[j, i];
}

// Vector-Matrix transformation
def *(V: [] real, M: [] real) where V.rank == 1 && M.rank == 2 {
    var R: [M.domain.dim(2)] real = 0.0;

    if V.domain.dim(1) != M.domain.dim(1) then
        halt("*: Vector and matrix dims must match");

    for (i, j) in M.domain do
        R[i] += V[j] * M[j, i];

    return R;
}

// Matrix-Vector multiplication
def *(M: [] real, V: [] real) where V.rank == 1 && M.rank == 2 {
    var R: [M.domain.dim(1)] real = 0.0;

    if V.domain.dim(1) != M.domain.dim(2) then
        halt("*: Vector and matrix dims must match");

    for (i, j) in M.domain do
        R[i] += V[j] * M[i, j];

    return R;
}

// Frobenius Norm -- 2-norm for Vectors
def normf(A) where A.rank == 1 {
    return sqrt(+ reduce A**2);
}

// Inner product of A and B
def inner(A, B) where A.rank == 1 && B.rank == 1 {
    return + reduce A*B;
}

config param eps            = 1e-16;
config param truncateAtEps  = false;

// Truncate values to 0 if they are below machine precision
def truncate(x) {
    if truncateAtEps then
      if x <= eps then return 0.0;
    return x;
}

// Copy matrix B's transpose into matrix A
def transposeCopy(A: [] real, B: [] real) where A.rank == 2 && B.rank == 2 {
    forall (i, j) in A.domain do
        A[i, j] = B[j, i];
}

// Vector-Matrix transformation
def *(V: [] real, M: [] real) where V.rank == 1 && M.rank == 2 {
    var R: [M.domain.dim(2)] real = 0.0;

    if V.domain.dim(1) != M.domain.dim(1) then
        halt("*: Vector and matrix dims must match");

    for (i, j) in M.domain do
        R[i] += V[j] * M[j, i];

    return R;
}

// Matrix-Vector multiplication
def *(M: [] real, V: [] real) where V.rank == 1 && M.rank == 2 {
    var R: [M.domain.dim(1)] real = 0.0;

    if V.domain.dim(1) != M.domain.dim(2) then
        halt("*: Vector and matrix dims must match");

    for (i, j) in M.domain do
        R[i] += V[j] * M[i, j];

    return R;
}

// Frobenius Norm -- 2-norm for Vectors
def normf(A) where A.rank == 1 {
    return sqrt(+ reduce A**2);
}

// Inner product of A and B
def inner(A, B) where A.rank == 1 && B.rank == 1 {
    return + reduce A*B;
}


Allow running the kernel using either remote tasks or remote references.

[Reviewed by Brad.]

This adds a config variable, which controls whether the computational
kernel is done using remote tasks ("on" statements) so that the table
updates are computed locally, or direct remote references for the table
updates.  Some communication layer implementations perform better with
one, and some with the other.  For now at least, the default is to use
direct remote references with the uGNI-based communication layer, and
remote tasks with the rest.

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

//
// Use standard modules for Block distributions and Timing routines
//
use BlockDist, Time;

//
// Use the user modules for computing HPCC problem sizes and for
// defining RA's random stream of values
//
use HPCCProblemSize, RARandomStream;

//
// The number of tables as well as the element and index types of
// that table
//
const numTables = 1;
type elemType = randType,
     indexType = randType;

//
// Configuration constants defining log2(problem size) -- n -- and
// the number of updates -- N_U
//
config const n = computeProblemSize(numTables, elemType,
                                    returnLog2=true, retType=indexType),
             N_U = 2**(n+2);

//
// Constants defining the problem size (m) and a bit mask for table
// indexing
//
const m = 2**n,
      indexMask = m-1;

//
// Configuration constant defining the number of errors to allow (as a
// fraction of the number of updates, N_U)
//
config const errorTolerance = 1e-2;

//
// Configuration constants to control what's printed -- benchmark
// parameters, input and output arrays, and/or statistics
//
config const printParams = true,
             printArrays = false,
             printStats = true;

//
// Use "on" statement in kernel, or remote memory references?
//
config const useOn = (CHPL_COMM != "ugni");

//
// TableDist is a 1D block distribution for domains storing indices
// of type "indexType", and it is computed by blocking the bounding
// box 0..m-1 across the set of locales.  UpdateDist is a similar
// distribution that is computed by blocking the indices 0..N_U-1
// across the locales.
//
const TableDist = new dmap(new Block(boundingBox=[0..m-1])),
      UpdateDist = new dmap(new Block(boundingBox=[0..N_U-1]));

//
// TableSpace describes the index set for the table.  It is a 1D
// domain storing indices of type indexType, it is distributed
// according to TableDist, and it contains the indices 0..m-1.
// Updates is an index set describing the set of updates to be made.
// It is distributed according to UpdateDist and contains the
// indices 0..N_U-1.
//
const TableSpace: domain(1, indexType) dmapped TableDist = [0..m-1],
      Updates: domain(1, indexType) dmapped UpdateDist = [0..N_U-1];

//
// The program entry point
//
proc main() {
  //
  // T is the distributed table itself, storing a variable of type
  // elemType for each index in TableSpace.
  //
  var T: [TableSpace] elemType;

  printConfiguration();   // print the problem size, number of trials, etc.

  //
  // In parallel, initialize the table such that each position
  // contains its index.  "[i in TableSpace]" is shorthand for "forall
  // i in TableSpace"
  //
  [i in TableSpace] T(i) = i;

  const startTime = getCurrentTime();              // capture the start time

  //
  // The main computation: Iterate over the set of updates and the
  // stream of random values in a parallel, zippered manner, dropping
  // the update index on the ground and storing the random value
  // in r.  Depending on the value of the useOn config variable, we
  // either use an on-clause to force the table updates to be executed
  // on the locales which own the table elements in question, or we
  // do the updates using direct GETs and PUTs.  Either way could be
  // faster, depending on the characteristics of the paricular comm
  // layer in use.  Compute the update using r both to compute the
  // index and as the update value.
  //
  if (useOn) then
    forall ( , r) in (Updates, RAStream()) do
      on TableDist.idxToLocale(r & indexMask) do {
        const myR = r;
        local {
          T(myR & indexMask) ^= myR;
        }
      }
  else
    forall ( , r) in (Updates, RAStream()) do
      T(r & indexMask) ^= r;

  const execTime = getCurrentTime() - startTime;   // capture the elapsed time

  const validAnswer = verifyResults(T);            // verify the updates
  printResults(validAnswer, execTime);             // print the results
}

//
// Print the problem size and number of updates
//
proc printConfiguration() {
  if (printParams) {
    if (printStats) then printLocalesTasks();
    printProblemSize(elemType, numTables, m);
    writeln("Number of updates = ", N_U, "\n");
  }
}

//
// Verify that the computation is correct
//
proc verifyResults(T) {
  //
  // Print the table, if requested
  //
  if (printArrays) then writeln("After updates, T is: ", T, "\n");

  //
  // Reverse the updates by recomputing them, this time using an
  // atomic statement to ensure no conflicting updates
  //
  if (useOn) then
    forall ( , r) in (Updates, RAStream()) do
      on TableDist.idxToLocale(r & indexMask) do {
        const myR = r;
        local {
          atomic T(myR & indexMask) ^= myR;
        }
     }
  else
   forall ( , r) in (Updates, RAStream()) do
      atomic T(r & indexMask) ^= r;

  //
  // Print the table again after the updates have been reversed
  //
  if (printArrays) then writeln("After verification, T is: ", T, "\n");

  //
  // Compute the number of table positions that weren't reverted
  // correctly.  This is an indication of the number of conflicting
  // updates.
  //
  const numErrors = + reduce [i in TableSpace] (T(i) != i);
  if (printStats) then writeln("Number of errors is: ", numErrors, "\n");

  //
  // Return whether or not the number of errors was within the benchmark's
  // tolerance.
  //
  return numErrors <= (errorTolerance * N_U);
}

//
// Print out success/failure, the execution time, and the GUPS value
//
proc printResults(successful, execTime) {
  writeln("Validation: ", if successful then "SUCCESS" else "FAILURE");
  if (printStats) {
    writeln("Execution time = ", execTime);
    writeln("Performance (GUPS) = ", (N_U / execTime) * 1e-9);
  }
}

//
// Use standard modules for Block distributions and Timing routines
//
use BlockDist, Time;

//
// Use the user modules for computing HPCC problem sizes and for
// defining RA's random stream of values
//
use HPCCProblemSize, RARandomStream;

//
// The number of tables as well as the element and index types of
// that table
//
const numTables = 1;
type elemType = randType,
     indexType = randType;

//
// Configuration constants defining log2(problem size) -- n -- and
// the number of updates -- N_U
//
config const n = computeProblemSize(numTables, elemType,
                                    returnLog2=true, retType=indexType),
             N_U = 2**(n+2);

//
// Constants defining the problem size (m) and a bit mask for table
// indexing
//
const m = 2**n,
      indexMask = m-1;

//
// Configuration constant defining the number of errors to allow (as a
// fraction of the number of updates, N_U)
//
config const errorTolerance = 1e-2;

//
// Configuration constants to control what's printed -- benchmark
// parameters, input and output arrays, and/or statistics
//
config const printParams = true,
             printArrays = false,
             printStats = true;

//
// TableDist is a 1D block distribution for domains storing indices
// of type "indexType", and it is computed by blocking the bounding
// box 0..m-1 across the set of locales.  UpdateDist is a similar
// distribution that is computed by blocking the indices 0..N_U-1
// across the locales.
//
const TableDist = new dmap(new Block(boundingBox=[0..m-1])),
      UpdateDist = new dmap(new Block(boundingBox=[0..N_U-1]));

//
// TableSpace describes the index set for the table.  It is a 1D
// domain storing indices of type indexType, it is distributed
// according to TableDist, and it contains the indices 0..m-1.
// Updates is an index set describing the set of updates to be made.
// It is distributed according to UpdateDist and contains the
// indices 0..N_U-1.
//
const TableSpace: domain(1, indexType) dmapped TableDist = [0..m-1],
      Updates: domain(1, indexType) dmapped UpdateDist = [0..N_U-1];

//
// The program entry point
//
proc main() {
  //
  // T is the distributed table itself, storing a variable of type
  // elemType for each index in TableSpace.
  //
  var T: [TableSpace] elemType;

  printConfiguration();   // print the problem size, number of trials, etc.

  //
  // In parallel, initialize the table such that each position
  // contains its index.  "[i in TableSpace]" is shorthand for "forall
  // i in TableSpace"
  //
  [i in TableSpace] T(i) = i;

  const startTime = getCurrentTime();              // capture the start time

  //
  // The main computation: Iterate over the set of updates and the
  // stream of random values in a parallel, zippered manner, dropping
  // the update index on the ground and storing the random value
  // in r.  Use an on-clause to force the table update to be executed on
  // the locale which owns the table element in question to minimize
  // communications.  Compute the update using r both to compute the
  // index and as the update value.
  //
  forall ( , r) in (Updates, RAStream()) do
    on TableDist.idxToLocale(r & indexMask) do {
      const myR = r;
      local {
        T(myR & indexMask) ^= myR;
      }
   }

  const execTime = getCurrentTime() - startTime;   // capture the elapsed time

  const validAnswer = verifyResults(T);            // verify the updates
  printResults(validAnswer, execTime);             // print the results
}

//
// Print the problem size and number of updates
//
proc printConfiguration() {
  if (printParams) {
    if (printStats) then printLocalesTasks();
    printProblemSize(elemType, numTables, m);
    writeln("Number of updates = ", N_U, "\n");
  }
}

//
// Verify that the computation is correct
//
proc verifyResults(T) {
  //
  // Print the table, if requested
  //
  if (printArrays) then writeln("After updates, T is: ", T, "\n");

  //
  // Reverse the updates by recomputing them, this time using an
  // atomic statement to ensure no conflicting updates
  //
  forall ( , r) in (Updates, RAStream()) do
    on TableDist.idxToLocale(r & indexMask) do {
      const myR = r;
      local {
        atomic T(myR & indexMask) ^= myR;
      }
   }

  //
  // Print the table again after the updates have been reversed
  //
  if (printArrays) then writeln("After verification, T is: ", T, "\n");

  //
  // Compute the number of table positions that weren't reverted
  // correctly.  This is an indication of the number of conflicting
  // updates.
  //
  const numErrors = + reduce [i in TableSpace] (T(i) != i);
  if (printStats) then writeln("Number of errors is: ", numErrors, "\n");

  //
  // Return whether or not the number of errors was within the benchmark's
  // tolerance.
  //
  return numErrors <= (errorTolerance * N_U);
}

//
// Print out success/failure, the execution time, and the GUPS value
//
proc printResults(successful, execTime) {
  writeln("Validation: ", if successful then "SUCCESS" else "FAILURE");
  if (printStats) {
    writeln("Execution time = ", execTime);
    writeln("Performance (GUPS) = ", (N_U / execTime) * 1e-9);
  }
}