Merge pull request #2587 from devitocodes/hotfix-parlang-lowering-3

compiler: Improve quality of generated code with Bundles and MPI

Author: FabioLuporini

devito/ir/iet/nodes.py

@@ -102,10 +102,16 @@ def ccode(self):
 
     @property
     def view(self):
-        """A representation of the IET rooted in ``self``."""
+        """A high-level representation of the IET rooted in `self`."""
         from devito.ir.iet.visitors import printAST
         return printAST(self)
 
+    @property
+    def view_cir(self):
+        from devito.ir.iet.visitors import CGen
+        from devito.passes.iet.languages.CIR import CIRPrinter
+        return str(CGen(printer=CIRPrinter).visit(self))
+
     @property
     def children(self):
         """Return the traversable children."""
@@ -148,7 +154,7 @@ def writes(self):
         return ()
 
     def _signature_items(self):
-        return (str(self),)
+        return (self.view_cir,)
 
 
 class ExprStmt:

devito/ir/support/basic.py

@@ -8,7 +8,8 @@
 from devito.ir.support.utils import AccessMode, extrema
 from devito.ir.support.vector import LabeledVector, Vector
 from devito.symbolics import (compare_ops, retrieve_indexed, retrieve_terminals,
-                              q_constant, q_affine, q_routine, search, uxreplace)
+                              q_constant, q_comp_acc, q_affine, q_routine, search,
+                              uxreplace)
 from devito.tools import (Tag, as_mapper, as_tuple, is_integer, filter_sorted,
                           flatten, memoized_meth, memoized_generator)
 from devito.types import (ComponentAccess, Dimension, DimensionTuple, Fence,
@@ -529,9 +530,16 @@ def __hash__(self):
             (self.source, self.sink, self.source.timestamp == self.sink.timestamp)
         )
 
-    @property
+    @cached_property
     def function(self):
-        return self.source.function
+        if q_comp_acc(self.source.access) and not q_comp_acc(self.sink.access):
+            # E.g., `source=ab[x].x` and `sink=ab[x]` -> `a(x)`
+            return self.source.access.function_access
+        elif not q_comp_acc(self.source.access) and q_comp_acc(self.sink.access):
+            # E.g., `source=ab[x]` and `sink=ab[x].y` -> `b(x)`
+            return self.sink.access.function_access
+        else:
+            return self.source.function
 
     @property
     def findices(self):
@@ -955,7 +963,7 @@ def reads_gen(self):
     @memoized_generator
     def reads_smart_gen(self, f):
         """
-        Generate all read access to a given function.
+        Generate all read accesses to a given function.
 
         StencilDimensions, if any, are replaced with their extrema.
 

devito/ir/support/properties.py

@@ -85,6 +85,13 @@ def __init__(self, name, val=None):
 is used for iteration spaces that are larger than the data space.
 """
 
+INBOUND_IF_RELAXED = Property('inbound-if-relaxed')
+"""
+Similar to INBOUND, but devised for IncrDimensions whose iteration space is
+_not_ larger than the data space and, as such, still require full lowering
+(see the `relax_incr_dimensions` pass for more info).
+"""
+
 PREFETCHABLE = Property('prefetchable')
 """
 A Dimension along which prefetching is feasible and beneficial.
@@ -295,6 +302,9 @@ def block(self, dims, kind='default'):
     def inbound(self, dims):
         return self.add(dims, INBOUND)
 
+    def inbound_if_relaxed(self, dims):
+        return self.add(dims, INBOUND_IF_RELAXED)
+
     def init_core_shm(self, dims):
         properties = self.add(dims, INIT_CORE_SHM)
         properties = properties.drop(properties={INIT_HALO_LEFT_SHM,
@@ -327,7 +337,8 @@ def is_affine(self, dims):
         return any(AFFINE in self.get(d, ()) for d in as_tuple(dims))
 
     def is_inbound(self, dims):
-        return any(INBOUND in self.get(d, ()) for d in as_tuple(dims))
+        return any({INBOUND, INBOUND_IF_RELAXED}.intersection(self.get(d, set()))
+                   for d in as_tuple(dims))
 
     def is_sequential(self, dims):
         return any(SEQUENTIAL in self.get(d, ()) for d in as_tuple(dims))

devito/mpi/halo_scheme.py

@@ -30,15 +30,18 @@ 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))
+        return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims,
+                     self.bundle))
 
     def union(self, other):
         """
