Merge pull request #2246 from devitocodes/subdim-mpi-fix

dsl: Patch domain decomposition bug with SubDomains

devito/data/decomposition.py

@@ -324,12 +324,16 @@ def index_glb_to_loc(self, *args, rel=True):
                 rel_ofs = self.glb_min + abs_ofs - base
                 if abs_ofs >= base and abs_ofs <= top:
                     return rel_ofs
+                elif abs_ofs > top:
+                    return top + 1
                 else:
                     return None
             else:
                 rel_ofs = abs_ofs - (self.glb_max - top)
                 if abs_ofs >= self.glb_max - top and abs_ofs <= self.glb_max - base:
                     return rel_ofs
+                elif abs_ofs > self.glb_max - base:
+                    return self.glb_max - base + 1
                 else:
                     return None
         else:

devito/mpi/distributed.py

@@ -14,6 +14,7 @@
 from devito.parameters import configuration
 from devito.tools import EnrichedTuple, as_tuple, ctypes_to_cstr, filter_ordered
 from devito.types import CompositeObject, Object
+from devito.types.utils import DimensionTuple
 
 
 # Do not prematurely initialize MPI
@@ -126,7 +127,8 @@ def glb_shape(self):
     @property
     def shape(self):
         """The calling MPI rank's local shape."""
-        return tuple(len(i) for i in self.glb_numb)
+        return DimensionTuple(*[len(i) for i in self.glb_numb],
+                              getters=self.dimensions)
 
     @property
     def dimensions(self):

devito/types/grid.py

@@ -13,6 +13,7 @@
 from devito.types.args import ArgProvider
 from devito.types.basic import Scalar
 from devito.types.dense import Function
+from devito.types.utils import DimensionTuple
 from devito.types.dimension import (Dimension, SpaceDimension, TimeDimension,
                                     Spacing, SteppingDimension, SubDimension)
 
@@ -37,7 +38,7 @@ def __init__(self, shape=None, dimensions=None, dtype=None):
     @property
     def shape(self):
         """Shape of the physical domain."""
-        return self._shape
+        return DimensionTuple(*self._shape, getters=self.dimensions)
 
     @property
     def dimensions(self):
@@ -226,12 +227,13 @@ def origin_ioffset(self):
         """Offset index of the local (per-process) origin from the domain origin."""
         grid_origin = [min(i) for i in self.distributor.glb_numb]
         assert len(grid_origin) == len(self.spacing)
-        return tuple(grid_origin)
+        return DimensionTuple(*grid_origin, getters=self.dimensions)
 
     @property
     def origin_offset(self):
         """Physical offset of the local (per-process) origin from the domain origin."""
-        return tuple(i*h for i, h in zip(self.origin_ioffset, self.spacing))
+        return DimensionTuple(*[i*h for i, h in zip(self.origin_ioffset, self.spacing)],
+                              getters=self.dimensions)
 
     @property
     def time_dim(self):

tests/test_subdomains.py

@@ -162,6 +162,47 @@ def define(self, dimensions):
 
         assert u0.data.all() == u1.data.all() == u2.data.all() == u3.data.all()
 
+    sd_specs = [('middle', 1, 7), ('middle', 2, 3), ('middle', 7, 1),
+                ('middle', 5, 5), ('right', 3), ('right', 7), ('left', 3),
+                ('left', 7)]
+
+    @pytest.mark.parametrize('spec', sd_specs)
+    @pytest.mark.parallel(mode=[2, 3])
+    def test_subdomains_mpi(self, spec):
+
+        class sd0(SubDomain):
+            name = 'd0'
+
+            def __init__(self, spec):
+                super().__init__()
+                self.spec = spec
+
+            def define(self, dimensions):
+                x = dimensions[0]
+                return {x: self.spec}
+        s_d0 = sd0(spec)
+
+        grid = Grid(shape=(11,), extent=(10.,), subdomains=(s_d0,))
+        x = grid.dimensions[0]
+        xd0 = grid.subdomains['d0'].dimensions[0]
+        f = Function(name='f', grid=grid)
+
+        Operator(Eq(f, f+1, subdomain=grid.subdomains['d0']))()
+
+        # Sets ones on a global array according to the subdomains specified
+        # then slices this according to the indices on this rank and compares
+        # to the operator output.
+        check = np.zeros(grid.shape)
+
+        mM_map = {x.symbolic_min: 0, x.symbolic_max: grid.shape[0]-1}
+        t_map = {k: v for k, v in xd0._thickness_map.items() if v is not None}
+        start = int(xd0.symbolic_min.subs({**mM_map, **t_map}))
+        stop = int(xd0.symbolic_max.subs({**mM_map, **t_map})+1)
+
+        check[slice(start, stop)] = 1
+
+        assert np.all(check[grid.distributor.glb_slices[x]] == f.data)
+
 
 class TestMultiSubDomain(object):