From 9794266b25273230edc65cbed38885e32496db11 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Fri, 28 Jun 2024 09:13:33 +0100
Subject: [PATCH 01/11] dsl: Initial reworking of MultiSubDimension and
 SubDomainSet

---
 devito/passes/clusters/implicit.py |  3 +++
 devito/types/grid.py               | 25 +++++++++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 560e293b5a..f5b39a51ad 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -70,6 +70,7 @@ class LowerExplicitMSD(LowerMSD):
     _q_syncs_in_key = True
 
     def callback(self, clusters, prefix):
+        print("prefix", prefix)
         try:
             dim = prefix[-1].dim
         except IndexError:
@@ -85,6 +86,7 @@ def callback(self, clusters, prefix):
         tip = None
         processed = []
         for c in clusters:
+            # FIXME: This needs to be reworked to be dimension-centric
             try:
                 dd = c.ispace[idx].dim
                 d = msdim(dd)
@@ -95,6 +97,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
+            # FIXME: What is inside this, and could it simply point to the dimension?
             mapper, dims = lower_msd(d.msd, c)
             if not dims:
                 # An Implicit MSD
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 523d044e2b..c659497459 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -589,6 +589,20 @@ def symbolic_max(self):
     def symbolic_size(self):
         return self.parent.symbolic_size
 
+    @property
+    def implicit_dim(self):
+        try:
+            return self._implicit_dim
+        except AttributeError:
+            return None
+
+    @property
+    def functions(self):
+        try:
+            return self._functions
+        except AttributeError:
+            return None
+
 
 class MultiSubDomain(AbstractSubDomain):
 
@@ -781,7 +795,8 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
         # compiler can use to generate code
         n = counter - npresets
         assert n >= 0
-        self._implicit_dimension = i_dim = Dimension(name='n%d' % n)
+
+        i_dim = Dimension(name='n%d' % n)
         functions = []
         for j in range(len(self._local_bounds)):
             index = floor(j/2)
@@ -801,7 +816,13 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
                 f.data[:] = self._local_bounds[j]
 
             functions.append(f)
-        self._functions = as_tuple(functions)
+
+        # FIXME: Might require a better data structure here
+        functions = as_tuple(functions)
+
+        for d in self.dimensions:
+            d._implicit_dimension = i_dim
+            d._functions = functions
 
     @property
     def n_domains(self):

From 1e7a02943d6d0557dafd29e9a06af52e98832910 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Fri, 28 Jun 2024 15:24:58 +0100
Subject: [PATCH 02/11] compiler: Drop SubDomainSet earlier in compilation for
 more graceful lowering

