Reduce the number of roundtrips for DistArray.get_localshapes to zero.

distarray/dist/distarray.py

@@ -262,8 +262,9 @@ def inner_fill(arr, value):
 
     def _reduce(self, local_reduce_name, axes=None, dtype=None, out=None):
 
-        if any(0 in localshape for localshape in self.get_localshapes()):
-            raise NotImplementedError("Reduction not implemented for empty LocalArrays")
+        if any(0 in localshape for localshape in self.localshapes()):
+            raise NotImplementedError("Reduction not implemented for empty "
+                                      "LocalArrays")
 
         if out is not None:
             _raise_nie()
@@ -338,10 +339,8 @@ def get(key):
             return key.copy()
         return self.context.apply(get, args=(self.key,), targets=self.targets)
 
-    def get_localshapes(self):
-        def get(key):
-            return key.local_shape
-        return self.context.apply(get, args=(self.key,), targets=self.targets)
+    def localshapes(self):
+        return self.distribution.localshapes()
 
     # Binary operators
 

distarray/dist/maps.py

@@ -38,7 +38,8 @@
                                       positivify,
                                       _start_stop_block,
                                       normalize_dim_dict,
-                                      normalize_reduction_axes)
+                                      normalize_reduction_axes,
+                                      shapes_from_dim_data_per_rank)
 
 
 def _dedup_dim_dicts(dim_dicts):
@@ -191,6 +192,7 @@ def __init__(self, size, grid_size):
             msg = "grid_size for NoDistMap must be 1 (given %s)"
             raise ValueError(msg % grid_size)
         self.size = size
+        self.grid_size = grid_size
 
     def owners(self, idx):
         return [0] if 0 <= idx < self.size else []
@@ -669,3 +671,6 @@ def reduce(self, axes):
                                        dist=reduced_dist,
                                        grid_shape=reduced_grid_shape,
                                        targets=reduced_targets)
+
+    def localshapes(self):
+        return shapes_from_dim_data_per_rank(self.get_dim_data_per_rank())

distarray/dist/tests/test_decorators.py

@@ -203,7 +203,7 @@ def test_local_sum(self):
         dd = self.local_sum(self.da)
         if self.ntargets == 1:
             dd = [dd]
-        lshapes = self.da.get_localshapes()
+        lshapes = self.da.localshapes()
         expected = []
         for lshape in lshapes:
             expected.append(lshape[0] * lshape[1] * (2 * numpy.pi))
@@ -234,7 +234,7 @@ def test_local_add_mixed(self):
 
     @unittest.skip('Locally adding ndarrays not supported.')
     def test_local_add_ndarray(self):
-        shp = self.da.get_localshapes()[0]
+        shp = self.da.localshapes()[0]
         ndarr = numpy.empty(shp)
         ndarr.fill(33)
         dk = self.local_add_ndarray(self.da, 11, ndarr)

distarray/dist/tests/test_distarray.py

@@ -365,9 +365,10 @@ class TestDistArrayCreationSubSet(ContextTestCase):
     def test_create_target_subset(self):
         shape = (100, 100)
         subtargets = self.context.targets[::2]
-        distribution = Distribution.from_shape(self.context, shape=shape, targets=subtargets)
+        distribution = Distribution.from_shape(self.context, shape=shape,
+                                               targets=subtargets)
         darr = self.context.ones(distribution)
-        lss = darr.get_localshapes()
+        lss = darr.localshapes()
         self.assertEqual(len(lss), len(subtargets))
 
         ddpr = distribution.get_dim_data_per_rank()

distarray/metadata_utils.py

@@ -283,3 +283,87 @@ def normalize_reduction_axes(axes, ndim):
     else:
         axes = tuple(positivify(a, ndim) for a in axes)
     return axes
