compiler: Add abstraction of efuncs w ComponentAccesses

Author: FabioLuporini

devito/mpi/routines.py

@@ -740,7 +740,7 @@ def _make_halowait(self, f, hse, key, wait, msg=None):
 
         parameters = list(f.handles) + list(fixed.values()) + [nb, msg]
 
-        return Callable('halowait%d' % key, iet, 'void', parameters, ('static',))
+        return HaloWait(f'halowait{key}', iet, parameters)
 
     def _call_halowait(self, name, f, hse, msg):
         nb = f.grid.distributor._obj_neighborhood
@@ -779,7 +779,7 @@ def _make_region(self, hs, key):
     def _make_msg(self, f, hse, key):
         # Only retain the halos required by the Diag scheme
         halos = sorted(i for i in hse.halos if isinstance(i.dim, tuple))
-        return MPIMsgEnriched('msg%d' % key, f, halos)
+        return MPIMsgEnriched(f'msg{key}', f, halos)
 
     def _make_sendrecv(self, *args, **kwargs):
         return
@@ -868,7 +868,7 @@ def _make_halowait(self, f, hse, key, *args, msg=None):
         ncomms = Symbol(name='ncomms')
         iet = Iteration([waitsend, waitrecv, scatter], dim, ncomms - 1)
         parameters = f.handles + tuple(fixed.values()) + (msg, ncomms)
-        return Callable('halowait%d' % key, iet, 'void', parameters, ('static',))
+        return HaloWait(f'halowait{key}', iet, parameters)
 
     def _call_halowait(self, name, f, hse, msg):
         args = f.handles + tuple(hse.loc_indices.values()) + (msg, msg.npeers)
@@ -1050,9 +1050,11 @@ def __init__(self, name, body, parameters, bufg, bufs):
 
 
 class HaloUpdate(MPICallable):
+    pass
 
-    def __init__(self, name, body, parameters):
-        super().__init__(name, body, parameters)
+
+class HaloWait(MPICallable):
+    pass
 
 
 class Remainder(ElementalFunction):
@@ -1254,12 +1256,14 @@ class MPIMsgEnriched(MPIMsg):
     _C_field_ofsg = 'ofsg'
     _C_field_from = 'fromrank'
     _C_field_to = 'torank'
+    _C_field_components = 'components'
 
     fields = MPIMsg.fields + [
         (_C_field_ofss, POINTER(c_int)),
         (_C_field_ofsg, POINTER(c_int)),
         (_C_field_from, c_int),
-        (_C_field_to, c_int)
+        (_C_field_to, c_int),
+        (_C_field_components, POINTER(c_int)),
     ]
 
     def _arg_defaults(self, allocator, alias=None, args=None):
@@ -1298,6 +1302,17 @@ def _arg_defaults(self, allocator, alias=None, args=None):
                     ofss.append(f._offset_owned[dim].left)
             entry.ofss = (c_int*len(ofss))(*ofss)
 
+            # Track the component accesses for packing/unpacking as numbers
+            # representing the field being accessed (that is: .x -> 0, .y -> 1,
+            # .z -> 2, .w -> 3), if any
+            if isinstance(self.target, BundleView):
+                ncomp = self.target.ncomp
+                component_indices = self.target.component_indices
+                entry.components = (c_int*ncomp)(*component_indices)
+            elif self.target.is_Bundle:
+                ncomp = self.target.ncomp
+                entry.components = (c_int*ncomp)(*range(ncomp))
+
         return {self.name: self.value}
 
 

devito/passes/iet/definitions.py

@@ -177,12 +177,14 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
         decl = Definition(obj)
 
+        sizeof_dtypeN = SizeOf(obj.indexed._C_typedata)
+        sizeof_dtype1 = SizeOf(obj.c0.indexed._C_typedata)
+
         # NOTE: the `arity` is calculated such as `sizeof(float3)/sizeof(float)`
         # for portability reasons (since we don't know the size of compound
         # types a priori)
         arity_param = Symbol(name='arity', dtype=size_t)
