Add some new sparse methods based on feedback (#26280)

[reviewed by @lydia-duncan — thanks!]

While at CLSAC, I had the chance to meet up with Rich Vuduc, who had
done a little algorithmic tinkering with sparse arrays in Chapel this
past year in his spare time. He mentioned some helper routines that he
had wanted, and ended up writing himself, digging into Chapel's
internals. In this PR, I've added versions of three of the main routines
he was missing to the relevant domain maps:

* A 'getCoordinates()' method on default/COO sparse arrays that returns
a [a reference to] an array of tuples representing the row/col
coordinates of the nonzeroes (for a 2D sparse array, or a k*tuple for kD
arrays).

* A zero-argument 'getLocalSubarray()' query that returns [a reference
to] an array representing the current locale's local chunk of a
block-distributed sparse array.

* A 'localSubarrays()' iterator that iterates over all of a
block-distributed sparse array's target locales, yielding the local
subarray on the corresponding locale. So this could be used to spin up
SPMD-style parallelism with a task per locale and a reference to that
locale's block. There's also a serial version that iterates over the
locales, yielding their blocks serially (mostly for debugging and
completeness; seems less likely to be used in practice when perf
matters)

While here, this also adds `ref` overloads for some of the existing
utility routines that only provided `const ref` versions before, to lean
on ref intent overloading and permit their contents to be modified
directly. This was another thing that came out of the conversation with
Rich.

I'm not sure these are the best/right names, but since sparse is
unstable, also didn't labor over them very much, thinking it would be
best to get them into users' hands and get feedback on their utility.

modules/dists/SparseBlockDist.chpl

@@ -553,14 +553,30 @@ class SparseBlockArr: BaseSparseArr(?) {
     return true;
   }
 
+  proc getLocalSubarray(localeRow, localeCol) ref {
+    return this.locArr[localeRow, localeCol]!.myElems;
+  }
+
   proc getLocalSubarray(localeRow, localeCol) const ref {
     return this.locArr[localeRow, localeCol]!.myElems;
   }
 
+  proc getLocalSubarray(localeIdx) ref {
+    return this.locArr[localeIdx]!.myElems;
+  }
+
   proc getLocalSubarray(localeIdx) const ref {
     return this.locArr[localeIdx]!.myElems;
   }
 
+  proc getLocalSubarray() ref {
+    return myLocArr!.myElems;
+  }
+
+  proc getLocalSubarray() const ref {
+    return myLocArr!.myElems;
+  }
+
   proc setLocalSubarray(locNonzeroes, loc: locale = here) {
     if loc != here then
       halt("setLocalSubarray() doesn't currently support remote updates");
@@ -571,6 +587,22 @@ class SparseBlockArr: BaseSparseArr(?) {
     else
       myBlock.data = locNonzeroes.data;
   }
+
+  iter localSubarrays() ref {
+    for locIdx in dom.dist.targetLocDom {
+      on locArr[locIdx] {
+        yield getLocalSubarray(locIdx);
+      }
+    }
+  }
+
+  iter localSubarrays(param tag: iterKind) ref where tag == iterKind.standalone {
+    coforall locIdx in dom.dist.targetLocDom {
+      on locArr[locIdx] {
+        yield getLocalSubarray(locIdx);
+      }
+    }
+  }
 }
 
 //

modules/internal/ChapelDistribution.chpl

@@ -636,6 +636,11 @@ module ChapelDistribution {
       }
     }
 
+    /* make the domain "just the right size" */
+    inline proc _fit(size: int) {
+      nnzDom = {0..<size};
+    }
+
     // This method assumes nnz is updated according to the size
     // requested. So, a bulk addition into a sparse domain should: (1)
     // calculate new nnz and update it, (2) call this method, (3) add

modules/internal/DefaultSparse.chpl

@@ -407,6 +407,16 @@ module DefaultSparse {
     override proc dsiSupportsAutoLocalAccess() param {
       return defaultSparseSupportsAutoLocalAccess;
     }
+
+    proc getCoordinates() const ref : [] rank*idxType {
+      _fit(_nnz);  // shrink the coordinate array to be "just the right size"
+      return _indices;
+    }
+
+    proc getCoordinates() ref : [] rank*idxType {
+      _fit(_nnz);  // shrink the coordinate array to be "just the right size"
+      return _indices;
+    }
   }
 
 