+
+
+# functions for getting a size from a dim_data for each dist_type
+# n
+def non_dist_size(dim_data):
+    return dim_data['size']
+
+
+# b
+def block_size(dim_data):
+    stop = dim_data['stop']
+    start = dim_data['start']
+    return stop - start
+
+
+# choose cyclic or block cyclic based on blocks size. This is necessary
+# becuse they have the same dist type character.
+def c_or_bc_chooser(dim_data):
+    block_size = dim_data.get('block_size', 1)
+    if block_size == 1:
+        return cyclic_size(dim_data)
+    elif block_size > 1:
+        return block_cyclic_size(dim_data)
+    else:
+        raise ValueError("block_size %s is invalid" % block_size)
+
+
+# c
+def cyclic_size(dim_data):
+    global_size = dim_data['size']
+    grid_rank = dim_data.get('proc_grid_rank', 0)
+    grid_size = dim_data.get('proc_grid_size', 1)
+    return (global_size - 1 - grid_rank) // grid_size + 1
+
+
+# c
+def block_cyclic_size(dim_data):
+    global_size = dim_data['size']
+    block_size = dim_data.get('block_size', 1)
+    grid_size = dim_data.get('proc_grid_size', 1)
+    grid_rank = dim_data.get('proc_grid_rank', 0)
+
+    global_nblocks, partial = divmod(global_size, block_size)
+    local_partial = partial if grid_rank == 0 else 0
+    local_nblocks = (global_nblocks - 1 - grid_rank) // grid_size + 1
+    return local_nblocks * block_size + local_partial
+
+
+# u
+def unstructured_size(dim_data):
+    return len(dim_data.get('indices', None))
+
+
+def size_from_dim_data(dim_data):
+    """
+    Get a size from a dim_data.
+    """
+    return size_chooser(dim_data['dist_type'])(dim_data)
+
+
+def size_chooser(dist_type):
+    """
+    Get a function from a dist_type.
+    """
+    chooser = {'n': non_dist_size,
+               'b': block_size,
+               'c': c_or_bc_chooser,
+               'u': unstructured_size}
+    return chooser[dist_type]
+
+
+def shapes_from_dim_data_per_rank(ddpr):  # ddpr = dim_data_per_rank
+    """
+    Given a dim_data_per_rank object, return the shapes of the localarrays.
+    This requires no communication.
+    """
+    # create the list of shapes
+    shape_list = []
+    for rank_dd in ddpr:
+        shape = []
+        for dd in rank_dd:
+            shape.append(size_from_dim_data(dd))
+        shape_list.append(tuple(shape))
+    return shape_list

distarray/tests/test_metadata_utils.py

@@ -6,6 +6,7 @@
 
 import unittest
 from distarray import metadata_utils
+from distarray.dist import Distribution, Context
 
 
 class TestMakeGridShape(unittest.TestCase):
@@ -26,5 +27,69 @@ def test_negative_index(self):
         self.assertEqual(result, 8)
 
 
+class TestGridSizes(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.context = Context()
+
+    def test_dist_sizes(self):
+        dist = Distribution.from_shape(self.context, (2, 3, 4),
+                                       dist=('n', 'b', 'c'))
+        ddpr = dist.get_dim_data_per_rank()
+        shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
+        if len(self.context.view) == 4:
+            self.assertEqual(shapes, [(2, 2, 2), (2, 2, 2), (2, 1, 2),
+                                      (2, 1, 2)])
+
+    def test_n_size(self):
+        dim_dict = {'dist_type': 'n',
+                    'size': 42,
+                    'proc_grid_size': 1,
+                    'proc_grid_rank': 0}
+
+        dist = Distribution(self.context, (dim_dict,))
+        ddpr = dist.get_dim_data_per_rank()
+        shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
+        self.assertEqual(shapes, [(42,)])
+
+    def test_b_size(self):
+        dim_dict = {'dist_type': 'b',
+                    'size': 42,
+                    'bounds': [0, 20, 42],
+                    'proc_grid_size': 2,
+                    'proc_grid_rank': 0,
+                    'start': 0,
+                    'stop': 42}
+        dist = Distribution(self.context, (dim_dict,))
+        ddpr = dist.get_dim_data_per_rank()
+        shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
+        self.assertEqual(shapes, [(20,), (22,)])
+
+    def test_c_size(self):
+        dim_dict = {'dist_type': 'c',
+                    'size': 42,
+                    'proc_grid_size': 2,
+                    'proc_grid_rank': 0,
+                    'start': 0}
+        dist = Distribution(self.context, (dim_dict,))
+        ddpr = dist.get_dim_data_per_rank()
+        shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
+        self.assertEqual(shapes, [(21,), (21,)])
+
+    def test_bc_size(self):
+        dim_dict = {'dist_type': 'b',
+                    'size': 42,
+                    'block_size': 2,
+                    'bounds': [0, 20, 42],
+                    'proc_grid_size': 2,
+                    'proc_grid_rank': 0,
+                    'start': 0}
+        dist = Distribution(self.context, (dim_dict,))
+        ddpr = dist.get_dim_data_per_rank()
+        shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
+        self.assertEqual(shapes, [(20,), (22,)])
+
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)