-        arity_arg = (SizeOf(obj.indexed._C_typedata) /
-                     SizeOf(obj.c0.indexed._C_typedata))
+        arity_arg = sizeof_dtypeN / sizeof_dtype1
         ndims_param = Symbol(name='ndims', dtype=size_t)
         ndims_arg = obj.ndim
         shape_param = Array(name=f'{obj.name}_shape', dtype=np.uint64,
@@ -217,7 +219,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
             dim,
             ndims_param - 1,
         ))
-        init.append(DummyExpr(ffp3, ffp2*arity_param))
+        init.append(DummyExpr(ffp3, ffp2*arity_param*sizeof_dtype1))
         init.append(DummyExpr(ffp4, arity_param))
 
         # Allocate the underlying host data
@@ -260,12 +262,14 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         """
         decl = Definition(obj)
 
+        sizeof_dtypeN = SizeOf(obj.indexed._C_typedata)
+        sizeof_dtype1 = SizeOf(obj.c0.indexed._C_typedata)
+
         # NOTE: the `arity` is calculated such as `sizeof(float3)/sizeof(float)`
         # for portability reasons (since we don't know the size of compound
         # types a priori)
         arity_param = Symbol(name='arity', dtype=size_t)
-        arity_arg = (SizeOf(obj.indexed._C_typedata) /
-                     SizeOf(obj.c0.indexed._C_typedata))
+        arity_arg = sizeof_dtypeN / sizeof_dtype1
         ndims_param = Symbol(name='ndims', dtype=size_t)
         ndims_arg = obj.ndim
         shape_param = Array(name=f'{obj.name}_shape', dtype=np.uint64,
@@ -299,7 +303,7 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
             dim,
             ndims_param - 1,
         ))
-        init.append(DummyExpr(ffp3, ffp2*arity_param))
+        init.append(DummyExpr(ffp3, ffp2*arity_param*sizeof_dtype1))
         init.append(DummyExpr(ffp4, arity_param))
 
         # Free all of the allocated data

devito/passes/iet/engine.py

@@ -1,18 +1,26 @@
 from collections import OrderedDict, defaultdict
 from functools import partial, singledispatch, wraps
 
-from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable,
-                           FindNodes, FindSymbols, MapNodes, MetaCall, Transformer,
-                           EntryFunction, ThreadCallable, Uxreplace,
-                           derive_parameters)
+import numpy as np
+from sympy import Mul
+
+from devito.ir.iet import (
+    Call, ExprStmt, Expression, Iteration, SyncSpot, AsyncCallable, FindNodes,
+    FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction,
+    ThreadCallable, Uxreplace, derive_parameters
+)
 from devito.ir.support import SymbolRegistry
 from devito.mpi.distributed import MPINeighborhood
+from devito.mpi.routines import Gather, Scatter, HaloUpdate, HaloWait, MPIMsg
 from devito.passes import needs_transfer
-from devito.symbolics import FieldFromComposite, FieldFromPointer
+from devito.symbolics import (FieldFromComposite, FieldFromPointer, IndexedPointer,
+                              search)
 from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass
-from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension,
-                          ModuloDimension, Indirection, Pointer, SharedData,
-                          ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard)
+from devito.types import (
+    Array, Bundle, ComponentAccess, CompositeObject, Lock, IncrDimension,
+    ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Symbol, Temp,
+    NPThreads, NThreadsBase, Wildcard
+)
 from devito.types.args import ArgProvider
 from devito.types.dense import DiscreteFunction
 from devito.types.dimension import AbstractIncrDimension, BlockDimension
@@ -147,6 +155,7 @@ def apply(self, func, **kwargs):
         # Minimize code size
         if len(efuncs) > len(self.efuncs):
             efuncs = reuse_compounds(efuncs, self.sregistry)
+            efuncs = abstract_component_accesses(efuncs)
             efuncs = reuse_efuncs(self.root, efuncs, self.sregistry)
 
         self.efuncs = efuncs
@@ -316,6 +325,80 @@ def _(i, sregistry=None):
     return i._rebuild(pname=pname, cfields=cfields, ncfields=ncfields, function=None)
 
 
+def abstract_component_accesses(efuncs):
+    """
+    Generalise `efuncs` by replacing ComponentAccesses with pointer arithmetic
+    where possible and useful.
+
+    This pass is only applied to selected Callables:
+
+        1) HaloUpdate/Gather
+        2) HaloWait/Scatter
+    """
+    processed = dict(efuncs)
+
+    for k, efunc in efuncs.items():
+        if not isinstance(efunc, (HaloUpdate, HaloWait)):
+            continue
+
+        calls = FindNodes((Gather, Scatter)).visit(efunc)
+        if len(calls) != 1:
+            continue
+        call_copy_buffer, = calls
+        copy_buffer = efuncs[call_copy_buffer.name]
+
+        # Retrieve expected objects. If this fails, it means it's a structure
+        # we don't recognize (e.g., an unsupported MPI scheme?), so we just
+        # give up and move on
+        try:
+            f, = [i for i in copy_buffer.parameters if i.is_Bundle]
+            msg, = [i for i in efunc.parameters if isinstance(i, MPIMsg)]
+            dim, = FindSymbols('dimensions').visit(efunc)
+        except ValueError:
+            continue
+
+        exprs = FindNodes(Expression).visit(copy_buffer)
+        compaccs = set().union(*[search(i.expr, ComponentAccess) for i in exprs])
+        if not compaccs or len(compaccs) == f.ncomp:
+            continue
+
+        # Sorted for deterministic codegen
+        compaccs = sorted(compaccs, key=lambda i: i.index)
+
+        # Access the same entry as `ca` but via pointer arithmetic instead
+        arity_param = Symbol(name='arity', dtype=np.uint32, is_const=True)
+        compoff_params = [Symbol(name=f'c{i}', dtype=np.uint32, is_const=True)
+                          for i in range(len(compaccs))]
+
+        f_flatten = f.func(name='flat_data', components=f.c0)
+
+        subs = {}
+        for ca, o in zip(compaccs, compoff_params):
+            indices = [Mul(arity_param, i, evaluate=False) for i in ca.indices]
+            indices[-1] += o
+            subs[ca] = f_flatten.indexed[indices]
+
+        # Transform the `copy_buffer` to use a flatten representation of the Bundle
+        copy_buffer1 = Uxreplace(subs).visit(copy_buffer)
+        parameters = [*copy_buffer1.parameters, arity_param, *compoff_params]
+        parameters[parameters.index(f)] = f_flatten
+        copy_buffer1 = copy_buffer1._rebuild(parameters=parameters)
+
+        # Update the `efunc` to use the new `copy_buffer1` version
+        arity_arg = FieldFromPointer(f._C_field_arity, f._C_symbol)
+        compoff_args = [IndexedPointer(
+            FieldFromComposite(msg._C_field_components, IndexedPointer(msg, dim)), i
+        ) for i in range(len(compoff_params))]
+
+        arguments = [*call_copy_buffer.arguments, arity_arg, *compoff_args]
+        call_copy_buffer1 = call_copy_buffer._rebuild(arguments=arguments)
+        efunc1 = Transformer({call_copy_buffer: call_copy_buffer1}).visit(efunc)
+
+        processed.update({k: efunc1, copy_buffer.name: copy_buffer1})
+
+    return processed
+
+
 def reuse_efuncs(root, efuncs, sregistry=None):
     """
     Generalise `efuncs` so that syntactically identical Callables may be dropped,

devito/types/array.py

@@ -567,6 +567,10 @@ def parent(self):
     def handles(self):
         return (self.parent,)
 
+    @property
+    def component_indices(self):
+        return tuple(self.parent.components.index(i) for i in self.components)
+
 
 class ComponentAccess(Expr, Pickable):