test/distributions/sparseBlock/localCoords.chpl

@@ -0,0 +1,64 @@
+use BlockDist;
+
+config const n = 10,
+             useTupleIndexing = true;
+
+const D = {0..<n, 0..<n} dmapped new blockDist({0..<n, 0..<n});
+const DS: sparse subdomain(D);
+
+var A: [DS] real;
+
+// build up the global sparse domain and array by having each locale set
+// its own local blocks
+coforall loc in DS.targetLocales() do on loc {
+  // compute this locale's contribution to the index set, creating a
+  // diagonal pattern per locale
+  const myInds = DS.parentDom.localSubdomain();
+  var locSpsInds: sparse subdomain(myInds);
+
+  for i in 0..<min(myInds.dim(0).size, myInds.dim(1).size) {
+    locSpsInds += myInds.low + (i,i);
+  }
+  //  writeln(here.id, ": ", locSpsInds);
+
+  // now do the same for its local non-zeroes
+  var locVals: [locSpsInds] real;
+  for i in 0..<min(myInds.dim(0).size, myInds.dim(1).size) {
+    const low = myInds.low;
+    locVals[low+(i,i)] = low(0)+i + (low(1)+i)/10.0;
+  }
+  //  writeln(here.id, ": ", locVals);
+
+  // Take those local portions and plug them into the global domain/array
+  DS.setLocalSubdomain(locSpsInds);
+  A.setLocalSubarray(locVals);
+}
+
+
+// check the results
+
+for block in A.localSubarrays() {
+  writeln("locale ", here.id, " owns:\n", block.domain.getCoordinates());
+}
+writeln();
+
+// perform some more updates to add the local antidiagonal, and make
+// sure those occurred
+
+coforall loc in DS.targetLocales() do on loc {
+  // compute this locale's contribution to the index set, creating a
+  // diagonal pattern per locale
+  const myInds = DS.parentDom.localSubdomain();
+  var locSpsInds = A.getLocalSubarray().domain;
+
+  for i in 0..<min(myInds.dim(0).size, myInds.dim(1).size) {
+    const newInd = myInds.low + (myInds.dim(0).size-(i+1), i);
+    locSpsInds += newInd;
+  }
+
+  DS.setLocalSubdomain(locSpsInds);
+}
+
+for block in A.localSubarrays() {
+  writeln("locale ", here.id, " owns:\n", block.domain.getCoordinates());
+}

test/distributions/sparseBlock/localCoords.comm-none.good

@@ -0,0 +1,5 @@
+locale 0 owns:
+(0, 0) (1, 1) (2, 2) (3, 3) (4, 4) (5, 5) (6, 6) (7, 7) (8, 8) (9, 9)
+
+locale 0 owns:
+(0, 0) (0, 9) (1, 1) (1, 8) (2, 2) (2, 7) (3, 3) (3, 6) (4, 4) (4, 5) (5, 4) (5, 5) (6, 3) (6, 6) (7, 2) (7, 7) (8, 1) (8, 8) (9, 0) (9, 9)

test/distributions/sparseBlock/localCoords.good

@@ -0,0 +1,17 @@
+locale 0 owns:
+(0, 0) (1, 1) (2, 2) (3, 3) (4, 4)
+locale 1 owns:
+(0, 5) (1, 6) (2, 7) (3, 8) (4, 9)
+locale 2 owns:
+(5, 0) (6, 1) (7, 2) (8, 3) (9, 4)
+locale 3 owns:
+(5, 5) (6, 6) (7, 7) (8, 8) (9, 9)
+
+locale 0 owns:
+(0, 0) (0, 4) (1, 1) (1, 3) (2, 2) (3, 1) (3, 3) (4, 0) (4, 4)
+locale 1 owns:
+(0, 5) (0, 9) (1, 6) (1, 8) (2, 7) (3, 6) (3, 8) (4, 5) (4, 9)
+locale 2 owns:
+(5, 0) (5, 4) (6, 1) (6, 3) (7, 2) (8, 1) (8, 3) (9, 0) (9, 4)
+locale 3 owns:
+(5, 5) (5, 9) (6, 6) (6, 8) (7, 7) (8, 6) (8, 8) (9, 5) (9, 9)

