<commit_before>//
// Use standard modules for vector and matrix Norms, Random numbers
// and Timing routines
//
use Norm, Random, Time;
use UtilMath;

//
// Use the user module for computing HPCC problem sizes
//
//use HPCCProblemSize;

//
// Use the distributions we need for this computation
//
use DimensionalDist2D, ReplicatedDim, BlockCycDim;
use ReplicatedDist;

//
// The number of matrices and the element type of those matrices
// indexType can be int or int(64), elemType can be real or complex
//
const numMatrices = 1;
type indexType = int,
     elemType = real;

//
// Configuration constants indicating the problem size (n) and the
// block size (blkSize)
//
config const n = 62,
             blkSize = 8;

//
// Configuration constant used for verification thresholds
//
config const epsilon = 2.0e-15;

//
// Configuration constants to indicate whether or not to use a
// pseudo-random seed (based on the clock) or a fixed seed; and to
// specify the fixed seed explicitly
//
config const useRandomSeed = false,
             seed = if useRandomSeed then SeedGenerator.currentTime else 31415;

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

// Perform verification?
config const verify = true;

// These are solely to make the testing system happy given the COMPOPTS file.
// To be removed once COMPOPTS becomes a non-issue.
config var reproducible = false, verbose = false;

config param maxBlkSize = 200;
if blkSize > maxBlkSize then
  halt("maxBlkSize ", maxBlkSize,
       " is exceeded by the actual blkSize ", blkSize);

// needed for performance
compilerAssert(CHPL_NETWORK_ATOMICS == "none",
               "export CHPL_NETWORK_ATOMICS=none");

