New Issue: Prototype partial sum reductions on Block-distributed arrays

Notifications Issues
1

20787, "bradcray", "Prototype partial sum reductions on Block-distributed arrays", "2022-10-07T00:00:42Z"

Partial reductions are a feature that's been intended for Chapel since forever (e.g., Partial Reductions · Issue #8951 · chapel-lang/chapel · GitHub, but also since the dawn of the project), yet never completed, essentially waiting for a user or use-case that wanted them badly enough. @sdbachman looks like such a user, which has revived past conversations about partial reductions lately.

My current best thought about a first step here is to add some methods that we could call manually on Block-distributed arrays as a starting point, avoiding the language design aspects for the time being. Specifically, I'm imagining us supporting:

proc BlockArr.partSumReduce(resultArr: [?resDom]) { /* implement partial reductions here */ }

proc BlockArr.partSumReduce(resultDom: domain) {
  var resultArr: [resultDom] eltType;
  this.partSumReduce(resultArr);
  return resultArr;
}

where, for the short-term, I think we should require resDom to be the same rank as this, to share its distribution/alignment, and to have each dimension either be a singleton index (e.g., 1..1) or the same as this in each dimension. I also think we should start with a 2D-only implementation to begin with, and then generalize it to multiple dimensions as we go.

I'm imagining that the "takes an array" version of the routine would

  • determine which dimensions are being collapsed by inspecting the domains' dimensions
  • declare a local array representing this locale's local contribution to the global result
  • loop over its local elements, reducing into the local array
  • communicate with other locales in its row/column/plane/pencil/etc. to combine results into one of their temporary arrays
  • copy the result into the appropriate sub-array of the resultArr (or perhaps that sub-array can be the place where the temporary arrays are accumulated from the start?)

In doing this, we should be sure to:

  • iterate over the local data in parallel, and in a cache-friendly way
  • only communicate ~1 time per locale to communicate the results, and only with locales that are in our row/column

Here are some rough sketches of how the first of these bullets can be accomplished, with some hard-coded problem sizes and local arrays as a demonstration:

   var A: [1..m, 1..n] real,
    colRes: [1..m] real,
    rowRes: [1..n] real;

// compute + partial reduction along rows, resulting in a column of values
forall i in 1..m do
  for j in 1..n do
    colRes[i] += A[i,j];

// compute + partial reduction along cols, resulting in row of values
forall i in 1..m with (+ reduce rowRes) do
  for j in 1..n do
    rowRes[j] += A[i,j];

Note that this second case can't be a simple forall due to the potential for RRWW races on rowRes[j]. We could parallelize the inner loop, but that would add a lot of tasking overhead. We could swap the order of the loops, but that would traverse A in a less cache-friendly manner. So we use a task-local reduction variable of the temporary array's size instead. We could also use a vector of atomic values, but I suspect that'd be too expensive.