dsl: Start halo update fix. Halo exchange now aware of subdomain-rank…

… overlap

devito/mpi/distributed.py

@@ -363,8 +363,7 @@ def glb_to_rank(self, index):
             ret[r] = filter_ordered(inds[0])
         return ret
 
-    @property
-    def neighborhood(self):
+    def neighborhood(self, subdomain=None):
         """
         A mapper ``M`` describing the calling MPI rank's neighborhood in the
         decomposed grid. Let
@@ -383,20 +382,27 @@ def neighborhood(self):
               ``d0``, ``s1`` the DataSide along ``d1``, and so on. This can be
               useful to retrieve the diagonal neighbours (e.g., ``M[(LEFT, LEFT)]``
               gives the top-left neighbour in a 2D grid).
+
+        An optional `subdomain` parameter allows a `SubDomain` to be passed. This
+        is used to check the calling MPI rank's neighborhood in the decomposed grid
+        against the footprint of the subdomain on this grid to identify no-ops.
         """
+        if subdomain:
+            crosses = subdomain._crosses
         # Set up horizontal neighbours
         shifts = {d: self.comm.Shift(i, 1) for i, d in enumerate(self.dimensions)}
         ret = {}
         for d, (src, dest) in shifts.items():
             ret[d] = {}
-            ret[d][LEFT] = src
-            ret[d][RIGHT] = dest
+            ret[d][LEFT] = MPI.PROC_NULL if subdomain and not crosses[d][LEFT] else src
+            ret[d][RIGHT] = MPI.PROC_NULL if subdomain and not crosses[d][RIGHT] else dest
 
         # Set up diagonal neighbours
         for i in product([LEFT, CENTER, RIGHT], repeat=self.ndim):
             neighbor = [c + s.val for c, s in zip(self.mycoords, i)]
 
-            if any(c < 0 or c >= s for c, s in zip(neighbor, self.topology)):
+            if any(c < 0 or c >= s for c, s in zip(neighbor, self.topology)) \
+               or subdomain and not crosses[i]:
                 ret[i] = MPI.PROC_NULL
             else:
                 ret[i] = self.comm.Get_cart_rank(neighbor)