---
 devito/passes/clusters/implicit.py | 52 +++++++++++++--------------
 devito/types/dimension.py          |  1 +
 devito/types/grid.py               | 58 +++++++++++-------------------
 3 files changed, 46 insertions(+), 65 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index f5b39a51ad..4d4fb26f50 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -3,8 +3,6 @@
 """
 
 from collections import defaultdict
-from functools import singledispatch
-from math import floor
 
 import numpy as np
 
@@ -12,7 +10,6 @@
 from devito.symbolics import retrieve_dimensions
 from devito.tools import Bunch, frozendict, timed_pass
 from devito.types import Eq, Symbol
-from devito.types.grid import MultiSubDimension, SubDomainSet
 
 __all__ = ['generate_implicit']
 
@@ -70,7 +67,6 @@ class LowerExplicitMSD(LowerMSD):
     _q_syncs_in_key = True
 
     def callback(self, clusters, prefix):
-        print("prefix", prefix)
         try:
             dim = prefix[-1].dim
         except IndexError:
@@ -86,7 +82,6 @@ def callback(self, clusters, prefix):
         tip = None
         processed = []
         for c in clusters:
-            # FIXME: This needs to be reworked to be dimension-centric
             try:
                 dd = c.ispace[idx].dim
                 d = msdim(dd)
@@ -96,9 +91,12 @@ def callback(self, clusters, prefix):
                 processed.append(c)
                 continue
 
+            msdims = [msdim(i.dim) for i in c.ispace[idx:]]
+            msdims = [d for d in msdims if d is not None]
+
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            # FIXME: What is inside this, and could it simply point to the dimension?
-            mapper, dims = lower_msd(d.msd, c)
+            mapper, dims = lower_msd(msdims)
+
             if not dims:
                 # An Implicit MSD
                 processed.append(c)
@@ -172,8 +170,12 @@ def callback(self, clusters, prefix):
             except IndexError:
                 continue
 
+            msdims = [i.dim for i in ispace]
+            msdims = [d for d in msdims if d is not None]
+
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd(d.msd, c)
+            mapper, dims = lower_msd(msdims)
+            # mapper, dims = lower_msd(d.msd, c)
             if dims:
                 # An Explicit MSD
                 continue
@@ -218,31 +220,27 @@ def callback(self, clusters, prefix):
 def msdim(d):
     try:
         for i in d._defines:
-            if isinstance(i, MultiSubDimension):
+            if i.is_MultiSub:
                 return i
     except AttributeError:
         pass
     return None
 
 
-@singledispatch
-def lower_msd(msd, cluster):
-    # Retval: (dynamic thickness mapper, iteration dimensions)
-    return (), ()
-
-
-@lower_msd.register(SubDomainSet)
-def _(msd, *args):
-    ret = {}
-    for j in range(len(msd._local_bounds)):
-        index = floor(j/2)
-        side = j % 2
-        d = msd.dimensions[index]
-        f = msd._functions[j]
-
-        ret[(d.root, side)] = f.indexify()
-
-    return frozendict(ret), (msd._implicit_dimension,)
+def lower_msd(msdims):
+    # FIXME: Variable names are horrible here. Suggestions encouraged
+    mapper = {}
+    dims = set()
+    for d in msdims:
+        # Pull out the parent MultiSubDimension if blocked etc
+        msds = [d for d in d._defines if d.is_MultiSub]
+        assert len(msds) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
+        msd = msds.pop()
+
+        mapper.update({(d.root, i): f.indexify()
+                       for i, f in enumerate(msd.functions)})
+        dims.add(msd.implicit_dimension)
+    return mapper, tuple(dims)
 
 
 def make_implicit_exprs(mapper, sregistry):
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index ec476e558b..92f476e0a2 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -103,6 +103,7 @@ class Dimension(ArgProvider):
     is_NonlinearDerived = False
     is_AbstractSub = False
     is_Sub = False
+    is_MultiSub = False
     is_Conditional = False
     is_Stepping = False
     is_Stencil = False
diff --git a/devito/types/grid.py b/devito/types/grid.py
index c659497459..23961eba8d 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -1,7 +1,6 @@
 from abc import ABC
 from collections import namedtuple
 from functools import cached_property
-from math import floor
 
 import numpy as np
 from sympy import prod
@@ -553,19 +552,9 @@ class MultiSubDimension(AbstractSubDimension):
     A special Dimension to be used in MultiSubDomains.
     """
 
-    __rargs__ = ('name', 'parent', 'msd')
+    is_MultiSub = True
 
-    def __init_finalize__(self, name, parent, msd):
-        # NOTE: a MultiSubDimension stashes a reference to the originating
-        # MultiSubDomain.  This creates a circular pattern as the `msd` itself
-        # carries references to its MultiSubDimensions. This is currently
-        # necessary because during compilation we drop the MultiSubDomain
-        # early, but when the MultiSubDimensions are processed we still need it
-        # to create the implicit equations. Untangling this is definitely
-        # possible, but not straightforward
-        self.msd = msd
-
-        super().__init_finalize__(name, parent)
+    __rargs__ = ('name', 'parent')
 
     def __hash__(self):
         # There is no possibility for two MultiSubDimensions to ever hash the
@@ -590,9 +579,9 @@ def symbolic_size(self):
         return self.parent.symbolic_size
 
     @property
