added test for set/get by name

added test for set/get by name

use NumSuch,
    Charcoal;

class NumSuchTest : UnitTest {
  var nv: int = 8,
      D: domain(2) = {1..nv, 1..nv},
      SD: sparse subdomain(D),
      X: [SD] real;

  var vn: [1..0] string;

  proc setUp() {
    vn.push_back("star lord");
    vn.push_back("gamora");
    vn.push_back("groot");
    vn.push_back("drax");
    vn.push_back("rocket");
    vn.push_back("mantis");
    vn.push_back("yondu");

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;
  }

  proc init() {
    super.init();
    this.initDone();
  }

  proc testIndexSort() {
    var p: bool = true;
    const Arr = [7, 10, 23, 1];
    const Idx = [2.2, 3.3, 1.1, 4.4];

    var a: [1..4] real = for x in indexSort(arr=Arr, idx=Idx) do x;
    var expected: [1..4] real = [23, 7, 10, 1];
    this.results.push_back(assertArrayEquals(msg="IndexSort", expected=expected, actual=a));

    var b: [1..4] real = for x in indexSort(arr=Arr, idx=Idx, reverse=true) do x;
    var e: [1..4] real = [1, 10, 7, 23];
    this.results.push_back(assertArrayEquals(msg="IndexSort reversed", expected=e, actual=b));
  }

  proc testNamedMatrix() {
    var nm = new NamedMatrix(X=X);
    //this.results.push_back(assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=nm.nnz():int ));
    this.results.push_back(assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=11:int ));
    this.results.push_back(assertRealEquals(msg="X.sparsity", expected=0.171875:real, actual=nm.sparsity() ));
  }

  proc testSetRowNames() {
    var nm = new NamedMatrix(X=X);
    try {
      nm.setRowNames(vn);
      this.results.push_back(
        assertThrowsError(msg="Set Row Names wrong size NOT ENFORCED"
        , passed=false, new Error())
      );
    } catch err: DimensionMatchError {
      this.results.push_back(assertThrowsError(msg="Set Row Names wrong size ENFORCED"
        ,  passed=true, err=err));
    } catch err: Error {
      this.results.push_back(
        assertThrowsError(msg="Set Row Names wrong size ENFORCED"
        ,  passed=true, err=err)
      );
    }

    vn.push_back("nebula");
    try {
      nm.setRowNames(vn);
      this.results.push_back(
        assertIntEquals(msg="Set Row Names right size ENFORCED",  expected=vn.size, actual=nm.nrows() )
      );
    } catch {
      this.results.push_back(
        assertIntEquals(msg="Set Row Names right size failed",  expected=8, actual=nm.nrows())
      );
    }
  }

  proc testSetColNames() {
    var nm = new NamedMatrix(X=X);
    //vn.push_back("nebula");
    try {
      nm.setColNames(vn);
      this.results.push_back(
        assertIntEquals(msg="Set Col Names right size ENFORCED",  expected=vn.size, actual=nm.ncols() )
      );
    } catch {
      this.results.push_back(
        assertIntEquals(msg="Set Col Names right size",  expected=vn.size, actual=nm.ncols())
      );
    }
  }

  proc testNamedMatrixInitWithNames() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      this.results.push_back(
        assertIntEquals(msg="Initialize X with names only (rows)", expected=vn.size, actual=nm.nrows())
      );
      this.results.push_back(
        assertIntEquals(msg="Initialize X with names only (cols)", expected=vn.size, actual=nm.ncols())
      );
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc testSetByName() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      nm.set("star lord", "yondu", 17.0);
      this.results.push_back(
        assertRealEquals(msg="Set/Get by name", expected=17.0, actual=nm.get("star lord", "yondu"))
      );
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc run() {
    super.run();
    testIndexSort();
    testNamedMatrix();
    testSetRowNames();
    testSetColNames();
    testNamedMatrixInitWithNames();
    testSetByName();
    return 0;
  }

}

proc main(args: [] string) : int {
  var t = new NumSuchTest();
  var ret = t.run();
  t.report();
  return ret;
}

use NumSuch,
    Charcoal;

class NumSuchTest : UnitTest {
  var nv: int = 8,
      D: domain(2) = {1..nv, 1..nv},
      SD: sparse subdomain(D),
      X: [SD] real;

  var vn: [1..0] string;

  proc setUp() {
    vn.push_back("star lord");
    vn.push_back("gamora");
    vn.push_back("groot");
    vn.push_back("drax");
    vn.push_back("rocket");
    vn.push_back("mantis");
    vn.push_back("yondu");

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;
  }

  proc init() {
    super.init();
    this.initDone();
  }

  proc testIndexSort() {
    var p: bool = true;
    const Arr = [7, 10, 23, 1];
    const Idx = [2.2, 3.3, 1.1, 4.4];

    var a: [1..4] real = for x in indexSort(arr=Arr, idx=Idx) do x;
    var expected: [1..4] real = [23, 7, 10, 1];
    this.results.push_back(assertArrayEquals(msg="IndexSort", expected=expected, actual=a));

    var b: [1..4] real = for x in indexSort(arr=Arr, idx=Idx, reverse=true) do x;
    var e: [1..4] real = [1, 10, 7, 23];
    this.results.push_back(assertArrayEquals(msg="IndexSort reversed", expected=e, actual=b));
  }

  proc testNamedMatrix() {
    var nm = new NamedMatrix(X=X);
    //this.results.push_back(assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=nm.nnz():int ));
    this.results.push_back(assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=11:int ));
    this.results.push_back(assertRealEquals(msg="X.sparsity", expected=0.171875:real, actual=nm.sparsity() ));
  }

  proc testSetRowNames() {
    var nm = new NamedMatrix(X=X);
    try {
      nm.setRowNames(vn);
      this.results.push_back(
        assertThrowsError(msg="Set Row Names wrong size NOT ENFORCED"
        , passed=false, new Error())
      );
    } catch err: DimensionMatchError {
      this.results.push_back(assertThrowsError(msg="Set Row Names wrong size ENFORCED"
        ,  passed=true, err=err));
    } catch err: Error {
      this.results.push_back(
        assertThrowsError(msg="Set Row Names wrong size ENFORCED"
        ,  passed=true, err=err)
      );
    }

    vn.push_back("nebula");
    try {
      nm.setRowNames(vn);
      this.results.push_back(
        assertIntEquals(msg="Set Row Names right size ENFORCED",  expected=vn.size, actual=nm.nrows() )
      );
    } catch {
      this.results.push_back(
        assertIntEquals(msg="Set Row Names right size failed",  expected=8, actual=nm.nrows())
      );
    }
  }

  proc testSetColNames() {
    var nm = new NamedMatrix(X=X);
    //vn.push_back("nebula");
    try {
      nm.setColNames(vn);
      this.results.push_back(
        assertIntEquals(msg="Set Col Names right size ENFORCED",  expected=vn.size, actual=nm.ncols() )
      );
    } catch {
      this.results.push_back(
        assertIntEquals(msg="Set Col Names right size",  expected=vn.size, actual=nm.ncols())
      );
    }
  }

  proc testNamedMatrixInitWithNames() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      this.results.push_back(
        assertIntEquals(msg="Initialize X with names only (rows)", expected=vn.size, actual=nm.nrows())
      );
      this.results.push_back(
        assertIntEquals(msg="Initialize X with names only (cols)", expected=vn.size, actual=nm.ncols())
      );
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc run() {
    super.run();
    testIndexSort();
    testNamedMatrix();
    testSetRowNames();
    testSetColNames();
    testNamedMatrixInitWithNames();
    return 0;
  }

}

proc main(args: [] string) : int {
  var t = new NumSuchTest();
  var ret = t.run();
  t.report();
  return ret;
}


Move fence to top of loop, comment changes. Responds to Greg's request.

// ChapelLocale.chpl
//
pragma "no use ChapelStandard"
module ChapelLocale {

  //
  // An abstract class. Specifies the required locale interface.
  // Each locale implementation must inherit from this class.
  //
  class locale {
    //- Constructor
    proc locale() { }

    //------------------------------------------------------------------------{
    //- Fields and accessors defined for all locale types (not overridable)
    //-

    // Every locale has a parent, except for the root locale.
    // The parent of the root locale is nil (by definition).
    const parent : locale;

    // To be removed from the required interface once legacy code is adjusted.
    const numCores: int;

    proc id : int return chpl_id();  // just the node part
    proc localeid : chpl_localeID_t return chpl_localeid(); // full locale id
    proc name return chpl_name();

    // This many tasks are running on this locale.
    //
    // Note handling runningTaskCounter <= 0 in runningTaskCnt().  The
    // references to it are done via here.runningTasks(), and currently
    // "here" means the requested locale, not the actual or concrete
    // one.  But this may be "any", and the running task count in it may
    // be zero, which in turn leads to overindexing the locales array
    // after the subtraction of 1 in locale.runningTasks(), below.  In
    // the long run we probably need to go through all the "here" refs
    // in the module code and differentiate more carefully which need
    // the requested locale and which need the concrete one.  But for
    // now, we make the assumption that all requests for the number of
    // running tasks want the count from the locale the calling task is
    // running on, so the minimum possible value must be 1.
    var runningTaskCounter : atomic int;

    inline proc runningTaskCntSet(val : int) {
      runningTaskCounter.write(val, memory_order_relaxed);
    }

