mpi: Fix Bundles support and improve codegen

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,15 +374,13 @@ def _make_copy(self, f, hse, key, swap=False):
             name = 'scatter%s' % key
 
         if isinstance(f, Bag):
-            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])))
+            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/passes/iet/definitions.py

@@ -19,7 +19,7 @@
                               SizeOf, VOID, pow_to_mul, unevaluate)
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap,
-                          DeviceRM, Eq, Symbol)
+                          DeviceRM, Eq, Symbol, size_t)
 
 __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage']
 
@@ -168,49 +168,56 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
         decl = Definition(obj)
 
+        arity_param = Symbol(name='arity', dtype=size_t)
+        arity_arg = SizeOf(obj.indexed._C_typedata)
+        ffp0 = FieldFromPointer(obj._C_field_data, obj._C_symbol)
+        ffp1 = FieldFromPointer(obj._C_field_shape, obj._C_symbol)
+        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
+        ffp3 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
+
         # Allocate the Array struct
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)
-        allocs = [self.langbb['host-alloc'](memptr, alignment, nbytes)]
+        alloc0 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
-        nbytes_param = Symbol(name='nbytes', dtype=np.uint64, is_const=True)
-        nbytes_arg = SizeOf(obj.indexed._C_typedata)*obj.size
+        # Allocate the shape array
+        memptr = VOID(Byref(ffp1), '**')
+        nbytes = SizeOf(obj._C_size_type)*obj.ndim
+        alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
+
+        # Initialize the Array struct
+        init = [*[DummyExpr(IndexedPointer(ffp1, i), s)
+                  for i, s in enumerate(obj.c0.symbolic_shape)],
+                DummyExpr(ffp2, obj.size),
+                DummyExpr(ffp3, ffp2*arity_param)]
 
         # Allocate the underlying host data
-        ffp0 = FieldFromPointer(obj._C_field_data, obj._C_symbol)
         memptr = VOID(Byref(ffp0), '**')
-        allocs.append(self.langbb['host-alloc-pin'](memptr, alignment, nbytes_param))
-
-        # Initialize the Array struct
-        ffp1 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
-        init0 = DummyExpr(ffp1, nbytes_param)
-        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
-        init1 = DummyExpr(ffp2, 0)
+        alloc2 = self.langbb['host-alloc-pin'](memptr, alignment, ffp3)
 
+        # Free all of the allocated data
         frees = [self.langbb['host-free-pin'](ffp0),
+                 self.langbb['host-free'](ffp1),
                  self.langbb['host-free'](obj._C_symbol)]
 
         # Allocate the underlying device data, if required by the backend
-        alloc, free = self._make_dmap_allocfree(obj, nbytes_param)
-
-        # Chain together all allocs and frees
-        allocs = as_tuple(allocs) + as_tuple(alloc)
-        frees = as_tuple(free) + as_tuple(frees)
+        alloc_dmap, free_dmap = self._make_dmap_allocfree(obj, ffp3)
 
         ret = Return(obj._C_symbol)
 
         # Wrap everything in a Callable so that we can reuse the same code
         # for equivalent Array structs
         name = self.sregistry.make_name(prefix='alloc')
-        body = (decl, *allocs, init0, init1, ret)
+        body = (decl, alloc0, alloc1, *init, alloc2, *as_tuple(alloc_dmap), ret)
         efunc0 = make_callable(name, body, retval=obj)
         args = list(efunc0.parameters)
-        args[args.index(nbytes_param)] = nbytes_arg
+        args[args.index(arity_param)] = arity_arg
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees
         name = self.sregistry.make_name(prefix='free')
+        frees = as_tuple(free_dmap) + as_tuple(frees)
         efunc1 = make_callable(name, frees)
         free = Call(name, efunc1.parameters)
 