-    def implicit_dim(self):
+    def implicit_dimension(self):
         try:
-            return self._implicit_dim
+            return self._implicit_dimension
         except AttributeError:
             return None
 
@@ -754,7 +743,7 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
 
         # Create the SubDomainSet SubDimensions
         self._dimensions = tuple(
-            MultiSubDimension('%si%d' % (d.name, counter), d, self)
+            MultiSubDimension('%si%d' % (d.name, counter), d)
             for d in grid.dimensions
         )
 
@@ -797,32 +786,25 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
         assert n >= 0
 
         i_dim = Dimension(name='n%d' % n)
-        functions = []
-        for j in range(len(self._local_bounds)):
-            index = floor(j/2)
-            d = self.dimensions[index]
-            if j % 2 == 0:
-                fname = "%s_%s" % (self.name, d.min_name)
-            else:
-                fname = "%s_%s" % (self.name, d.max_name)
-            f = Function(name=fname, grid=self._grid, shape=(self._n_domains,),
-                         dimensions=(i_dim,), dtype=np.int32)
-
-            # Check if shorthand notation has been provided:
-            if isinstance(self._local_bounds[j], int):
-                f.data[:] = np.full((self._n_domains,), self._local_bounds[j],
-                                    dtype=np.int32)
-            else:
-                f.data[:] = self._local_bounds[j]
 
-            functions.append(f)
+        for i, d in enumerate(self.dimensions):
+            fnames = ["%s_%s" % (self.name, minmax)
+                      for minmax in (d.min_name, d.max_name)]
+            f = tuple(Function(name=fname, grid=self._grid, shape=(self._n_domains,),
+                               dimensions=(i_dim,), dtype=np.int32)
+                      for fname in fnames)
 
-        # FIXME: Might require a better data structure here
-        functions = as_tuple(functions)
+            # Check if shorthand notation has been provided:
+            for j in range(2):
+                idx = 2*i + j
+                if isinstance(self._local_bounds[idx], int):
+                    f[j].data[:] = np.full((self._n_domains,), self._local_bounds[idx],
+                                           dtype=np.int32)
+                else:
+                    f[j].data[:] = self._local_bounds[idx]
 
-        for d in self.dimensions:
             d._implicit_dimension = i_dim
-            d._functions = functions
+            d._functions = f
 
     @property
     def n_domains(self):

From 9deb1a3255dbbb440ecba36034176a17be9d84f0 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Fri, 28 Jun 2024 16:16:41 +0100
Subject: [PATCH 03/11] misc: Tidy up

---
 devito/passes/clusters/implicit.py | 24 +++++++++---------------
 devito/types/grid.py               |  6 +++---
 2 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 4d4fb26f50..9c36aa41e6 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -91,11 +91,9 @@ def callback(self, clusters, prefix):
                 processed.append(c)
                 continue
 
-            msdims = [msdim(i.dim) for i in c.ispace[idx:]]
-            msdims = [d for d in msdims if d is not None]
-
-            # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd(msdims)
+            # Get all MultiSubDimensions in the cluster and get the dynamic thickness
+            # mapper for the associated MultiSubDomain
+            mapper, dims = lower_msd([msdim(i.dim) for i in c.ispace[idx:]])
 
             if not dims:
                 # An Implicit MSD
@@ -170,11 +168,8 @@ def callback(self, clusters, prefix):
             except IndexError:
                 continue
 
-            msdims = [i.dim for i in ispace]
-            msdims = [d for d in msdims if d is not None]
-
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd(msdims)
+            mapper, dims = lower_msd([i.dim for i in ispace])
             # mapper, dims = lower_msd(d.msd, c)
             if dims:
                 # An Explicit MSD
@@ -228,17 +223,16 @@ def msdim(d):
 
 
 def lower_msd(msdims):
-    # FIXME: Variable names are horrible here. Suggestions encouraged
+    msdims = [d for d in msdims if d is not None]
     mapper = {}
     dims = set()
     for d in msdims:
         # Pull out the parent MultiSubDimension if blocked etc