    inline proc runningTaskCntAdd(val : int) {
      runningTaskCounter.add(val, memory_order_relaxed);
    }

    inline proc runningTaskCntSub(val : int) {
      runningTaskCounter.sub(val, memory_order_relaxed);
    }

    inline proc runningTaskCnt() {
      var rtc = runningTaskCounter.read(memory_order_relaxed);
      return if (rtc <= 0) then 1 else rtc;
    }
    //------------------------------------------------------------------------}

    //------------------------------------------------------------------------{
    //- User Interface Methods (overridable)
    //-

    // These are dynamically dispatched, so they can be overridden in
    // concrete classes.
    proc chpl_id() : int {
      _throwPVFCError();
      return -1;
    }

    proc chpl_localeid() : chpl_localeID_t {
      _throwPVFCError();
      extern const c_sublocid_none: chpl_sublocID_t;
      return chpl_buildLocaleID(-1:chpl_nodeID_t, c_sublocid_none);
    }

    proc chpl_name() : string {
      _throwPVFCError();
      return "";
    }

    // A useful default definition is provided (not pure virtual).
    proc readWriteThis(f) {
      f <~> name;
    }

    proc getChildCount() : int {
      _throwPVFCError();
      return 0;
    }

// Part of the required locale interface.
// Commented out because presently iterators are statically bound.
//    iter getChildIndices() : int {
//      for idx in this.getChildSpace do
//        yield idx;
//    }

    proc addChild(loc:locale)
    {
      _throwPVFCError();
    }

    proc getChild(idx:int) : locale {
      _throwPVFCError();
      return this;
    }

// Part of the required locale interface.
// Commented out because presently iterators are statically bound.
//    iter getChildren() : locale  {
//    _throwPVFCError();
//      yield 0;
//    }

    //------------------------------------------------------------------------}
  }

  class AbstractLocaleModel : locale {
    // This will be used for interfaces that will be common to all
    // (non-RootLocale) locale models
  }

  // rootLocale is declared to be of type locale rather than
  // RootLocale because doing so would make locales depend on arrays
  // which depend on locales, etc.
  //
  // If we had a way to express domain and array return types abstractly,
  // that problem would go away. <hilde>
  //
  // The rootLocale is private to each locale.  It cannot be
  // initialized until LocaleModel is initialized.  To disable this
  // replication, set replicateRootLocale to false.
  pragma "private" var rootLocale : locale = nil;
  config param replicateRootLocale = true;

  // The rootLocale needs to be initalized on all locales prior to
  // initializing the Locales array.  Unfortunately, the rootLocale
  // cannot be properly replicated until DefaultRectangular can be
  // initialized (as the private copies of the defaultDist are
  // needed).  We remedy this by initializing a rootLocale on locale 0
  // (called origRootLocale), and setting all locales' rootLocale to
  // origRootLocale.  Later, after DefaultRectangular can be
  // initialized, we create local copies of the rootLocale (and the
  // Locales array).  This is being done in the InitPrivateGlobals
  // module.
  //
  var origRootLocale : locale = nil;

  class AbstractRootLocale : locale {
    // These functions are used to establish values for Locales[] and
    // LocaleSpace -- an array of locales and its correponding domain
    // which are used as the default set of targetLocales in many
    // distributions.
    proc getDefaultLocaleSpace() {
      _throwPVFCError();
      const emptyLocaleSpace: domain(1) = {1..0};
      return emptyLocaleSpace;
    }

    proc getDefaultLocaleArray() {
      _throwPVFCError();
      const emptyLocaleSpace: domain(1) = {1..0};
      const emptyLocales: [emptyLocaleSpace] locale;
      return emptyLocales;
    }

    proc localeIDtoLocale(id : chpl_localeID_t) : locale {
      _throwPVFCError();
      return this;
    }

    // These iterators are to be used by RootLocale:init() to
    // initialize the LocaleModel.  The calling loop body cannot
    // contain any non-local code, since the rootLocale is not yet
    // initialized.
    iter initOnLocales() {
      if numLocales > 1 then
        halt("The locales must be initialized in parallel");
      for locIdx in (origRootLocale:RootLocale).getDefaultLocaleSpace() {
        yield locIdx;
        rootLocale = origRootLocale;
      }
    }

    // Since we are going on to all the locales here, we use the
    // opportunity to initialize any global private variables we
    // either need (e.g., defaultDist) or can do at this point in
    // initialization (e.g., rootLocale).
    iter initOnLocales(param tag: iterKind)
      where tag==iterKind.leader {
      // Simple locales barrier, see implementation below for notes
      var b: localesBarrier;
      var flags: [1..#numLocales-1] localesSignal;
      extern const c_sublocid_any: chpl_sublocID_t;
      coforall locIdx in 0..#numLocales /*ref(b)*/ {
        on __primitive("chpl_on_locale_num",
                       chpl_buildLocaleID(locIdx:chpl_nodeID_t,
                                          c_sublocid_any)) {
          chpl_defaultDistInitPrivate();
          yield locIdx;
          b.wait(locIdx, flags);
          chpl_rootLocaleInitPrivate(locIdx);
        }
      }
    }

    iter initOnLocales(param tag: iterKind, followThis)
      where tag==iterKind.follower {
      yield followThis;
    }
  }

  //
  // Simple locales barrier
  //
  // This sits outside of the abstract root locale definition above
  // because the compiler cannot resolve the constructor (known issues
  // with nested classes/records).  In addition, we cannot have the
  // flags array in the record, because the initCopy function needs
  // the dataPar* configs declared in ChapelDistribution.
  //
  // Each non-0 locale increments the count, and then waits on a
  // *local* atomic.  This is done by creating an array of type
  // class localesSignal, one per locale, and allocating each locale's
  // copy before updating the count.  Locale 0 waits for the others to
  // arrive and then set the signals to true.  We can't do anything
  // too complicated this early on, so we are using a for loop to
  // broadcast that we are done.
  class localesSignal {
    var s: atomic bool;
  }
  record localesBarrier {
    proc wait(locIdx, flags) {
      if locIdx==0 {
        // locale 0 has nothing else to do, so check flags
        while (true) {
          // This fence ensures that writes to the count variables
          // are available to this task. (Note that they aren't
          // atomic if they're 128-bit writes though - so we
          // have some risk of getting part of a wide pointer).
          // Without this fence, there is a race condition on
          // a weakly-ordered architecture.
          atomic_fence();
          var count = 0;
          for f in flags do
            if f then count += 1;
          if count==numLocales-1 then break;
          // Give time to other tasks/threads/processes
          // like we do in waitFor
          chpl_task_yield();
        }
        // Let the others go
        for f in flags do
          f.s.testAndSet();
      } else {
        var f = new localesSignal();
        // expose my flag to locale 0
        flags[locIdx] = f;
        // wait (locally) for locale 0 to set my flag
        f.s.waitFor(true);
        // clean up
        delete f;
      }
    }
  }

  // This function is called in the LocaleArray module to initialize
  // the rootLocale.  It sets up the origRootLocale and also includes
  // set up of the each locale's LocaleModel via RootLocale:init().
  //
  // The init() function must use the initOnLocales() iterator above
  // to iterate in parallel over the locales to set up the LocaleModel
  // object.
  proc chpl_init_rootLocale() {
    origRootLocale = new RootLocale();
    (origRootLocale:RootLocale).init();
  }

  // This function sets up a private copy of rootLocale by replicating
  // origRootLocale and resets the Locales array to point to the local
  // copy on all but locale 0 (which is done in LocalesArray.chpl as
  // part of the declaration).
  proc chpl_rootLocaleInitPrivate(locIdx) {
    // Even when not replicating the rootLocale, we must temporarily
    // set the rootLocale to the original version on locale 0, because
    // the initialization below needs to get/set locale ids.
    rootLocale = origRootLocale;
    if replicateRootLocale && locIdx!=0 {
      // Create a new local rootLocale
      var newRootLocale = new RootLocale();
      // We don't want to be doing unnecessary ref count updates here
      // as they require additional tasks.  We know we don't need them
      // so tell the compiler to not insert them.
      pragma "no copy" pragma "no auto destroy"
      const origLocales => (origRootLocale:RootLocale).getDefaultLocaleArray();
      var origRL = origLocales._value.theData;
      var newRL = newRootLocale.getDefaultLocaleArray()._value.theData;
      // We must directly implement a bulk copy here, as the mechanisms
      // for doing so via a whole array assignment are not initialized
      // yet and copying element-by-element via a for loop is is costly.
      __primitive("chpl_comm_get",
                  __primitive("array_get", newRL, 0),
                  0 /* locale 0 */,
                  __primitive("array_get", origRL, 0), numLocales);
      // Set the rootLocale to the local copy
      rootLocale = newRootLocale;
    }
    if locIdx!=0 {
      // We mimic a private Locales array alias by using the move
      // primitive.
      __primitive("move", Locales,
                  (rootLocale:RootLocale).getDefaultLocaleArray());
    }
  }

  // We need a temporary value for "here" before the architecture is defined.
  // This is due to the fact that "here" is used for memory and task control
  // in setting up the architecture itself.
  // Its type should probably be renamed dummyLocale or something
  // representative.
  // The dummy locale provides system-default tasking and memory management.
  const dummyLocale = new locale();

