diff --git a/distarray/dist/distarray.py b/distarray/dist/distarray.py index fa079e07..0cfe79be 100644 --- a/distarray/dist/distarray.py +++ b/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 diff --git a/distarray/dist/maps.py b/distarray/dist/maps.py index bf09e618..46195202 100644 --- a/distarray/dist/maps.py +++ b/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()) diff --git a/distarray/dist/tests/test_decorators.py b/distarray/dist/tests/test_decorators.py index bbf4359a..e6ea707e 100644 --- a/distarray/dist/tests/test_decorators.py +++ b/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) diff --git a/distarray/dist/tests/test_distarray.py b/distarray/dist/tests/test_distarray.py index 357f6d52..4ec083f7 100644 --- a/distarray/dist/tests/test_distarray.py +++ b/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() diff --git a/distarray/metadata_utils.py b/distarray/metadata_utils.py index f3c3eef9..1bf3d6e0 100644 --- a/distarray/metadata_utils.py +++ b/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 diff --git a/distarray/tests/test_metadata_utils.py b/distarray/tests/test_metadata_utils.py index bab24861..fe8760c2 100644 --- a/distarray/tests/test_metadata_utils.py +++ b/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)