@@ -222,35 +229,50 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         """
         decl = Definition(obj)
 
+        arity_param = Symbol(name='arity', dtype=size_t)
+        arity_arg = SizeOf(obj.indexed._C_typedata)
+        ffp1 = FieldFromPointer(obj._C_field_shape, obj._C_symbol)
+        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
+        ffp3 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
+
         # Allocate the Bundle struct
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)
-        alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
+        alloc0 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
-        nbytes_param = Symbol(name='nbytes', dtype=np.uint64, is_const=True)
-        nbytes_arg = SizeOf(obj.indexed._C_typedata)*obj.size
+        # Allocate the shape array
+        memptr = VOID(Byref(ffp1), '**')
+        nbytes = SizeOf(obj._C_size_type)*obj.ndim
+        alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         # Initialize the Bundle struct
-        ffp1 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
-        init0 = DummyExpr(ffp1, nbytes_param)
-        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
-        init1 = DummyExpr(ffp2, 0)
+        init = [*[DummyExpr(IndexedPointer(ffp1, i), s)
+                  for i, s in enumerate(obj.c0.symbolic_shape)],
+                DummyExpr(ffp2, obj.size),
+                DummyExpr(ffp3, ffp2*arity_param)]
 
-        free = self.langbb['host-free'](obj._C_symbol)
+        # Free all of the allocated data
+        frees = [self.langbb['host-free'](ffp1),
+                 self.langbb['host-free'](obj._C_symbol)]
 
         ret = Return(obj._C_symbol)
 
         # Wrap everything in a Callable so that we can reuse the same code
         # for equivalent Bundle structs
         name = self.sregistry.make_name(prefix='alloc')
-        body = (decl, alloc, init0, init1, ret)
+        body = (decl, alloc0, alloc1, *init, ret)
         efunc0 = make_callable(name, body, retval=obj)
         args = list(efunc0.parameters)
-        args[args.index(nbytes_param)] = nbytes_arg
+        args[args.index(arity_param)] = arity_arg
         alloc = Call(name, args, retobj=obj)
 
-        storage.update(obj, site, allocs=alloc, frees=free, efuncs=efunc0)
+        # Same story for the frees
+        name = self.sregistry.make_name(prefix='free')
+        efunc1 = make_callable(name, frees)
+        free = Call(name, efunc1.parameters)
+
+        storage.update(obj, site, allocs=alloc, frees=free, efuncs=(efunc0, efunc1))
 
     def _alloc_object_array_on_low_lat_mem(self, site, obj, storage):
         """

devito/types/array.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, Structure, c_void_p, c_ulong, c_uint64
+from ctypes import POINTER, Structure, c_void_p, c_uint64
 from functools import cached_property
 
 import numpy as np
@@ -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):
@@ -60,6 +60,13 @@ def shape_allocated(self):
     def is_const(self):
         return self._is_const
 
+    @property
+    def c0(self):
+        # ArrayBasic can be used as a base class for tensorial objects (that is,
+        # arrays whose components are AbstractFunctions). This property enables
+        # treating the two cases uniformly in some lowering passes
+        return self
+
 
 class Array(ArrayBasic):
 
@@ -202,19 +209,27 @@ def _make_pointer(self, dim):
         return PointerArray(name='p%s' % self.name, dimensions=dim, array=self)
 
 
-class ArrayMapped(Array):
+class MappedArrayMixin:
 
     _C_structname = 'array'
     _C_field_data = 'data'
     _C_field_dmap = 'dmap'
-    _C_field_nbytes = 'nbytes'
+    _C_field_shape = 'shape'
     _C_field_size = 'size'
+    _C_field_nbytes = 'nbytes'
+
+    _C_size_type = c_uint64
 
     _C_ctype = POINTER(type(_C_structname, (Structure,),
                             {'_fields_': [(_C_field_data, c_restrict_void_p),
-                                          (_C_field_nbytes, c_ulong),
                                           (_C_field_dmap, c_void_p),
-                                          (_C_field_size, c_uint64)]}))
+                                          (_C_field_shape, POINTER(_C_size_type)),
+                                          (_C_field_size, _C_size_type),
+                                          (_C_field_nbytes, _C_size_type)]}))
+
+
+class ArrayMapped(MappedArrayMixin, Array):
+    pass
 
 
 class ArrayObject(ArrayBasic):
