2014-06 APUG talk

distarray/dist/distarray.py

@@ -17,10 +17,12 @@
 import operator
 from itertools import product
 from functools import reduce
+from collections import Sequence
 
 import numpy as np
 
 import distarray
+from distarray.metadata_utils import sanitize_indices
 from distarray.dist.maps import Distribution
 from distarray.utils import _raise_nie
 from distarray.metadata_utils import normalize_reduction_axes
@@ -32,7 +34,6 @@
 # Code
 # ---------------------------------------------------------------------------
 
-
 class DistArray(object):
 
     __array_priority__ = 20.0
@@ -84,7 +85,8 @@ def get_dim_datas_and_dtype(arr):
 
         # has context, get dist and dtype
         elif (distribution is None) and (dtype is None):
-            res = context.apply(get_dim_datas_and_dtype, args=(key,))
+            res = context.apply(get_dim_datas_and_dtype, args=(key,),
+                                targets=targets)
             dim_datas = [i[0] for i in res]
             dtypes = [i[1] for i in res]
             da._dtype = dtypes[0]
@@ -95,7 +97,8 @@ def get_dim_datas_and_dtype(arr):
         # has context and dtype, get dist
         elif (distribution is None) and (dtype is not None):
             da._dtype = dtype
-            dim_datas = context.apply(getattr, args=(key, 'dim_data'))
+            dim_datas = context.apply(getattr, args=(key, 'dim_data'),
+                                      targets=targets)
             da.distribution = Distribution.from_dim_data_per_rank(context,
                                                                   dim_datas,
                                                                   targets)
@@ -128,10 +131,31 @@ def __repr__(self):
             (self.shape, self.targets)
         return s
 
+    def _process_return_value(self, result, return_proxy, index, targets):
+
+        if return_proxy:
+            # proxy returned as result of slice
+            # slicing shouldn't alter the dtype
+            result = result[0]
+            return DistArray.from_localarrays(key=result,
+                                              context=self.context,
+                                              targets=targets,
+                                              dtype=self.dtype)
+
+        elif isinstance(result, Sequence):
+            somethings = [i for i in result if i is not None]
+            if len(somethings) == 0:
+                # using checked_getitem and all return None
+                raise IndexError("Index %r is is not present." % (index,))
+            if len(somethings) == 1:
+                return somethings[0]
+            else:
+                return result
+        else:
+            assert False  # impossible is nothing
+
+
     def __getitem__(self, index):
-        #TODO: FIXME: major performance improvements possible here,
-        # especially for special cases like `index == slice(None)`.
-        # This would dramatically improve tondarray's performance.
 
         # to be run locally
         def checked_getitem(arr, index):
@@ -141,38 +165,35 @@ def checked_getitem(arr, index):
         def raw_getitem(arr, index):
             return arr.global_index[index]
 
-        if isinstance(index, int) or isinstance(index, slice):
-            tuple_index = (index,)
-            return self.__getitem__(tuple_index)
-
-        elif isinstance(index, tuple):
-            targets = self.distribution.owning_targets(index)
-
-            args = (self.key, index)
-            if self.distribution.has_precise_index:
-                result = self.context.apply(raw_getitem, args=args,
-                                            targets=targets)
-            else:
-                result = self.context.apply(checked_getitem, args=args,
-                                            targets=targets)
-            result = [i for i in result if i is not None]
-            if len(result) != 1:
-                raise IndexError("Getting more than one result (%s) is not "
-                                 "supported yet." % (result,))
-            elif result is None:
-                raise IndexError("Index %r is out of bounds" % (index,))
-            else:
-                return result[0]
-        else:
-            raise TypeError("Invalid index type.")
+        # to be run locally
+        def get_slice(arr, index, ddpr, comm):
+            from distarray.local.maps import Distribution
+            local_distribution = Distribution(comm=comm,
+                                              dim_data=ddpr[comm.Get_rank()])
+            result = arr.global_index.get_slice(index, local_distribution)
+            return proxyize(result)
+
+        return_type, index = sanitize_indices(index, ndim=self.ndim,
+                                              shape=self.shape)
+        return_proxy = (return_type == 'view')
+        targets = self.distribution.owning_targets(index)
+
+        args = [self.key, index]
+        if self.distribution.has_precise_index:
+            if return_proxy:  # returning a new DistArray view
+                new_distribution = self.distribution.slice(index)
+                ddpr = new_distribution.get_dim_data_per_rank()
+                args.extend([ddpr, new_distribution.comm])
+                local_fn = get_slice
+            else:  # returning a value
+                local_fn = raw_getitem
+        else:  # returning a value from unstructured
+            local_fn = checked_getitem
+
+        result = self.context.apply(local_fn, args=args, targets=targets)
+        return self._process_return_value(result, return_proxy, index, targets)
 
     def __setitem__(self, index, value):
