Support for reading zarr stores with a custom locale distribution (ch…

…apel-lang#25803)

This PR adds support for customizing the distribution of Zarr stores
across locales with an `targetLocale` argument. This will help with
getting reasonable distributions for stores that have small sizes in the
first dimension.

Tested locally and on chapdl with none and gasnet comms

Reviewed by: @benharsh

modules/packages/Zarr.chpl

@@ -347,7 +347,7 @@ module Zarr {
     :arg dimCount: Dimensionality of the zarr array
 
   */
-  proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) throws {
+  proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1, targetLocales: [] locale = Locales) throws {
     var md = getMetadata(directoryPath);
     validateMetadata(md, dtype, dimCount);
     // Size and shape tuples
@@ -365,7 +365,7 @@ module Zarr {
 
     // Initialize the distributed domain and array
     const undistD : domain(dimCount) = totalRanges;
-    const Dist = new blockDist(boundingBox=undistD);
+    const Dist = new blockDist(boundingBox=undistD, targetLocales=targetLocales);
     const D = Dist.createDomain(undistD);
     var A: [D] dtype;
 

test/library/packages/Zarr/.gitignore

@@ -5,3 +5,8 @@ Test3D
 c-blosc
 *.dylib
 *.so*
+*.dll*
+LocalIOStore_*
+TestLocal
+TestTargetLocale
+

test/library/packages/Zarr/CLEANFILES

@@ -5,3 +5,5 @@ Test1D
 Test2D
 Test3D
 LocalIOStore_*
+TestLocal
+TestTargetLocale

test/library/packages/Zarr/ZarrTargetLocale.chpl

@@ -0,0 +1,39 @@
+use Zarr;
+use Random;
+use BlockDist;
+use IO;
+
+
+const D = {0..<10, 0..<10};
+var A: [D] real(32);
+fillRandom(A);
+
+// write the store that we will read in a distributed fashion
+writeZarrArrayLocal("TestTargetLocale", A, (3,3));
+
+// read with the default
+var defaultArray = readZarrArray("TestTargetLocale", real(32), 2);
+
+// read with custom locale distribution
+const myTargetLocales = reshape(Locales, {0..0, 0..#numLocales});
+var targetedArray = readZarrArray("TestTargetLocale", real(32), 2, targetLocales=myTargetLocales);
+
+on Locales[0] {
+  var defaultLocale0 = defaultArray.localSubdomain();
+  assert(defaultLocale0 == {0..<5, 0..<10}, "Locale 0 subdomain incorrect for default locale distribution: %?".format(defaultLocale0));
+  var targetLocale0 = targetedArray.localSubdomain();
+  assert(targetLocale0 == {0..<10, 0..<5}, "Locale 0 subdomain incorrect for targeted locale distribution: %?".format(targetLocale0));
+}
+
+on Locales[1] {
+  var defaultLocale1 = defaultArray.localSubdomain();
+  assert(defaultLocale1 == {5..<10, 0..<10}, "Locale 1 subdomain incorrect for default locale distribution: %?".format(defaultLocale1));
+  var targetLocale1 = targetedArray.localSubdomain();
+  assert(targetLocale1 == {0..<10, 5..<10}, "Locale 1 subdomain incorrect for targeted locale distribution: %?".format(targetLocale1));
+}
+
+for i in D {
+  assert(defaultArray[i] == A[i], "Default locale distribution read incorrect at index %?".format(i));
+  assert(targetedArray[i] == A[i], "Targeted locale distribution read incorrect at index %?".format(i));
+}
+writeln("Pass");

test/library/packages/Zarr/ZarrTargetLocale.good

@@ -0,0 +1 @@
+Pass

test/library/packages/Zarr/ZarrTargetLocale.numlocales

@@ -0,0 +1 @@
+2

test/library/packages/Zarr/ZarrTargetLocale.skipif

@@ -0,0 +1 @@
+CHPL_COMM==none