  extern proc chpl_task_getRequestedSubloc(): chpl_sublocID_t;
  extern var chpl_nodeID: int(32);
  // Return the locale ID of the current locale
  inline proc here_id {
    return chpl_buildLocaleID(chpl_nodeID,chpl_task_getRequestedSubloc());
  }
  // Return the current locale
  inline proc here {
    return chpl_localeID_to_locale(here_id);
  }

  // Returns a wide pointer to the locale with the given id.
  proc chpl_localeID_to_locale(id : chpl_localeID_t) : locale {
    if rootLocale then
      return (rootLocale:AbstractRootLocale).localeIDtoLocale(id);
    else
      // For code prior to rootLocale initialization
      return dummyLocale;
  }

  pragma "insert line file info"
  extern proc chpl_memhook_malloc_pre(number:int, size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_malloc_post(ptr:opaque, number:int,
                                       size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_realloc_pre(ptr:object, size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_realloc_post(newPtr:opaque, ptr:object,
                                        size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_free_pre(ptr:opaque): void;
  extern proc chpl_memhook_md_num(): int(16);

  proc chpl_getPrivatizedCopy(type objectType, objectPid:int): objectType
    return __primitive("chpl_getPrivatizedClass", nil:objectType, objectPid);

//########################################################################{
//# Locale diagnostics
//#

  proc locale.totalThreads() {
    var totalThreads: int;
    extern proc chpl_task_getNumThreads() : int;
    on this do totalThreads = chpl_task_getNumThreads();
    return totalThreads;
  }

  proc locale.idleThreads() {
    var idleThreads: int;
    extern proc chpl_task_getNumIdleThreads() : int;
    on this do idleThreads = chpl_task_getNumIdleThreads();
    return idleThreads;
  }

  proc locale.queuedTasks() {
    var queuedTasks: int;
    extern proc chpl_task_getNumQueuedTasks() : int;
    on this do queuedTasks = chpl_task_getNumQueuedTasks();
    return queuedTasks;
  }

  proc locale.runningTasks() {
    return this.runningTaskCnt();
  }

  proc locale.blockedTasks() {
    var blockedTasks: int;
    extern proc chpl_task_getNumBlockedTasks() : int;
    on this do blockedTasks = chpl_task_getNumBlockedTasks();
    return blockedTasks;
  }

//########################################################################}

  //
  // Increment and decrement the task count.
  //
  // Elsewhere in the module code we adjust the running task count
  // directly, but at least for now the runtime also needs to be
  // able to do so.  These functions support that.
  //
  pragma "insert line file info"
  export
  proc chpl_taskRunningCntInc() {
    here.runningTaskCntAdd(1);
  }

  pragma "insert line file info"
  export
  proc chpl_taskRunningCntDec() {
    here.runningTaskCntSub(1);
  }
}

// ChapelLocale.chpl
//
pragma "no use ChapelStandard"
module ChapelLocale {

  //
  // An abstract class. Specifies the required locale interface.
  // Each locale implementation must inherit from this class.
  //
  class locale {
    //- Constructor
    proc locale() { }

    //------------------------------------------------------------------------{
    //- Fields and accessors defined for all locale types (not overridable)
    //-

    // Every locale has a parent, except for the root locale.
    // The parent of the root locale is nil (by definition).
    const parent : locale;

    // To be removed from the required interface once legacy code is adjusted.
    const numCores: int;

    proc id : int return chpl_id();  // just the node part
    proc localeid : chpl_localeID_t return chpl_localeid(); // full locale id
    proc name return chpl_name();

    // This many tasks are running on this locale.
    //
    // Note handling runningTaskCounter <= 0 in runningTaskCnt().  The
    // references to it are done via here.runningTasks(), and currently
    // "here" means the requested locale, not the actual or concrete
    // one.  But this may be "any", and the running task count in it may
    // be zero, which in turn leads to overindexing the locales array
    // after the subtraction of 1 in locale.runningTasks(), below.  In
    // the long run we probably need to go through all the "here" refs
    // in the module code and differentiate more carefully which need
    // the requested locale and which need the concrete one.  But for
    // now, we make the assumption that all requests for the number of
    // running tasks want the count from the locale the calling task is
    // running on, so the minimum possible value must be 1.
    var runningTaskCounter : atomic int;

    inline proc runningTaskCntSet(val : int) {
      runningTaskCounter.write(val, memory_order_relaxed);
    }

    inline proc runningTaskCntAdd(val : int) {
      runningTaskCounter.add(val, memory_order_relaxed);
    }

    inline proc runningTaskCntSub(val : int) {
      runningTaskCounter.sub(val, memory_order_relaxed);
    }

    inline proc runningTaskCnt() {
      var rtc = runningTaskCounter.read(memory_order_relaxed);
      return if (rtc <= 0) then 1 else rtc;
    }
    //------------------------------------------------------------------------}

    //------------------------------------------------------------------------{
    //- User Interface Methods (overridable)
    //-

    // These are dynamically dispatched, so they can be overridden in
    // concrete classes.
    proc chpl_id() : int {
      _throwPVFCError();
      return -1;
    }

    proc chpl_localeid() : chpl_localeID_t {
      _throwPVFCError();
      extern const c_sublocid_none: chpl_sublocID_t;
      return chpl_buildLocaleID(-1:chpl_nodeID_t, c_sublocid_none);
    }

    proc chpl_name() : string {
      _throwPVFCError();
      return "";
    }

    // A useful default definition is provided (not pure virtual).
    proc readWriteThis(f) {
      f <~> name;
    }

    proc getChildCount() : int {
      _throwPVFCError();
      return 0;
    }

// Part of the required locale interface.
// Commented out because presently iterators are statically bound.
//    iter getChildIndices() : int {
//      for idx in this.getChildSpace do
//        yield idx;
//    }

    proc addChild(loc:locale)
    {
      _throwPVFCError();
    }

    proc getChild(idx:int) : locale {
      _throwPVFCError();
      return this;
    }

// Part of the required locale interface.
// Commented out because presently iterators are statically bound.
//    iter getChildren() : locale  {
//    _throwPVFCError();
//      yield 0;
//    }

    //------------------------------------------------------------------------}
  }

  class AbstractLocaleModel : locale {
    // This will be used for interfaces that will be common to all
    // (non-RootLocale) locale models
  }

  // rootLocale is declared to be of type locale rather than
  // RootLocale because doing so would make locales depend on arrays
  // which depend on locales, etc.
  //
  // If we had a way to express domain and array return types abstractly,
  // that problem would go away. <hilde>
  //
  // The rootLocale is private to each locale.  It cannot be
  // initialized until LocaleModel is initialized.  To disable this
  // replication, set replicateRootLocale to false.
  pragma "private" var rootLocale : locale = nil;
  config param replicateRootLocale = true;

  // The rootLocale needs to be initalized on all locales prior to
  // initializing the Locales array.  Unfortunately, the rootLocale
  // cannot be properly replicated until DefaultRectangular can be
  // initialized (as the private copies of the defaultDist are
  // needed).  We remedy this by initializing a rootLocale on locale 0
  // (called origRootLocale), and setting all locales' rootLocale to
  // origRootLocale.  Later, after DefaultRectangular can be
  // initialized, we create local copies of the rootLocale (and the
  // Locales array).  This is being done in the InitPrivateGlobals
  // module.
  //
  var origRootLocale : locale = nil;

  class AbstractRootLocale : locale {
    // These functions are used to establish values for Locales[] and
    // LocaleSpace -- an array of locales and its correponding domain
    // which are used as the default set of targetLocales in many
    // distributions.
    proc getDefaultLocaleSpace() {
      _throwPVFCError();
      const emptyLocaleSpace: domain(1) = {1..0};
      return emptyLocaleSpace;
    }

    proc getDefaultLocaleArray() {
      _throwPVFCError();
      const emptyLocaleSpace: domain(1) = {1..0};
      const emptyLocales: [emptyLocaleSpace] locale;
      return emptyLocales;
    }

    proc localeIDtoLocale(id : chpl_localeID_t) : locale {
      _throwPVFCError();
      return this;
    }

    // These iterators are to be used by RootLocale:init() to
    // initialize the LocaleModel.  The calling loop body cannot
    // contain any non-local code, since the rootLocale is not yet
    // initialized.
    iter initOnLocales() {
      if numLocales > 1 then
        halt("The locales must be initialized in parallel");
      for locIdx in (origRootLocale:RootLocale).getDefaultLocaleSpace() {
        yield locIdx;
        rootLocale = origRootLocale;
      }
    }

    // Since we are going on to all the locales here, we use the
    // opportunity to initialize any global private variables we
    // either need (e.g., defaultDist) or can do at this point in
    // initialization (e.g., rootLocale).
    iter initOnLocales(param tag: iterKind)
      where tag==iterKind.leader {
      // Simple locales barrier, see implementation below for notes
      var b: localesBarrier;
      var flags: [1..#numLocales-1] localesSignal;
      extern const c_sublocid_any: chpl_sublocID_t;
      coforall locIdx in 0..#numLocales /*ref(b)*/ {
        on __primitive("chpl_on_locale_num",
                       chpl_buildLocaleID(locIdx:chpl_nodeID_t,
                                          c_sublocid_any)) {
          chpl_defaultDistInitPrivate();
          yield locIdx;
          b.wait(locIdx, flags);
          chpl_rootLocaleInitPrivate(locIdx);
        }
      }
    }

    iter initOnLocales(param tag: iterKind, followThis)
      where tag==iterKind.follower {
      yield followThis;
    }
  }

  //
  // Simple locales barrier
  //
  // This sits outside of the abstract root locale definition above
  // because the compiler cannot resolve the constructor (known issues
  // with nested classes/records).  In addition, we cannot have the
  // flags array in the record, because the initCopy function needs
  // the dataPar* configs declared in ChapelDistribution.
  //
  // Each non-0 locale increments the count, and then waits on a
  // *local* atomic.  This is done by creating an array of type
  // class localesSignal, one per locale, and allocating each locale's
  // copy before updating the count.  Locale 0 waits for the others to
  // arrive and then set the signals to true.  We can't do anything
  // too complicated this early on, so we are using a for loop to
  // broadcast that we are done.
  class localesSignal {
    var s: atomic bool;
  }
  record localesBarrier {
    proc wait(locIdx, flags) {
      if locIdx==0 {
        // locale 0 has nothing else to do, so check flags
        while (true) {
          var count = 0;
          for f in flags do
            if f then count += 1;
          if count==numLocales-1 then break;
          atomic_fence(); // race condition without fence.
          chpl_task_yield();
        }
        // Let the others go
        for f in flags do
          f.s.testAndSet();
      } else {
        var f = new localesSignal();
        // expose my flag to locale 0
        flags[locIdx] = f;
        // wait (locally) for locale 0 to set my flag
        f.s.waitFor(true);
        // clean up
        delete f;
      }
    }
  }

  // This function is called in the LocaleArray module to initialize
  // the rootLocale.  It sets up the origRootLocale and also includes
  // set up of the each locale's LocaleModel via RootLocale:init().
  //
  // The init() function must use the initOnLocales() iterator above
  // to iterate in parallel over the locales to set up the LocaleModel
  // object.
  proc chpl_init_rootLocale() {
    origRootLocale = new RootLocale();
    (origRootLocale:RootLocale).init();
  }

  // This function sets up a private copy of rootLocale by replicating
  // origRootLocale and resets the Locales array to point to the local
  // copy on all but locale 0 (which is done in LocalesArray.chpl as
  // part of the declaration).
  proc chpl_rootLocaleInitPrivate(locIdx) {
    // Even when not replicating the rootLocale, we must temporarily
    // set the rootLocale to the original version on locale 0, because
    // the initialization below needs to get/set locale ids.
    rootLocale = origRootLocale;
    if replicateRootLocale && locIdx!=0 {
      // Create a new local rootLocale
      var newRootLocale = new RootLocale();
      // We don't want to be doing unnecessary ref count updates here
      // as they require additional tasks.  We know we don't need them
      // so tell the compiler to not insert them.
      pragma "no copy" pragma "no auto destroy"
      const origLocales => (origRootLocale:RootLocale).getDefaultLocaleArray();
      var origRL = origLocales._value.theData;
      var newRL = newRootLocale.getDefaultLocaleArray()._value.theData;
      // We must directly implement a bulk copy here, as the mechanisms
      // for doing so via a whole array assignment are not initialized
      // yet and copying element-by-element via a for loop is is costly.
      __primitive("chpl_comm_get",
                  __primitive("array_get", newRL, 0),
                  0 /* locale 0 */,
                  __primitive("array_get", origRL, 0), numLocales);
      // Set the rootLocale to the local copy
      rootLocale = newRootLocale;
    }
    if locIdx!=0 {
      // We mimic a private Locales array alias by using the move
      // primitive.
      __primitive("move", Locales,
                  (rootLocale:RootLocale).getDefaultLocaleArray());
    }
  }

  // We need a temporary value for "here" before the architecture is defined.
  // This is due to the fact that "here" is used for memory and task control
  // in setting up the architecture itself.
  // Its type should probably be renamed dummyLocale or something
  // representative.
  // The dummy locale provides system-default tasking and memory management.
  const dummyLocale = new locale();

  extern proc chpl_task_getRequestedSubloc(): chpl_sublocID_t;
  extern var chpl_nodeID: int(32);
  // Return the locale ID of the current locale
  inline proc here_id {
    return chpl_buildLocaleID(chpl_nodeID,chpl_task_getRequestedSubloc());
  }
  // Return the current locale
  inline proc here {
    return chpl_localeID_to_locale(here_id);
  }

  // Returns a wide pointer to the locale with the given id.
  proc chpl_localeID_to_locale(id : chpl_localeID_t) : locale {
    if rootLocale then
      return (rootLocale:AbstractRootLocale).localeIDtoLocale(id);
    else
      // For code prior to rootLocale initialization
      return dummyLocale;
  }

  pragma "insert line file info"
  extern proc chpl_memhook_malloc_pre(number:int, size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_malloc_post(ptr:opaque, number:int,
                                       size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_realloc_pre(ptr:object, size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_realloc_post(newPtr:opaque, ptr:object,
                                        size:int, md:int(16)): void;
  pragma "insert line file info"
  extern proc chpl_memhook_free_pre(ptr:opaque): void;
  extern proc chpl_memhook_md_num(): int(16);

  proc chpl_getPrivatizedCopy(type objectType, objectPid:int): objectType
    return __primitive("chpl_getPrivatizedClass", nil:objectType, objectPid);

//########################################################################{
//# Locale diagnostics
//#

  proc locale.totalThreads() {
    var totalThreads: int;
    extern proc chpl_task_getNumThreads() : int;
    on this do totalThreads = chpl_task_getNumThreads();
    return totalThreads;
  }

  proc locale.idleThreads() {
    var idleThreads: int;
    extern proc chpl_task_getNumIdleThreads() : int;
    on this do idleThreads = chpl_task_getNumIdleThreads();
    return idleThreads;
  }

  proc locale.queuedTasks() {
    var queuedTasks: int;
    extern proc chpl_task_getNumQueuedTasks() : int;
    on this do queuedTasks = chpl_task_getNumQueuedTasks();
    return queuedTasks;
  }

  proc locale.runningTasks() {
    return this.runningTaskCnt();
  }

  proc locale.blockedTasks() {
    var blockedTasks: int;
    extern proc chpl_task_getNumBlockedTasks() : int;
    on this do blockedTasks = chpl_task_getNumBlockedTasks();
    return blockedTasks;
  }

//########################################################################}

  //
  // Increment and decrement the task count.
  //
  // Elsewhere in the module code we adjust the running task count
  // directly, but at least for now the runtime also needs to be
  // able to do so.  These functions support that.
  //
  pragma "insert line file info"
  export
  proc chpl_taskRunningCntInc() {
    here.runningTaskCntAdd(1);
  }

  pragma "insert line file info"
  export
  proc chpl_taskRunningCntDec() {
    here.runningTaskCntSub(1);
  }
}


Use a for- instead of while-loop in isPrime() function.

/*
 * Math helper functions.
 */

// requires CHPL_GMP=gmp or system
use GMP;

// factorial(n) / (factorial(k) * factorial(n - k));
proc choose(n: int, k: int): BigInt {
  var result = factorial(n);
  result.div_q(Round.ZERO, result, factorial(k));
  result.div_q(Round.ZERO, result, factorial(n - k));
  return result;
}

// Yields proper divisors of n. Proper divisors of n are numbers less than n
// which divide evenly into n.
iter divisors(n) {
  var current = 1;

  while current <= n / 2 {
    yield current;

    do {
      current += 1;
    } while n % current != 0;
  }
}

// n!
proc factorial(n: int): BigInt {
  if n < 0 {
    halt("Error! non-positive factorial!?");
  } else if n == 0 {
    return new BigInt(1);
  } else {
    var result = factorial(n - 1);
    result.mul_ui(result, n: c_ulong);
    return result;
  }
}

// Print the Fibonacci number less than and equal to 4 million.
iter fibonacci(): BigInt {
  var current = new BigInt(1),
    next = new BigInt(1),
    tmp: BigInt;
  while true {
    yield current;

    current.add(current, next);
    current <=> next;
  }
}

// This is a fairly optimized version of calculating if a number is prime or
// not.
// See: http://stackoverflow.com/a/15285588
proc isPrime(number) {
  if number == 2 || number == 3 {
    return true;
  } else if number < 2 || number % 2 == 0 {
    return false;
  } else if number < 9 {
    return true;
  } else if number % 3 == 0 {
    return false;
  }

  const num = number.safeCast(uint),
    r = (number ** 0.5): uint;
  for f in 5..r by 6 {
    if num % f == 0 {
      return false;
    } else if num % (f + 2) == 0 {
      return false;
    }
  }
  return true;
}

// Yield successive permutations of the elements.
//
// Inspired by: https://docs.python.org/2/library/itertools.html#itertools.permutations
iter permutations(elements: [] string, length=-1): string {
  var n = elements.size;
  var r = length;

  // Default for r is length.
  if r == -1 then
    r = n;

  // If permutations are longer than number of elements, no permutations are
  // possible.
  if r > n then
    return;

  var indices: [{0..n-1}] int = [i in 0..n-1] i,
    cycles: [{0..n-1}] int = [c in n-r+1..n by -1] c;

  var perm = "";
  for i in indices[0.. # r] do
    perm += elements[i+1];
  yield perm;

  while true {
    var done = true;
    for i in 0..r-1 by -1 {
      cycles[i] -= 1;
      if cycles[i] == 0 {
        var rightInd = indices[i+1..],
          leftInd = indices[i..i];
        indices[i.. # rightInd.size] = rightInd;
        indices[(i + rightInd.size)..] = leftInd;
        cycles[i] = n - i;
      } else {
        indices[i] <=> indices[indices.size - cycles[i]];

        perm = "";
        for i in indices[0.. # r] do
          perm += elements[i+1];
        yield perm;

        done = false;
        break;
      }
    }

    // Did not yield a permutation, so we're done.
    if done then
      return;
  }
}

// Sieve of Eratosthenes
// Generates first k prime numbers.
//
// Inspired by: http://stackoverflow.com/a/568618
// And:         http://code.activestate.com/recipes/117119/#c4
iter primes(k: uint, param debug=false): uint {
  if k >= 1 then
    yield 2;

  // If only one prime requested, stop here.
  if k == 1 then
    return;

  var i: uint = 1,
    q: uint = 3;
  var D: domain(uint),
    A: [D] uint;

  while true {
    if ! D.member(q) {
      if debug then
        assert(isPrime(q));

      yield q;

      i += 1;
      if i == k then
        break;

      A[q * q] = 2 * q;
    } else {
      const p = A[q];
      var pq = p + q;
      while D.member(pq) do
        pq += p;
      A[pq] = p;
      D -= q;
    }
    q += 2;
  }
}

// Sieve of Eratosthenes up to max value.
// Generates prime numbers less than or equal to `m`.
iter primesUpTo(m: uint, param debug=false): uint {
  if m < 2 then
    halt("no primes less than 2");

  if m >= 2 then
    yield 2;

  if m == 2 then
    return;

  var q: uint = 3,
    D: domain(uint),
    A: [D] uint;

  while q <= m {
    if ! D.member(q) {
      if debug then
        assert(isPrime(q));

      yield q;

      A[q * q] = 2 * q;
    } else {
      const p = A[q];
      var pq = p + q;
      while D.member(pq) do
        pq += p;
      A[pq] = p;
      D -= q;
    }
    q += 2;
  }
}

/* // Yield the first `k` prime numbers as BigInts. */
/* iter bigPrimes(k: int(64)): BigInt { */
/*   var i = 0, */
/*     q = new BigInt(2); */
/*   var D: domain(BigInt); */
/*   var A: [D] domain(BigInt); */

/*   while true { */
/*     if ! D.member(q) { */
/*       yield q; */
/*       i += 1; */
/*       var qSquared = new BigInt(); */
/*       qSquared.mul(q, q); */
/*       D += qSquared; */
/*       A[qSquared] += q; */
/*     } else { */
/*       for p in A[q] { */
/*         var pq = new BigInt(); */
/*         pq.add(p, q); */
/*         D += pq; */
/*         A[pq] = p; */
/*       } */
/*       D.remove(q); */
/*     } */
/*     if i == k { */
/*       break; */
/*     } */
/*     q += 1; */
/*   } */
/* } */

// Returns sum of individual digits of n. n is serialized as string, then each
// character is cast to an int, and added together.
proc sumDigits(n) {
  var nStr = n: string,
    result = 0;
  for i in 1..nStr.length {
    result += nStr.substring(i): int;
  }
  return result;
}

// Returns sum of all divisors of n.
proc sumDivisors(n) {
  var sum = 0;
  for i in divisors(n) {
    sum += i;
  }
  return sum;
}

/*
 * Math helper functions.
 */

// requires CHPL_GMP=gmp or system
use GMP;

// factorial(n) / (factorial(k) * factorial(n - k));
proc choose(n: int, k: int): BigInt {
  var result = factorial(n);
  result.div_q(Round.ZERO, result, factorial(k));
  result.div_q(Round.ZERO, result, factorial(n - k));
  return result;
}

// Yields proper divisors of n. Proper divisors of n are numbers less than n
// which divide evenly into n.
iter divisors(n) {
  var current = 1;

  while current <= n / 2 {
    yield current;

    do {
      current += 1;
    } while n % current != 0;
  }
}

// n!
proc factorial(n: int): BigInt {
  if n < 0 {
    halt("Error! non-positive factorial!?");
  } else if n == 0 {
    return new BigInt(1);
  } else {
    var result = factorial(n - 1);
    result.mul_ui(result, n: c_ulong);
    return result;
  }
}

// Print the Fibonacci number less than and equal to 4 million.
iter fibonacci(): BigInt {
  var current = new BigInt(1),
    next = new BigInt(1),
    tmp: BigInt;
  while true {
    yield current;

    current.add(current, next);
    current <=> next;
  }
}

// This is a fairly optimized version of calculating if a number is prime or
// not.
// See: http://stackoverflow.com/a/15285588
proc isPrime(number) {
  if number == 2 || number == 3 {
    return true;
  } else if number < 2 || number % 2 == 0 {
    return false;
  } else if number < 9 {
    return true;
  } else if number % 3 == 0 {
    return false;
  }

  const num = number.safeCast(uint);
  var r = (number ** 0.5): uint,
    f: uint = 5;
  while (f <= r) {
    if num % f == 0 {
      return false;
    } else if num % (f + 2) == 0 {
      return false;
    }
    f += 6;
  }
  return true;
}

// Yield successive permutations of the elements.
//
// Inspired by: https://docs.python.org/2/library/itertools.html#itertools.permutations
iter permutations(elements: [] string, length=-1): string {
  var n = elements.size;
  var r = length;

  // Default for r is length.
  if r == -1 then
    r = n;

  // If permutations are longer than number of elements, no permutations are
  // possible.
  if r > n then
    return;

  var indices: [{0..n-1}] int = [i in 0..n-1] i,
    cycles: [{0..n-1}] int = [c in n-r+1..n by -1] c;

  var perm = "";
  for i in indices[0.. # r] do
    perm += elements[i+1];
  yield perm;

  while true {
    var done = true;
    for i in 0..r-1 by -1 {
      cycles[i] -= 1;
      if cycles[i] == 0 {
        var rightInd = indices[i+1..],
          leftInd = indices[i..i];
        indices[i.. # rightInd.size] = rightInd;
        indices[(i + rightInd.size)..] = leftInd;
        cycles[i] = n - i;
      } else {
        indices[i] <=> indices[indices.size - cycles[i]];

        perm = "";
        for i in indices[0.. # r] do
          perm += elements[i+1];
        yield perm;

        done = false;
        break;
      }
    }

    // Did not yield a permutation, so we're done.
    if done then
      return;
  }
}

// Sieve of Eratosthenes
// Generates first k prime numbers.
//
// Inspired by: http://stackoverflow.com/a/568618
// And:         http://code.activestate.com/recipes/117119/#c4
iter primes(k: uint, param debug=false): uint {
  if k >= 1 then
    yield 2;

  // If only one prime requested, stop here.
  if k == 1 then
    return;

  var i: uint = 1,
    q: uint = 3;
  var D: domain(uint),
    A: [D] uint;

  while true {
    if ! D.member(q) {
      if debug then
        assert(isPrime(q));

      yield q;

      i += 1;
      if i == k then
        break;

      A[q * q] = 2 * q;
    } else {
      const p = A[q];
      var pq = p + q;
      while D.member(pq) do
        pq += p;
      A[pq] = p;
      D -= q;
    }
    q += 2;
  }
}

// Sieve of Eratosthenes up to max value.
// Generates prime numbers less than or equal to `m`.
iter primesUpTo(m: uint, param debug=false): uint {
  if m < 2 then
    halt("no primes less than 2");

  if m >= 2 then
    yield 2;

  if m == 2 then
    return;

  var q: uint = 3,
    D: domain(uint),
    A: [D] uint;

  while q <= m {
    if ! D.member(q) {
      if debug then
        assert(isPrime(q));

      yield q;

      A[q * q] = 2 * q;
    } else {
      const p = A[q];
      var pq = p + q;
      while D.member(pq) do
        pq += p;
      A[pq] = p;
      D -= q;
    }
    q += 2;
  }
}

/* // Yield the first `k` prime numbers as BigInts. */
/* iter bigPrimes(k: int(64)): BigInt { */
/*   var i = 0, */
/*     q = new BigInt(2); */
/*   var D: domain(BigInt); */
/*   var A: [D] domain(BigInt); */

/*   while true { */
/*     if ! D.member(q) { */
/*       yield q; */
/*       i += 1; */
/*       var qSquared = new BigInt(); */
/*       qSquared.mul(q, q); */
/*       D += qSquared; */
/*       A[qSquared] += q; */
/*     } else { */
/*       for p in A[q] { */
/*         var pq = new BigInt(); */
/*         pq.add(p, q); */
/*         D += pq; */
/*         A[pq] = p; */
/*       } */
/*       D.remove(q); */
/*     } */
/*     if i == k { */
/*       break; */
/*     } */
/*     q += 1; */
/*   } */
/* } */

// Returns sum of individual digits of n. n is serialized as string, then each
// character is cast to an int, and added together.
proc sumDigits(n) {
  var nStr = n: string,
    result = 0;
  for i in 1..nStr.length {
    result += nStr.substring(i): int;
  }
  return result;
}

// Returns sum of all divisors of n.
proc sumDivisors(n) {
  var sum = 0;
  for i in divisors(n) {
    sum += i;
  }
  return sum;
}


Replace stray tabs with spaces for my own readability.

// MPlot.chpl
//
// Plotting routines for the Mandelbrot exercise.
//

//
// Types of image files: Black & White, Greyscale, or Color.  The integer
// values chosen here represent the header encoding for PBM, PGM, and
// PPM files, respectively.
//
enum imageType { bw=1, grey, color};

//
// User-specifiable image type to produce; defaults to color
//
config const imgType = imageType.color;

//
// Base filename for output file
//
config const filename = "mandelbrot";

//
// Maximum color depth for image file
//
config const maxColor = 15;

//
// This routine will plot a rectangular array with any coordinate
// system to the output file specified by the filename config const.
//
proc plot(NumSteps:[]) where NumSteps.rank == 2 {
  //
  // An associative domain+array mapping from the image type enum to
  // file extensions.  Note that associative domains over enumerated
  // types are fully-populated by default.
  //
  const extensionSpace: domain(imageType);
  const extensions: [extensionSpace] string = (".pbm", ".pgm", ".ppm");

  //
  // Compute the full output filename and open the file and a writer
  // to it.
  //
  const outfilename = filename + extensions(imgType);
  const outfile = open(outfilename, iomode.cw).writer();

  //
  // Plot the image to the file (could also pass stdout in as the file...)
  //
  plotToFile(NumSteps, outfile);

  //
  // Close file and tell user what it was called
  //
  outfile.close();
  writeln("Wrote output to ", outfilename);
}

//
// This is a helper routine that plots to an outfile 'channel'; this
// channel could simply be stdout or some other channel instead of an
// actual file channel.
//
proc plotToFile(NumSteps: [?Dom], outfile) {
  //
  // Capture the number of rows and columns in the array to be plotted
  //
  const rows = Dom.dim(1).length,
        cols = Dom.dim(2).length;

  //
  // Write the image header corresponding to the file type
  //
  outfile.writeln("P", imgType:int);
  outfile.writeln(rows, " ", cols);

  //
  // For file types other than greyscale, we have to write the number
  // of colors we're using
  //
  if (imgType != imageType.bw) then
    outfile.writeln(maxColor);

  //
  // compute the maximum number of steps that were taken, just in case
  // it wasn't the user-supplied cutoff.
  //
  const maxSteps = max reduce NumSteps;

  //
  // Write the output data.  Though verbose, we use three loop nests
  // here to avoid extra conditionals in the inner loop.
  //
  select (imgType) {
    when imageType.bw {
      for i in Dom.dim(1) {
        for j in Dom.dim(2) {
          outfile.write(if NumSteps[i,j] then 0 else 1, " ");
        }
        outfile.writeln();
      }
    }

    when imageType.grey {
      for i in Dom.dim(1) {
        for j in Dom.dim(2) {
          outfile.write((maxColor*NumSteps[i,j])/maxSteps, " ");
        }
        outfile.writeln();
      }
    }

    when imageType.color {
      for i in Dom.dim(1) {
        for j in Dom.dim(2) {
          outfile.write((maxColor*NumSteps[i,j])/maxSteps, " ", 0, " ", 0, " ");
        }
        outfile.writeln();
      }
    }
  }
}


I spoke too soon earlier, having believed that if
logN==10 worked other problem sizes would work as
well.  Turned out there were still some subtle bugs
that only showed up at larger and smaller problem
sizes.  I believe things are now working (which is
to say, the error keeps pace with John's for
logN=3..17 on a ramp of input values and is less
than the tolerance for a random input).

The bugs that this version fixes are:
  - strided degenerate sequences are correct again
    (thanks to Steve's fix)
  - two loops were strided incorrectly and are now
    fixed -- one by m2 rather than 2; the other by
    1 rather than 2.
  - the code that manually cleans up "odd" problem
    sizes is now correct.

Now I should be in a position to clean this code up
tomorrow.

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

use Random;
use Timers;

// BLC: clean up complex arithmetic

config const logN = 5;
config const debug = false;
config const printTiming = false;

// BLC: This call needs to go away
def resetA(A: [?D] complex) {
  /*
  [i in D] {
    A(i).real = 2*i;
    A(i).imag = A(i).real + 1.0;
  }
  */
}

def writeA(A) {
  write("a[] =");
  for i in A.domain do write(A(i):" %g %g");
  writeln();
}

def main() {
  const N = 1 << logN;
  const gflop = 5.0 * N * logN / 1000000000.0;

  const EPS = 2.0 ** -51.0;
  const THRESHOLD = 16.0;

  const D = [0..N-1];
  const DW = [0..N/4-1];

  var A: [D] complex;
  var B: [D] complex;
  var W: [DW] complex;

  [i in D] {
    A(i).real = 2*i;
    A(i).imag = A(i).real + 1.0;
  }

  fillRandomVec(A);
  if debug then writeA(A);

  B = A;                       // save A for verification step

  twiddles(W);
  W = bitReverse(W);

  // BLC: would like this to be:
  //  A.imag = -A.imag;            // conjugate data
  A = conjg(A);
  if debug then writeA(A);

  var fftTimer: Timer;
  fftTimer.start();

  A = bitReverse(A);

  // BLC: This line needs to go away
  resetA(A);

  // BLC: Why was John timing the second version?
  if debug then writeA(A);
  dfft(A, W);
  if debug then writeA(A);
  fftTimer.stop();

  // BLC: originally wrote this as:
  //  A.real =  A.real / N;        // conjugate and scale data
  //  A.imag = -A.imag / N;
  // but it doesn't currently work.  Maybe that's OK because the
  // rewrite is pretty clean

  // BLC: This line only works if /(complex,real) is implemented
  // correctly:
  A = conjg(A) / N;

  if debug then writeA(A);

  A = bitReverse(A);
  dfft(A, W);
  if debug then writeA(A);

  var time = fftTimer.value;

  if debug then {
    var C: [D] float = sqrt((B.real - A.real)**2 + (B.imag - A.imag)**2);
    writeln("error[] =", C);
  }
  var maxerr = max reduce sqrt((B.real - A.real)**2 + (B.imag - A.imag)**2);

  maxerr = maxerr / logN / EPS;

  var status = if (maxerr < THRESHOLD) then "SUCCESS" else "FAILURE";
  writeln(status, ", error = ", maxerr);
  writeln("\n");

  writeln("N      = ", N);
  if (printTiming) {
    writeln("Time   = ", time);
    writeln("GFlops = ", gflop / time);
  }
}

def twiddles(W: [?WD] complex) {
  const n = WD(1).length;
  const delta = 2.0 * atan(1.0) / n;

  W(0) = 1.0;
  W(n/2) = let cosDeltaN = cos(delta * n/2)
            in (cosDeltaN, cosDeltaN):complex;
  forall i in 1..n/2-1 {
    const x = cos(delta*i);
    const y = sin(delta*i);
    W(i)     = (x, y):complex;
    W(n - i) = (y, x):complex;
  }
}

// Check what the NSA supports for bit reversal
// rename this?
def bitReverse(W: [?WD] complex) {  // BLC: would be nice to drop complex?
  const n = WD(1).length;
  const reverse = lg(n);
  var V: [WD] complex;
  /* BLC: could we do this as a permutation instead?
  var P: [i in WD] index(WD) = i;
  bitReverse(P);
  V(P) = W;
  */
  forall i in WD {  // BLC: could this be a uint domain?
    // BLC: what does this bitReverse function do with high-order bits?
    // BLC: could this be rewritten as one line?
    var ndx = bitReverse(i:uint(64), numBits=reverse);
    V(ndx:int) = W(i); // BLC: unfortunate cast
  }
  // BLC: W = V would be preferable
  return V;
}

def dfft(A: [?AD] complex, W) {
  /*
  if (debug) {
    write("w[] =");
    for i in W.domain do write(W(i):" %g %g");
    writeln();
  }
  */

  if debug then writeln("cft1st();");
  cft1st(A, W);

  var l = 4;
  var lasti = 2;
  for i in 4..logN/2 by 2 {
    if debug then writeln("cftmd1(l=", l, ", A, W);");
    cftmd1(l, A, W);
    l *= 4;
    lasti = i;
  }
  //  writeln("lasti+2 = ", lasti+2, " logN-1 = ", logN-1);
  for i in lasti+2..logN-1 by 2 {
    if debug then writeln("i=", i, "  cftmd2(l=", l, ", A, W);");
    cftmd2(l, A, W);
    if debug then writeln("i=", i, "  returned from cftmd2();");
    l *= 4;
  }

  const n = AD(1).length;
  if ((l << 2) == n) {
    if debug then writeln("l << 2 == n");
    forall j in 0..l-1 {
      resetA(A);
      butterfly(1.0, 1.0, 1.0, A[j..j+3*l by l]);
    }
  } else {
    if debug then writeln("l << 2 != n");
    forall j in 0..l-1 {
      var a = A(j);
      var b = A(j+l);
      A(j) = a + b;
      A(j+l) = a - b;
    }
  }
}

def cft1st(A, W) {
  const n = A.domain(1).length;
  var x0 = A(0) + A(1);
  var x1 = A(0) - A(1);
  var x2 = A(2) + A(3);
  var x3 = A(2) - A(3);

  const wk1r = W(1).real;

  if debug then writeA(A);

  A(0) = x0 + x2;
  A(2) = x0 - x2;
  A(1) = (x1.real - x3.imag, x1.imag + x3.real):complex;
  A(3) = (x1.real + x3.imag, x1.imag - x3.real):complex;

  x0 = A(4) + A(5);
  x1 = A(4) - A(5);
  x2 = A(6) + A(7);
  x3 = A(6) - A(7);
  A(4) = x0 + x2;
  A(6) = (x2.imag - x0.imag, x0.real - x2.real):complex;
  x0 = (x1.real - x3.imag, x1.imag + x3.real):complex;
  A(5) = wk1r * (x0.real - x0.imag, x0.real + x0.imag):complex;
  x0 = (x3.imag + x1.real, x3.real - x1.imag):complex;
  A(7) = wk1r * (x0.imag - x0.real, x0.imag + x0.real):complex;

  if debug then writeln("  // computes first 8 complexes manually");

  if debug then writeA(A);

  var k1 = 1;
  forall j in 8..n-1 by 8 {
    var wk2 = W(k1);
    var wk1 = W(2*k1);
    var wk3 = (wk1.real - 2* wk2.imag * wk1.imag,
               2 * wk2.imag * wk1.real - wk1.imag):complex;

    resetA(A);
    butterfly(wk1, wk2, wk3, A[j..j+3]);

    //    writeln("accessing: ", 2*k1+1);
    wk1 = W(2*k1+1);
    wk3 = (wk1.real - 2*wk2.real * wk1.imag,
           2*wk2.real * wk1.real - wk1.imag):complex;
    wk2 = (-wk2.imag, wk2.real): complex;
    resetA(A);
    butterfly(wk1, wk2, wk3, A[j+4..j+7]);

    k1 += 1;
  }
}

def cftmd0(l, A, W) {
  var wk1r = W(1).real;
  const m = l << 2;

  //  if debug then writeA(A);

  forall j in 0..l-1 {
    resetA(A);
    butterfly(1.0, 1.0, 1.0, A[j..j+3*l by l]);
  }

  forall j in m..l+m-1 {
    resetA(A);
    butterfly((wk1r, wk1r):complex, 1.0i, (-wk1r, wk1r):complex,
              A[j..j+3*l by l]);
  }
}

def cftmd1(l, A, W) {
  const m = l << 2;
  const m2 = 2*m;
  const n = A.domain(1).length;

  if debug then writeln("  cftmd0(l=", l, ", A, W);");
  cftmd0(l, A, W);
  if debug then writeln("  returned from cftmd0();");
  var k1 = 1;
  forall k in m2..n-1 by m2 {
    var wk2 = W[k1];
    var wk1 = W[2*k1];
    var wk3 = (wk1.real - 2 * wk2.imag * wk1.imag,
               2 * wk2.imag * wk1.real - wk1.imag):complex;
    for j in k..l+k-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    wk1 = W[2*k1+1];
    wk3 = (wk1.real - 2 * wk2.real * wk1.imag,
           2 * wk2.real * wk1.real - wk1.imag):complex;

    for j in k+m..l+k+m-1 {
      resetA(A);
      butterfly(wk1, (-wk2.imag, wk2.real):complex, wk3, A[j..j+3*l by l]);
    }

    k1 += 1;
  }
}

def cftmd2(l, A, W) {
  var m = l << 2;
  var m2 = 2*m;
  const n = A.domain(1).length;

  if debug then writeln("  cftmd0(l=", l, ", A, W);");
  cftmd0(l, A, W);
  if debug then writeln("  returned from cftmd0();");
  if (m2 >= n) return;
  if (m2 >= n / 8) {
    if debug then writeln("  cftmd21(l=", l, ", A, W);");
    cftmd21(l, A, W);
    if debug then writeln("  returned from cftmd21();");
    return;
  }

  forall j in 0..l-1 {
    var k1 = 1;  // BLC: zip this in
    forall k in m2..n-1 by m2 {
      var wk2 = W[k1];
      var wk1 = W[k1 + k1];
      var wk3 = (wk1.real - 2*wk2.imag * wk1.imag,
                 2 * wk2.imag * wk1.real - wk1.imag): complex;
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j+k..j+k+3*l by l]);

      k1 += 1;
    }

    k1 = 1;  // BLC: zip this in
    forall k in m2..n-1 by m2 {
      var wk2 = W[k1];
      var wk1 = W[2*k1 + 1];
      var wk3 = (wk1.real - 2*wk2.real * wk1.imag,
                 2*wk2.real * wk1.real - wk1.imag): complex;
      wk2 = (-wk2.imag, wk2.real): complex;

      resetA(A);
      butterfly(wk1, wk2, wk3, A[j+k+m..j+k+m+3*l by l]);

      k1 += 1;
    }
  }
}

def cftmd21(l, A, W) {
  const n = A.domain(1).length;
  var m = l << 2;
  var m2 = 2*m;
  var m3 = 3*m;

  var k1 = 1;
  for k in m2..n-1 by m2 {
    var wk2 = W[k1];
    var wk1 = W[2*k1];
    var wk3 = (wk1.real - 2*wk2.imag * wk1.imag,
               2* wk2.imag * wk1.real - wk1.imag): complex;

    forall j in k..k+l-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    wk1 = W[2*k1 + 1];
    wk3 = (wk1.real - 2*wk2.real * wk1.imag,
           2*wk2.real * wk1.real - wk1.imag): complex;
    wk2 = (-wk2.imag, wk2.real): complex;

    forall j in k+m..k+m+l-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    k1 += 1;
  }
}

def butterfly(wk1: complex, wk2: complex, wk3: complex,
              inout A:[1..4] complex) {
  var x0 = A[1] + A[2];
  var x1 = A[1] - A[2];
  var x2 = A[3] + A[4];
  var x3 = A[3] - A[4];

  if (debug) {
    writeln("    a=", A[1]:"{%g,%g}", ", b=", A[2]:"{%g,%g}", ", c=", A[3]:"{%g,%g}", ", d=", A[4]:"{%g,%g}");
    writeln("      wk1=", wk1:"%g + %gi");
    writeln("      wk2=", wk2:"%g + %gi");
    writeln("      wk3=", wk3:"%g + %gi");

    writeln();

    writeln("      x0=", x0:"%g %g");
    writeln("      x1=", x1:"%g %g");
    writeln("      x2=", x2:"%g %g");
    writeln("      x3=", x3:"%g %g");
  }

  A[1] = x0 + x2;
  x0 -= x2;
  if debug then writeln("      x0=", x0:"%g %g");
  A[3] = wk2 * x0;
  x0 = (x1.real - x3.imag, x1.imag + x3.real):complex;
  if debug then writeln("      x0=", x0:"%g %g");
  A[2] = wk1 * x0;
  x0 = (x1.real + x3.imag, x1.imag - x3.real):complex;
  if debug then writeln("      x0=", x0:"%g %g");
  A[4] = wk3 * x0;

  if (debug) {
    writeln("    a=", A[1]:"{%g,%g}", ", b=", A[2]:"{%g,%g}", ", c=", A[3]:"{%g,%g}", ", d=", A[4]:"{%g,%g}\n");
  }
}

use Random;
use Timers;

// BLC: clean up complex arithmetic

config const logN = 5;
config const debug = false;

// BLC: This call needs to go away
def resetA(A: [?D] complex) {
  /*
  [i in D] {
    A(i).real = 2*i;
    A(i).imag = A(i).real + 1.0;
  }
  */
}

def writeA(A) {
  write("a[] =");
  for i in A.domain do write(A(i):" %g %g");
  writeln();
}

def main() {
  const N = 1 << logN;
  const gflop = 5.0 * N * logN / 1000000000.0;

  const EPS = 2.0 ** -51.0;
  const THRESHOLD = 16.0;

  const D = [0..N-1];
  const DW = [0..N/4-1];

  var A: [D] complex;
  var B: [D] complex;
  var W: [DW] complex;

  [i in D] {
    A(i).real = 2*i;
    A(i).imag = A(i).real + 1.0;
  }
  writeA(A);

  //  fillRandomVec(A);

  B = A;                       // save A for verification step

  twiddles(W);
  W = bitReverse(W);

  // BLC: would like this to be:
  //  A.imag = -A.imag;            // conjugate data
  A = conjg(A);
  writeA(A);

  var fftTimer: Timer;
  fftTimer.start();

  A = bitReverse(A);

  // BLC: This line needs to go away
  resetA(A);

  if debug then writeA(A);

  // BLC: Why was John timing the second version?
  writeA(A);
  dfft(A, W);
  writeA(A);
  fftTimer.stop();

  if (debug) then halt();

  // BLC: originally wrote this as:
  //  A.real =  A.real / N;        // conjugate and scale data
  //  A.imag = -A.imag / N;
  // but it doesn't currently work.  Maybe that's OK because the
  // rewrite is pretty clean

  // BLC: This line only works if /(complex,real) is implemented
  // correctly:
  A = conjg(A) / N;

  writeA(A);

  A = bitReverse(A);
  dfft(A, W);
  writeA(A);

  var time = fftTimer.value;

  var maxerr = max reduce sqrt((B.real - A.real)**2 + (B.imag - A.imag)**2);

  maxerr = maxerr / logN / EPS;

  var status = if (maxerr < THRESHOLD) then "SUCCESS" else "FAILURE";
  writeln(status, ", error = ", maxerr);
  writeln("\n");

  writeln("N      = ", N);
  writeln("Time   = ", time);
  writeln("GFlops = ", gflop / time);
}

def twiddles(W: [?WD] complex) {
  const n = WD(1).length;
  const delta = 2.0 * atan(1.0) / n;

  W(0) = 1.0;
  W(n/2) = let cosDeltaN = cos(delta * n/2)
            in (cosDeltaN, cosDeltaN):complex;
  forall i in 1..n/2-1 {
    const x = cos(delta*i);
    const y = sin(delta*i);
    W(i)     = (x, y):complex;
    W(n - i) = (y, x):complex;
  }
}

// Check what the NSA supports for bit reversal
// rename this?
def bitReverse(W: [?WD] complex) {  // BLC: would be nice to drop complex?
  const n = WD(1).length;
  const reverse = lg(n);
  var V: [WD] complex;
  /* BLC: could we do this as a permutation instead?
  var P: [i in WD] index(WD) = i;
  bitReverse(P);
  V(P) = W;
  */
  forall i in WD {  // BLC: could this be a uint domain?
    // BLC: what does this bitReverse function do with high-order bits?
    // BLC: could this be rewritten as one line?
    var ndx = bitReverse(i:uint(64), numBits=reverse);
    V(ndx:int) = W(i); // BLC: unfortunate cast
  }
  // BLC: W = V would be preferable
  return V;
}

def dfft(A: [?AD] complex, W) {
  if (debug) {
    write("w[] =");
    for i in W.domain do write(W(i):" %g %g");
    writeln();
  }

  if debug then writeln("cft1st();");
  cft1st(A, W);

  var l = 4;
  var lasti = 2;
  for i in 4..logN/2 by 2 {
    if debug then writeln("cftmd1(l=", l, ", A, W);");
    cftmd1(l, A, W);
    l *= 4;
    lasti = i;
  }
  for i in lasti+2..logN-1 by 2 {
    if debug then writeln("i=", i, "  cftmd2(l=", l, ", A, W);");
    cftmd2(l, A, W);
    l *= 4;
  }

  const n = AD(1).length;
  if ((l << 2) == n) {
    if debug then writeln("l << 2 == n");
    forall j in 0..l-1 {
      resetA(A);
      butterfly(1.0, 1.0, 1.0, A[j..j+3*l by l]);
    }
  } else {
    if debug then writeln("l << 2 != n");
    forall j in 0..l-1 {
      var x0 = A(j).real;
      var x1 = A(j).imag;
      A(j).real = x0+x1;
      A(j).imag = x0-x1;
    }
  }
}

def cft1st(A, W) {
  const n = A.domain(1).length;
  var x0 = A(0) + A(1);
  var x1 = A(0) - A(1);
  var x2 = A(2) + A(3);
  var x3 = A(2) - A(3);

  const wk1r = W(1).real;

  A(0) = x0 + x2;
  A(2) = x0 - x2;
  A(1) = (x1.real - x3.imag, x1.imag + x3.real):complex;
  A(3) = (x1.real + x3.imag, x1.imag - x3.real):complex;

  x0 = A(4) + A(5);
  x1 = A(4) - A(5);
  x2 = A(6) + A(7);
  x3 = A(6) - A(7);
  A(4) = x0 + x2;
  A(6) = (x2.imag - x0.imag, x0.real - x2.real):complex;
  x0 = (x1.real - x3.imag, x1.imag + x3.real):complex;
  A(5) = wk1r * (x0.real - x0.imag, x0.real + x0.imag):complex;
  x0 = (x3.imag + x1.real, x3.real - x1.imag):complex;
  A(7) = wk1r * (x0.imag - x0.real, x0.imag + x0.real):complex;

  if debug then writeln("  // computes first 8 complexes manually");

  var k1 = 1;
  forall j in 8..n-1 by 8 {
    var wk2 = W(k1);
    var wk1 = W(2*k1);
    var wk3 = (wk1.real - 2* wk2.imag * wk1.imag,
               2 * wk2.imag * wk1.real - wk1.imag):complex;

    resetA(A);
    butterfly(wk1, wk2, wk3, A[j..j+3]);

    //    writeln("accessing: ", 2*k1+1);
    wk1 = W(2*k1+1);
    wk3 = (wk1.real - 2*wk2.real * wk1.imag,
           2*wk2.real * wk1.real - wk1.imag):complex;
    wk2 = (-wk2.imag, wk2.real): complex;
    resetA(A);
    butterfly(wk1, wk2, wk3, A[j+4..j+7]);

    k1 += 1;
  }
}

def cftmd0(l, A, W) {
  var wk1r = W(1).real;
  const m = l << 2;

  if debug then writeA(A);

  forall j in 0..l-1 {
    resetA(A);
    butterfly(1.0, 1.0, 1.0, A[j..j+3*l by l]);
  }

  forall j in m..l+m-1 {
    resetA(A);
    butterfly((wk1r, wk1r):complex, 1.0i, (-wk1r, wk1r):complex,
              A[j..j+3*l by l]);
  }
}

def cftmd1(l, A, W) {
  const m = l << 2;
  const m2 = 2*m;
  const n = A.domain(1).length;

  if debug then writeln("  cftmd0(l=", l, ", A, W);");
  cftmd0(l, A, W);
  var k1 = 1;
  forall k in m2..n-1 by m2 {
    var wk2 = W[k1];
    var wk1 = W[2*k1];
    var wk3 = (wk1.real - 2 * wk2.imag * wk1.imag,
               2 * wk2.imag * wk1.real - wk1.imag):complex;
    for j in k..l+k-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    wk1 = W[2*k1+1];
    wk3 = (wk1.real - 2 * wk2.real * wk1.imag,
           2 * wk2.real * wk1.real - wk1.imag):complex;

    for j in k+m..l+k+m-1 {
      resetA(A);
      butterfly(wk1, (-wk2.imag, wk2.real):complex, wk3, A[j..j+3*l by l]);
    }

    k1 += 1;
  }
}

def cftmd2(l, A, W) {
  var m = l << 2;
  var m2 = 2*m;
  const n = A.domain(1).length;

  if debug then writeln("  cftmd0(l=", l, ", A, W);");
  cftmd0(l, A, W);
  if (m2 >= n) return;
  if (m2 >= n / 8) {
    if debug then writeln("  cftmd21(l=", l, ", A, W);");
    cftmd21(l, A, W);
    return;
  }

  forall j in 0..l-1 by 2 {
    var k1 = 1;  // BLC: zip this in
    forall k in m2..n-1 by m2 {
      var wk2 = W[k1];
      var wk1 = W[k1 + k1];
      var wk3 = (wk1.real - 2*wk2.imag * wk1.imag,
                 2 * wk2.imag * wk1.real - wk1.imag): complex;
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j+k..j+k+3*l by l]);

      k1 += 1;
    }

    k1 = 1;  // BLC: zip this in
    forall k in m2..n-1 by 2 {
      var wk2 = W[k1];
      var wk1 = W[2*k1 + 1];
      var wk3 = (wk1.real - 2*wk2.real * wk1.imag,
                 2*wk2.real * wk1.real - wk1.imag): complex;
      wk2 = (-wk2.imag, wk2.real): complex;

      resetA(A);
      butterfly(wk1, wk2, wk3, A[j+k+m..j+k+m+3*l by l]);

      k1 += 1;
    }
  }
}

def cftmd21(l, A, W) {
  const n = A.domain(1).length;
  var m = l << 2;
  var m2 = 2*m;
  var m3 = 3*m;

  var k1 = 1;
  for k in m2..n-1 by m2 {
    var wk2 = W[k1];
    var wk1 = W[2*k1];
    var wk3 = (wk1.real - 2*wk2.imag * wk1.imag,
               2* wk2.imag * wk1.real - wk1.imag): complex;

    forall j in k..k+l-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    wk1 = W[2*k1 + 1];
    wk3 = (wk1.real - 2*wk2.real * wk1.imag,
           2*wk2.real * wk1.real - wk1.imag): complex;
    wk2 = (-wk2.imag, wk2.real): complex;

    forall j in k+m..k+m+l-1 {
      resetA(A);
      butterfly(wk1, wk2, wk3, A[j..j+3*l by l]);
    }

    k1 += 1;
  }
}

def butterfly(wk1: complex, wk2: complex, wk3: complex,
              inout A:[1..4] complex) {
  var x0 = A[1] + A[2];
  var x1 = A[1] - A[2];
  var x2 = A[3] + A[4];
  var x3 = A[3] - A[4];

  if (debug) {
    writeln("    a=", A[1]:"{%g,%g}", ", b=", A[2]:"{%g,%g}", ", c=", A[3]:"{%g,%g}", ", d=", A[4]:"{%g,%g}");
    writeln("      wk1=", wk1:"%g + %gi");
    writeln("      wk2=", wk2:"%g + %gi");
    writeln("      wk3=", wk3:"%g + %gi");
  }

  A[1] = x0 + x2;
  x0 -= x2;
  A[3] = wk2 * x0;
  x0 = (x1.real - x3.imag, x1.imag + x3.real):complex;
  A[2] = wk1 * x0;
  x0 = (x1.real + x3.imag, x1.imag - x3.real):complex;
  A[4] = wk3 * x0;

  if (debug) {
    writeln("    a=", A[1]:"{%g,%g}", ", b=", A[2]:"{%g,%g}", ", c=", A[3]:"{%g,%g}", ", d=", A[4]:"{%g,%g}\n");
  }
}