mpi: Support halo exchanges from individual Bundle components

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -30,12 +30,14 @@ class HaloLabel(Tag):
 class HaloSchemeEntry(EnrichedTuple):
 
     __rargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')
+    __rkwargs__ = ('bundle',)
 
-    def __new__(cls, loc_indices, loc_dirs, halos, dims, getters=None):
+    def __new__(cls, loc_indices, loc_dirs, halos, dims, bundle=None, getters=None):
+        getters = cls.__rargs__ + cls.__rkwargs__
         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)
+                 frozenset(halos), frozenset(dims), bundle]
+        kwargs = dict(zip(getters, items))
+        return super().__new__(cls, *items, getters=getters, **kwargs)
 
     def __hash__(self):
         return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims))
@@ -168,7 +170,7 @@ def union(self, halo_schemes):
                 elif not v.loc_indices or hse.loc_indices == v.loc_indices:
                     loc_indices, loc_dirs = hse.loc_indices, hse.loc_dirs
                 else:
-                    # The `loc_dirs` must match otherwise it'd be a symptom there's
+                    # These must match otherwise it'd be a symptom there's
                     # something horribly broken elsewhere!
                     assert hse.loc_dirs == v.loc_dirs
                     assert list(hse.loc_indices) == list(v.loc_indices)
@@ -185,7 +187,11 @@ def union(self, halo_schemes):
                 halos = hse.halos | v.halos
                 dims = hse.dims | v.dims
 
-                fmapper[k] = HaloSchemeEntry(loc_indices, loc_dirs, halos, dims)
+                assert hse.bundle is v.bundle
+
+                fmapper[k] = HaloSchemeEntry(
+                    loc_indices, loc_dirs, halos, dims, bundle=hse.bundle
+                )
 
             # Compute the `honored` union
             for d, v in i.honored.items():
@@ -641,8 +647,12 @@ def _uxreplace_dispatch_haloscheme(hs0, rule):
         for i, v in rule.items():
             if i is f:
                 # Yes!
-                g = v
-                hse = hse0
+                if v.is_Bundle:
+                    g = f
+                    hse = hse0._rebuild(bundle=v)
+                else:
+                    g = v
+                    hse = hse0
 
             elif i.is_Indexed and i.function is f and v.is_Indexed:
                 # Yes, but through an Indexed, hence the `loc_indices` may now

devito/mpi/routines.py

@@ -17,7 +17,7 @@
 from devito.symbolics import (Byref, CondNe, FieldFromPointer, FieldFromComposite,
                               IndexedPointer, Macro, cast, subs_op_args)
 from devito.tools import (as_mapper, dtype_to_mpitype, dtype_len, infer_datasize,
-                          flatten, generator, is_integer, split)
+                          flatten, generator, is_integer)
 from devito.types import (Array, Bag, Dimension, Eq, Symbol, LocalObject,
                           CompositeObject, CustomDimension)
 
@@ -292,19 +292,21 @@ def _make_bundles(self, hs):
 
         mapper = as_mapper(halo_scheme.fmapper, lambda i: halo_scheme.fmapper[i])
         for hse, components in mapper.items():
-            # We recast everything as Bags for simplicity -- worst case scenario
-            # all Bags only have one component. Existing Bundles are preserved
             halo_scheme = halo_scheme.drop(components)
-            bundles, candidates = split(tuple(components), lambda i: i.is_Bundle)
-            for b in bundles:
-                halo_scheme = halo_scheme.add(b, hse)
 
+            # Existing Bundles are preserved
+            if hse.bundle and set(components) == set(hse.bundle.components):
+                halo_scheme = halo_scheme.add(hse.bundle, hse)
+                continue
+
+            # We recast everything else as Bags for simplicity -- worst case
+            # scenario all Bags only have one component.
             try:
-                name = "bag_%s" % "".join(f.name for f in candidates)
-                bag = Bag(name=name, components=candidates)
+                name = "bag_%s" % "".join(f.name for f in components)
+                bag = Bag(name=name, components=components)
                 halo_scheme = halo_scheme.add(bag, hse)
             except ValueError:
-                for i in candidates:
+                for i in components:
                     name = "bag_%s" % i.name
                     bag = Bag(name=name, components=i)
                     halo_scheme = halo_scheme.add(bag, hse)
@@ -363,10 +365,19 @@ def _make_copy(self, f, hse, key, swap=False):
         else:
             swap = lambda i, j: (j, i)
             name = 'scatter%s' % key
+
         if isinstance(f, Bag):
-            for i, c in enumerate(f.components):
-                eqns.append(Eq(*swap(buf[[i] + bdims], c[findices])))
+            if hse.bundle is not None:
+                # `f` is the only component of `hse.bundle` that is
+                # being communicated
+                assert f.ncomp == 1
+                i = hse.bundle.components.index(f.c0)
+                eqns.append(Eq(*swap(buf[[0] + bdims], hse.bundle[[i] + findices])))
+            else:
+                for i, c in enumerate(f.components):
+                    eqns.append(Eq(*swap(buf[[i] + bdims], c[findices])))
         else:
+            assert f.is_Bundle
             for i in range(f.ncomp):
                 eqns.append(Eq(*swap(buf[[i] + bdims], f[[i] + findices])))
 

devito/passes/iet/mpi.py

@@ -198,8 +198,8 @@ def _drop_if_unwritten(iet, options=None, **kwargs):
     writes = {i.write for i in FindNodes(Expression).visit(iet)}
     mapper = {}
     for hs in FindNodes(HaloSpot).visit(iet):
-        for f in hs.fmapper:
-            if f not in writes and key(f):
+        for f, v in hs.fmapper.items():
+            if not writes.intersection({f, v.bundle}) and key(f):
                 mapper[hs] = mapper.get(hs, hs.halo_scheme).drop(f)
 
     # Post-process analysis