types: Remove increments from SubDimension names

devito/ir/equations/algorithms.py

@@ -1,13 +1,16 @@
 from collections.abc import Iterable
+from itertools import groupby
 
 from sympy import sympify
 
-from devito.symbolics import retrieve_indexed, uxreplace, retrieve_dimensions
+from devito.symbolics import (retrieve_indexed, uxreplace, dxreplace,
+                              retrieve_dimensions)
 from devito.tools import Ordering, as_tuple, flatten, filter_sorted, filter_ordered
+from devito.ir.support import pull_dims
 from devito.types import Dimension, IgnoreDimSort
 from devito.types.basic import AbstractFunction
 
-__all__ = ['dimension_sort', 'lower_exprs']
+__all__ = ['dimension_sort', 'lower_exprs', 'separate_dimensions']
 
 
 def dimension_sort(expr):
@@ -147,3 +150,51 @@ def lower_exprs(expressions, **kwargs):
     else:
         assert len(processed) == 1
         return processed.pop()
+
+
+def separate_dimensions(expressions):
+    """
+    Rename Dimensions with clashing names within expressions.
+    """
+    resolutions = {}
+    count = {}  # Keep track of increments on dim names
+    processed = []
+    for e in expressions:
+        # Just want dimensions which appear in the expression
+        # Dimensions in indices
+        dims = set().union(*[i.function.dimensions
+                             for i in retrieve_indexed(e)])
+        # Dimensions in conditions and ConditionalDimension parents
+        dims = dims.union(*[pull_dims(d.condition, flag=False)
+                            for d in dims if d.is_Conditional
+                            and d.condition is not None],
+                          set(d.parent for d in dims if d.is_Conditional))
+        # Sort for groupby
+        dims = sorted(dims, key=lambda x: x.name)
+
+        # TODO: Needs to check for dimensions in factors too
+
+        # Group dimensions by matching names
+        groups = tuple(tuple(g) for n, g in groupby(dims, key=lambda x: x.name))
+        clashes = tuple(g for g in groups if len(g) > 1)
+
+        subs = {}
+        for c in clashes:
+            # Preferentially rename dims that aren't root dims
+            ddims = tuple(d for d in c if not d.is_Root)
+
+            for d in ddims:
+                if d in resolutions.keys():
+                    subs[d] = resolutions[d]
+                else:
+                    try:
+                        subs[d] = d._rebuild('%s%s' % (d.name, count[d.name]))
+                        count[d.name] += 1
+                    except KeyError:
+                        subs[d] = d._rebuild('%s0' % d.name)
+                        count[d.name] = 1
+                    resolutions[d] = subs[d]
+
+        processed.append(dxreplace(e, subs))
+
+    return processed

devito/ir/equations/equation.py

@@ -8,6 +8,7 @@
 from devito.symbolics import IntDiv, uxreplace
 from devito.tools import Pickable, Tag, frozendict
 from devito.types import Eq, Inc, ReduceMax, ReduceMin
+from devito.symbolics.manipulation import _dxreplace_registry
 
 __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax']
 
@@ -220,8 +221,11 @@ def writes(self):
     def xreplace(self, rules):
         return LoweredEq(self.lhs.xreplace(rules), self.rhs.xreplace(rules), **self.state)
 
-    def func(self, *args):
-        return self._rebuild(*args, evaluate=False)
+    def func(self, *args, **kwargs):
+        return self._rebuild(*args, **kwargs, evaluate=False)
+
+
+_dxreplace_registry.register(LoweredEq)
 
 
 class ClusterizedEq(IREq):

devito/ir/support/guards.py

@@ -8,6 +8,7 @@
 
 from devito.ir.support.space import Forward, IterationDirection
 from devito.symbolics import CondEq, CondNe
+from devito.symbolics.manipulation import _uxreplace_dispatch, _uxreplace
 from devito.tools import Pickable, as_tuple, frozendict, split
 from devito.types import Dimension
 
@@ -57,6 +58,12 @@ def _args_rebuild(self):
         return (self.d,)
 
 
+@_uxreplace_dispatch.register(GuardFactor)
+def _(expr, rule, mode='ux'):
+    d, changed = _uxreplace(expr.d, rule, mode=mode)
+    return expr.__class__(d), changed
+
+
 class GuardFactorEq(GuardFactor, CondEq):
     pass
 

devito/ir/support/space.py

@@ -11,6 +11,7 @@
 from devito.tools import (Ordering, Stamp, as_list, as_tuple, filter_ordered,
                           flatten, frozendict, is_integer, toposort)
 from devito.types import Dimension, ModuloDimension