@@ -343,7 +358,7 @@ def array(self):
         return self._array
 
 
-class Bundle(ArrayBasic):
+class Bundle(MappedArrayMixin, ArrayBasic):
 
     """
     Tensor symbol representing an unrolled vector of AbstractFunctions.
@@ -417,7 +432,6 @@ def __halo_setup__(self, components=(), **kwargs):
 
     @property
     def c0(self):
-        # Shortcut for self.components[0]
         return self.components[0]
 
     # Class attributes overrides
@@ -456,18 +470,26 @@ def ncomp(self):
     def initvalue(self):
         return None
 
-    # Overrides defaulting to self.c0's behaviour
-
+    # Defaulting to self.c0's behaviour
     for i in ('_mem_internal_eager', '_mem_internal_lazy', '_mem_local',
               '_mem_mapped', '_mem_host', '_mem_stack', '_mem_constant',
               '_mem_shared', '_mem_shared_remote', '__padding_dtype__',
               '_size_domain', '_size_halo', '_size_owned', '_size_padding',
-              '_size_nopad', '_size_nodomain', '_offset_domain',
-              '_offset_halo', '_offset_owned', '_dist_dimensions',
-              '_C_get_field', 'grid', 'symbolic_shape',
+              '_size_nopad', '_size_nodomain', '_offset_domain', '_offset_halo',
+              '_offset_owned', '_dist_dimensions', '_C_get_field', 'grid',
               *AbstractFunction.__properties__):
         locals()[i] = property(lambda self, v=i: getattr(self.c0, v))
 
+    # Other overrides
+
+    @cached_property
+    def symbolic_shape(self):
+        from devito.symbolics import FieldFromPointer, IndexedPointer  # noqa
+        ffp = FieldFromPointer(self._C_field_shape, self._C_symbol)
+        ret = [s if is_integer(s) else IndexedPointer(ffp, i)
+               for i, s in enumerate(super().symbolic_shape)]
+        return DimensionTuple(*ret, getters=self.dimensions)
+
     @property
     def _mem_heap(self):
         return not any([self._mem_stack, self._mem_shared, self._mem_shared_remote])
@@ -490,17 +512,12 @@ def __getitem__(self, index):
             raise ValueError("Expected %d or %d indices, got %d instead"
                              % (self.ndim, self.ndim + 1, len(index)))
 
-    _C_structname = ArrayMapped._C_structname
-    _C_field_data = ArrayMapped._C_field_data
-    _C_field_nbytes = ArrayMapped._C_field_nbytes
-    _C_field_dmap = ArrayMapped._C_field_dmap
-    _C_field_size = ArrayMapped._C_field_size
-
     @property
     def _C_ctype(self):
         if self._mem_mapped:
-            return ArrayMapped._C_ctype
+            return super()._C_ctype
         else:
+            #TODO DROP???
             return POINTER(dtype_to_ctype(self.dtype))
 
 
@@ -518,6 +535,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')

devito/types/misc.py

@@ -14,7 +14,8 @@
 
 __all__ = ['Timer', 'Pointer', 'VolatileInt', 'FIndexed', 'Wildcard', 'Fence',
            'Global', 'Hyperplane', 'Indirection', 'Temp', 'TempArray', 'Jump',
-           'nop', 'WeakFence', 'CriticalRegion', 'Auto', 'AutoRef', 'auto']
+           'nop', 'WeakFence', 'CriticalRegion', 'Auto', 'AutoRef', 'auto',
+           'size_t']
 
 
 class Timer(CompositeObject):
@@ -344,7 +345,9 @@ def closing(self):
 """
 
 
-# *** CXX support types
+# *** C/CXX support types
+
+size_t = CustomDtype('size_t')
 
 # NOTE: In C++, `auto` is a type specifier more than a type itself, but
 # it's a distinction we can afford to ignore, at least for now