-        #TODO: FIXME: major performance improvements possible here.
-        # Especially when `index == slice(None)` and value is an
-        # ndarray, since for block and cyclic, we can generate slices of
-        # `value` and assign to local arrays. This would dramatically
-        # improve the fromndarray method's performance.
-
         # to be run locally
         def checked_setitem(arr, index, value):
             return arr.global_index.checked_setitem(index, value)
@@ -181,26 +202,57 @@ def checked_setitem(arr, index, value):
         def raw_setitem(arr, index, value):
             arr.global_index[index] = value
 
-        if isinstance(index, int) or isinstance(index, slice):
-            tuple_index = (index,)
-            return self.__setitem__(tuple_index, value)
-
-        elif isinstance(index, tuple):
-            targets = self.distribution.owning_targets(index)
-            args = (self.key, index, value)
-            if self.distribution.has_precise_index:
-                self.context.apply(raw_setitem, args=args, targets=targets)
+        # to be run locally
+        def set_slice(arr, index, value, value_slices):
+            local_slice = value_slices[arr.comm_rank]
+            arr.global_index[index] = value[local_slice]
+
+        set_type, index = sanitize_indices(index, ndim=self.ndim,
+                                           shape=self.shape)
+
+        targets = self.distribution.owning_targets(index)
+        args = [self.key, index, value]
+        if self.distribution.has_precise_index:
+            if set_type == 'value':
+                local_fn = raw_setitem
+            elif set_type == 'view':
+                args[-1] = np.asarray(args[-1])  # convert to array
+                # this could be made more efficient
+                # we only need the bounds computed by distribution.slice
+                new_distribution = self.distribution.slice(index)
+                if args[-1].shape != new_distribution.shape:
+                    msg = "Slice shape does not equal rvalue shape."
+                    raise ValueError(msg)
+                ddpr = new_distribution.get_dim_data_per_rank()
+                def bounds_slice(dd):
+                    if dd['dist_type'] == 'b':
+                        return slice(dd['start'], dd['stop'])
+                    elif dd['dist_type'] == 'n':
+                        return slice(0, dd['size'])
+                    else:
+                        msg = "Function only works for 'n' and 'b' 'dist_type's"
+                        raise TypeError(msg)
+                value_slices =  [tuple(bounds_slice(dd) for dd in dim_data)
+                                 for dim_data in ddpr]
+                # but we need a data structure indexable by a target's rank
+                # assume contigious range of targets here
+                value_slices_per_target = [None] * len(self.targets)
+                value_slices_per_target[targets[0]:targets[-1]] = value_slices
+                args.append(value_slices_per_target)
+                local_fn = set_slice
             else:
-                result = self.context.apply(checked_setitem, args=args,
-                                            targets=targets)
-                result = [i for i in result if i is not None]
-                if len(result) > 1:
-                    raise IndexError("Setting more than one result (%s) is "
-                                     "not supported yet." % (result,))
-                elif result == []:
-                    raise IndexError("Index %s is out of bounds" % (index,))
-        else:
-            raise TypeError("Invalid index type.")
+                assert False
+            self.context.apply(local_fn, args=args, targets=targets)
+
+        else:  # setting unstructured elements
+            local_fn = checked_setitem
+            result = self.context.apply(local_fn, args=args, targets=targets)
+            result = [i for i in result if i is not None]
+            if len(result) > 1:
+                raise IndexError("Setting more than one result (%s) is "
+                                 "not supported yet." % (result,))
+            elif result == []:
+                raise IndexError("Index %s is out of bounds" % (index,))
 
     @property
     def context(self):
@@ -246,7 +298,8 @@ def tondarray(self):
         """Returns the distributed array as an ndarray."""
         arr = np.empty(self.shape, dtype=self.dtype)
         local_name = self.context._generate_key()
-        self.context._execute('%s = %s.copy()' % (local_name, self.key), targets=self.targets)
+        self.context._execute('%s = %s.copy()' % (local_name, self.key),
+                              targets=self.targets)
         local_arrays = self.context._pull(local_name, targets=self.targets)
         for local_array in local_arrays:
             maps = (list(ax_map.global_iter) for ax_map in