//
// The program entry point, currently module initialization
//

  //
  // Compute targetLocales - required for Dimensional.
  // We hard-code 2 dimensions.
  //
  var targetIds: domain(2);
  var targetLocales: [targetIds] locale;
  setupTargetLocalesArray(targetIds, targetLocales, Locales);

  // Here are the dimensions of our grid of locales.
  const tl1 = targetIds.dim(1).length,
        tl2 = targetIds.dim(2).length;

  if tl1 != tl2 then
    halt("backwardSub() is implemented only for a square locale grid");

  printBriefConfiguration();

  const distReplicated = new ReplicatedDist();

  // Replicate 'targetLocales' to reduce comm.
  // 'targetLocales' itself could be replicated instead,
  // but this way the changes are smaller.
  const targetIdsRepl = targetIds dmapped new dmap(distReplicated);
  var targetLocalesRepl: [targetIdsRepl] locale;
  coforall (l, id) in zip(Locales, LocaleSpace) do on l {
    targetLocalesRepl._value.localArrs[id].arrLocalRep = targetLocales;
  }

  // Create individual dimension descriptors
  const
    // block-cyclic for 1st dimension
    bdim1 = new BlockCyclicDim(lowIdx=1, blockSize=blkSize, numLocales=tl1),
    // replicated for 1st dimension
    rdim1 = new ReplicatedDim(tl1),

    // block-cyclic for 2nd dimension
    bdim2 = new BlockCyclicDim(lowIdx=1, blockSize=blkSize, numLocales=tl2),
    // replicated for 2nd dimension
    rdim2 = new ReplicatedDim(tl2);

  const
    dist1b2b = new DimensionalDist2D(targetLocales, bdim1, bdim2, "dist1b2b"),
    dist1b2r = new DimensionalDist2D(targetLocales, bdim1, rdim2, "dist1b2r"),
    dist1r2b = new DimensionalDist2D(targetLocales, rdim1, bdim2, "dist1r2b");

  //
  // MatVectSpace is a 2D domain of type indexType that represents the
  // n x n matrix adjacent to the column vector b.  MatrixSpace is a
  // subdomain that is created by slicing into MatVectSpace,
  // inheriting all of its rows and its low column bound.  As our
  // standard distribution library is filled out, MatVectSpace will be
  // distributed using a BlockCyclic(blkSize) distribution.
  //
  // We use 'AbD' instead of 'MatVectSpace' throughout.
  //
  const MVS = {1..n, 1..n+1};
  const AbD = MVS dmapped new dmap(dist1b2b);

  var Ab : [AbD] elemType,           // the matrix A and vector b
      piv: [1..n] indexType;         // a vector of pivot values


  //
  // Create the 1-d replicated arrays for schurComplement().
  //
  const replAD = {1..n, 1..blkSize}   dmapped new dmap(dist1b2r),
        replBD = {1..blkSize, 1..n+1} dmapped new dmap(dist1r2b);

  var replA: [replAD] elemType,
      replB: [replBD] elemType;


  const replKD = {0..0, 1..n+1} dmapped new dmap(dist1r2b);
  var   replK: [replKD] elemType;

  const replUD = {0..#blkSize, 0..#blkSize} dmapped new dmap(distReplicated);
        // can't use targetLocales because of "consistency" req. - that's OK
  var   replU: [replUD] elemType;


  // For panelSolve().
  record psRedResultT {
    var row: indexType;
    var elmx, absmx: elemType;
    proc init() { absmx = min(absmx.type); }
    // updateE: // caller to ensure exclusive access
    proc updateE(r: indexType, e: elemType) { updateE(r, e, abs(e)); }
    proc updateE(alt: psRedResultT) { updateE(alt.row, alt.elmx, alt.absmx); }
    proc updateE(r: indexType, e: elemType, a: elemType) {
      if a > absmx { absmx = a; elmx = e; row = r; }
    }
  }
  var psRedLocalResults: [0..#tl1] psRedResultT;

  //
  // Additional data structures for backwardSub().
  //
  type cpuAtomicElemType = // presently unused
    if elemType == real(64) then atomic_real64
    else compilerError("need to define cpuAtomicElemType");
  type cpuAtomicCountType = atomic_int64;
  const replPD = {0..#blkSize} dmapped new dmap(distReplicated);

  // the working 'x' vector - we reuse replK
  //
  // Implementation consequences: (A) Need to zero it out in backwardSub().
  // (B) If we allocated replX separately from replK, we need to
  // replicate X where backwardSub currently invokes replicateK().
  // 
  var replX => replK;

  // Allocate these arrays just once.
  // Partial sums local to each node.
  var bsLocalPartSums: [replPD] elemType;

  // Partial sums from another node in the row
  // brought to the current diagonal node, with an "available" flag.
  param spsNA=0, spsAV=1, spsNBL=2;
  record bsSetofPartSums {
    var values: [0..#blkSize] elemType;
    // spsNA - not available yet or has already been read
    // spsAV - available
    // spsNBL - no applicable blocks on that node
    var avail: cpuAtomicCountType; // initially 0
  }

  // (Only needed on diagonal nodes, for simplicity create everywhere.)
  const bsOthersPartSumsD = {0..#tl2} dmapped new dmap(distReplicated);
  var bsOthersPartSums: [bsOthersPartSumsD] bsSetofPartSums;

  initAB();

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

  LUFactorize(n, piv);                 // compute the LU factorization

  var x => backwardSub(n);  // perform the back substitution

  var execTime = getCurrentTime() - startTime;  // store the elapsed time
  if execTime < 0 then execTime += 24*3600;          // adjust for date change
  printTime(execTime);

  //
  // Validate the answer and print the results
  const validAnswer = verifyResults(x);
  printResults(validAnswer, execTime);

//
// blocked LU factorization with pivoting for matrix augmented with
// vector of RHS values.
//
proc LUFactorize(n: indexType,
                piv: [1..n] indexType) {
  
  // Initialize the pivot vector to represent the initially unpivoted matrix.
  piv = 1..n;

  /* The following diagram illustrates how we partition the matrix.
     Each iteration of the loop increments a variable blk by blkSize;
     point (blk, blk) is the upper-left location of the currently
     unfactored matrix (the dotted region represents the areas
     factored in prior iterations).  The unfactored matrix is
     partioned into four subdomains: tl, tr, bl, and br, and an
     additional domain (not shown), l, that is the union of tl and bl.

                (point blk, blk)
       +-------//------------------+
       |......//...................|
       |.....//....................|
       |....+-----+----------------|
       |....|     |                |
       |....| tl  |      tr        |
       |....|     |                |
       |....+-----+----------------|
       |....|     |                |
       |....|     |                |
       |....| bl  |      br        |
       |....|     |                |
       |....|     |                |
       +----+-----+----------------+
  */
  for blk in 1..n by blkSize {
    if printStats then writeln("processing block ", blk);

    // MVS was AbD
    const tl = MVS[blk..#blkSize, blk..#blkSize],
          tr = MVS[blk..#blkSize, blk+blkSize..],
          bl = MVS[blk+blkSize.., blk..#blkSize],
          br = MVS[blk+blkSize.., blk+blkSize..],
          l  = MVS[blk.., blk..#blkSize];

    //
    // Now that we've sliced and diced Ab properly, do the blocked-LU
    // computation:
    //
    panelSolve(l, piv);
    updateBlockRow(tl, tr);
    
    //
    // update trailing submatrix (if any)
    //
    schurComplement(blk, bl, tr, br);
  }
}

//
// Distributed matrix-multiply for HPL. The idea behind this algorithm is that
// some point the matrix will be partioned as shown in the following diagram:
//
//    [1]----+-----+-----+-----+
//     |     |bbbbb|bbbbb|bbbbb|  Solve for the dotted region by
//     |     |bbbbb|bbbbb|bbbbb|  multiplying the 'a' and 'b' region.
//     |     |bbbbb|bbbbb|bbbbb|  The 'a' region is a block column, the
//     +----[2]----+-----+-----+  'b' region is a block row.
//     |aaaaa|.....|.....|.....|
//     |aaaaa|.....|.....|.....|  The 'a' region was 'bl' in the calling
//     |aaaaa|.....|.....|.....|  function but called AD here.  Similarly,
//     +-----+-----+-----+-----+  'b' was 'tr' in the calling code, but BD
//     |aaaaa|.....|.....|.....|  here.
//     |aaaaa|.....|.....|.....|  
//     |aaaaa|.....|.....|.....|
//     +-----+-----+-----+-----+
//
// Every locale with a block of data in the dotted region updates
// itself by multiplying the neighboring a-region block to its left
// with the neighboring b-region block above it and subtracting its
// current data from the result of this multiplication. To ensure that
// all locales have local copies of the data needed to perform this
// multiplication we copy the data A and B data into the replA and
// replB arrays, which will use a dimensional (block-cyclic,
// replicated-block) distribution (or vice-versa) to ensure that every
// locale only stores one copy of each block it requires for all of
// its rows/columns.
//
proc schurComplement(blk, AD, BD, Rest) {

  // Prevent replication of unequal-sized slices
  if Rest.numIndices == 0 then return;

  //
  // Copy data into replicated arrays so every processor has a local copy
  // of the data it will need to perform a local matrix-multiply.
  //
  replicateA(blk, AD.dim(2));
  //
  replicateB(blk, BD.dim(1));

    const low1 = Rest.dim(1).low,
          low2 = Rest.dim(2).low;
    coforall lid1 in 0..#tl1 do
      coforall lid2 in 0..#tl2 do
        on targetLocalesRepl[lid1, lid2] do
          local {
            const myStarts1 = low1..n by blkSize*tl1 align 1+blkSize*lid1;
            const myStarts2 = low2..n+1 by blkSize*tl2 align 1+blkSize*lid2;
            const blkRange = 1..blkSize;
            forall j1 in myStarts1 {
              const outerRange = j1..min(j1+blkSize-1, n);
              var h2 => replA._value.dsiLocalSlice1((outerRange, blkRange));
              forall j2 in myStarts2 {
                const innerRange = j2..min(j2+blkSize-1, n+1);
                var h1 => Ab._value.dsiLocalSlice1((outerRange, innerRange)),
                    h3 => replB._value.dsiLocalSlice1((blkRange, innerRange));
                for a in outerRange {
                  const
                    h2dd = h2._value.data,
                    h2off = hoistOffset(h2, a, blkRange);
                  for w in blkRange  {
                    const h2aw = h2dd(h2off+w); // h2[a,w];
                    // The code below hoists loop-invariant computations in
                    // dsiAccess() out of the 'for innerRange' loop by hand.
                    // We are using tuples instead of objects because
                    // the generated code is more efficient, which is
                    // also a needed compiler optimization.
                    const
                      h1dd  = h1._value.data,
                      h1off = hoistOffset(h1, a, innerRange),
                      h3dd  = h3._value.data,
                      h3off = hoistOffset(h3, w, innerRange);
                    for b in innerRange do
                      // Ab[a,b] -= replA[a,w] * replB[w,b];
                      h1dd(h1off+b) -= h2aw * h3dd(h3off+b);
                  } // for w
                } // for a
              } // forall j2
            } // forall j1
          } // local
}

/////////////////////////////////

// NB works only when storage indices for sliceDim1,2 are contiguous
// Each sliceDim1,2 can be either a range or an int

proc DimensionalArr.dsiLocalSlice1((sliceDim1, sliceDim2)) {
  // might be more elegant to replace the explicit arg tuple with 'slice'
  proc toScalar(slice)
    return if isIntegralType(slice.type) then slice else slice.low;
  proc toRange(slice)
    return if isIntegralType(slice.type) then slice..slice else slice;
  proc origScalar(param dim) param
    return isIntegralType(if dim==1 then sliceDim1.type else sliceDim2.type);
  proc toOrig(param dim, arg)
    return if origScalar(dim) then toScalar(arg) else toRange(arg);

  const dom = this.dom;
  const dist = dom.dist;
  // todo: cache (l1, l2) in privatized copies when possible
  // (i.e. if privatization is supported and there is no oversubscription)
  // Assuming dsiLocalSlice is guaranteed to be local to 'here'.
  const l1 = dist.di1.dsiIndexToLocale1d(toScalar(sliceDim1)),
        l2 = dist.di2.dsiIndexToLocale1d(toScalar(sliceDim2)),
        locAdesc = this.localAdescs[l1, l2],
        r1 = if dom.dom1.dsiStorageUsesUserIndices()
             then if origScalar(1) then sliceDim1 else dom.whole.dim(1)(sliceDim1)
             else toOrig(1, locAdesc.locDom.doml1.dsiLocalSliceStorageIndices1d(dom.dom1, toRange(sliceDim1))),
        r2 = if dom.dom2.dsiStorageUsesUserIndices()
             then if origScalar(2) then sliceDim2 else dom.whole.dim(2)(sliceDim2)
             else toOrig(2, locAdesc.locDom.doml2.dsiLocalSliceStorageIndices1d(dom.dom2, toRange(sliceDim2)));

  const reindexExpr =
    if origScalar(1) then
      if origScalar(2) then compilerError("dsiLocalSlice1 is not supported for two scalars")
      else (sliceDim2,)
    else
      if origScalar(2) then (sliceDim1,)
      else (sliceDim1, sliceDim2);
  var result: [(...reindexExpr)] => locAdesc.myStorageArr[r1, r2];
  return result;
}

inline proc hoistOffset(A: [], i1, slice2) {
  const d = A._value;
  compilerAssert(d.rank == 2);
  const result = d.origin + i1 * d.blk(1) - d.factoredOffs;
  return result;
}

/////////////////////////////////

//
// do unblocked-LU decomposition within the specified panel, update the
// pivot vector accordingly
//
proc panelSolve(
               panel: domain,
               piv: [] indexType) {

  const blk = panel.dim(1).low;
  const dim2 = panel.dim(2);
  const cornerLocale = targetLocaleCorner(blk);
  var tStart, tReduce, tSwap, tRepl, tComp: real;
  var tSwapMsg: string;

  for k in dim2[..n] {             // iterate through the columns

/* psReduce() does the following:

    const col = panel[k.., k..k];
    
    // Find the pivot, the element with the largest absolute value.
    const (_, (pivotRow, _)) = maxloc reduce zip(abs(Ab(col)), col);

    // Capture the pivot value explicitly (note that result of maxloc
    // is absolute value, so it can't be used directly).
    // k <= pivotRow <= n
    //
    const pivotVal = Ab[pivotRow, k];
*/
    const (pivotRow, pivotVal) = psReduce(blk, k);
    if (pivotVal == 0) then
      halt("Matrix cannot be factorized");

    // Swap the current row with the pivot row and update the pivot vector
    // to reflect that
   if k == pivotRow {
    // nothing
   } else {
    var wasLocal: bool; // fyi only
    //Ab[k..k, ..] <=> Ab[pivotRow..pivotRow, ..];
    psSwap(k, pivotRow, wasLocal);
    piv[k] <=> piv[pivotRow];
   }

/* replicate + psCompute() do the following:

    // divide all values below and in the same col as the pivot by
    // the pivot value
    Ab[k+1.., k..k] /= pivotVal;
    
    // update all other values below the pivot
    forall (i,j) in panel[k+1.., k+1..] do
      Ab[i,j] -= Ab[i,k] * Ab[k,j];
*/

    // replicate
    on cornerLocale {
      local
        for j in dim2[k+1..] do
          replK[0,j] = Ab[k,j];
      replicateK(blk);
    }

    // compute
    psCompute(panel, blk, k, pivotVal);
  }
}

proc psReduce(blk, k) {
  const lid2 = targetLocalesIndexForAbIndex(2, k);
  const rlrLoc = psRedLocalResults.locale;

  coforall lid1 in 0..#tl1 {
    on targetLocalesRepl[lid1, lid2] {
      var locResult: psRedResultT;
      locResult.init();

      local {
        // guard updates to locResult and maxRes within the forall
        var upd$: sync bool;
        var maxRes: atomic real;
        maxRes.write(min(real));
        // Starts of all the blocks in the panel that are local to this node.
        const panelStarts = blk..n by blkSize*tl1 align 1+blkSize*lid1;
          forall iStart in panelStarts {
            const iRange = max(iStart, k)..min(iStart + blkSize - 1, n);
            var myResult: psRedResultT;
            myResult.init();
            var locAB => Ab._value.dsiLocalSlice1((iRange, k));
            for i in iRange do myResult.updateE(i, locAB[i]);
            if myResult.absmx > maxRes.read() { // only lock if we might update
              upd$ = true;  // lock<commit_msg>Added a missing 'with' clause to hpl.hpcc2012.chpl<commit_after>//
// Use standard modules for vector and matrix Norms, Random numbers
// and Timing routines
//
use Norm, Random, Time;
use UtilMath;

//
// Use the user module for computing HPCC problem sizes
//
//use HPCCProblemSize;

//
// Use the distributions we need for this computation
//
use DimensionalDist2D, ReplicatedDim, BlockCycDim;
use ReplicatedDist;

//
// The number of matrices and the element type of those matrices
// indexType can be int or int(64), elemType can be real or complex
//
const numMatrices = 1;
type indexType = int,
     elemType = real;

//
// Configuration constants indicating the problem size (n) and the
// block size (blkSize)
//
config const n = 62,
             blkSize = 8;

//
// Configuration constant used for verification thresholds
//
config const epsilon = 2.0e-15;

//
// Configuration constants to indicate whether or not to use a
// pseudo-random seed (based on the clock) or a fixed seed; and to
// specify the fixed seed explicitly
//
config const useRandomSeed = false,
             seed = if useRandomSeed then SeedGenerator.currentTime else 31415;

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

// Perform verification?
config const verify = true;

// These are solely to make the testing system happy given the COMPOPTS file.
// To be removed once COMPOPTS becomes a non-issue.
config var reproducible = false, verbose = false;

config param maxBlkSize = 200;
if blkSize > maxBlkSize then
  halt("maxBlkSize ", maxBlkSize,
       " is exceeded by the actual blkSize ", blkSize);

// needed for performance
compilerAssert(CHPL_NETWORK_ATOMICS == "none",
               "export CHPL_NETWORK_ATOMICS=none");

//
// The program entry point, currently module initialization
//

  //
  // Compute targetLocales - required for Dimensional.
  // We hard-code 2 dimensions.
  //
  var targetIds: domain(2);
  var targetLocales: [targetIds] locale;
  setupTargetLocalesArray(targetIds, targetLocales, Locales);

  // Here are the dimensions of our grid of locales.
  const tl1 = targetIds.dim(1).length,
        tl2 = targetIds.dim(2).length;

  if tl1 != tl2 then
    halt("backwardSub() is implemented only for a square locale grid");

  printBriefConfiguration();

  const distReplicated = new ReplicatedDist();

  // Replicate 'targetLocales' to reduce comm.
  // 'targetLocales' itself could be replicated instead,
  // but this way the changes are smaller.
  const targetIdsRepl = targetIds dmapped new dmap(distReplicated);
  var targetLocalesRepl: [targetIdsRepl] locale;
  coforall (l, id) in zip(Locales, LocaleSpace) do on l {
    targetLocalesRepl._value.localArrs[id].arrLocalRep = targetLocales;
  }

  // Create individual dimension descriptors
  const
    // block-cyclic for 1st dimension
    bdim1 = new BlockCyclicDim(lowIdx=1, blockSize=blkSize, numLocales=tl1),
    // replicated for 1st dimension
    rdim1 = new ReplicatedDim(tl1),

    // block-cyclic for 2nd dimension
    bdim2 = new BlockCyclicDim(lowIdx=1, blockSize=blkSize, numLocales=tl2),
    // replicated for 2nd dimension
    rdim2 = new ReplicatedDim(tl2);

  const
    dist1b2b = new DimensionalDist2D(targetLocales, bdim1, bdim2, "dist1b2b"),
    dist1b2r = new DimensionalDist2D(targetLocales, bdim1, rdim2, "dist1b2r"),
    dist1r2b = new DimensionalDist2D(targetLocales, rdim1, bdim2, "dist1r2b");

  //
  // MatVectSpace is a 2D domain of type indexType that represents the
  // n x n matrix adjacent to the column vector b.  MatrixSpace is a
  // subdomain that is created by slicing into MatVectSpace,
  // inheriting all of its rows and its low column bound.  As our
  // standard distribution library is filled out, MatVectSpace will be
  // distributed using a BlockCyclic(blkSize) distribution.
  //
  // We use 'AbD' instead of 'MatVectSpace' throughout.
  //
  const MVS = {1..n, 1..n+1};
  const AbD = MVS dmapped new dmap(dist1b2b);

  var Ab : [AbD] elemType,           // the matrix A and vector b
      piv: [1..n] indexType;         // a vector of pivot values


  //
  // Create the 1-d replicated arrays for schurComplement().
  //
  const replAD = {1..n, 1..blkSize}   dmapped new dmap(dist1b2r),
        replBD = {1..blkSize, 1..n+1} dmapped new dmap(dist1r2b);

  var replA: [replAD] elemType,
      replB: [replBD] elemType;


  const replKD = {0..0, 1..n+1} dmapped new dmap(dist1r2b);
  var   replK: [replKD] elemType;

  const replUD = {0..#blkSize, 0..#blkSize} dmapped new dmap(distReplicated);
        // can't use targetLocales because of "consistency" req. - that's OK
  var   replU: [replUD] elemType;


  // For panelSolve().
  record psRedResultT {
    var row: indexType;
    var elmx, absmx: elemType;
    proc init() { absmx = min(absmx.type); }
    // updateE: // caller to ensure exclusive access
    proc updateE(r: indexType, e: elemType) { updateE(r, e, abs(e)); }
    proc updateE(alt: psRedResultT) { updateE(alt.row, alt.elmx, alt.absmx); }
    proc updateE(r: indexType, e: elemType, a: elemType) {
      if a > absmx { absmx = a; elmx = e; row = r; }
    }
  }
  var psRedLocalResults: [0..#tl1] psRedResultT;

  //
  // Additional data structures for backwardSub().
  //
  type cpuAtomicElemType = // presently unused
    if elemType == real(64) then atomic_real64
    else compilerError("need to define cpuAtomicElemType");
  type cpuAtomicCountType = atomic_int64;
  const replPD = {0..#blkSize} dmapped new dmap(distReplicated);

  // the working 'x' vector - we reuse replK
  //
  // Implementation consequences: (A) Need to zero it out in backwardSub().
  // (B) If we allocated replX separately from replK, we need to
  // replicate X where backwardSub currently invokes replicateK().
  // 
  var replX => replK;

  // Allocate these arrays just once.
  // Partial sums local to each node.
  var bsLocalPartSums: [replPD] elemType;

  // Partial sums from another node in the row
  // brought to the current diagonal node, with an "available" flag.
  param spsNA=0, spsAV=1, spsNBL=2;
  record bsSetofPartSums {
    var values: [0..#blkSize] elemType;
    // spsNA - not available yet or has already been read
    // spsAV - available
    // spsNBL - no applicable blocks on that node
    var avail: cpuAtomicCountType; // initially 0
  }

  // (Only needed on diagonal nodes, for simplicity create everywhere.)
  const bsOthersPartSumsD = {0..#tl2} dmapped new dmap(distReplicated);
  var bsOthersPartSums: [bsOthersPartSumsD] bsSetofPartSums;

  initAB();

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

  LUFactorize(n, piv);                 // compute the LU factorization

  var x => backwardSub(n);  // perform the back substitution

  var execTime = getCurrentTime() - startTime;  // store the elapsed time
  if execTime < 0 then execTime += 24*3600;          // adjust for date change
  printTime(execTime);

  //
  // Validate the answer and print the results
  const validAnswer = verifyResults(x);
  printResults(validAnswer, execTime);

//
// blocked LU factorization with pivoting for matrix augmented with
// vector of RHS values.
//
proc LUFactorize(n: indexType,
                piv: [1..n] indexType) {
  
  // Initialize the pivot vector to represent the initially unpivoted matrix.
  piv = 1..n;

  /* The following diagram illustrates how we partition the matrix.
     Each iteration of the loop increments a variable blk by blkSize;
     point (blk, blk) is the upper-left location of the currently
     unfactored matrix (the dotted region represents the areas
     factored in prior iterations).  The unfactored matrix is
     partioned into four subdomains: tl, tr, bl, and br, and an
     additional domain (not shown), l, that is the union of tl and bl.

                (point blk, blk)
       +-------//------------------+
       |......//...................|
       |.....//....................|
       |....+-----+----------------|
       |....|     |                |
       |....| tl  |      tr        |
       |....|     |                |
       |....+-----+----------------|
       |....|     |                |
       |....|     |                |
       |....| bl  |      br        |
       |....|     |                |
       |....|     |                |
       +----+-----+----------------+
  */
  for blk in 1..n by blkSize {
    if printStats then writeln("processing block ", blk);

    // MVS was AbD
    const tl = MVS[blk..#blkSize, blk..#blkSize],
          tr = MVS[blk..#blkSize, blk+blkSize..],
          bl = MVS[blk+blkSize.., blk..#blkSize],
          br = MVS[blk+blkSize.., blk+blkSize..],
          l  = MVS[blk.., blk..#blkSize];

    //
    // Now that we've sliced and diced Ab properly, do the blocked-LU
    // computation:
    //
    panelSolve(l, piv);
    updateBlockRow(tl, tr);
    
    //
    // update trailing submatrix (if any)
    //
    schurComplement(blk, bl, tr, br);
  }
}

//
// Distributed matrix-multiply for HPL. The idea behind this algorithm is that
// some point the matrix will be partioned as shown in the following diagram:
//
//    [1]----+-----+-----+-----+
//     |     |bbbbb|bbbbb|bbbbb|  Solve for the dotted region by
//     |     |bbbbb|bbbbb|bbbbb|  multiplying the 'a' and 'b' region.
//     |     |bbbbb|bbbbb|bbbbb|  The 'a' region is a block column, the
//     +----[2]----+-----+-----+  'b' region is a block row.
//     |aaaaa|.....|.....|.....|
//     |aaaaa|.....|.....|.....|  The 'a' region was 'bl' in the calling
//     |aaaaa|.....|.....|.....|  function but called AD here.  Similarly,
//     +-----+-----+-----+-----+  'b' was 'tr' in the calling code, but BD
//     |aaaaa|.....|.....|.....|  here.
//     |aaaaa|.....|.....|.....|  
//     |aaaaa|.....|.....|.....|
//     +-----+-----+-----+-----+
//
// Every locale with a block of data in the dotted region updates
// itself by multiplying the neighboring a-region block to its left
// with the neighboring b-region block above it and subtracting its
// current data from the result of this multiplication. To ensure that
// all locales have local copies of the data needed to perform this
// multiplication we copy the data A and B data into the replA and
// replB arrays, which will use a dimensional (block-cyclic,
// replicated-block) distribution (or vice-versa) to ensure that every
// locale only stores one copy of each block it requires for all of
// its rows/columns.
//
proc schurComplement(blk, AD, BD, Rest) {

  // Prevent replication of unequal-sized slices
  if Rest.numIndices == 0 then return;

  //
  // Copy data into replicated arrays so every processor has a local copy
  // of the data it will need to perform a local matrix-multiply.
  //
  replicateA(blk, AD.dim(2));
  //
  replicateB(blk, BD.dim(1));

    const low1 = Rest.dim(1).low,
          low2 = Rest.dim(2).low;
    coforall lid1 in 0..#tl1 do
      coforall lid2 in 0..#tl2 do
        on targetLocalesRepl[lid1, lid2] do
          local {
            const myStarts1 = low1..n by blkSize*tl1 align 1+blkSize*lid1;
            const myStarts2 = low2..n+1 by blkSize*tl2 align 1+blkSize*lid2;
            const blkRange = 1..blkSize;
            forall j1 in myStarts1 {
              const outerRange = j1..min(j1+blkSize-1, n);
              var h2 => replA._value.dsiLocalSlice1((outerRange, blkRange));
              forall j2 in myStarts2 {
                const innerRange = j2..min(j2+blkSize-1, n+1);
                var h1 => Ab._value.dsiLocalSlice1((outerRange, innerRange)),
                    h3 => replB._value.dsiLocalSlice1((blkRange, innerRange));
                for a in outerRange {
                  const
                    h2dd = h2._value.data,
                    h2off = hoistOffset(h2, a, blkRange);
                  for w in blkRange  {
                    const h2aw = h2dd(h2off+w); // h2[a,w];
                    // The code below hoists loop-invariant computations in
                    // dsiAccess() out of the 'for innerRange' loop by hand.
                    // We are using tuples instead of objects because
                    // the generated code is more efficient, which is
                    // also a needed compiler optimization.
                    const
                      h1dd  = h1._value.data,
                      h1off = hoistOffset(h1, a, innerRange),
                      h3dd  = h3._value.data,
                      h3off = hoistOffset(h3, w, innerRange);
                    for b in innerRange do
                      // Ab[a,b] -= replA[a,w] * replB[w,b];
                      h1dd(h1off+b) -= h2aw * h3dd(h3off+b);
                  } // for w
                } // for a
              } // forall j2
            } // forall j1
          } // local
}

/////////////////////////////////

// NB works only when storage indices for sliceDim1,2 are contiguous
// Each sliceDim1,2 can be either a range or an int

proc DimensionalArr.dsiLocalSlice1((sliceDim1, sliceDim2)) {
  // might be more elegant to replace the explicit arg tuple with 'slice'
  proc toScalar(slice)
    return if isIntegralType(slice.type) then slice else slice.low;
  proc toRange(slice)
    return if isIntegralType(slice.type) then slice..slice else slice;
  proc origScalar(param dim) param
    return isIntegralType(if dim==1 then sliceDim1.type else sliceDim2.type);
  proc toOrig(param dim, arg)
    return if origScalar(dim) then toScalar(arg) else toRange(arg);

  const dom = this.dom;
  const dist = dom.dist;
  // todo: cache (l1, l2) in privatized copies when possible
  // (i.e. if privatization is supported and there is no oversubscription)
  // Assuming dsiLocalSlice is guaranteed to be local to 'here'.
  const l1 = dist.di1.dsiIndexToLocale1d(toScalar(sliceDim1)),
        l2 = dist.di2.dsiIndexToLocale1d(toScalar(sliceDim2)),
        locAdesc = this.localAdescs[l1, l2],
        r1 = if dom.dom1.dsiStorageUsesUserIndices()
             then if origScalar(1) then sliceDim1 else dom.whole.dim(1)(sliceDim1)
             else toOrig(1, locAdesc.locDom.doml1.dsiLocalSliceStorageIndices1d(dom.dom1, toRange(sliceDim1))),
        r2 = if dom.dom2.dsiStorageUsesUserIndices()
             then if origScalar(2) then sliceDim2 else dom.whole.dim(2)(sliceDim2)
             else toOrig(2, locAdesc.locDom.doml2.dsiLocalSliceStorageIndices1d(dom.dom2, toRange(sliceDim2)));

  const reindexExpr =
    if origScalar(1) then
      if origScalar(2) then compilerError("dsiLocalSlice1 is not supported for two scalars")
      else (sliceDim2,)
    else
      if origScalar(2) then (sliceDim1,)
      else (sliceDim1, sliceDim2);
  var result: [(...reindexExpr)] => locAdesc.myStorageArr[r1, r2];
  return result;
}

inline proc hoistOffset(A: [], i1, slice2) {
  const d = A._value;
  compilerAssert(d.rank == 2);
  const result = d.origin + i1 * d.blk(1) - d.factoredOffs;
  return result;
}

/////////////////////////////////

//
// do unblocked-LU decomposition within the specified panel, update the
// pivot vector accordingly
//
proc panelSolve(
               panel: domain,
               piv: [] indexType) {

  const blk = panel.dim(1).low;
  const dim2 = panel.dim(2);
  const cornerLocale = targetLocaleCorner(blk);
  var tStart, tReduce, tSwap, tRepl, tComp: real;
  var tSwapMsg: string;

  for k in dim2[..n] {             // iterate through the columns

/* psReduce() does the following:

    const col = panel[k.., k..k];
    
    // Find the pivot, the element with the largest absolute value.
    const (_, (pivotRow, _)) = maxloc reduce zip(abs(Ab(col)), col);

    // Capture the pivot value explicitly (note that result of maxloc
    // is absolute value, so it can't be used directly).
    // k <= pivotRow <= n
    //
    const pivotVal = Ab[pivotRow, k];
*/
    const (pivotRow, pivotVal) = psReduce(blk, k);
    if (pivotVal == 0) then
      halt("Matrix cannot be factorized");

    // Swap the current row with the pivot row and update the pivot vector
    // to reflect that
   if k == pivotRow {
    // nothing
   } else {
    var wasLocal: bool; // fyi only
    //Ab[k..k, ..] <=> Ab[pivotRow..pivotRow, ..];
    psSwap(k, pivotRow, wasLocal);
    piv[k] <=> piv[pivotRow];
   }

/* replicate + psCompute() do the following:

    // divide all values below and in the same col as the pivot by
    // the pivot value
    Ab[k+1.., k..k] /= pivotVal;
    
    // update all other values below the pivot
    forall (i,j) in panel[k+1.., k+1..] do
      Ab[i,j] -= Ab[i,k] * Ab[k,j];
*/

    // replicate
    on cornerLocale {
      local
        for j in dim2[k+1..] do
          replK[0,j] = Ab[k,j];
      replicateK(blk);
    }

    // compute
    psCompute(panel, blk, k, pivotVal);
  }
}

proc psReduce(blk, k) {
  const lid2 = targetLocalesIndexForAbIndex(2, k);
  const rlrLoc = psRedLocalResults.locale;

  coforall lid1 in 0..#tl1 {
    on targetLocalesRepl[lid1, lid2] {
      var locResult: psRedResultT;
      locResult.init();

      local {
        // guard updates to locResult and maxRes within the forall
        var upd$: sync bool;
        var maxRes: atomic real;
        maxRes.write(min(real));
        // Starts of all the blocks in the panel that are local to this node.
        const panelStarts = blk..n by blkSize*tl1 align 1+blkSize*lid1;
          forall iStart in panelStarts with (ref locResult) {
            const iRange = max(iStart, k)..min(iStart + blkSize - 1, n);
            var myResult: psRedResultT;
            myResult.init();
            var locAB => Ab._value.dsiLocalSlice1((iRange, k));
            for i in iRange do myResult.updateE(i, locAB[i]);
            if myResult.absmx > maxRes.read() { // only lock if we might update
              upd$ = true;  // lock
{'fraction_non_alphanumeric': 0.12938080842010974, 'fraction_numerical': 0.014891949389765983, 'mean_word_length': 2.2184684684684686, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5821598', 'n_tokens_mistral': 11636, 'n_tokens_neox': 11156, 'n_words': 3646}