@@ -414,7 +420,7 @@ def _obj_neighborhood(self):
         A CompositeObject describing the calling MPI rank's neighborhood
         in the decomposed grid.
         """
-        return MPINeighborhood(self.neighborhood)
+        return MPINeighborhood(self.neighborhood())
 
     def _rebuild(self, shape=None, dimensions=None, comm=None):
         return Distributor(shape or self.shape, dimensions or self.dimensions,

devito/mpi/routines.py

@@ -1250,7 +1250,7 @@ def _arg_defaults(self, allocator, alias=None, args=None):
         super()._arg_defaults(allocator, alias, args=args)
 
         f = alias or self.target.c0
-        neighborhood = f.grid.distributor.neighborhood
+        neighborhood = f.grid.distributor.neighborhood()
 
         for i, halo in enumerate(self.halos):
             entry = self.value[i]

devito/types/basic.py

@@ -1194,7 +1194,9 @@ def _dist_dimensions(self):
         if self._distributor is None:
             return ()
         else:
+            # FIXME: Generating the correct halo size breaks things when subdomains are off rank
             dims = [d.parent if d.is_Sub else d for d in self.dimensions]
+            # dims = self.dimensions
             return tuple(d for d in dims if d in self._distributor.dimensions)
 
     @cached_property

devito/types/dense.py

@@ -322,6 +322,8 @@ def _size_outhalo(self):
         right = [max(i.loc_abs_max+j-i.glb_max, 0) if i and not i.loc_empty else 0
                  for i, j in zip(self._decomposition, self._size_inhalo.right)]
 
+        print("Decomposition", self._decomposition)
+
         sizes = tuple(Size(i, j) for i, j in zip(left, right))
 
         if self._distributor.is_parallel and (any(left) > 0 or any(right)) > 0:
@@ -752,9 +754,19 @@ def _halo_exchange(self):
             raise RuntimeError("`%s` cannot perform a halo exchange as it has "
                                "no Grid attached" % self.name)
 
-        neighborhood = self._distributor.neighborhood
+        # TODO: Will need a function to check the distributor neighbourhood against where
+        # the subdomain crosses the edges of the rank
+        if self.grid and self.grid.is_SubDomain:
+            neighborhood = self._distributor.neighborhood(subdomain=self.grid)
+        else:
+            neighborhood = self._distributor.neighborhood()
         comm = self._distributor.comm
 
+        print("Rank: %s" % comm.Get_rank(), self.shape,
+              self._size_inhalo, self._size_outhalo)
+
+        print("Neighbourhood", neighborhood)
+
         for d in self._dist_dimensions:
             for i in [LEFT, RIGHT]:
                 # Get involved peers
@@ -767,9 +779,11 @@ def _halo_exchange(self):
 
                 # Setup recv buffer
                 shape = self._data_in_region(HALO, d, i.flip()).shape
+                print("Rank: %s" % comm.Get_rank(), d, i, dest, source, shape)
                 recvbuf = np.ndarray(shape=shape, dtype=self.dtype)
 
                 # Communication
+                # FIXME: This line fails
                 comm.Sendrecv(sendbuf, dest=dest, recvbuf=recvbuf, source=source)
 
                 # Scatter received data

devito/types/equation.py

@@ -70,7 +70,7 @@ def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=No
                 **kwargs):
         kwargs['evaluate'] = False
         obj = sympy.Eq.__new__(cls, lhs, rhs, **kwargs)
-        obj._subdomain = cls.__subdomain_setup__(lhs, rhs, **kwargs)
+        obj._subdomain = cls.__subdomain_setup__(lhs, rhs, subdomain)
         obj._substitutions = coefficients
         obj._implicit_dims = as_tuple(implicit_dims)
 
@@ -126,8 +126,7 @@ def _flatten(self):
             return [self]
 
     @classmethod
-    def __subdomain_setup__(cls, lhs, rhs, **kwargs):
-        subdomain = kwargs.get('subdomain')
+    def __subdomain_setup__(cls, lhs, rhs, subdomain):
         subdomains = {subdomain} if subdomain else set()
 
         # Get all the functions in the LHS and RHS

devito/types/grid.py

@@ -1,15 +1,16 @@
 from abc import ABC
 from collections import namedtuple
+from itertools import product
 from math import floor
 
 import numpy as np
 from sympy import prod, Interval
 from cached_property import cached_property
 
-from devito.data import LEFT, RIGHT
+from devito.data import LEFT, CENTER, RIGHT
 from devito.logger import warning
 from devito.mpi import Distributor, MPI
-from devito.tools import ReducerMap, as_tuple
+from devito.tools import ReducerMap, as_tuple, frozendict
 from devito.types.args import ArgProvider
 from devito.types.basic import Scalar
 from devito.types.dense import Function
@@ -581,9 +582,6 @@ def __subdomain_finalize__(self, grid=None, **kwargs):
         self._dimensions = tuple(sub_dimensions)
         self._dtype = self.grid.dtype
 
-        # TODO: Rebuild the distributor with the SubDomain dimensions
-        # TODO: Figure out what shape this should have
-
         dist_interval = {dim: Interval(s.start, s.stop-1)
                          for dim, s in self.distributor.glb_slices.items()}
 
@@ -608,14 +606,31 @@ def __subdomain_finalize__(self, grid=None, **kwargs):
                            else dist_interval[dim].intersect(sdim_interval[dim]))
                           for dim in self.grid.dimensions)
 
+        # Keep track of the edges of the rank crossed by the SubDomain
+        # Format is {dimension: {side: crosses}} or {(sides,): crosses}
+        crosses = {}
         for dim, sdim in zip(self.grid.dimensions, self.dimensions):
             if sdim.is_Sub:
-                if sdim.local:
-                    in_rank = dist_interval[dim].issuperset(sdim_interval[dim])
-                    off_rank = dist_interval[dim].isdisjoint(sdim_interval[dim])
-                    if not in_rank and not off_rank:
-                        raise ValueError("SubDimension %s is local and cannot be"
-                                         " decomposed across MPI ranks" % dim)
+                in_rank = dist_interval[dim].issuperset(sdim_interval[dim])
+                off_rank = dist_interval[dim].isdisjoint(sdim_interval[dim])
+                if in_rank or off_rank:
+                    crosses[dim] = {LEFT: False, RIGHT: False}
+                elif sdim.local:
+                    raise ValueError("SubDimension %s is local and cannot be"
+                                     " decomposed across MPI ranks" % dim)
+                else:
+                    crosses[dim] = {LEFT: sdim_interval[dim].left
+                                    < dist_interval[dim].left,
+                                    RIGHT: sdim_interval[dim].right
+                                    > dist_interval[dim].right}
+            else:
+                crosses[dim] = {LEFT: True, RIGHT: True}
+
+        for i in product([LEFT, CENTER, RIGHT], repeat=len(self.shape)):
+            crosses[i] = all(crosses[d][s] for d, s in zip(self.grid.dimensions, i)
+                             if s in crosses[d])  # Skip over CENTER
+
+        self._crosses = frozendict(crosses)
 
         if any([i.is_empty for i in intervals]):
             # SubDomain does not overlap this rank

tests/test_mpi.py

@@ -109,8 +109,8 @@ def test_neighborhood_horizontal_2d(self):
             7: {x: {LEFT: 4, RIGHT: PN}, y: {LEFT: 6, RIGHT: 8}},
             8: {x: {LEFT: 5, RIGHT: PN}, y: {LEFT: 7, RIGHT: PN}},
         }
-        assert expected[distributor.myrank][x] == distributor.neighborhood[x]
-        assert expected[distributor.myrank][y] == distributor.neighborhood[y]
+        assert expected[distributor.myrank][x] == distributor.neighborhood()[x]
+        assert expected[distributor.myrank][y] == distributor.neighborhood()[y]
 
     @pytest.mark.parallel(mode=9)
     def test_neighborhood_diagonal_2d(self):
@@ -140,7 +140,7 @@ def test_neighborhood_diagonal_2d(self):
             7: {(LEFT, LEFT): 3, (LEFT, RIGHT): 5, (RIGHT, LEFT): PN, (RIGHT, RIGHT): PN},
             8: {(LEFT, LEFT): 4, (LEFT, RIGHT): PN, (RIGHT, LEFT): PN, (RIGHT, RIGHT): PN}  # noqa
         }
-        assert all(expected[distributor.myrank][i] == distributor.neighborhood[i]
+        assert all(expected[distributor.myrank][i] == distributor.neighborhood()[i]
                    for i in [(LEFT, LEFT), (LEFT, RIGHT), (RIGHT, LEFT), (RIGHT, RIGHT)])
 
     @pytest.mark.parallel(mode=[2, 4])

tests/test_subdomains.py

@@ -892,8 +892,8 @@ class TestSubdomainFunctionsParallel:
     #                                ('middle', 2, 3), ('middle', 1, 7),
     #                                None])
     @pytest.mark.parallel(mode=[(2, 'full')])
-    @pytest.mark.parametrize('x', [('left', 3)])
-    @pytest.mark.parametrize('y', [('left', 3)])
+    @pytest.mark.parametrize('x', [('middle', 2, 3)])
+    @pytest.mark.parametrize('y', [('middle', 2, 3)])
     def test_function_data_shape_mpi(self, x, y):
         """
         Check that defining a Function on a subset of a Grid results in arrays
