WIP WIP

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -184,11 +184,14 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         shape_param = Array(name=f'{obj.name}_shape', dtype=np.uint64,
                             dimensions=(Dimension(name='d'),), scope='rvalue')
         shape_arg = ListInitializer(obj.c0.symbolic_shape, dtype=shape_param.dtype)
+        sizeofelem_param = Symbol(name='sizeofelem', dtype=size_t)
+        sizeofelem_arg = 1  #TODO
 
         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)
+        ffp4 = FieldFromPointer(obj._C_field_sizeofelem, obj._C_symbol)
 
         # Allocate the Array struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -213,6 +216,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
             ndims_param - 1,
         ))
         init.append(DummyExpr(ffp3, ffp2*arity_param))
+        init.append(DummyExpr(ffp4, sizeofelem_param))
 
         # Allocate the underlying host data
         memptr = VOID(Byref(ffp0), '**')
@@ -237,6 +241,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         args[args.index(arity_param)] = arity_arg
         args[args.index(ndims_param)] = ndims_arg
         args[args.index(shape_param)] = shape_arg
+        args[args.index(sizeofelem_param)] = sizeofelem_arg  #TODO
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees
@@ -261,10 +266,14 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         shape_param = Array(name=f'{obj.name}_shape', dtype=np.uint64,
                             dimensions=(Dimension(name='d'),), scope='rvalue')
         shape_arg = ListInitializer(obj.c0.symbolic_shape, dtype=shape_param.dtype)
+        sizeofelem_param = Symbol(name='sizeofelem', dtype=size_t)
+        sizeofelem_arg = 1  #TODO
+        from IPython import embed; embed()
 
         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)
+        ffp4 = FieldFromPointer(obj._C_field_sizeofelem, obj._C_symbol)
 
         # Allocate the Bundle struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -305,6 +314,7 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         args[args.index(arity_param)] = arity_arg
         args[args.index(ndims_param)] = ndims_arg
         args[args.index(shape_param)] = shape_arg
+        args[args.index(sizeofelem_param)] = sizeofelem_arg  #TODO
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees

devito/passes/iet/engine.py

@@ -1,18 +1,25 @@
 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 devito.ir.iet import (
+    Call, DummyExpr, ExprStmt, Expression, List, 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 CopyBuffer
 from devito.passes import needs_transfer
-from devito.symbolics import FieldFromComposite, FieldFromPointer
+from devito.symbolics import (FieldFromComposite, FieldFromPointer, Byref, Cast,
+                              Deref, 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 +154,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(self.root, efuncs, self.sregistry)
             efuncs = reuse_efuncs(self.root, efuncs, self.sregistry)
 
         self.efuncs = efuncs
@@ -316,6 +324,79 @@ def _(i, sregistry=None):
     return i._rebuild(pname=pname, cfields=cfields, ncfields=ncfields, function=None)
 
 
+def abstract_component_accesses(root, efuncs, sregistry=None):
+    """
+    Generalise `efuncs` by replacing ComponentAccesses with pointer arithmetic
+    where possible and useful.
+    """
+    candidates = (CopyBuffer,)
+
+    # Transform the candidate efuncs
+    processed = {}
+    for k, efunc in efuncs.items():
+        if not isinstance(efunc, candidates):
+            continue
+
+        found = defaultdict(set)
+        for e in FindNodes(Expression).visit(efunc):
+            for v in search(e.expr, ComponentAccess):
+                found[e].add(v)
+
+        # Is it a supported `efunc` structure?
+        if len(found) != 1:
+            continue
+        expr, compaccs = found.popitem()
+        if len(compaccs) != 1:
+            # Pointless -- either none or a whole-Bundle access, so not worth
+            continue
+
+        ca, = compaccs
+        f = ca.function
+
+        dtype = f.c0.indexed._C_ctype
+
+        # Access the same entry as `ca` but via pointer arithmetic instead
+        from IPython import embed; embed()
+        base = Symbol(name='base', dtype=dtype)
+        ptr = Symbol(name='ptr', dtype=dtype)
+        index = Symbol(name='index', dtype=np.uint32, is_const=True)
+
+        body = List(body=[
+            DummyExpr(base, Cast(Byref(f.indexed[ca.indices]), dtype=dtype,
+                                 reinterpret=True)),
+            DummyExpr(ptr, base + index),
+            Uxreplace({ca: Deref(ptr)}).visit(expr)
+        ])
+
+        efunc1 = Transformer({expr: body}).visit(efunc)
+        efunc1 = efunc1._rebuild(parameters=(*efunc1.parameters, index))
+
+        processed[k] = (efunc1, ca)
+
+    if not processed:
+        return efuncs
+
+    # Update the Call sites
+    mapper = {}
+    dag = create_call_graph(root.name, efuncs)
+    for k, (efunc, ca) in processed.items():
+        mapper[k] = efunc
+
+        for i in dag.downstream(k):
+            caller = efuncs[i]
+
+            calls = [c for c in FindNodes(Call).visit(caller) if c.name == k]
+            subs = {c: c._rebuild(arguments=(*c.arguments, ca.index))
+                    for c in calls}
+
+            mapper[i] = Transformer(subs).visit(caller)
+
+    # Update the efuncs
+    efuncs = {**dict(efuncs), **mapper}
+
+    return efuncs
+
+
 def reuse_efuncs(root, efuncs, sregistry=None):
     """
     Generalise `efuncs` so that syntactically identical Callables may be dropped,

devito/types/array.py

@@ -223,6 +223,7 @@ class MappedArrayMixin:
     _C_field_shape = 'shape'
     _C_field_size = 'size'
     _C_field_nbytes = 'nbytes'
+    _C_field_sizeofelem = 'sizeofelem'
 
     _C_size_type = c_uint64
 
@@ -231,7 +232,8 @@ class MappedArrayMixin:
                                           (_C_field_dmap, c_void_p),
                                           (_C_field_shape, POINTER(_C_size_type)),
                                           (_C_field_size, _C_size_type),
-                                          (_C_field_nbytes, _C_size_type)]}))
+                                          (_C_field_nbytes, _C_size_type),
+                                          (_C_field_sizeofelem, _C_size_type)]}))
 
 
 class ArrayMapped(MappedArrayMixin, Array):