compiler: Drop SubDomainSet earlier during compilation for more graceful lowering

devito/passes/clusters/implicit.py

@@ -4,15 +4,15 @@
 
 from collections import defaultdict
 from functools import singledispatch
-from math import floor
 
 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.grid import MultiSubDimension, SubDomainSet
+from devito.types.dimension import BlockDimension
+from devito.types.grid import MultiSubDimension
 
 __all__ = ['generate_implicit']
 
@@ -94,8 +94,10 @@ def callback(self, clusters, prefix):
                 processed.append(c)
                 continue
 
-            # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd(d.msd, c)
+            # 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}, c)
+
             if not dims:
                 # An Implicit MSD
                 processed.append(c)
@@ -170,7 +172,7 @@ def callback(self, clusters, prefix):
                 continue
 
             # Get the dynamic thickness mapper for the given MultiSubDomain
-            mapper, dims = lower_msd(d.msd, c)
+            mapper, dims = lower_msd(ispace.itdims, c)
             if dims:
                 # An Explicit MSD
                 continue
@@ -181,8 +183,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 = {}
@@ -215,31 +218,44 @@ 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 (), ()
+def _lower_msd(dim, cluster):
+    # Retval: (dynamic thickness mapper, iteration dimension)
+    return {}, None
+
 
+@_lower_msd.register(MultiSubDimension)
+def _(dim, cluster):
+    i_dim = dim.implicit_dimension
+    mapper = {(dim.root, i): dim.functions[i_dim, mM]
+              for i, mM in enumerate(dim.bounds_indices)}
+    return mapper, i_dim
 
-@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()
+@_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()
+    return _lower_msd(msd, cluster)
 
-    return frozendict(ret), (msd._implicit_dimension,)
+
+def lower_msd(msdims, cluster):
+    mapper = {}
+    dims = set()
+    for d in msdims:
+        dmapper, ddim = _lower_msd(d, cluster)
+        mapper.update(dmapper)
+        dims.add(ddim)
+    return frozendict(mapper), tuple(dims - {None})
 
 
 def make_implicit_exprs(mapper, sregistry):

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

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
@@ -16,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']
 
@@ -553,19 +552,16 @@ 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
+    __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
 
+    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):
         # There is no possibility for two MultiSubDimensions to ever hash the
@@ -738,12 +734,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, self)
-            for d in grid.dimensions
-        )
-
         # Compute the SubDomainSet shapes
         global_bounds = []
         for i in self._global_bounds:
@@ -774,34 +764,34 @@ 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
         # compiler can use to generate code
         n = counter - npresets
         assert n >= 0
-        self._implicit_dimension = 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)
 
+        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):
             # 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]
+            for j in range(2):
+                idx = 2*i + j
+                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))
 
-            functions.append(f)
-        self._functions = as_tuple(functions)
+        self._dimensions = tuple(dimensions)
 
     @property
     def n_domains(self):

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):
         """
@@ -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