@@ -903,21 +903,36 @@ def test_function_data_shape_mpi(self, x, y):
         reduced_domain = ReducedDomain(x, y, grid=grid)
         f = Function(name='f', grid=reduced_domain, space_order=2)
 
-        g = Function(name='g', grid=grid)
-        Operator(Eq(g, g+1, subdomain=reduced_domain))()
+        g = Function(name='g', grid=grid, space_order=2)
+        eq = Eq(g, g+1, subdomain=reduced_domain)
+        Operator(eq)()
 
         slices = tuple(grid.distributor.glb_slices[dim]
                        for dim in grid.dimensions)
+
         check_data = g.data[slices]
 
+        print("g data")
+        print(g.data)
+        print("Check data")
+        print(check_data)
+        print("Data")
+        print(f.data)
+        print("Data with halo")
+        print(f.data_with_halo)
+        for d in grid.dimensions:
+            print('Inhalo %s:' % d.name, f._size_inhalo[d])
+            print('Outhalo %s:' % d.name, f._size_outhalo[d])
+
         assert np.count_nonzero(check_data) == f.data.size
 
         if np.count_nonzero(check_data) != 0:
             coords = np.nonzero(check_data)
             shape = tuple(np.amax(c)-np.amin(c)+1 for c in coords)
             assert f.data.shape == shape
-            assert f.data_with_halo.shape == tuple(i+4 for i in f.data.shape)
-            assert f._distributor.shape == grid.shape
-            for d in grid.dimensions:
-                assert all([i == 2 for i in f._size_inhalo[d]])
-                assert all([i == 2 for i in f._size_outhalo[d]])
+            # assert f.data_with_halo.shape == tuple(i+4 for i in f.data.shape)
+            # assert f._distributor.shape == grid.shape
+            # for d in grid.dimensions:
+            #     assert all([i == 2 for i in f._size_inhalo[d]])
+            #     assert all([i == 2 for i in f._size_outhalo[d]])
+        assert False