api: prevent sympy warnings with diag

Author: mloubout

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

@@ -309,34 +309,37 @@ def _(d, mapper, rebuilt, sregistry):
             rebuilt[idim0] = idim1 = idim0._rebuild(name=iname)
 
         kwargs['implicit_dimension'] = idim1
-        fkwargs['dimensions'] = (idim1,) + d.functions.dimensions[1:]
+        fkwargs['dimensions'] = (idim1,) + as_tuple(f.dimensions[1:] for f in d.functions)
 
-    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}."
+                )
+
+            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 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

@@ -178,6 +178,17 @@ def _has_enough_large_blockable_dims(self, cluster, d, nested=False):
                     if (cluster.properties.is_parallel_relaxed(i) and
                         not self._has_short_trip_count(i))]) >= 3
 
+    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 = set().union(*[d.dim._defines for d in prefix if not d.dim.is_Time])
+        return any(p_dims & set(f.dimensions) for f in d.functions)
+
     def callback(self, clusters, prefix):
         if not prefix:
             return clusters
@@ -190,6 +201,8 @@ def callback(self, clusters, prefix):
                 return clusters
 
             if is_on_device(c.functions, self.gpu_fit):
+                if self._is_non_uniform(prefix[:-1], d):
+                    return clusters
                 if self._has_short_trip_count(d):
                     if self._has_atomic_blockable_dim(c, d):
                         # Optimization: minimize number of parallel reductions

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