From 914dfa0bfa5a120b9884cef4c63ea9fca1480d49 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Mon, 15 Jul 2024 17:10:23 +0000
Subject: [PATCH 1/9] compiler: Revamp MultiSubDimension lowering

---
 devito/ir/equations/algorithms.py  | 92 +++++++++++++++++++++++++++++-
 devito/operator/operator.py        | 10 ++--
 devito/passes/clusters/implicit.py | 60 ++++---------------
 devito/types/dimension.py          | 74 ++++++++++++++----------
 devito/types/grid.py               | 39 +++++++------
 5 files changed, 171 insertions(+), 104 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index ab1d80af8b..cad931afad 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -1,11 +1,17 @@
 from collections.abc import Iterable
+from collections import defaultdict
+from functools import singledispatch
+
+import numpy as np
 
 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, Symbol,
+                          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 +152,85 @@ 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.
+    """
+    # {root Dimension -> {SubDimension -> concrete SubDimension}}
+    mapper = defaultdict(dict)
+
+    _concretize_subdims(exprs, mapper)
+    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):
+    return {}
+
+
+@_concretize_subdims.register(list)
+@_concretize_subdims.register(tuple)
+def _(v, mapper):
+    for i in v:
+        _concretize_subdims(i, mapper)
+
+
+@_concretize_subdims.register(Eq)
+def _(expr, mapper):
+    for d in expr.free_symbols:
+        _concretize_subdims(d, mapper)
+
+
+@_concretize_subdims.register(SubDimension)
+def _(d, mapper):
+    # TODO: to be implemented as soon as we drop the counter machinery in
+    # Grid.__subdomain_finalize__
+    return {}
+
+
+@_concretize_subdims.register(ConditionalDimension)
+def _(d, mapper):
+    # 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!
+    return {}
+
+
+@_concretize_subdims.register(MultiSubDimension)
+def _(d, mapper):
+    if any(d.thickness):
+        # TODO: for now Grid.__subdomain_finalize__ creates the thickness, but
+        # soon it will be done here instead
+        return
+    else:
+        pd = d.parent
+
+        subs = mapper[pd]
+
+        ltkn = Symbol(name="%s_ltkn%d" % (pd.name, len(subs)), dtype=np.int32,
+                      is_const=True, nonnegative=True)
+        rtkn = Symbol(name="%s_rtkn%d" % (pd.name, len(subs)), dtype=np.int32,
+                      is_const=True, nonnegative=True)
+
+        left = pd.symbolic_min + ltkn
+        right = pd.symbolic_max - rtkn
+
+        thickness = ((ltkn, None), (rtkn, None))
+
+        subs[d] = d._rebuild(d.name, pd, left, right, thickness)
diff --git a/devito/operator/operator.py b/devito/operator/operator.py
index 363c2507e3..f1fac6d592 100644
--- a/devito/operator/operator.py
+++ b/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)
diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 715a97fc27..675f1effc6 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -5,12 +5,10 @@
 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
 
@@ -18,7 +16,7 @@
 
 
 @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,7 +221,7 @@ 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]
+    mapper = {(dim, i): dim.functions[i_dim, mM]
               for i, mM in enumerate(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(list(d._thickness_map)[side], v) for (d, side), v in mapper.items()]
 
 
 def reduce(m0, m1, edims, prefix):
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 92f476e0a2..7c3fc05fb1 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -539,17 +539,56 @@ def _arg_check(self, *args, **kwargs):
 class AbstractSubDimension(DerivedDimension):
 
     """
-    Symbol defining a convex iteration sub-space derived from a ``parent``
+    Symbol defining a convex iteration sub-space derived from a `parent`
     Dimension.
 
     Notes
     -----
