FIX FIX FIX (now generate much better)

Author: FabioLuporini

devito/mpi/routines.py

@@ -18,8 +18,8 @@
                               IndexedPointer, Macro, cast, subs_op_args)
 from devito.tools import (as_mapper, dtype_to_mpitype, dtype_len, infer_datasize,
                           flatten, generator, is_integer)
-from devito.types import (Array, Bag, Dimension, Eq, Symbol, LocalObject,
-                          CompositeObject, CustomDimension)
+from devito.types import (Array, Bag, BundleView, Dimension, Eq, Symbol,
+                          LocalObject, CompositeObject, CustomDimension)
 
 __all__ = ['HaloExchangeBuilder', 'ReductionBuilder', 'mpi_registry']
 
@@ -295,8 +295,15 @@ def _make_bundles(self, hs):
             halo_scheme = halo_scheme.drop(components)
 
             # Existing Bundles are preserved
-            if hse.bundle and set(components) == set(hse.bundle.components):
-                halo_scheme = halo_scheme.add(hse.bundle, hse)
+            if hse.bundle:
+                if set(components) == set(hse.bundle.components):
+                    halo_scheme = halo_scheme.add(hse.bundle, hse)
+                else:
+                    name = f'bundleview_{hse.bundle.name}'
+                    bundle_view = BundleView(
+                        name=name, components=components, parent=hse.bundle
+                    )
+                    halo_scheme = halo_scheme.add(bundle_view, hse)
                 continue
 
             # We recast everything else as Bags for simplicity -- worst case
@@ -367,13 +374,13 @@ def _make_copy(self, f, hse, key, swap=False):
             name = 'scatter%s' % key
 
         if isinstance(f, Bag):
-            for i, c0 in enumerate(f.components):
-                if hse.bundle is not None:
-                    indices = [hse.bundle.components.index(c0), *findices]
-                    rhs = hse.bundle[indices]
-                else:
-                    rhs = c0[findices]
-                eqns.append(Eq(*swap(buf[[i] + bdims], rhs)))
+            for i, c in enumerate(f.components):
+                eqns.append(Eq(*swap(buf[[i] + bdims], c[findices])))
+        elif isinstance(f, BundleView):
+            assert f.parent is hse.bundle
+            for i, c in enumerate(f.components):
+                indices = [f.parent.components.index(c), *findices]
+                eqns.append(Eq(*swap(buf[[i] + bdims], f.parent[indices])))
         else:
             assert f.is_Bundle
             for i in range(f.ncomp):

devito/types/array.py

@@ -10,7 +10,7 @@
 from devito.types.utils import CtypesFactory, DimensionTuple
 
 __all__ = ['Array', 'ArrayMapped', 'ArrayObject', 'PointerArray', 'Bundle',
-           'ComponentAccess', 'Bag']
+           'ComponentAccess', 'Bag', 'BundleView']
 
 
 class ArrayBasic(AbstractFunction, LocalType):
@@ -518,6 +518,31 @@ def handles(self):
         return self.components
 
 
+class BundleView(Bundle):
+
+    """
+    A BundleView is like a Bundle but it doesn't represent a concrete object
+    in the generated code. It's used by the compiler to represent a subset
+    of the components of a Bundle.
+    """
+
+    __rkwargs__ = Bundle.__rkwargs__ + ('parent',)
+
+    def __new__(cls, *args, parent=None, **kwargs):
+        obj = super().__new__(cls, *args, **kwargs)
+        obj._parent = parent
+
+        return obj
+
+    @property
+    def parent(self):
+        return self._parent
+
+    @property
+    def handles(self):
+        return (self.parent,)
+
+
 class ComponentAccess(Expr, Pickable):
 
     _component_names = ('x', 'y', 'z', 'w')