+from devito.symbolics.manipulation import _uxreplace_dispatch
 
 __all__ = ['NullInterval', 'Interval', 'IntervalGroup', 'IterationSpace',
            'IterationInterval', 'DataSpace', 'Forward', 'Backward', 'Any',
@@ -1000,4 +1001,43 @@ def nonderived_directions(self):
         return {k: v for k, v in self.directions.items() if not k.is_Derived}
 
 
+@_uxreplace_dispatch.register(Interval)
+def _(expr, rule, mode='ux'):
+    changed = False
+    dim, flag = _uxreplace_dispatch(expr.dim, rule, mode=mode)
+    changed |= flag
+    lower, flag = _uxreplace_dispatch(expr.lower, rule, mode=mode)
+    changed |= flag
+    upper, flag = _uxreplace_dispatch(expr.upper, rule, mode=mode)
+    changed |= flag
+
+    return Interval(dim, lower=lower, upper=upper, stamp=expr.stamp), changed
+
+
+@_uxreplace_dispatch.register(IterationSpace)
+def _(expr, rule, mode='ux'):
+    intervals = []
+    relations = set()
+    changed = False
+    for interval in expr.intervals:
+        i, flag = _uxreplace_dispatch(interval, rule, mode=mode)
+        intervals.append(i)
+        changed |= flag
+    for relation in expr.intervals.relations:
+        r, flag = _uxreplace_dispatch(relation, rule, mode=mode)
+        relations.add(r)
+        changed |= flag
+    si, flag = _uxreplace_dispatch(expr.sub_iterators, rule, mode=mode)
+    changed |= flag
+    di, flag = _uxreplace_dispatch(expr.directions, rule, mode=mode)
+    changed |= flag
+
+    intervals = sorted(intervals, key=lambda x: x.dim.name)
+    intervals = IntervalGroup(intervals, relations=frozenset(relations),
+                              mode=expr.intervals.mode)
+
+    return IterationSpace(intervals, sub_iterators=si,
+                          directions=di), changed
+
+
 null_ispace = IterationSpace([])

devito/ir/support/utils.py

@@ -187,7 +187,7 @@ def detect_accesses(exprs):
     for e in as_tuple(exprs):
         other_dims.update(i for i in e.free_symbols if isinstance(i, Dimension))
         other_dims.update(e.implicit_dims)
-    other_dims = filter_sorted(other_dims)
+    other_dims = sorted(other_dims, key=lambda x: x.name)
     mapper[None] = Stencil([(i, 0) for i in other_dims])
 
     return mapper

devito/mpi/halo_scheme.py

@@ -9,7 +9,7 @@
 from devito import configuration
 from devito.data import CORE, OWNED, LEFT, CENTER, RIGHT
 from devito.ir.support import Forward, Scope
-from devito.symbolics.manipulation import _uxreplace_registry
+from devito.symbolics.manipulation import _uxreplace_registry, _dxreplace_registry
 from devito.tools import (Reconstructable, Tag, as_tuple, filter_ordered, flatten,
                           frozendict, is_integer)
 from devito.types import Grid
@@ -590,7 +590,7 @@ def __eq__(self, other):
     func = Reconstructable._rebuild
 
 
-def _uxreplace_dispatch_haloscheme(hs0, rule):
+def _uxreplace_dispatch_haloscheme(hs0, rule, mode='ux'):
     changed = False
     hs = hs0
     for f, hse0 in hs0.fmapper.items():
@@ -636,3 +636,6 @@ def _uxreplace_dispatch_haloscheme(hs0, rule):
 
 _uxreplace_registry.register(HaloTouch,
                              {HaloScheme: _uxreplace_dispatch_haloscheme})
+
+_dxreplace_registry.register(HaloTouch,
+                             {HaloScheme: _uxreplace_dispatch_haloscheme})

devito/operator/operator.py

@@ -9,7 +9,7 @@
 from devito.data import default_allocator
 from devito.exceptions import InvalidOperator
 from devito.logger import debug, info, perf, warning, is_log_enabled_for
-from devito.ir.equations import LoweredEq, lower_exprs
+from devito.ir.equations import LoweredEq, lower_exprs, separate_dimensions
 from devito.ir.clusters import ClusterGroup, clusterize
 from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols, MetaCall,
                            derive_parameters, iet_build)
@@ -330,6 +330,9 @@ def _lower_exprs(cls, expressions, **kwargs):
 
         processed = [LoweredEq(i) for i in expressions]
 
+        # Resolve clashing dimension names
+        processed = separate_dimensions(processed)
+
         return processed
 
     # Compilation -- Cluster level