diff --git a/devito/passes/clusters/implicit.py b/devito/passes/clusters/implicit.py index 560e293b5a..590b0a0377 100644 --- a/devito/passes/clusters/implicit.py +++ b/devito/passes/clusters/implicit.py @@ -4,7 +4,6 @@ from collections import defaultdict from functools import singledispatch -from math import floor import numpy as np @@ -12,7 +11,8 @@ 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,7 +218,7 @@ 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 @@ -223,23 +226,36 @@ def msdim(d): @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): 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 523d044e2b..d49cf95ad7 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 @@ -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): diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index 196df33348..5ee2076fd9 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): """ @@ -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