Merge pull request #2582 from devitocodes/fixup-mpi-bundles

mpi: Fix correctness of halo exchanges for Bundles

Author: mloubout

devito/mpi/halo_scheme.py

@@ -31,17 +31,32 @@ class HaloSchemeEntry(EnrichedTuple):
 
     __rargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')
 
-    def __init__(self, loc_indices, loc_dirs, halos, dims, getters=None):
-        self.loc_indices = frozendict(loc_indices)
-        self.loc_dirs = frozendict(loc_dirs)
-        self.halos = frozenset(halos)
-        self.dims = frozenset(dims)
+    def __new__(cls, loc_indices, loc_dirs, halos, dims, getters=None):
+        items = [frozendict(loc_indices), frozendict(loc_dirs),
+                 frozenset(halos), frozenset(dims)]
+        kwargs = dict(zip(cls.__rargs__, items))
+        return super().__new__(cls, *items, getters=cls.__rargs__, **kwargs)
 
     def __hash__(self):
-        return hash((self.loc_indices,
-                     self.loc_dirs,
-                     self.halos,
-                     self.dims))
+        return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims))
+
+    def union(self, other):
+        """
+        Return a new HaloSchemeEntry that is the union of this and `other`.
+        The `loc_indices` and `loc_dirs` must be the same, otherwise an
+        exception is raised.
+        """
+        if self.loc_indices != other.loc_indices or \
+           self.loc_dirs != other.loc_dirs:
+            raise HaloSchemeException(
+                "Inconsistency found while building a HaloScheme"
+            )
+
+        halos = self.halos | other.halos
+        dims = self.dims | other.dims
+
+        return HaloSchemeEntry(self.loc_indices, self.loc_dirs, halos, dims,
+                               getters=self.getters)
 
 
 Halo = namedtuple('Halo', 'dim side')
@@ -410,9 +425,11 @@ def add(self, f, hse):
         """
         Create a new HaloScheme that contains all entries in `self` plus the one
         passed in input. If `f` already exists in `self`, the old value is
-        overridden.
+        augmented with `hse` (i.e., the halos are unioned).
         """
-        fmapper = dict(self.fmapper.items())
+        fmapper = dict(self.fmapper)
+        if f in fmapper:
+            hse = fmapper[f].union(hse)
         fmapper[f] = hse
         return HaloScheme.build(fmapper, self.honored)
 

devito/mpi/routines.py

@@ -1217,7 +1217,7 @@ def _arg_defaults(self, allocator, alias, args=None):
 
     def _arg_values(self, args=None, **kwargs):
         # Any will do
-        for f in self.target.handles:
+        for f in self.target.components:
             try:
                 alias = kwargs[f.name]
                 break

devito/tools/data_structures.py

@@ -82,7 +82,8 @@ def __getitem_hook__(self, key):
     def __getnewargs_ex__(self):
         # Bypass default reconstruction logic since this class spawns
         # objects with varying number of attributes
-        return tuple(self), dict(self.__dict__)
+        sdict = {k: v for k, v in self.__dict__.items() if k not in self.getters}
+        return tuple(self), sdict
 
     def get(self, key, val=None):
         return self.getters.get(key, val)

devito/tools/dtypes_lowering.py

@@ -218,7 +218,10 @@ def dtype_len(dtype):
     """
     try:
         # Vector types
-        return len(dtype)
+        dlen = len(dtype)
+        if 'padding0' in dtype.fields:
+            dlen -= 1
+        return dlen
     except TypeError:
         return 1