-    This is an abstract class. The actual implementations are SubDimension
-    and MultiSubDimension.
+    This is just the abstract base class for various types of SubDimensions.
     """
 
     is_AbstractSub = True
 
+    __rargs__ = (DerivedDimension.__rargs__ +
+                 ('symbolic_min', 'symbolic_max', 'thickness'))
+    __rkwargs__ = ()
+
+    def __init_finalize__(self, name, parent, left, right, thickness, **kwargs):
+        super().__init_finalize__(name, parent)
+        self._interval = sympy.Interval(left, right)
+        self._thickness = Thickness(*thickness)
+
+    @cached_property
+    def symbolic_min(self):
+        return self._interval.left
+
+    @cached_property
+    def symbolic_max(self):
+        return self._interval.right
+
+    @cached_property
+    def symbolic_size(self):
+        # The size must be given as a function of the parent's symbols
+        return self.symbolic_max - self.symbolic_min + 1
+
+    @property
+    def thickness(self):
+        return self._thickness
+
+    @cached_property
+    def _thickness_map(self):
+        return dict(self.thickness)
+
+    @property
+    def ltkn(self):
+        # Shortcut for the left thickness symbol
+        return self.thickness.left[0]
+
+    @property
+    def rtkn(self):
+        # Shortcut for the right thickness symbol
+        return self.thickness.right[0]
+
 
 class SubDimension(AbstractSubDimension):
 
@@ -601,14 +640,10 @@ class SubDimension(AbstractSubDimension):
 
     is_Sub = True
 
-    __rargs__ = (DerivedDimension.__rargs__ +
-                 ('symbolic_min', 'symbolic_max', 'thickness', 'local'))
-    __rkwargs__ = ()
+    __rargs__ = AbstractSubDimension.__rargs__ + ('local',)
 
     def __init_finalize__(self, name, parent, left, right, thickness, local, **kwargs):
-        super().__init_finalize__(name, parent)
-        self._interval = sympy.Interval(left, right)
-        self._thickness = Thickness(*thickness)
+        super().__init_finalize__(name, parent, left, right, thickness)
         self._local = local
 
     @classmethod
@@ -645,27 +680,10 @@ def middle(cls, name, parent, thickness_left, thickness_right, local=False):
                    thickness=((lst, thickness_left), (rst, thickness_right)),
                    local=local)
 
-    @cached_property
-    def symbolic_min(self):
-        return self._interval.left
-
-    @cached_property
-    def symbolic_max(self):
-        return self._interval.right
-
-    @cached_property
-    def symbolic_size(self):
-        # The size must be given as a function of the parent's symbols
-        return self.symbolic_max - self.symbolic_min + 1
-
     @property
     def local(self):
         return self._local
 
-    @property
-    def thickness(self):
-        return self._thickness
-
     @property
     def is_left(self):
         return self.thickness.right[1] is None
@@ -687,10 +705,6 @@ def bound_symbols(self):
     def _maybe_distributed(self):
         return not self.local
 
-    @cached_property
-    def _thickness_map(self):
-        return dict(self.thickness)
-
     @cached_property
     def _offset_left(self):
         # The left extreme of the SubDimension can be related to either the
diff --git a/devito/types/grid.py b/devito/types/grid.py
index d49cf95ad7..c98da4d61c 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -10,7 +10,7 @@
 from devito.mpi import Distributor, MPI
 from devito.tools import ReducerMap, as_tuple
 from devito.types.args import ArgProvider
-from devito.types.basic import Scalar
+from devito.types.basic import Scalar, Symbol
 from devito.types.dense import Function
 from devito.types.utils import DimensionTuple
 from devito.types.dimension import (Dimension, SpaceDimension, TimeDimension,
@@ -556,9 +556,10 @@ class MultiSubDimension(AbstractSubDimension):
 
     __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
 
-    def __init_finalize__(self, name, parent, functions=None, bounds_indices=None,
+    def __init_finalize__(self, name, parent, left, right, thickness,
+                          functions=None, bounds_indices=None,
                           implicit_dimension=None):
-        super().__init_finalize__(name, parent)
+        super().__init_finalize__(name, parent, left, right, thickness)
         self.functions = functions
         self.bounds_indices = bounds_indices
         self.implicit_dimension = implicit_dimension
@@ -573,18 +574,6 @@ def __hash__(self):
     def bound_symbols(self):
         return self.parent.bound_symbols
 
-    @cached_property
-    def symbolic_min(self):
-        return self.parent.symbolic_min
-
-    @cached_property
-    def symbolic_max(self):
-        return self.parent.symbolic_max
-
-    @cached_property
-    def symbolic_size(self):
-        return self.parent.symbolic_size
-
 
 class MultiSubDomain(AbstractSubDomain):
 
@@ -786,10 +775,22 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
                 sd_func.data[:, idx] = self._local_bounds[idx]
 
             dname = '%si%d' % (d.name, counter)
-            dimensions.append(MultiSubDimension(dname, d,
-                                                functions=sd_func,
-                                                bounds_indices=(2*i, 2*i+1),
-                                                implicit_dimension=i_dim))
+
+            ltkn = Symbol(name="%s_ltkn%d" % (d.name, counter), dtype=np.int32,
+                          is_const=True, nonnegative=True)
+            rtkn = Symbol(name="%s_rtkn%d" % (d.name, counter), dtype=np.int32,
+                          is_const=True, nonnegative=True)
+
+            left = d.symbolic_min + ltkn
+            right = d.symbolic_max - rtkn
+
+            thickness = ((ltkn, None), (rtkn, None))
+
+            dimensions.append(MultiSubDimension(
+                dname, d, left, right, thickness,
+                functions=sd_func, bounds_indices=(2*i, 2*i+1),
+                implicit_dimension=i_dim
+            ))
 
         self._dimensions = tuple(dimensions)
 

From 46cc75a7d5ec5b73250c8fb768174fa439e1e060 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 16 Jul 2024 15:13:24 +0100
Subject: [PATCH 2/9] compiler: Move MultiSubDimension to dimensions

---
 devito/types/dimension.py | 31 ++++++++++++++++++++++++++++++-
 devito/types/grid.py      | 31 +------------------------------
 2 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 7c3fc05fb1..c4d556ae7b 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -15,7 +15,7 @@
 from devito.types.constant import Constant
 
 __all__ = ['Dimension', 'SpaceDimension', 'TimeDimension', 'DefaultDimension',
-           'CustomDimension', 'SteppingDimension', 'SubDimension',
+           'CustomDimension', 'SteppingDimension', 'SubDimension', 'MultiSubDimension',
            'ConditionalDimension', 'ModuloDimension', 'IncrDimension',
            'BlockDimension', 'StencilDimension', 'VirtualDimension',
            'Spacing', 'dimensions']
@@ -790,6 +790,35 @@ def _arg_values(self, interval, grid=None, **kwargs):
         return {i.name: v for i, v in zip(self._thickness_map, (ltkn, rtkn))}
 
 
+class MultiSubDimension(AbstractSubDimension):
+
+    """
+    A special Dimension to be used in MultiSubDomains.
+    """
+
+    is_MultiSub = True
+
+    __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
+
+    def __init_finalize__(self, name, parent, left, right, thickness,
+                          functions=None, bounds_indices=None,
+                          implicit_dimension=None):
+        super().__init_finalize__(name, parent, left, right, thickness)
+        self.functions = functions
+        self.bounds_indices = bounds_indices
+        self.implicit_dimension = implicit_dimension
+
+    def __hash__(self):
+        # There is no possibility for two MultiSubDimensions to ever hash the
+        # same, since a MultiSubDimension carries a reference to a MultiSubDomain,
+        # which is unique
+        return id(self)
+
+    @cached_property
+    def bound_symbols(self):
+        return self.parent.bound_symbols
+
+
 class ConditionalDimension(DerivedDimension):
 
     """
