Merge pull request #2411 from devitocodes/revamp-msd-lowering

compiler: Revamp MultiSubDimension lowering

devito/ir/equations/algorithms.py

@@ -1,11 +1,15 @@
 from collections.abc import Iterable
+from collections import defaultdict
+from functools import singledispatch
 
 from devito.symbolics import retrieve_indexed, uxreplace, retrieve_dimensions
 from devito.tools import Ordering, as_tuple, flatten, filter_sorted, filter_ordered
-from devito.types import Dimension, IgnoreDimSort
+from devito.types import (Dimension, Eq, IgnoreDimSort, SubDimension,
+                          ConditionalDimension)
 from devito.types.basic import AbstractFunction
+from devito.types.grid import MultiSubDimension
 
-__all__ = ['dimension_sort', 'lower_exprs']
+__all__ = ['dimension_sort', 'lower_exprs', 'concretize_subdims']
 
 
 def dimension_sort(expr):
@@ -146,3 +150,86 @@ def _lower_exprs(expressions, subs):
     else:
         assert len(processed) == 1
         return processed.pop()
+
+
+def concretize_subdims(exprs, **kwargs):
+    """
+    Given a list of expressions, return a new list where all user-defined
+    SubDimensions have been replaced by their concrete counterparts.
+
+    A concrete SubDimension binds objects that are guaranteed to be unique
+    across `exprs`, such as the thickness symbols.
+    """
+    sregistry = kwargs.get('sregistry')
+
+    # {root Dimension -> {SubDimension -> concrete SubDimension}}
+    mapper = defaultdict(dict)
+
+    _concretize_subdims(exprs, mapper, sregistry)
+    if not mapper:
+        return exprs
+
+    subs = {}
+    for i in mapper.values():
+        subs.update(i)
+
+    processed = [uxreplace(e, subs) for e in exprs]
+
+    return processed
+
+
+@singledispatch
+def _concretize_subdims(a, mapper, sregistry):
+    pass
+
+
+@_concretize_subdims.register(list)
+@_concretize_subdims.register(tuple)
+def _(v, mapper, sregistry):
+    for i in v:
+        _concretize_subdims(i, mapper, sregistry)
+
+
+@_concretize_subdims.register(Eq)
+def _(expr, mapper, sregistry):
+    for d in expr.free_symbols:
+        _concretize_subdims(d, mapper, sregistry)
+
+
+@_concretize_subdims.register(SubDimension)
+def _(d, mapper, sregistry):
+    # TODO: to be implemented as soon as we drop the counter machinery in
+    # Grid.__subdomain_finalize__
+    pass
+
+
+@_concretize_subdims.register(ConditionalDimension)
+def _(d, mapper, sregistry):
+    # TODO: to be implemented as soon as we drop the counter machinery in
+    # Grid.__subdomain_finalize__
+    # TODO: call `_concretize_subdims(d.parent, mapper)` as the parent might be
+    # a SubDimension!
+    pass
+
+
+@_concretize_subdims.register(MultiSubDimension)
+def _(d, mapper, sregistry):
+    if not d.is_abstract:
+        # TODO: for now Grid.__subdomain_finalize__ creates the thickness, but
+        # soon it will be done here instead
+        return
+
+    pd = d.parent
+
+    subs = mapper[pd]
+
+    if d in subs:
+        # Already have a substitution for this dimension
+        return
+
+    name = sregistry.make_name(prefix=d.name)
+    ltkn, rtkn = MultiSubDimension._symbolic_thickness(name)
+
+    thickness = (ltkn, rtkn)
+
+    subs[d] = d._rebuild(d.name, pd, thickness)

devito/operator/operator.py

@@ -12,7 +12,7 @@
 from devito.data import default_allocator
 from devito.exceptions import InvalidOperator, ExecutionError
 from devito.logger import debug, info, perf, warning, is_log_enabled_for, switch_log_level
-from devito.ir.equations import LoweredEq, lower_exprs
+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)
@@ -342,6 +342,10 @@ def _lower_exprs(cls, expressions, **kwargs):
         # "True" lowering (indexification, shifting, ...)
         expressions = lower_exprs(expressions, **kwargs)
 
+        # Turn user-defined SubDimensions into concrete SubDimensions,
+        # in particular uniqueness across expressions is ensured
+        expressions = concretize_subdims(expressions, **kwargs)
+
         processed = [LoweredEq(i) for i in expressions]
 
         return processed