test/distributions/sparseBlock/localSubarraysIter.chpl

@@ -0,0 +1,71 @@
+use BlockDist, CompressedSparseLayout;
+
+config const n = 10,
+             useTupleIndexing = true;
+
+config param csr = true;
+
+const D = {0..<n, 0..<n} dmapped if csr then new blockDist({0..<n, 0..<n},
+                                               sparseLayoutType=csrLayout)
+                                        else new blockDist({0..<n, 0..<n});
+const DS: sparse subdomain(D);
+
+// check that targetLocales on domains is working as expected
+writeln("DS target locales:");
+writeln(DS.targetLocales());
+writeln();
+
+var A: [DS] real;
+
+// build up the global sparse domain and array by having each locale set
+// its own local blocks
+coforall loc in DS.targetLocales() do on loc {
+  // compute this locale's contribution to the index set, creating a
+  // diagonal pattern per locale
+  const myInds = DS.parentDom.localSubdomain();
+  var locSpsInds: if csr then sparse subdomain(myInds) dmapped new csrLayout()
+                         else sparse subdomain(myInds);
+
+  for i in 0..<min(myInds.dim(0).size, myInds.dim(1).size) {
+    locSpsInds += myInds.low + (i,i);
+  }
+  //  writeln(here.id, ": ", locSpsInds);
+
+  // now do the same for its local non-zeroes
+  var locVals: [locSpsInds] real;
+  for i in 0..<min(myInds.dim(0).size, myInds.dim(1).size) {
+    const low = myInds.low;
+    locVals[low+(i,i)] = low(0)+i + (low(1)+i)/10.0;
+  }
+  //  writeln(here.id, ": ", locVals);
+
+  // Take those local portions and plug them into the global domain/array
+  DS.setLocalSubdomain(locSpsInds);
+  A.setLocalSubarray(locVals);
+}
+
+
+// check the results
+
+
+for block in A.localSubarrays() {
+  writeln("locale ", here.id, " owns:\n", block);
+}
+
+forall block in A.localSubarrays() {
+  forall a in block {
+    a *= 10;
+  }
+}
+
+/* Equivalent to the following loop, but already tested above
+for block in A.localSubarrays() {
+  writeln("locale ", here.id, " owns:\n", block);
+}
+*/
+
+for loc in Locales {
+  on loc {
+    writeln("locale ", here.id, " owns:\n", A.getLocalSubarray());
+  }
+}

test/distributions/sparseBlock/localSubarraysIter.comm-none.good

@@ -0,0 +1,27 @@
+DS target locales:
+LOCALE0
+
+locale 0 owns:
+0.0
+1.1
+2.2
+3.3
+4.4
+5.5
+6.6
+7.7
+8.8
+9.9
+
+locale 0 owns:
+0.0
+11.0
+22.0
+33.0
+44.0
+55.0
+66.0
+77.0
+88.0
+99.0
+

test/distributions/sparseBlock/localSubarraysIter.good

@@ -0,0 +1,60 @@
+DS target locales:
+LOCALE0 LOCALE1
+LOCALE2 LOCALE3
+
+locale 0 owns:
+0.0
+1.1
+2.2
+3.3
+4.4
+
+locale 1 owns:
+0.5
+1.6
+2.7
+3.8
+4.9
+
+locale 2 owns:
+5.0
+6.1
+7.2
+8.3
+9.4
+
+locale 3 owns:
+5.5
+6.6
+7.7
+8.8
+9.9
+
+locale 0 owns:
+0.0
+11.0
+22.0
+33.0
+44.0
+
+locale 1 owns:
+5.0
+16.0
+27.0
+38.0
+49.0
+
+locale 2 owns:
+50.0
+61.0
+72.0
+83.0
+94.0
+
+locale 3 owns:
+55.0
+66.0
+77.0
+88.0
+99.0
+