compiler: Remove counters from subdomains and automatically generate unique names

devito/ir/equations/algorithms.py

@@ -1,13 +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.symbolics import (retrieve_indexed, uxreplace, retrieve_dimensions,
+                              retrieve_functions)
 from devito.tools import Ordering, as_tuple, flatten, filter_sorted, filter_ordered
 from devito.types import (Dimension, Eq, IgnoreDimSort, SubDimension,
                           ConditionalDimension)
 from devito.types.basic import AbstractFunction
-from devito.types.grid import MultiSubDimension
+from devito.types.dimension import MultiSubDimension
+from devito.data.allocators import DataReference
+from devito.logger import warning
 
 __all__ = ['dimension_sort', 'lower_exprs', 'concretize_subdims']
 
@@ -162,74 +164,139 @@ def concretize_subdims(exprs, **kwargs):
     """
     sregistry = kwargs.get('sregistry')
 
-    # {root Dimension -> {SubDimension -> concrete SubDimension}}
-    mapper = defaultdict(dict)
+    mapper = {}
+    rebuilt = {}  # Rebuilt implicit dims etc which are shared between dimensions
 
-    _concretize_subdims(exprs, mapper, sregistry)
+    _concretize_subdims(exprs, mapper, rebuilt, sregistry)
     if not mapper:
         return exprs
 
-    subs = {}
-    for i in mapper.values():
-        subs.update(i)
+    # There may be indexed Arrays defined on SubDimensions in the expressions
+    # These must have their dimensions replaced and their .function attribute
+    # reset to prevent recovery of the original SubDimensions
+    functions = {f for f in retrieve_functions(exprs) if f.is_Array}
+    for f in functions:
+        dimensions = tuple(mapper.get(d, d) for d in f.dimensions)
+        if dimensions != f.dimensions:
+            # A dimension has been rebuilt, so build a mapper for Indexed
+            mapper[f.indexed] = f._rebuild(dimensions=dimensions, function=None).indexed
 
-    processed = [uxreplace(e, subs) for e in exprs]
+    processed = [uxreplace(e, mapper) for e in exprs]
 
     return processed
 
 
 @singledispatch
-def _concretize_subdims(a, mapper, sregistry):
+def _concretize_subdims(a, mapper, rebuilt, sregistry):
     pass
 
 
 @_concretize_subdims.register(list)
 @_concretize_subdims.register(tuple)
-def _(v, mapper, sregistry):
+def _(v, mapper, rebuilt, sregistry):
     for i in v:
-        _concretize_subdims(i, mapper, sregistry)
+        _concretize_subdims(i, mapper, rebuilt, sregistry)
 
 
 @_concretize_subdims.register(Eq)
-def _(expr, mapper, sregistry):
+def _(expr, mapper, rebuilt, sregistry):
     for d in expr.free_symbols:
-        _concretize_subdims(d, mapper, sregistry)
+        _concretize_subdims(d, mapper, rebuilt, sregistry)
+
+    # Subdimensions can be hiding in implicit dims
+    _concretize_subdims(expr.implicit_dims, mapper, rebuilt, 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
+def _(d, mapper, rebuilt, sregistry):
+    if d in mapper:
+        # Already have a substitution for this dimension
+        return
 
+    tkns = [tkn._rebuild(name=sregistry.make_name(prefix=tkn.name))
+            for tkn in d.tkns]
+    tkns_subs = {tkn0: tkn1 for tkn0, tkn1 in zip(d.tkns, tkns)}
+    left, right = [mM.subs(tkns_subs) for mM in (d.symbolic_min, d.symbolic_max)]
+    thickness = tuple((v, d._thickness_map[k]) for k, v in tkns_subs.items())
 
-@_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
+    mapper[d] = d._rebuild(symbolic_min=left, symbolic_max=right, thickness=thickness)
 
 
-@_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
+@_concretize_subdims.register(ConditionalDimension)
+def _(d, mapper, rebuilt, sregistry):
+    if d in mapper:
+        # Already have a substitution for this dimension
         return
 
-    pd = d.parent
+    _concretize_subdims(d.parent, mapper, rebuilt, sregistry)
+
+    kwargs = {}
+
+    # Parent may be a subdimension
+    if d.parent in mapper:
+        kwargs['parent'] = mapper[d.parent]
+
+    # Condition may contain subdimensions
+    if d.condition is not None:
+        for v in d.condition.free_symbols:
+            _concretize_subdims(v, mapper, rebuilt, sregistry)
 
-    subs = mapper[pd]
+        if any(v in mapper for v in d.condition.free_symbols):
+            # Substitute into condition
+            kwargs['condition'] = d.condition.subs(mapper)
 
-    if d in subs:
+    if kwargs:
+        # Rebuild if parent or condition need replacing
+        mapper[d] = d._rebuild(**kwargs)
+
+
+@_concretize_subdims.register(MultiSubDimension)
+def _(d, mapper, rebuilt, sregistry):
+    if d in mapper:
         # Already have a substitution for this dimension
         return
 
-    name = sregistry.make_name(prefix=d.name)
-    ltkn, rtkn = MultiSubDimension._symbolic_thickness(name)
+    abstract_tkns = MultiSubDimension._symbolic_thickness(d.parent.name)
+    concrete_tkns = tuple(tkn._rebuild(name=sregistry.make_name(prefix=tkn.name))
+                          for tkn in abstract_tkns)
+
+    kwargs = {'thickness': concrete_tkns}
+    fkwargs = {}
+
+    idim0 = d.implicit_dimension
+    if idim0 is not None:
+        if idim0 in rebuilt:
+            idim1 = rebuilt[idim0]
+        else:
+            iname = sregistry.make_name(prefix='n')
+            rebuilt[idim0] = idim1 = idim0._rebuild(name=iname)
+
+        kwargs['implicit_dimension'] = idim1
+        fkwargs.update({'dimensions': (idim1,) + d.functions.dimensions[1:],
+                        'halo': None,
+                        'padding': None})
+
+    if d.functions in rebuilt:
+        functions = rebuilt[d.functions]
+    else:
+        # Increment every instance of this name after the first encountered
+        fname = sregistry.make_name(prefix=d.functions.name, increment_first=False)
+        # Warn the user if name has been changed, since this will affect overrides
+        if fname != d.functions.name:
+            fkwargs['name'] = fname
+            warning("%s <%s> renamed as '%s'" %
+                    (str(d.functions), id(d.functions), fname))
+
+        fkwargs['function'] = None
+
+        # Data in MultiSubDimension function may not have been touched at this point,
+        # in which case do not use an allocator, as there is nothing to allocate, and
+        # doing so will petrify the `_data` attribute as None.
+        if d.functions._data is not None:
+            fkwargs.update({'allocator': DataReference(d.functions._data)})
+
+        rebuilt[d.functions] = functions = d.functions._rebuild(**fkwargs)
 
-    thickness = (ltkn, rtkn)
+    kwargs['functions'] = functions
 
-    subs[d] = d._rebuild(d.name, pd, thickness)
+    mapper[d] = d._rebuild(**kwargs)

devito/ir/support/symregistry.py

@@ -30,7 +30,7 @@ def __init__(self):
         # passes, to maximize symbol (especially Dimension) reuse
         self.caches = {}
 
-    def make_name(self, prefix=None):
+    def make_name(self, prefix=None, increment_first=True):
         # By default we're creating a new symbol
         if prefix is None:
             prefix = self._symbol_prefix
@@ -40,6 +40,10 @@ def make_name(self, prefix=None):
         except KeyError:
             counter = self.counters.setdefault(prefix, generator())
 
+            # Only increment symbol names after the first encountered
+            if not increment_first:
+                return prefix
+
         return "%s%d" % (prefix, counter())
 
     def make_npthreads(self, size):

devito/types/basic.py

@@ -347,6 +347,7 @@ def _filter_assumptions(cls, **kwargs):
         for i in list(kwargs):
             if i in _assume_rules.defined_facts:
                 assumptions[i] = kwargs.pop(i)
+
         return assumptions, kwargs
 
     def __new__(cls, *args, **kwargs):

devito/types/dimension.py

@@ -666,7 +666,7 @@ def __init_finalize__(self, name, parent, left, right, thickness, local,
 
     @classmethod
     def left(cls, name, parent, thickness, local=True):
-        lst, rst = cls._symbolic_thickness(name)
+        lst, rst = cls._symbolic_thickness(parent.name)
         return cls(name, parent,
                    left=parent.symbolic_min,
                    right=parent.symbolic_min+lst-1,
@@ -675,7 +675,7 @@ def left(cls, name, parent, thickness, local=True):
 
     @classmethod
     def right(cls, name, parent, thickness, local=True):
-        lst, rst = cls._symbolic_thickness(name)
+        lst, rst = cls._symbolic_thickness(parent.name)
         return cls(name, parent,
                    left=parent.symbolic_max-rst+1,
                    right=parent.symbolic_max,
@@ -684,7 +684,7 @@ def right(cls, name, parent, thickness, local=True):
 
     @classmethod
     def middle(cls, name, parent, thickness_left, thickness_right, local=False):
-        lst, rst = cls._symbolic_thickness(name)
+        lst, rst = cls._symbolic_thickness(parent.name)
         return cls(name, parent,
                    left=parent.symbolic_min+lst,
                    right=parent.symbolic_max-rst,

devito/types/grid.py

@@ -182,8 +182,8 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
 
         # Initialize SubDomains
         subdomains = tuple(i for i in (Domain(), Interior(), *as_tuple(subdomains)))
-        for counter, i in enumerate(subdomains):
-            i.__subdomain_finalize__(self, counter=counter)
+        for i in subdomains:
+            i.__subdomain_finalize__(self)
         self._subdomains = subdomains
 
         self._origin = as_tuple(origin or tuple(0. for _ in self.shape))
@@ -502,7 +502,6 @@ def __subdomain_finalize__(self, grid, **kwargs):
         # Create the SubDomain's SubDimensions
         sub_dimensions = []
         sdshape = []
-        counter = kwargs.get('counter', 0) - 1
         for k, v, s in zip(self.define(grid.dimensions).keys(),
                            self.define(grid.dimensions).values(), grid.shape):
             if isinstance(v, Dimension):
@@ -514,8 +513,7 @@ def __subdomain_finalize__(self, grid, **kwargs):
                     side, thickness_left, thickness_right = v
                     if side != 'middle':
                         raise ValueError("Expected side 'middle', not `%s`" % side)
-                    sub_dimensions.append(SubDimension.middle('i%d%s' %
-                                                              (counter, k.name),
+                    sub_dimensions.append(SubDimension.middle('i%s' % k.name,
                                                               k, thickness_left,
                                                               thickness_right))
                     thickness = s-thickness_left-thickness_right
@@ -526,16 +524,14 @@ def __subdomain_finalize__(self, grid, **kwargs):
                         if s-thickness < 0:
                             raise ValueError("Maximum thickness of dimension %s "
                                              "is %d, not %d" % (k.name, s, thickness))
-                        sub_dimensions.append(SubDimension.left('i%d%s' %
-                                                                (counter, k.name),
+                        sub_dimensions.append(SubDimension.left('i%s' % k.name,
                                                                 k, thickness))
                         sdshape.append(thickness)
                     elif side == 'right':
                         if s-thickness < 0:
                             raise ValueError("Maximum thickness of dimension %s "
                                              "is %d, not %d" % (k.name, s, thickness))
-                        sub_dimensions.append(SubDimension.right('i%d%s' %
-                                                                 (counter, k.name),
+                        sub_dimensions.append(SubDimension.right('i%s' % k.name,
                                                                  k, thickness))
                         sdshape.append(thickness)
                     else:
@@ -701,7 +697,7 @@ def __init__(self, **kwargs):
         except AttributeError:
             pass
 
-    def __subdomain_finalize__(self, grid, counter=0, **kwargs):
+    def __subdomain_finalize__(self, grid, **kwargs):
         self._grid = grid
         self._dtype = grid.dtype
 
@@ -740,11 +736,12 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
 
         # Associate the `_local_bounds` to suitable symbolic objects that the
         # compiler can use to generate code
-        n = counter - npresets
-        assert n >= 0
 
-        i_dim = Dimension(name='n%d' % n)
-        d_dim = DefaultDimension(name='d%d' % n, default_value=2*grid.dim)
+        # Dimensions with identical names hash the same, hence tag them with the
+        # SubDomainSet ID to make them unique so they can be used to key a dictionary
+        # of replacements without risking overwriting.
+        i_dim = Dimension('n_%s' % str(id(self)))
+        d_dim = DefaultDimension(name='d', default_value=2*grid.dim)
         sd_func = Function(name=self.name, grid=self._grid,
                            shape=(self._n_domains, 2*grid.dim),
                            dimensions=(i_dim, d_dim), dtype=np.int32)
@@ -756,12 +753,8 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
                 idx = 2*i + j
                 sd_func.data[:, idx] = self._local_bounds[idx]
 
-            dname = '%si%d' % (d.name, counter)
-
-            thickness = MultiSubDimension._symbolic_thickness(dname)
-
             dimensions.append(MultiSubDimension(
-                dname, d, thickness, functions=sd_func,
+                'i%s' % d.name, d, None, functions=sd_func,
                 bounds_indices=(2*i, 2*i+1), implicit_dimension=i_dim
             ))