@@ -367,8 +371,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
               as parallelism.
             * Optimize Clusters for performance
         """
-        sregistry = kwargs['sregistry']
-
         # Build a sequence of Clusters from a sequence of Eqs
         clusters = clusterize(expressions, **kwargs)
 
@@ -385,7 +387,7 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
             pass
 
         # Generate implicit Clusters from higher level abstractions
-        clusters = generate_implicit(clusters, sregistry=sregistry)
+        clusters = generate_implicit(clusters)
 
         # Lower all remaining high order symbolic objects
         clusters = lower_index_derivatives(clusters, **kwargs)

devito/passes/clusters/implicit.py

@@ -5,20 +5,18 @@
 from collections import defaultdict
 from functools import singledispatch
 
-import numpy as np
-
 from devito.ir import SEQUENTIAL, Queue, Forward
 from devito.symbolics import retrieve_dimensions
 from devito.tools import Bunch, frozendict, timed_pass
-from devito.types import Eq, Symbol
+from devito.types import Eq
 from devito.types.dimension import BlockDimension
 from devito.types.grid import MultiSubDimension
 
 __all__ = ['generate_implicit']
 
 
 @timed_pass()
-def generate_implicit(clusters, sregistry):
+def generate_implicit(clusters):
     """
     Create and add implicit expressions from high-level abstractions.
 
@@ -29,17 +27,14 @@ def generate_implicit(clusters, sregistry):
 
         * MultiSubDomains attached to input equations.
     """
-    clusters = LowerExplicitMSD(sregistry).process(clusters)
-    clusters = LowerImplicitMSD(sregistry).process(clusters)
+    clusters = LowerExplicitMSD().process(clusters)
+    clusters = LowerImplicitMSD().process(clusters)
 
     return clusters
 
 
 class LowerMSD(Queue):
-
-    def __init__(self, sregistry):
-        super().__init__()
-        self.sregistry = sregistry
+    pass
 
 
 class LowerExplicitMSD(LowerMSD):
@@ -105,7 +100,7 @@ def callback(self, clusters, prefix):
                 processed.append(c)
                 continue
 
-            exprs, thickness = make_implicit_exprs(mapper, self.sregistry)
+            exprs = make_implicit_exprs(mapper)
 
             ispace = c.ispace.insert(dim, dims)
 
@@ -121,7 +116,6 @@ def callback(self, clusters, prefix):
 
             # The Cluster performing the actual computation, enriched with
             # the thicknesses
-            ispace = inject_thickness(ispace, thickness)
             processed.append(c.rebuild(ispace=ispace))
 
         return processed
@@ -190,12 +184,9 @@ def callback(self, clusters, prefix):
                                                mapper, edims, prefix)
 
         # Turn the reduced mapper into a list of equations
-        mapper = {}
         processed = []
         for bunch in found.values():
-            exprs, thickness = make_implicit_exprs(bunch.mapper, self.sregistry)
-
-            mapper.update({c: thickness for c in bunch.clusters})
+            exprs = make_implicit_exprs(bunch.mapper)
 
             # Only retain outer guards (e.g., along None) if any
             key = lambda i: i is None or i in prefix.prefix([pd])
@@ -206,13 +197,7 @@ def callback(self, clusters, prefix):
                 c.rebuild(exprs=exprs, ispace=prefix, guards=guards, syncs=syncs)
             )
 
-        # Add in the dynamic thickness
-        for c in clusters:
-            try:
-                ispace = inject_thickness(c.ispace, mapper[c])
-                processed.append(c.rebuild(ispace=ispace))
-            except KeyError:
-                processed.append(c)
+        processed.extend(clusters)
 
         return processed
 
@@ -236,8 +221,8 @@ def _lower_msd(dim, cluster):
 @_lower_msd.register(MultiSubDimension)
 def _(dim, cluster):
     i_dim = dim.implicit_dimension
-    mapper = {(dim.root, i): dim.functions[i_dim, mM]
-              for i, mM in enumerate(dim.bounds_indices)}
+    mapper = {tkn: dim.functions[i_dim, mM]
+              for tkn, mM in zip(dim.tkns, dim.bounds_indices)}
     return mapper, i_dim
 
 
@@ -260,31 +245,8 @@ def lower_msd(msdims, cluster):
     return frozendict(mapper), tuple(dims - {None})
 
 
-def make_implicit_exprs(mapper, sregistry):
-    exprs = []
-    thickness = defaultdict(lambda: [None, None])
-    for (d, side), v in mapper.items():
-        tkn = 'l' if side == 0 else 'r'
-        name = sregistry.make_name('%s_%stkn' % (d.name, tkn))
-        s = Symbol(name=name, dtype=np.int32, is_const=True, nonnegative=True)
-
-        exprs.append(Eq(s, v))
-        thickness[d][side] = s
-
-    return exprs, frozendict(thickness)
-
-
-def inject_thickness(ispace, thickness):
-    for i in ispace.itdims:
-        if i.is_Block and i._depth > 1:
-            # The thickness should be injected once only!
-            continue
-        try:
-            v0, v1 = thickness[i.root]
-            ispace = ispace.translate(i, v0, -v1)
-        except KeyError:
-            pass
-    return ispace
+def make_implicit_exprs(mapper):
+    return [Eq(k, v) for k, v in mapper.items()]
 
 
 def reduce(m0, m1, edims, prefix):