@@ -47,7 +50,8 @@ def union(self, other):
         exception is raised.
         """
         if self.loc_indices != other.loc_indices or \
-           self.loc_dirs != other.loc_dirs:
+           self.loc_dirs != other.loc_dirs or \
+           self.bundle is not other.bundle:
             raise HaloSchemeException(
                 "Inconsistency found while building a HaloScheme"
             )
@@ -56,7 +60,7 @@ def union(self, other):
         dims = self.dims | other.dims
 
         return HaloSchemeEntry(self.loc_indices, self.loc_dirs, halos, dims,
-                               getters=self.getters)
+                               bundle=self.bundle, getters=self.getters)
 
 
 Halo = namedtuple('Halo', 'dim side')
@@ -168,7 +172,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 +189,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 +649,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,9 +17,9 @@
 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)
-from devito.types import (Array, Bag, Dimension, Eq, Symbol, LocalObject,
-                          CompositeObject, CustomDimension)
+                          flatten, generator, is_integer)
+from devito.types import (Array, Bag, BundleView, Dimension, Eq, Symbol,
+                          LocalObject, CompositeObject, CustomDimension)
 
 __all__ = ['HaloExchangeBuilder', 'ReductionBuilder', 'mpi_registry']
 
@@ -292,19 +292,28 @@ 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:
+                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
+            # 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 +372,17 @@ 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])))
+        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):
                 eqns.append(Eq(*swap(buf[[i] + bdims], f[[i] + findices])))
 
@@ -724,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
@@ -763,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
@@ -852,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)
@@ -1034,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):
@@ -1238,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):
@@ -1282,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/operator/operator.py

@@ -13,11 +13,12 @@
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
-from devito.logger import debug, info, perf, warning, is_log_enabled_for, switch_log_level
+from devito.logger import (debug, info, perf, warning, is_log_enabled_for,
+                           switch_log_level)
 from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
 from devito.ir.clusters import ClusterGroup, clusterize
-from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols, MetaCall,
-                           derive_parameters, iet_build)
+from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols,
+                           MetaCall, derive_parameters, iet_build)
 from devito.ir.support import AccessMode, SymbolRegistry
 from devito.ir.stree import stree_build
 from devito.operator.profiling import create_profile
@@ -26,8 +27,7 @@
 from devito.parameters import configuration
 from devito.passes import (Graph, lower_index_derivatives, generate_implicit,
                            generate_macros, minimize_symbols, unevaluate,
-                           error_mapper, is_on_device)
-from devito.passes.iet.dtypes import lower_dtypes
+                           error_mapper, is_on_device, lower_dtypes)
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple,
                           flatten, filter_sorted, frozendict, is_integer,
@@ -488,7 +488,7 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):
         # Extract the necessary macros from the symbolic objects
         generate_macros(graph, **kwargs)
 
-        # Add type specific metadata
+        # Target-specific lowering
         lower_dtypes(graph, **kwargs)
 
         # Target-independent optimizations

devito/passes/clusters/buffering.py

@@ -123,9 +123,7 @@ class InjectBuffers(Queue):
     def __init__(self, mapper, sregistry, options):
         super().__init__()
 
-        # Sort the mapper so that we always process the same Function in the
-        # same order, hence we get deterministic code generation
-        self.mapper = {i: mapper[i] for i in sorted(mapper, key=lambda i: i.name)}
+        self.mapper = mapper
 
         self.sregistry = sregistry
         self.options = options
@@ -302,6 +300,9 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
     # {candidate buffered Function -> [Clusters that access it]}
     bfmap = map_buffered_functions(clusters, key)
 
+    # Sort for deterministic code generation
+    bfmap = {i: bfmap[i] for i in sorted(bfmap, key=lambda i: i.name)}
+
     # {buffered Function -> Buffer}
     xds = {}
     mapper = {}
@@ -718,7 +719,7 @@ def offset_from_centre(d, indices):
             # `time/factor` -- the starting pointing at time_m or time_M
             v = indices[0]
             try:
-                p = sum(v.args[1:])
+                p = v.func(*[i for i in v.args if not is_integer(i)])
                 if not ((p - v).is_Integer or (p - v).is_Symbol):
                     raise ValueError
             except (IndexError, ValueError):