compiler: fix msd thickness lowering

devito/finite_differences/elementary.py

@@ -2,7 +2,6 @@
 from packaging.version import Version
 
 from devito.finite_differences.differentiable import DifferentiableFunction, diffify
-from devito.types.lazy import Evaluable
 
 
 class factorial(DifferentiableFunction, sympy.factorial):
@@ -89,13 +88,13 @@ def root(x):
     return diffify(sympy.root(x))
 
 
-class Min(sympy.Min, Evaluable):
+class Min(sympy.Min, DifferentiableFunction):
 
     def _evaluate(self, **kwargs):
         return self.func(*self._evaluate_args(**kwargs), evaluate=False)
 
 
-class Max(sympy.Max, Evaluable):
+class Max(sympy.Max, DifferentiableFunction):
 
     def _evaluate(self, **kwargs):
         return self.func(*self._evaluate_args(**kwargs), evaluate=False)

devito/finite_differences/operators.py

@@ -1,3 +1,6 @@
+from sympy import S
+
+
 def div(func, shift=None, order=None, method='FD', side=None, **kwargs):
     """
     Divergence of the input Function.
@@ -194,6 +197,6 @@ def diag(func, size=None):
     to = getattr(func, 'time_order', 0)
 
     tens_func = TensorTimeFunction if func.is_TimeDependent else TensorFunction
-    comps = [[func if i == j else 0 for i in range(dim)] for j in range(dim)]
+    comps = [[func if i == j else S.Zero for i in range(dim)] for j in range(dim)]
     return tens_func(name='diag', grid=func.grid, space_order=func.space_order,
                      components=comps, time_order=to, diagonal=True)

devito/ir/equations/algorithms.py

@@ -298,9 +298,10 @@ def _(d, mapper, rebuilt, sregistry):
     tkns = tuple(tkn._rebuild(name=sregistry.make_name(prefix=tkn.name))
                  for tkn in d.thickness)
     kwargs = {'thickness': tkns}
-    fkwargs = {}
+    fkwargs = {'function': None, 'halo': None, 'padding': None}
 
     idim0 = d.implicit_dimension
+    idim1 = None
     if idim0 is not None:
         if idim0 in rebuilt:
             idim1 = rebuilt[idim0]
@@ -309,33 +310,34 @@ def _(d, mapper, rebuilt, sregistry):
             rebuilt[idim0] = idim1 = idim0._rebuild(name=iname)
 
         kwargs['implicit_dimension'] = idim1
-        fkwargs['dimensions'] = (idim1,) + d.functions.dimensions[1:]
 
-    if d.functions in rebuilt:
-        functions = rebuilt[d.functions]
-    else:
-        # Increment every instance of this name after the first encountered
-        fname = sregistry.make_name(prefix=d.functions.name, increment_first=False)
-        # Warn the user if name has been changed, since this will affect overrides
-        if fname != d.functions.name:
-            fkwargs['name'] = fname
-            warning(
-                f"{str(d.functions)} <{id(d.functions)}> renamed as '{fname}'. "
-                "Consider assigning a unique name to {d.functions.name}."
-            )
-
-        fkwargs.update({'function': None,
-                        'halo': None,
-                        'padding': None})
-
-        # Data in MultiSubDimension function may not have been touched at this point,
-        # in which case do not use an allocator, as there is nothing to allocate, and
-        # doing so will petrify the `_data` attribute as None.
-        if d.functions._data is not None:
-            fkwargs.update({'allocator': DataReference(d.functions._data)})
-
-        rebuilt[d.functions] = functions = d.functions._rebuild(**fkwargs)
-
-    kwargs['functions'] = functions
+    functions = []
+    for f in as_tuple(d.functions):
+        if f in rebuilt:
+            functions.append(rebuilt[f])
+        else:
+            # Increment every instance of this name after the first encountered
+            fname = sregistry.make_name(prefix=f.name, increment_first=False)
+            # Warn the user if name has been changed, since this will affect overrides
+            if fname != f.name:
+                fkwargs['name'] = fname
+                warning(
+                    f"{str(f)} <{id(f)}> renamed as '{fname}'. "
+                    f"Consider assigning a unique name to {f.name}."
+                )
+
+            if idim1 is not None:
+                fkwargs['dimensions'] = (idim1,) + f.dimensions[1:]
+            # Data in MultiSubDimension function may not have been touched at this point,
+            # in which case do not use an allocator, as there is nothing to allocate, and
+            # doing so will petrify the `_data` attribute as None.
+            if f._data is not None:
+                fkwargs.update({'allocator': DataReference(f._data)})
+
+            rebuilt[f] = f._rebuild(**fkwargs)
+            functions.append(rebuilt[f])
+
+    kwargs['functions'] = tuple(functions)
 
     mapper[d] = d._rebuild(**kwargs)
+    mapper.update(zip(d.thickness, tkns, strict=True))

devito/passes/clusters/blocking.py

@@ -125,6 +125,17 @@ def _has_short_trip_count(self, d):
         # most a few tens of unit, so they wouldn't benefit from blocking
         return is_integer(d.symbolic_size)
 
+    def _is_non_uniform(self, prefix, d):
+        # A non-uniform Dimension is one that is not guaranteed to have the same
+        # bounds and cannot be blocked
+        if not d.is_MultiSub:
+            # MultiSubdimensions are the only ones that can have variable bounds
+            return False
+
+        # Check if the bound of the MultiSubdimension is variable for d
+        p_dims = {d for pd in prefix for d in pd.dim._defines if not d.is_Time}
+        return any(p_dims & set(f.dimensions) for f in d.functions)
+
 
 class AnalyzeBlocking(AnayzeBlockingBase):
 
@@ -136,6 +147,9 @@ def callback(self, clusters, prefix):
         if self._has_short_trip_count(d):
             return clusters
 
+        if self._is_non_uniform(prefix[:-1], d):
+            return clusters
+
         processed = []
         for c in clusters:
             if not {PARALLEL,
@@ -184,6 +198,9 @@ def callback(self, clusters, prefix):
 
         d = prefix[-1].dim
 
+        if self._is_non_uniform(prefix[:-1], d):
+            return clusters
+
         processed = []
         for c in clusters:
             if not c.properties.is_parallel_relaxed(d):
@@ -252,6 +269,9 @@ def callback(self, clusters, prefix):
         if self._has_short_trip_count(d):
             return clusters
 
+        if self._is_non_uniform(prefix[:-1], d):
+            return clusters
+
         # Pointless if there's no data reuse
         if all(not self._has_data_reuse(c) for c in clusters):
             return clusters

devito/passes/clusters/implicit.py

@@ -183,12 +183,14 @@ def key(tkn):
             if not mapper:
                 continue
 
-            found[d.functions].clusters.append(c)
-            found[d.functions].mapper = reduce(found[d.functions].mapper,
-                                               mapper, {dim}, prefix)
+            for df in d.functions:
+                found[df].clusters.append(c)
+                found[df].mapper = reduce(found[df].mapper,
+                                          mapper, {dim}, prefix)
 
         # Turn the reduced mapper into a list of equations
         processed = []
+
         for bunch in found.values():
             exprs = make_implicit_exprs(bunch.mapper)
 
@@ -226,7 +228,8 @@ def _lower_msd(dim, cluster):
 def _(dim, cluster):
     i_dim = dim.implicit_dimension
     mapper = {
-        tkn: dim.functions[i_dim, mM]
+        tkn: f[i_dim, mM]
+        for f in dim.functions
         for tkn, mM in zip(dim.tkns, dim.bounds_indices, strict=True)
     }
     return mapper, i_dim

devito/types/dimension.py

@@ -11,7 +11,7 @@
 from devito.deprecations import deprecations
 from devito.exceptions import InvalidArgument
 from devito.logger import debug
-from devito.tools import Pickable, is_integer, is_number, memoized_meth
+from devito.tools import Pickable, as_tuple, is_integer, is_number, memoized_meth
 from devito.types.args import ArgProvider
 from devito.types.basic import DataSymbol, Scalar, Symbol
 from devito.types.constant import Constant
@@ -823,7 +823,7 @@ def __init_finalize__(self, name, parent, thickness, functions=None,
                           bounds_indices=None, implicit_dimension=None):
 
         super().__init_finalize__(name, parent, thickness)
-        self.functions = functions
+        self.functions = as_tuple(functions)
         self.bounds_indices = bounds_indices
         self.implicit_dimension = implicit_dimension
 

devito/types/grid.py

@@ -39,7 +39,8 @@ class CartesianDiscretization:
 
     def __init__(self, shape=None, dimensions=None, dtype=None):
         self._shape = as_tuple(shape)
-        self._dimensions = as_tuple(dimensions)
+        dimensions = as_tuple(dimensions)
+        self._dimensions = DimensionTuple(*dimensions, getters=dimensions)
         self._dtype = dtype
 
     @property
@@ -653,7 +654,7 @@ def __subdomain_finalize_legacy__(self, grid):
                     sdshape.append(thickness)
 
         self._shape = tuple(sdshape)
-        self._dimensions = tuple(sub_dimensions)
+        self._dimensions = DimensionTuple(*sub_dimensions, getters=grid.dimensions)
         self._dtype = grid.dtype
 
     @property
@@ -914,7 +915,7 @@ def __subdomain_finalize_core__(self, grid):
                 bounds_indices=(2*i, 2*i+1), implicit_dimension=i_dim
             ))
 
-        self._dimensions = tuple(dimensions)
+        self._dimensions = DimensionTuple(*dimensions, getters=grid.dimensions)
         self._subfunction = sd_func
 
     def __subdomain_finalize__(self):

tests/test_subdomains.py

@@ -693,7 +693,7 @@ class Dummy(SubDomainSet):
         # 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]
+                    and i.expr.rhs.function in xi.functions]
         # The thickness expressions Eq(x_ltkn0, dummy[n0][0]), ...
         # should be scheduled once per dimension
         assert len(sdsexprs) == 6