MPI-idioms.chpl

use BlockDist;

proc main() throws {

  var D = {1..8};
  var distD = D dmapped Block(boundingBox=D);
  var A: [distD] int;

  // Test array (from http://mpitutorial.com/tutorials/mpi-reduce-and-allreduce/)
  A = [5, 1, 2, 3, 7, 8, 4, 2];

  var ret = MPI_Reduce(A);
  writeln(ret);
  // Should be [18, 14]

  MPI_Allreduce(A);
  writeln(A);
  // Should be [18, 14, 18, 14, 18, 14, 18, 14]
}


/* Sum ith element on each locale and place sum onto ith element on locale 0 */
proc MPI_Reduce(A: [?D]) throws {
  if D.size & numLocales != 0 then
    throw new owned IllegalArgumentError('Array size must be evenly divisible by numLocales');

  const N = D.size / numLocales;

  var results: [1..N] int;

  forall i in results.domain {
    results[i] = + reduce A[i.. by N];
  }

  return results;
}

/* Sum ith element on each locale and place sum onto ith element of each locale */
proc MPI_Allreduce(ref A: [?D]) throws {
  var results = MPI_Reduce(A);

  const N = D.size / numLocales;
  forall loc in Locales {
    const slice = (loc.id*N+1)..#N;
    // localeSlice asserts we are writing locally
    A.localSlice(slice) = results;
  }
}

Tested with -nl 4