Functions on subdomains (WIP)

devito/ir/equations/algorithms.py

@@ -89,9 +89,15 @@ def lower_exprs(expressions, **kwargs):
             # Some Relationals may be pure SymPy objects, thus lacking the subdomain
             dimension_map = {}
 
+        # Gather `Function`'s defined on a `SubDomain`
+        fosd = set([f for f in retrieve_functions(expr, mode='unique')
+                    if f.is_Function and f._subdomain])
+
         # Handle Functions (typical case)
-        mapper = {f: f.indexify(lshift=True, subs=dimension_map)
-                  for f in retrieve_functions(expr)}
+        mapper = {**{f: f.indexify(lshift=True, subs=dimension_map)
+                     for f in set(retrieve_functions(expr)).difference(fosd)},
+                  **{f: f.indexify(lshift=True, subs=f._subdomain._access_map)
+                     for f in fosd}}
 
         # Handle Indexeds (from index notation)
         for i in retrieve_indexed(expr):

devito/operations/interpolators.py

@@ -6,8 +6,8 @@
 
 from devito.finite_differences import Evaluable
 from devito.logger import warning
-from devito.symbolics import retrieve_function_carriers, indexify, INT
-from devito.tools import powerset, flatten, prod
+from devito.symbolics import retrieve_function_carriers, retrieve_functions, indexify, INT
+from devito.tools import powerset, flatten, prod, as_list
 from devito.types import Eq, Inc
 from devito.types.basic import Scalar
 from devito.types.dense import SubFunction
@@ -232,6 +232,21 @@ def callback():
             idx_subs, temps = self._interpolation_indices(variables, offset,
                                                           field_offset=field_offset)
 
+            # FIXME: Maybe needs a re-think wrt placement?
+            # NOTE: Probably other better approaches available
+            # Add idx_subs for `Function`'s on `SubDomain`'s
+            add_subs = []
+            fosd = set([f for f in retrieve_functions(expr, mode='unique')
+                        if f.is_Function and f._subdomain])
+            if len(fosd) > 1:
+                raise NotImplementedError
+            for f in fosd:
+                dims = f._subdomain.dimensions
+                for i in idx_subs:
+                    add_subs.append({k: v.subs({d: v for d, v in zip(dims, i.values())})
+                                     for k, v in f._subdomain._access_map.items()})
+                idx_subs = add_subs
+
             # Substitute coordinate base symbols into the interpolation coefficients
             args = [_expr.xreplace(v_sub) * b.xreplace(v_sub)
                     for b, v_sub in zip(self._interpolation_coeffs, idx_subs)]
@@ -278,11 +293,26 @@ def callback():
             # List of indirection indices for all adjacent grid points
             idx_subs, temps = self._interpolation_indices(variables, offset,
                                                           field_offset=field_offset)
+            add_subs = []
+            fosd = set([f for f in as_list(field) if f.is_Function and f._subdomain])
+            if len(fosd) > 1:
+                raise NotImplementedError
+            for f in fosd:
+                dims = f._subdomain.dimensions
+                for i in idx_subs:
+                    add_subs.append({k: v.subs({d: v for d, v in zip(dims, i.values())})
+                                     for k, v in f._subdomain._access_map.items()})
 
             # Substitute coordinate base symbols into the interpolation coefficients
-            eqns = [Inc(field.xreplace(vsub), _expr.xreplace(vsub) * b,
-                        implicit_dims=self.sfunction.dimensions)
-                    for b, vsub in zip(self._interpolation_coeffs, idx_subs)]
+            if add_subs:
+                eqns = [Inc(field.xreplace(asub), _expr.xreplace(vsub) * b,
+                            implicit_dims=self.sfunction.dimensions)
+                        for b, asub, vsub in zip(self._interpolation_coeffs,
+                                                 add_subs, idx_subs)]
+            else:
+                eqns = [Inc(field.xreplace(vsub), _expr.xreplace(vsub) * b,
+                            implicit_dims=self.sfunction.dimensions)
+                        for b, vsub in zip(self._interpolation_coeffs, idx_subs)]
 
             return temps + eqns
 

devito/types/dense.py

@@ -50,6 +50,7 @@ class DiscreteFunction(AbstractFunction, ArgProvider, Differentiable):
     # its key routines (e.g., solve)
     _iterable = False
 