diff --git a/devito/types/grid.py b/devito/types/grid.py
index c98da4d61c..44fa9941bb 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -15,7 +15,7 @@
 from devito.types.utils import DimensionTuple
 from devito.types.dimension import (Dimension, SpaceDimension, TimeDimension,
                                     Spacing, SteppingDimension, SubDimension,
-                                    AbstractSubDimension, DefaultDimension)
+                                    MultiSubDimension, DefaultDimension)
 
 __all__ = ['Grid', 'SubDomain', 'SubDomainSet']
 
@@ -546,35 +546,6 @@ def define(self, dimensions):
         raise NotImplementedError
 
 
-class MultiSubDimension(AbstractSubDimension):
-
-    """
-    A special Dimension to be used in MultiSubDomains.
-    """
-
-    is_MultiSub = True
-
-    __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
-
-    def __init_finalize__(self, name, parent, left, right, thickness,
-                          functions=None, bounds_indices=None,
-                          implicit_dimension=None):
-        super().__init_finalize__(name, parent, left, right, thickness)
-        self.functions = functions
-        self.bounds_indices = bounds_indices
-        self.implicit_dimension = implicit_dimension
-
-    def __hash__(self):
-        # There is no possibility for two MultiSubDimensions to ever hash the
-        # same, since a MultiSubDimension carries a reference to a MultiSubDomain,
-        # which is unique
-        return id(self)
-
-    @cached_property
-    def bound_symbols(self):
-        return self.parent.bound_symbols
-
-
 class MultiSubDomain(AbstractSubDomain):
 
     """

From 539d0dbeeb793d23531d33cdb37d01eeb59c1d71 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 16 Jul 2024 16:26:44 +0100
Subject: [PATCH 3/9] compiler: Move symbolic thickness

---
 devito/types/dimension.py | 14 +++++++-------
 devito/types/grid.py      |  7 ++-----
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index c4d556ae7b..26788c7356 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -558,6 +558,13 @@ def __init_finalize__(self, name, parent, left, right, thickness, **kwargs):
         self._interval = sympy.Interval(left, right)
         self._thickness = Thickness(*thickness)
 
+    @classmethod
+    def _symbolic_thickness(cls, name):
+        return (Scalar(name="%s_ltkn" % name, dtype=np.int32,
+                       is_const=True, nonnegative=True),
+                Scalar(name="%s_rtkn" % name, dtype=np.int32,
+                       is_const=True, nonnegative=True))
+
     @cached_property
     def symbolic_min(self):
         return self._interval.left
@@ -646,13 +653,6 @@ def __init_finalize__(self, name, parent, left, right, thickness, local, **kwarg
         super().__init_finalize__(name, parent, left, right, thickness)
         self._local = local
 
-    @classmethod
-    def _symbolic_thickness(cls, name):
-        return (Scalar(name="%s_ltkn" % name, dtype=np.int32,
-                       is_const=True, nonnegative=True),
-                Scalar(name="%s_rtkn" % name, dtype=np.int32,
-                       is_const=True, nonnegative=True))
-
     @classmethod
     def left(cls, name, parent, thickness, local=True):
         lst, rst = cls._symbolic_thickness(name)
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 44fa9941bb..dd2c890c09 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -10,7 +10,7 @@
 from devito.mpi import Distributor, MPI
 from devito.tools import ReducerMap, as_tuple
 from devito.types.args import ArgProvider
-from devito.types.basic import Scalar, Symbol
+from devito.types.basic import Scalar
 from devito.types.dense import Function
 from devito.types.utils import DimensionTuple
 from devito.types.dimension import (Dimension, SpaceDimension, TimeDimension,
@@ -747,10 +747,7 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
 
             dname = '%si%d' % (d.name, counter)
 
-            ltkn = Symbol(name="%s_ltkn%d" % (d.name, counter), dtype=np.int32,
-                          is_const=True, nonnegative=True)
-            rtkn = Symbol(name="%s_rtkn%d" % (d.name, counter), dtype=np.int32,
-                          is_const=True, nonnegative=True)
+            ltkn, rtkn = MultiSubDimension._symbolic_thickness(dname)
 
             left = d.symbolic_min + ltkn
             right = d.symbolic_max - rtkn

From bd716e016643b3702d30304d9579483e0730a8be Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 16 Jul 2024 17:06:41 +0100
Subject: [PATCH 4/9] compiler: Tweak MultiSubDimension lowering

---
 devito/passes/clusters/implicit.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 675f1effc6..a7c494cc04 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -221,8 +221,8 @@ def _lower_msd(dim, cluster):
 @_lower_msd.register(MultiSubDimension)
 def _(dim, cluster):
     i_dim = dim.implicit_dimension
-    mapper = {(dim, i): dim.functions[i_dim, mM]
-              for i, mM in enumerate(dim.bounds_indices)}
+    mapper = {t[0]: dim.functions[i_dim, mM]
+              for t, mM in zip(dim.thickness, dim.bounds_indices)}
     return mapper, i_dim
 
 
@@ -246,7 +246,7 @@ def lower_msd(msdims, cluster):
 
 
 def make_implicit_exprs(mapper):
-    return [Eq(list(d._thickness_map)[side], v) for (d, side), v in mapper.items()]
+    return [Eq(k, v) for k, v in mapper.items()]
 
 
 def reduce(m0, m1, edims, prefix):

From 801b0046f6f4b7b013e5d1dfd194d18facb1d813 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Wed, 17 Jul 2024 08:20:53 +0000
Subject: [PATCH 5/9] compiler: Polish MultiSubDimension lowering

---
 devito/ir/equations/algorithms.py  |  4 ++--
 devito/passes/clusters/implicit.py |  4 ++--
 devito/types/dimension.py          | 13 +++++++++----
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index cad931afad..c82a619742 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -200,7 +200,7 @@ def _(expr, mapper):
 def _(d, mapper):
     # TODO: to be implemented as soon as we drop the counter machinery in
     # Grid.__subdomain_finalize__
-    return {}
+    pass
 
 
 @_concretize_subdims.register(ConditionalDimension)
@@ -209,7 +209,7 @@ def _(d, mapper):
     # Grid.__subdomain_finalize__
     # TODO: call `_concretize_subdims(d.parent, mapper)` as the parent might be
     # a SubDimension!
-    return {}
+    pass
 
 
 @_concretize_subdims.register(MultiSubDimension)
diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index a7c494cc04..a2f7b0ad42 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -221,8 +221,8 @@ def _lower_msd(dim, cluster):
 @_lower_msd.register(MultiSubDimension)
 def _(dim, cluster):
     i_dim = dim.implicit_dimension
-    mapper = {t[0]: dim.functions[i_dim, mM]
-              for t, mM in zip(dim.thickness, dim.bounds_indices)}
+    mapper = {t: dim.functions[i_dim, mM]
+              for t, mM in zip(dim.tkns, dim.bounds_indices)}
     return mapper, i_dim
 
 
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 26788c7356..641d2898fb 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -15,10 +15,10 @@
 from devito.types.constant import Constant
 
 __all__ = ['Dimension', 'SpaceDimension', 'TimeDimension', 'DefaultDimension',
-           'CustomDimension', 'SteppingDimension', 'SubDimension', 'MultiSubDimension',
-           'ConditionalDimension', 'ModuloDimension', 'IncrDimension',
-           'BlockDimension', 'StencilDimension', 'VirtualDimension',
-           'Spacing', 'dimensions']
+           'CustomDimension', 'SteppingDimension', 'SubDimension',
+           'MultiSubDimension', 'ConditionalDimension', 'ModuloDimension',
+           'IncrDimension', 'BlockDimension', 'StencilDimension',
+           'VirtualDimension', 'Spacing', 'dimensions']
 
 
 Thickness = namedtuple('Thickness', 'left right')
@@ -596,6 +596,11 @@ def rtkn(self):
         # Shortcut for the right thickness symbol
         return self.thickness.right[0]
 
+    @property
+    def tkns(self):
+        # Shortcut for both thickness symbols
+        return self.ltkn, self.rtkn
+
 
 class SubDimension(AbstractSubDimension):
 

From 52d8e09e64a6b74cc846d0c815578a2c9afdc688 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Wed, 17 Jul 2024 08:55:56 +0000
Subject: [PATCH 6/9] compiler: Lift left/right logic into MultiSubDim.__init__

---
 devito/ir/equations/algorithms.py | 22 ++++++++-----------
 devito/types/dimension.py         | 35 ++++++++++++++++++++++++++-----
 devito/types/grid.py              | 12 +++--------
 3 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index c82a619742..c7dd44a80e 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -214,23 +214,19 @@ def _(d, mapper):
 
 @_concretize_subdims.register(MultiSubDimension)
 def _(d, mapper):
-    if any(d.thickness):
+    if not d.is_abstract:
         # TODO: for now Grid.__subdomain_finalize__ creates the thickness, but
         # soon it will be done here instead
         return
-    else:
-        pd = d.parent
-
-        subs = mapper[pd]
 
-        ltkn = Symbol(name="%s_ltkn%d" % (pd.name, len(subs)), dtype=np.int32,
-                      is_const=True, nonnegative=True)
-        rtkn = Symbol(name="%s_rtkn%d" % (pd.name, len(subs)), dtype=np.int32,
-                      is_const=True, nonnegative=True)
+    pd = d.parent
 
-        left = pd.symbolic_min + ltkn
-        right = pd.symbolic_max - rtkn
+    subs = mapper[pd]
 
-        thickness = ((ltkn, None), (rtkn, None))
+    ltkn = Symbol(name="%s_ltkn%d" % (pd.name, len(subs)), dtype=np.int32,
+                  is_const=True, nonnegative=True)
+    rtkn = Symbol(name="%s_rtkn%d" % (pd.name, len(subs)), dtype=np.int32,
+                  is_const=True, nonnegative=True)
+    thickness = (ltkn, rtkn)
 
-        subs[d] = d._rebuild(d.name, pd, left, right, thickness)
+    subs[d] = d._rebuild(d.name, pd, thickness)
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 641d2898fb..3b70151969 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -9,7 +9,7 @@
 from devito.data import LEFT, RIGHT
 from devito.exceptions import InvalidArgument
 from devito.logger import debug
-from devito.tools import Pickable, is_integer
+from devito.tools import Pickable, is_integer, flatten
 from devito.types.args import ArgProvider
 from devito.types.basic import Symbol, DataSymbol, Scalar
 from devito.types.constant import Constant
@@ -586,6 +586,10 @@ def thickness(self):
     def _thickness_map(self):
         return dict(self.thickness)
 
+    @property
+    def is_abstract(self):
+        return all(i is None for i in flatten(self.thickness))
+
     @property
     def ltkn(self):
         # Shortcut for the left thickness symbol
@@ -654,7 +658,8 @@ class SubDimension(AbstractSubDimension):
 
     __rargs__ = AbstractSubDimension.__rargs__ + ('local',)
 
-    def __init_finalize__(self, name, parent, left, right, thickness, local, **kwargs):
+    def __init_finalize__(self, name, parent, left, right, thickness, local,
+                          **kwargs):
         super().__init_finalize__(name, parent, left, right, thickness)
         self._local = local
 
@@ -803,11 +808,31 @@ class MultiSubDimension(AbstractSubDimension):
 
     is_MultiSub = True
 
+    __rargs__ = (DerivedDimension.__rargs__ + ('thickness',))
     __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
 
-    def __init_finalize__(self, name, parent, left, right, thickness,
-                          functions=None, bounds_indices=None,
-                          implicit_dimension=None):
+    def __init_finalize__(self, name, parent, thickness, functions=None,
+                          bounds_indices=None, implicit_dimension=None):
+        # Canonicalize thickness
+        if thickness is None:
+            # Using dummy left/right is the only thing we can do for such
+            # an abstract MultiSubDimension
+            thickness = ((None, None), (None, None))
+
+            left = sympy.S.NegativeInfinity
+            right = sympy.S.Infinity
+        elif isinstance(thickness, tuple):
+            if all(isinstance(i, Symbol) for i in thickness):
+                ltkn, rtkn = thickness
+                thickness = ((ltkn, None), (rtkn, None))
+            elif all(isinstance(i, tuple) and len(i) == 2 for i in thickness):
+                (ltkn, _), (rtkn, _) = thickness
+
+            left = parent.symbolic_min + ltkn
+            right = parent.symbolic_max - rtkn
+        else:
+            raise ValueError("MultiSubDimension expects a tuple of thicknesses")
+
         super().__init_finalize__(name, parent, left, right, thickness)
         self.functions = functions
         self.bounds_indices = bounds_indices
diff --git a/devito/types/grid.py b/devito/types/grid.py
index dd2c890c09..041f4156b1 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -747,17 +747,11 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
 
             dname = '%si%d' % (d.name, counter)
 
-            ltkn, rtkn = MultiSubDimension._symbolic_thickness(dname)
-
-            left = d.symbolic_min + ltkn
-            right = d.symbolic_max - rtkn
-
-            thickness = ((ltkn, None), (rtkn, None))
+            thickness = MultiSubDimension._symbolic_thickness(dname)
 
             dimensions.append(MultiSubDimension(
-                dname, d, left, right, thickness,
-                functions=sd_func, bounds_indices=(2*i, 2*i+1),
-                implicit_dimension=i_dim
+                dname, d, thickness, functions=sd_func,
+                bounds_indices=(2*i, 2*i+1), implicit_dimension=i_dim
             ))
 
         self._dimensions = tuple(dimensions)

From d3d260e82bf5dfd4d4b43998d18ddac05191aff5 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Wed, 17 Jul 2024 11:19:50 +0100
Subject: [PATCH 7/9] compiler: Use sregistry and _symbolic_thickness

---
 devito/ir/equations/algorithms.py | 37 +++++++++++++++++--------------
 devito/types/dimension.py         | 14 +++++++-----
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index c7dd44a80e..05cee3089d 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -2,11 +2,9 @@
 from collections import defaultdict
 from functools import singledispatch
 
-import numpy as np
-
 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, Eq, IgnoreDimSort, SubDimension, Symbol,
+from devito.types import (Dimension, Eq, IgnoreDimSort, SubDimension,
                           ConditionalDimension)
 from devito.types.basic import AbstractFunction
 from devito.types.grid import MultiSubDimension
@@ -162,10 +160,12 @@ def concretize_subdims(exprs, **kwargs):
     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)
+    _concretize_subdims(exprs, mapper, sregistry)
     if not mapper:
         return exprs
 
@@ -179,32 +179,32 @@ def concretize_subdims(exprs, **kwargs):
 
 
 @singledispatch
-def _concretize_subdims(a, mapper):
-    return {}
+def _concretize_subdims(a, mapper, sregistry):
+    pass
 
 
 @_concretize_subdims.register(list)
 @_concretize_subdims.register(tuple)
-def _(v, mapper):
+def _(v, mapper, sregistry):
     for i in v:
-        _concretize_subdims(i, mapper)
+        _concretize_subdims(i, mapper, sregistry)
 
 
 @_concretize_subdims.register(Eq)
-def _(expr, mapper):
+def _(expr, mapper, sregistry):
     for d in expr.free_symbols:
-        _concretize_subdims(d, mapper)
+        _concretize_subdims(d, mapper, sregistry)
 
 
 @_concretize_subdims.register(SubDimension)
-def _(d, mapper):
+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):
+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
@@ -213,7 +213,7 @@ def _(d, mapper):
 
 
 @_concretize_subdims.register(MultiSubDimension)
-def _(d, mapper):
+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
@@ -223,10 +223,13 @@ def _(d, mapper):
 
     subs = mapper[pd]
 
-    ltkn = Symbol(name="%s_ltkn%d" % (pd.name, len(subs)), dtype=np.int32,
-                  is_const=True, nonnegative=True)
-    rtkn = Symbol(name="%s_rtkn%d" % (pd.name, len(subs)), dtype=np.int32,
-                  is_const=True, nonnegative=True)
+    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)
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 3b70151969..8a30fd3182 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -559,11 +559,11 @@ def __init_finalize__(self, name, parent, left, right, thickness, **kwargs):
         self._thickness = Thickness(*thickness)
 
     @classmethod
-    def _symbolic_thickness(cls, name):
-        return (Scalar(name="%s_ltkn" % name, dtype=np.int32,
-                       is_const=True, nonnegative=True),
-                Scalar(name="%s_rtkn" % name, dtype=np.int32,
-                       is_const=True, nonnegative=True))
+    def _symbolic_thickness(cls, name, stype=Scalar):
+        return (stype(name="%s_ltkn" % name, dtype=np.int32,
+                      is_const=True, nonnegative=True),
+                stype(name="%s_rtkn" % name, dtype=np.int32,
+                      is_const=True, nonnegative=True))
 
     @cached_property
     def symbolic_min(self):
@@ -844,6 +844,10 @@ def __hash__(self):
         # which is unique
         return id(self)
 
+    @classmethod
+    def _symbolic_thickness(cls, name):
+        return super()._symbolic_thickness(name, stype=Symbol)
+
     @cached_property
     def bound_symbols(self):
         return self.parent.bound_symbols

From e34ecde2ab8f069f77fcd7c6400ba504001d6245 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Wed, 17 Jul 2024 12:13:23 +0100
Subject: [PATCH 8/9] misc: Change thickness iterator name

---
 devito/passes/clusters/implicit.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index a2f7b0ad42..33a05e8ecc 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -221,8 +221,8 @@ def _lower_msd(dim, cluster):
 @_lower_msd.register(MultiSubDimension)
 def _(dim, cluster):
     i_dim = dim.implicit_dimension
-    mapper = {t: dim.functions[i_dim, mM]
-              for t, mM in zip(dim.tkns, dim.bounds_indices)}
+    mapper = {tkn: dim.functions[i_dim, mM]
+              for tkn, mM in zip(dim.tkns, dim.bounds_indices)}
     return mapper, i_dim
 
 

From a35c7db734f0d5c91855a1aeb018972f9e62f2f7 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Wed, 17 Jul 2024 14:19:01 +0000
Subject: [PATCH 9/9] compiler: Strengthen AbstractSubDimension.init

---
 devito/types/dimension.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 8a30fd3182..17981587a1 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -828,8 +828,13 @@ def __init_finalize__(self, name, parent, thickness, functions=None,
             elif all(isinstance(i, tuple) and len(i) == 2 for i in thickness):
                 (ltkn, _), (rtkn, _) = thickness
 
-            left = parent.symbolic_min + ltkn
-            right = parent.symbolic_max - rtkn
+            try:
+                left = parent.symbolic_min + ltkn
+                right = parent.symbolic_max - rtkn
+            except TypeError:
+                # May end up here after a reconstruction
+                left = sympy.S.NegativeInfinity
+                right = sympy.S.Infinity
         else:
             raise ValueError("MultiSubDimension expects a tuple of thicknesses")