use distributed domains for slice indexing in [slice] commands (#3634)

Signed-off-by: Jeremiah Corrado <jeremiah.corrado@hpe.com>
Bears-R-Us · Aug 7, 2024 · 7c828f8 · 7c828f8
1 parent 8df0476
commit 7c828f8
Show file tree

Hide file tree

Showing 4 changed files with 143 additions and 148 deletions.
diff --git a/src/IndexingMsg.chpl b/src/IndexingMsg.chpl
@@ -215,13 +215,15 @@ module IndexingMsg
             rngs[dim] = convertSlice(starts[dim], stops[dim], strides[dim]);
             outSizes[dim] = rngs[dim].size;
         }
-        const sliceDom = {(...rngs)};
+
+        const sliceDom = makeDistDom({(...rngs)});
         var arraySlice = makeDistArray((...outSizes), t);
 
-        forall (elt,j) in zip(arraySlice, sliceDom) with (var agg = newSrcAggregator(t)) do
+        forall (elt,j) in zip(arraySlice, sliceDom) with (var agg = newSrcAggregator(t)) {
             agg.copy(elt,array[j]);
+        }
 
-        return new shared SymEntry(arraySlice, max_bits);
+        return new shared SymEntry(arraySlice, max_bits=max_bits);
     }
 
     /*
@@ -917,7 +919,7 @@ module IndexingMsg
         var rngs: d.rank*range(strides=strideKind.any);
         for param dim in 0..<d.rank do
             rngs[dim] = convertSlice(starts[dim], stops[dim], strides[dim]);
-        const sliceDom = {(...rngs)};
+        const sliceDom = makeDistDom({(...rngs)});
 
         if t == bigint {
             var val_mb = value;
@@ -948,7 +950,7 @@ module IndexingMsg
         var sliceRanges: array_nd * range(strides=strideKind.any);
         for param dim in 0..<array_nd do
             sliceRanges[dim] = convertSlice(starts[dim], stops[dim], strides[dim]);
-        const sliceDom = {(...sliceRanges)};
+        const sliceDom = makeDistDom({(...sliceRanges)});
 
         if array_dtype_a == bigint {
             var bb = b.a:bigint;

diff --git a/src/MultiTypeSymEntry.chpl b/src/MultiTypeSymEntry.chpl
@@ -246,7 +246,7 @@ module MultiTypeSymEntry
           :type a: [] ?etype
         */
         proc init(a: [?D] ?etype, max_bits=-1) where MyDmap != Dmap.defaultRectangular && a.isDefaultRectangular() {
-            this.init(D.size, etype, D.rank);
+            this.init((...D.shape), etype);
             this.tupShape = D.shape;
             this.a = a;
             this.shape = tupShapeString(this.tupShape);

diff --git a/src/SymArrayDmap.chpl b/src/SymArrayDmap.chpl
@@ -18,30 +18,23 @@ module SymArrayDmap {
         :type shape: int
     */
     proc makeDistDom(shape: int ...?N) {
-        var rngs: N*range;
-        for i in 0..#N do rngs[i] = 0..#shape[i];
-        const dom = {(...rngs)};
-
-        select MyDmap
-        {
-            when Dmap.defaultRectangular {
-                return dom;
-            }
-            when Dmap.blockDist {
-                if dom.size > 0 {
-                    return blockDist.createDomain(dom);
-                }
-                // fix the annoyance about boundingBox being empty
-                else {
-                  return dom dmapped new blockDist(boundingBox=dom.expand(1));
-                }
-            }
-            otherwise {
-                halt("Unsupported distribution " + MyDmap:string);
-            }
-        }
+      var rngs: N*range;
+      for i in 0..#N do rngs[i] = 0..#shape[i];
+      const dom = {(...rngs)};
+
+      return makeDistDom(dom);
     }
 
+    proc makeDistDom(dom: domain(?)) {
+      select MyDmap {
+        when Dmap.defaultRectangular {
+          return dom;
+        }
+        when Dmap.blockDist {
+          return blockDist.createDomain(dom);
+        }
+      }
+    }
 
     /*
         Makes an array of specified type over a distributed domain