+    is_Function = False
     is_Input = True
     is_DiscreteFunction = True
     is_Tensor = True
@@ -329,14 +330,17 @@ def _size_outhalo(self):
                 if not self._distributor.is_boundary_rank:
                     warning(warning_msg)
                 else:
-                    for i, j, k, l in zip(left, right, self._distributor.mycoords,
+                    for i, j, k, l in zip(left[-len(self._distributor.dimensions):],
+                                          right[-len(self._distributor.dimensions):],
+                                          self._distributor.mycoords,
                                           self._distributor.topology):
                         if l > 1 and ((j > 0 and k == 0) or (i > 0 and k == l-1)):
                             warning(warning_msg)
                             break
             except AttributeError:
                 pass
 
+        #from IPython import embed; embed()
         return DimensionTuple(*sizes, getters=self.dimensions, left=left, right=right)
 
     @property
@@ -722,6 +726,8 @@ def _halo_exchange(self):
             raise RuntimeError("`%s` cannot perform a halo exchange as it has "
                                "no Grid attached" % self.name)
 
+        if self.name == 'u':
+            from IPython import embed; embed()
         neighborhood = self._distributor.neighborhood
         comm = self._distributor.comm
 
@@ -961,6 +967,9 @@ def __init_finalize__(self, *args, **kwargs):
         # Dynamically add derivative short-cuts
         self._fd = generate_fd_shortcuts(self)
 
+        # Check if defined on a subdomain
+        self._subdomain = kwargs.get('subdomain', None)
+
         # Flag whether it is a parameter or a variable.
         # Used at operator evaluation to evaluate the Function at the
         # variable location (i.e. if the variable is staggered in x the
@@ -988,12 +997,16 @@ def _eval_at(self, func):
     @classmethod
     def __indices_setup__(cls, **kwargs):
         grid = kwargs.get('grid')
+        subdomain = kwargs.get('subdomain')
         dimensions = kwargs.get('dimensions')
         if grid is None:
             if dimensions is None:
                 raise TypeError("Need either `grid` or `dimensions`")
         elif dimensions is None:
-            dimensions = grid.dimensions
+            if subdomain is None:
+                dimensions = grid.dimensions
+            else:
+                dimensions = subdomain.dimensions
 
         # Staggered indices
         staggered = kwargs.get("staggered", None)
@@ -1015,15 +1028,21 @@ def is_Staggered(self):
     @classmethod
     def __shape_setup__(cls, **kwargs):
         grid = kwargs.get('grid')
+        subdomain = kwargs.get('subdomain')
         dimensions = kwargs.get('dimensions')
         shape = kwargs.get('shape', kwargs.get('shape_global'))
         if grid is None:
             if shape is None:
                 raise TypeError("Need either `grid` or `shape`")
         elif shape is None:
-            if dimensions is not None and dimensions != grid.dimensions:
+            if dimensions is not None and dimensions != (grid.dimensions or
+                                                         subdomain.dimensions):
                 raise TypeError("Need `shape` as not all `dimensions` are in `grid`")
-            shape = grid.shape_local
+            # TODO: Fix this:
+            if subdomain:
+                shape = subdomain.shape_local
+            else:
+                shape = grid.shape_local
         elif dimensions is None:
             raise TypeError("`dimensions` required if both `grid` and "
                             "`shape` are provided")
@@ -1096,6 +1115,35 @@ def __padding_setup__(self, **kwargs):
         else:
             raise TypeError("`padding` must be int or %d-tuple of ints" % self.ndim)
 
+    @cached_property
+    def _dist_dimensions(self):
+        """Tuple of MPI-distributed Dimensions."""
+        if self._distributor is None:
+            return ()
+        # FIXME: Tidy + fix for certain subdim cases
+        if self._subdomain:
+            dist_dimensions = set(self._distributor.dimensions + self._subdomain.dimensions)
+            return tuple(d for d in self.dimensions if d in  dist_dimensions)
+        else:
+            return tuple(d for d in self.dimensions if d in self._distributor.dimensions)
+
+    @cached_property
+    def _decomposition(self):
+        """
+        Tuple of Decomposition objects, representing the domain decomposition.
+        None is used as a placeholder for non-decomposed Dimensions.
+        """
+        #from IPython import embed; embed()
+        if self._distributor is None:
+            return (None,)*self.ndim
+        if self._subdomain:
+            # FIXME: Test mixed case
+            mapper = {d: self._subdomain._decomposition[d] for d in self._dist_dimensions}
+        else:
+            mapper = {d: self._distributor.decomposition[d] for d in self._dist_dimensions}
+        #temp = tuple(mapper.get(d) for d in self.dimensions)
+        return tuple(mapper.get(d) for d in self.dimensions)
+
     @property
     def space_order(self):
         """The space order."""
@@ -1301,8 +1349,9 @@ def __indices_setup__(cls, **kwargs):
     @classmethod
     def __shape_setup__(cls, **kwargs):
         grid = kwargs.get('grid')
+        subdomain = kwargs.get('subdomain')
         save = kwargs.get('save') or None  # Force to None if 0/False/None/...
-        shape = kwargs.get('shape')
+        shape = kwargs.get('shape', kwargs.get('shape_global'))
         time_order = kwargs.get('time_order', 1)
 
         if grid is None:
@@ -1312,7 +1361,10 @@ def __shape_setup__(cls, **kwargs):
                 raise TypeError("Ambiguity detected: provide either `grid` and `save` "
                                 "or just `shape` ")
         elif shape is None:
-            shape = list(grid.shape_local)
+            if subdomain:
+                shape = list(subdomain.shape_local)
+            else:
+                shape = list(grid.shape_local)
             if save is None:
                 shape.insert(cls._time_position, time_order + 1)
             elif isinstance(save, Buffer):

devito/types/grid.py

@@ -4,7 +4,7 @@
 from sympy import prod
 from math import floor
 
-from devito.data import LEFT, RIGHT
+from devito.data import LEFT, RIGHT, Decomposition
 from devito.mpi import Distributor
 from devito.tools import ReducerMap, as_tuple, memoized_meth
 from devito.types.args import ArgProvider
@@ -349,7 +349,7 @@ def __init__(self):
             raise ValueError("SubDomain requires a `name`")
         self._dimensions = None
 
-    def __subdomain_finalize__(self, dimensions, shape, **kwargs):
+    def __subdomain_finalize__(self, dimensions, shape, distributor=None, **kwargs):
         # Create the SubDomain's SubDimensions
         sub_dimensions = []
         sdshape = []
@@ -360,42 +360,127 @@ def __subdomain_finalize__(self, dimensions, shape, **kwargs):
                 sub_dimensions.append(v)
                 sdshape.append(s)
             else:
+                sd_name = 'i%d%s' % (counter, k.name)
                 try:
                     # Case ('middle', int, int)
                     side, thickness_left, thickness_right = v
                     if side != 'middle':
                         raise ValueError("Expected side 'middle', not `%s`" % side)
-                    sub_dimensions.append(SubDimension.middle('i%d%s' %
-                                                              (counter, k.name),
-                                                              k, thickness_left,
+                    sub_dimensions.append(SubDimension.middle(sd_name, k, thickness_left,
                                                               thickness_right))
-                    thickness = s-thickness_left-thickness_right
-                    sdshape.append(thickness)
+                    sdshape.append(s-thickness_left-thickness_right)
                 except ValueError:
                     side, thickness = v
                     if side == 'left':
                         if s-thickness < 0:
                             raise ValueError("Maximum thickness of dimension %s "
                                              "is %d, not %d" % (k.name, s, thickness))
-                        sub_dimensions.append(SubDimension.left('i%d%s' %
-                                                                (counter, k.name),
-                                                                k, thickness))
+                        sub_dimensions.append(SubDimension.left(sd_name, k,
+                                                                thickness))
                         sdshape.append(thickness)
                     elif side == 'right':
                         if s-thickness < 0:
                             raise ValueError("Maximum thickness of dimension %s "
                                              "is %d, not %d" % (k.name, s, thickness))
-                        sub_dimensions.append(SubDimension.right('i%d%s' %
-                                                                 (counter, k.name),
-                                                                 k, thickness))
+                        sub_dimensions.append(SubDimension.right(sd_name, k,
+                                                                 thickness))
                         sdshape.append(thickness)
                     else:
                         raise ValueError("Expected sides 'left|right', not `%s`" % side)
         self._dimensions = tuple(sub_dimensions)
 
-        # Compute the SubDomain shape
+        # Set the SubDomain shape
         self._shape = tuple(sdshape)
 
+        # Derive the local shape for `SubDomain`'s on distributed grids along with the
+        # memory access map for any `Function` defined on this `SubDomain`.
+        access_map = {}
+        shift = {}
+        shape_local = []
+        sub_decomposition = []
+        for counter, (dim, d, s) in enumerate(zip(sub_dimensions, distributor.decomposition, self._shape)):
+            if dim.is_Sub:
+                c_name = 'c_%s' % dim.name
+                shift[c_name] = Constant(name=c_name, dtype=np.int32)
+                if dim._local:
+                    if distributor and distributor.is_parallel:
+                        if dim.thickness.right[1] == 0:
+                            ls = len(d.loc_abs_numb)
+                            minc = d.glb_min
+                            maxc = d.glb_max - (ls-dim.thickness.left[1])
+                            l = d.index_glb_to_loc(0, LEFT)
+                            r = d.index_glb_to_loc(ls-dim.thickness.left[1], RIGHT)
+                        else:
+                            ls = len(d.loc_abs_numb)
+                            minc = d.glb_min + (ls-dim.thickness.right[1])
+                            maxc = d.glb_max
+                            l = d.index_glb_to_loc(ls-dim.thickness.right[1], LEFT)
+                            r = d.index_glb_to_loc(0, RIGHT)
+                        if d.loc_abs_min > maxc:
+                            shape_local.append(0)
+                        elif d.loc_abs_max < minc:
+                            shape_local.append(0)
+                        else:
+                            if l is None:
+                                l = 0
+                            if r is None:
+                                r = 0
+                            shape_local.append(ls-l-r)
+                        shift[c_name].data = l
+                        # FIXME: Looks like we can lift all these terms
+                        access_map.update({dim: dim-shift[c_name]})
+                    else:
+                        if dim.thickness.left[1] == 0:
+                            shift[c_name].data = (s-dim.thickness.right[1])
+                        else:
+                            shift[c_name].data = 0
+                        access_map.update({dim: dim-shift[c_name]})
+                        shape_local.append(s)
+                else:
+                    if distributor and distributor.is_parallel:
+                        tl = dim.thickness.left[1]
+                        tr = dim.thickness.right[1]
+                        ls = len(d.loc_abs_numb)
+                        minc = d.glb_min + tl
+                        maxc = d.glb_max - tr
+                        l = d.index_glb_to_loc(tl, LEFT)
+                        r = d.index_glb_to_loc(tr, RIGHT)
+                        if d.loc_abs_min > maxc:
+                            shape_local.append(0)
+                        elif d.loc_abs_max < minc:
+                            shape_local.append(0)
+                        else:
+                            if l is None:
+                                l = 0
+                            if r is None:
+                                r = 0
+                            shape_local.append(ls-l-r)
+                        shift[c_name].data = l
+                    else:
+                        shift[c_name].data = dim.thickness.left[1]
+                        shape_local.append(s)
+                    access_map.update({dim: dim-shift[c_name]})
+                # Create the 'sub-distributor'
+                send_min = np.array([max(minc, d.loc_abs_min)])
+                send_max = np.array([min(maxc, d.loc_abs_max)])
+                rec_min = np.zeros(distributor.comm.Get_size(),
+                                   dtype=send_min.dtype.type)
+                rec_max = np.zeros(distributor.comm.Get_size(),
+                                   dtype=send_max.dtype.type)
+                distributor.comm.Allgather(send_min, rec_min)
+                distributor.comm.Allgather(send_max, rec_max)
+                mins = sorted(set(rec_min))
+                maxs = sorted(set(rec_max))
+                decomp_arrays = [np.array(range(mi, mx+1)) for mi, mx in zip(mins, maxs)]
+                sub_decomposition.append(Decomposition(decomp_arrays,
+                                                       distributor.mycoords[counter]))
+            else:
+                shape_local.append(len(d.loc_abs_numb))
+                access_map.update({dim: dim})
+        self._access_map = access_map
+        self._shape_local = tuple(shape_local)
+        self._decomposition = tuple(sub_decomposition)
+
     def __eq__(self, other):
         if not isinstance(other, SubDomain):
             return False
@@ -421,6 +506,10 @@ def dimension_map(self):
     def shape(self):
         return self._shape
 
+    @property
+    def shape_local(self):
+        return self._shape_local
+
     def define(self, dimensions):
         """
         Parametrically describe the SubDomain w.r.t. a generic Grid.