-        msds = [d for d in d._defines if d.is_MultiSub]
-        assert len(msds) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
-        msd = msds.pop()
+        msd = [d for d in d._defines if d.is_MultiSub]
+        assert len(msd) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
+        msd = msd.pop()
 
-        mapper.update({(d.root, i): f.indexify()
-                       for i, f in enumerate(msd.functions)})
+        mapper.update({(d.root, i): f.indexify() for i, f in enumerate(msd.functions)})
         dims.add(msd.implicit_dimension)
     return mapper, tuple(dims)
 
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 23961eba8d..367420ad00 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -791,14 +791,14 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
             fnames = ["%s_%s" % (self.name, minmax)
                       for minmax in (d.min_name, d.max_name)]
             f = tuple(Function(name=fname, grid=self._grid, shape=(self._n_domains,),
-                               dimensions=(i_dim,), dtype=np.int32)
-                      for fname in fnames)
+                               dimensions=(i_dim,), dtype=np.int32) for fname in fnames)
 
             # Check if shorthand notation has been provided:
             for j in range(2):
                 idx = 2*i + j
                 if isinstance(self._local_bounds[idx], int):
-                    f[j].data[:] = np.full((self._n_domains,), self._local_bounds[idx],
+                    f[j].data[:] = np.full((self._n_domains,),
+                                           self._local_bounds[idx],
                                            dtype=np.int32)
                 else:
                     f[j].data[:] = self._local_bounds[idx]

From b5df33643ce2bf2c68d78db3fbf5c6b8ae2682e6 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Fri, 28 Jun 2024 17:09:36 +0100
Subject: [PATCH 04/11] dsl: Avoid modification of MultiSubDimensions

---
 devito/passes/clusters/implicit.py |  6 ++---
 devito/types/grid.py               | 42 ++++++++++++------------------
 2 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 9c36aa41e6..2c6d9fdfc2 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -93,7 +93,7 @@ def callback(self, clusters, prefix):
 
             # Get all MultiSubDimensions in the cluster and get the dynamic thickness
             # mapper for the associated MultiSubDomain
-            mapper, dims = lower_msd([msdim(i.dim) for i in c.ispace[idx:]])
+            mapper, dims = lower_msd({msdim(i.dim) for i in c.ispace[idx:]} - {None})
 
             if not dims:
                 # An Implicit MSD
@@ -169,8 +169,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd([i.dim for i in ispace])
-            # mapper, dims = lower_msd(d.msd, c)
+            mapper, dims = lower_msd({i.dim for i in ispace} - {None})
             if dims:
                 # An Explicit MSD
                 continue
@@ -223,7 +222,6 @@ def msdim(d):
 
 
 def lower_msd(msdims):
-    msdims = [d for d in msdims if d is not None]
     mapper = {}
     dims = set()
     for d in msdims:
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 367420ad00..f8caf2e7ce 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -555,6 +555,12 @@ class MultiSubDimension(AbstractSubDimension):
     is_MultiSub = True
 
     __rargs__ = ('name', 'parent')
+    __rkwargs__ = ('functions', 'implicit_dimension')
+
+    def __init_finalize__(self, name, parent, functions=None, implicit_dimension=None):
+        super().__init_finalize__(name, parent)
+        self.functions = functions
+        self.implicit_dimension = implicit_dimension
 
     def __hash__(self):
         # There is no possibility for two MultiSubDimensions to ever hash the
@@ -578,20 +584,6 @@ def symbolic_max(self):
     def symbolic_size(self):
         return self.parent.symbolic_size
 
-    @property
-    def implicit_dimension(self):
-        try:
-            return self._implicit_dimension
-        except AttributeError:
-            return None
-
-    @property
-    def functions(self):
-        try:
-            return self._functions
-        except AttributeError:
-            return None
-
 
 class MultiSubDomain(AbstractSubDomain):
 
@@ -741,12 +733,6 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
         self._grid = grid
         self._dtype = grid.dtype
 
-        # Create the SubDomainSet SubDimensions
-        self._dimensions = tuple(
-            MultiSubDimension('%si%d' % (d.name, counter), d)
-            for d in grid.dimensions
-        )
-
         # Compute the SubDomainSet shapes
         global_bounds = []
         for i in self._global_bounds:
@@ -777,7 +763,7 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
             self._local_bounds = self._global_bounds
 
         # Sanity check
-        if len(self._local_bounds) != 2*len(self.dimensions):
+        if len(self._local_bounds) != 2*len(grid.dimensions):
             raise ValueError("Left and right bounds must be supplied for each dimension")
 
         # Associate the `_local_bounds` to suitable symbolic objects that the
@@ -787,9 +773,11 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
 
         i_dim = Dimension(name='n%d' % n)
 
-        for i, d in enumerate(self.dimensions):
-            fnames = ["%s_%s" % (self.name, minmax)
-                      for minmax in (d.min_name, d.max_name)]
+        dimensions = []
+        for i, d in enumerate(grid.dimensions):
+            dname = '%si%d' % (d.name, counter)
+            fnames = ["%s_%s_%s" % (self.name, dname, mM)
+                      for mM in ('m', 'M')]
             f = tuple(Function(name=fname, grid=self._grid, shape=(self._n_domains,),
                                dimensions=(i_dim,), dtype=np.int32) for fname in fnames)
 
@@ -803,8 +791,10 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
                 else:
                     f[j].data[:] = self._local_bounds[idx]
 
-            d._implicit_dimension = i_dim
-            d._functions = f
+            dimensions.append(MultiSubDimension(dname, d, functions=f,
+                                                implicit_dimension=i_dim))
+
+        self._dimensions = tuple(dimensions)
 
     @property
     def n_domains(self):

From 032b95cab19f5fdb7f2c8342eaa19b2c058b2d52 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Fri, 28 Jun 2024 17:14:12 +0100
Subject: [PATCH 05/11] misc: Remove unnecessary set operations

---
 devito/passes/clusters/implicit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 2c6d9fdfc2..01e518acdb 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -169,7 +169,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd({i.dim for i in ispace} - {None})
+            mapper, dims = lower_msd({i.dim for i in ispace})
             if dims:
                 # An Explicit MSD
                 continue

From 6510bdbdd8394cce309e2deff21b0f0632dbeb49 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Mon, 1 Jul 2024 09:48:33 +0100
Subject: [PATCH 06/11] dsl: Remove unnecessary rargs

---
 devito/types/grid.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/devito/types/grid.py b/devito/types/grid.py
index f8caf2e7ce..b29ee79ab3 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -554,7 +554,6 @@ class MultiSubDimension(AbstractSubDimension):
 
     is_MultiSub = True
 
-    __rargs__ = ('name', 'parent')
     __rkwargs__ = ('functions', 'implicit_dimension')
 
     def __init_finalize__(self, name, parent, functions=None, implicit_dimension=None):

From db66a4f3b65ab116132e683064f0293dd8ec0850 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Mon, 1 Jul 2024 14:17:17 +0100
Subject: [PATCH 07/11] dsl: Reintroduce singledispatch

---
 devito/passes/clusters/implicit.py | 31 ++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 01e518acdb..c0c179a878 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -3,6 +3,7 @@
 """
 
 from collections import defaultdict
+from functools import singledispatch
 
 import numpy as np
 
@@ -10,6 +11,7 @@
 from devito.symbolics import retrieve_dimensions
 from devito.tools import Bunch, frozendict, timed_pass
 from devito.types import Eq, Symbol
+from devito.types.grid import MultiSubDimension
 
 __all__ = ['generate_implicit']
 
@@ -221,18 +223,31 @@ def msdim(d):
     return None
 
 
+@singledispatch
+def _lower_msd(dim):
+    # Retval: (dynamic thickness mapper, iteration dimensions)
+    return (), ()
+
+
+@_lower_msd.register(MultiSubDimension)
+def _(dim):
+    # Pull out the parent MultiSubDimension if blocked etc
+    msd = [d for d in dim._defines if d.is_MultiSub]
+    assert len(msd) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
+    msd = msd.pop()
+
+    mapper = {(dim.root, i): f.indexify() for i, f in enumerate(msd.functions)}
+    return mapper, msd.implicit_dimension
+
+
 def lower_msd(msdims):
     mapper = {}
     dims = set()
     for d in msdims:
-        # Pull out the parent MultiSubDimension if blocked etc
-        msd = [d for d in d._defines if d.is_MultiSub]
-        assert len(msd) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
-        msd = msd.pop()
-
-        mapper.update({(d.root, i): f.indexify() for i, f in enumerate(msd.functions)})
-        dims.add(msd.implicit_dimension)
-    return mapper, tuple(dims)
+        dmapper, ddim = _lower_msd(d)
+        mapper.update(dmapper)
+        dims.add(ddim)
+    return mapper, tuple(dims - {None})
 
 
 def make_implicit_exprs(mapper, sregistry):

From cff01df07bc13f22d0af548ec2cb3051e4ea3db9 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Mon, 1 Jul 2024 16:43:16 +0100
Subject: [PATCH 08/11] dsl: Encapsulate bounds with single Function

---
 devito/passes/clusters/implicit.py | 22 +++++++++++----------
 devito/types/grid.py               | 31 ++++++++++++++++--------------
 tests/test_subdomains.py           |  2 +-
 3 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index c0c179a878..6c1cab378a 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -95,7 +95,7 @@ def callback(self, clusters, prefix):
 
             # Get all MultiSubDimensions in the cluster and get the dynamic thickness
             # mapper for the associated MultiSubDomain
-            mapper, dims = lower_msd({msdim(i.dim) for i in c.ispace[idx:]} - {None})
+            mapper, dims = lower_msd({msdim(i.dim) for i in c.ispace[idx:]} - {None}, c)
 
             if not dims:
                 # An Implicit MSD
@@ -171,7 +171,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd({i.dim for i in ispace})
+            mapper, dims = lower_msd({i.dim for i in ispace}, c)
             if dims:
                 # An Explicit MSD
                 continue
@@ -224,30 +224,32 @@ def msdim(d):
 
 
 @singledispatch
-def _lower_msd(dim):
+def _lower_msd(dim, cluster):
     # Retval: (dynamic thickness mapper, iteration dimensions)
-    return (), ()
+    return {}, None
 
 
 @_lower_msd.register(MultiSubDimension)
-def _(dim):
+def _(dim, cluster):
     # Pull out the parent MultiSubDimension if blocked etc
     msd = [d for d in dim._defines if d.is_MultiSub]
     assert len(msd) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
     msd = msd.pop()
 
-    mapper = {(dim.root, i): f.indexify() for i, f in enumerate(msd.functions)}
-    return mapper, msd.implicit_dimension
+    i_dim = msd.implicit_dimension
+    mapper = {(dim.root, i): msd.functions[i_dim, mM]
+              for i, mM in enumerate(msd.bounds_indices)}
+    return mapper, i_dim
 
 
-def lower_msd(msdims):
+def lower_msd(msdims, cluster):
     mapper = {}
     dims = set()
     for d in msdims:
-        dmapper, ddim = _lower_msd(d)
+        dmapper, ddim = _lower_msd(d, cluster)
         mapper.update(dmapper)
         dims.add(ddim)
-    return mapper, tuple(dims - {None})
+    return frozendict(mapper), tuple(dims - {None})
 
 
 def make_implicit_exprs(mapper, sregistry):
diff --git a/devito/types/grid.py b/devito/types/grid.py
index b29ee79ab3..ab767932fd 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)
+                                    AbstractSubDimension, DefaultDimension)
 
 __all__ = ['Grid', 'SubDomain', 'SubDomainSet']
 
@@ -554,11 +554,13 @@ class MultiSubDimension(AbstractSubDimension):
 
     is_MultiSub = True
 
-    __rkwargs__ = ('functions', 'implicit_dimension')
+    __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
 
-    def __init_finalize__(self, name, parent, functions=None, implicit_dimension=None):
+    def __init_finalize__(self, name, parent, functions=None, bounds_indices=None,
+                          implicit_dimension=None):
         super().__init_finalize__(name, parent)
         self.functions = functions
+        self.bounds_indices = bounds_indices
         self.implicit_dimension = implicit_dimension
 
     def __hash__(self):
@@ -771,26 +773,27 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
         assert n >= 0
 
         i_dim = Dimension(name='n%d' % n)
+        d_dim = DefaultDimension(name='d%d' % n, 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)
 
         dimensions = []
         for i, d in enumerate(grid.dimensions):
-            dname = '%si%d' % (d.name, counter)
-            fnames = ["%s_%s_%s" % (self.name, dname, mM)
-                      for mM in ('m', 'M')]
-            f = tuple(Function(name=fname, grid=self._grid, shape=(self._n_domains,),
-                               dimensions=(i_dim,), dtype=np.int32) for fname in fnames)
-
             # Check if shorthand notation has been provided:
             for j in range(2):
                 idx = 2*i + j
                 if isinstance(self._local_bounds[idx], int):
-                    f[j].data[:] = np.full((self._n_domains,),
-                                           self._local_bounds[idx],
-                                           dtype=np.int32)
+                    sd_func.data[:, idx] = np.full((self._n_domains,),
+                                                   self._local_bounds[idx],
+                                                   dtype=np.int32)
                 else:
-                    f[j].data[:] = self._local_bounds[idx]
+                    sd_func.data[:, idx] = self._local_bounds[idx]
 
-            dimensions.append(MultiSubDimension(dname, d, functions=f,
+            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))
 
         self._dimensions = tuple(dimensions)
diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py
index 196df33348..432f13e230 100644
--- a/tests/test_subdomains.py
+++ b/tests/test_subdomains.py
@@ -305,7 +305,7 @@ class MySubdomains(SubDomainSet):
         # unique -- see issue #1474
         exprs = FindNodes(Expression).visit(op)
         reads = set().union(*[e.reads for e in exprs])
-        assert len(reads) == 7  # f, g, h, xi_n_m, xi_n_M, yi_n_m, yi_n_M
+        assert len(reads) == 4  # f, g, h, mydomains
 
     def test_multi_sets(self):
         """

From 9580e267530858bbd0b8c2953a0cc5fbbdf82312 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 2 Jul 2024 10:43:19 +0100
Subject: [PATCH 09/11] tests: Added additional check to SubDomainSet with
 blocking

---
 devito/passes/clusters/implicit.py | 5 +++--
 devito/types/grid.py               | 7 +------
 tests/test_subdomains.py           | 9 +++++++++
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 6c1cab378a..100ba4842a 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -182,8 +182,9 @@ def callback(self, clusters, prefix):
             if dim not in edims or not edims.issubset(prefix.dimensions):
                 continue
 
-            found[d.msd].clusters.append(c)
-            found[d.msd].mapper = reduce(found[d.msd].mapper, mapper, edims, prefix)
+            found[d.functions].clusters.append(c)
+            found[d.functions].mapper = reduce(found[d.functions].mapper,
+                                               mapper, edims, prefix)
 
         # Turn the reduced mapper into a list of equations
         mapper = {}
diff --git a/devito/types/grid.py b/devito/types/grid.py
index ab767932fd..d49cf95ad7 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -783,12 +783,7 @@ def __subdomain_finalize__(self, grid, counter=0, **kwargs):
             # Check if shorthand notation has been provided:
             for j in range(2):
                 idx = 2*i + j
-                if isinstance(self._local_bounds[idx], int):
-                    sd_func.data[:, idx] = np.full((self._n_domains,),
-                                                   self._local_bounds[idx],
-                                                   dtype=np.int32)
-                else:
-                    sd_func.data[:, idx] = self._local_bounds[idx]
+                sd_func.data[:, idx] = self._local_bounds[idx]
 
             dname = '%si%d' % (d.name, counter)
             dimensions.append(MultiSubDimension(dname, d,
diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py
index 432f13e230..5ee2076fd9 100644
--- a/tests/test_subdomains.py
+++ b/tests/test_subdomains.py
@@ -642,6 +642,15 @@ class Dummy(SubDomainSet):
         assert_structure(op, ['t,n0', 't,n0,xi20_blk0,yi20_blk0,x,y,z'],
                          't,n0,xi20_blk0,yi20_blk0,x,y,z')
 
+        xi, _, _ = dummy.dimensions
+        # Check that the correct number of thickness expressions are generated
+        sdsexprs = [i.expr for i in FindNodes(Expression).visit(op)
+                    if i.expr.rhs.is_Indexed
+                    and i.expr.rhs.function is xi.functions]
+        # The thickness expressions Eq(x_ltkn0, dummy[n0][0]), ...
+        # should be scheduled once per dimension
+        assert len(sdsexprs) == 6
+
     def test_sequential_implicit(self):
         """
         Make sure the implicit dimensions of the MultiSubDomain define a sequential

From 5bfc7d7c6f28345b20363db6786761c1c2f4c4fa Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 2 Jul 2024 13:08:55 +0100
Subject: [PATCH 10/11] compiler: Better lowering of blocked MultiSubDimensions

---
 devito/passes/clusters/implicit.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 100ba4842a..60b435dbde 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -11,6 +11,7 @@
 from devito.symbolics import retrieve_dimensions
 from devito.tools import Bunch, frozendict, timed_pass
 from devito.types import Eq, Symbol
+from devito.types.dimension import BlockDimension
 from devito.types.grid import MultiSubDimension
 
 __all__ = ['generate_implicit']
@@ -226,21 +227,25 @@ def msdim(d):
 
 @singledispatch
 def _lower_msd(dim, cluster):
-    # Retval: (dynamic thickness mapper, iteration dimensions)
+    # Retval: (dynamic thickness mapper, iteration dimension)
     return {}, None
 
 
 @_lower_msd.register(MultiSubDimension)
 def _(dim, cluster):
-    # Pull out the parent MultiSubDimension if blocked etc
+    i_dim = dim.implicit_dimension
+    mapper = {(dim.root, i): dim.functions[i_dim, mM]
+              for i, mM in enumerate(dim.bounds_indices)}
+    return mapper, i_dim
+
+
+@_lower_msd.register(BlockDimension)
+def _(dim, cluster):
+    # Pull out the parent MultiSubDimension
     msd = [d for d in dim._defines if d.is_MultiSub]
     assert len(msd) == 1  # Sanity check. MultiSubDimensions shouldn't be nested.
     msd = msd.pop()
-
-    i_dim = msd.implicit_dimension
-    mapper = {(dim.root, i): msd.functions[i_dim, mM]
-              for i, mM in enumerate(msd.bounds_indices)}
-    return mapper, i_dim
+    return _lower_msd(msd, cluster)
 
 
 def lower_msd(msdims, cluster):

From b306e07e1350ad4a30bc4b68a771876991fe29d0 Mon Sep 17 00:00:00 2001
From: Edward Caunt <ed@devitocodes.com>
Date: Tue, 2 Jul 2024 14:08:25 +0100
Subject: [PATCH 11/11] misc: Minor aestheic tweak

---
 devito/passes/clusters/implicit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py
index 60b435dbde..590b0a0377 100644
--- a/devito/passes/clusters/implicit.py
+++ b/devito/passes/clusters/implicit.py
@@ -172,7 +172,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd({i.dim for i in ispace}, c)
+            mapper, dims = lower_msd(ispace.itdims, c)
             if dims:
                 # An Explicit MSD
                 continue