diff --git a/devito/finite_differences/elementary.py b/devito/finite_differences/elementary.py index 55a0054c03..c6625e7488 100644 --- a/devito/finite_differences/elementary.py +++ b/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) diff --git a/devito/finite_differences/operators.py b/devito/finite_differences/operators.py index 156466324e..c8acb8b0a5 100644 --- a/devito/finite_differences/operators.py +++ b/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) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index b3a78c0ebe..653e9dd76f 100644 --- a/devito/ir/equations/algorithms.py +++ b/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)) diff --git a/devito/passes/clusters/blocking.py b/devito/passes/clusters/blocking.py index d04962a52c..708bbd7d08 100644 --- a/devito/passes/clusters/blocking.py +++ b/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 diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py index 36058f6393..8a42436d2d 100644 --- a/devito/passes/clusters/implicit.py +++ b/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 diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 4010472bef..43dc070e88 100644 --- a/devito/types/dimension.py +++ b/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 diff --git a/devito/types/grid.py b/devito/types/grid.py index e166dd1262..69888ac7b0 100644 --- a/devito/types/grid.py +++ b/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): diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 2fafd4e02b..1e32240656 100644 --- a/tests/test_subdomains.py +++ b/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