diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index de477b0292..4224b63134 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -156,6 +156,8 @@ def _lower_exprs(expressions, subs): # Add dimensions map to the mapper in case dimensions are used # as an expression, i.e. Eq(u, x, subdomain=xleft) mapper.update(dimension_map) + # Dimensions used as an expression are global indices, unlike loop indices + mapper.update(_make_global_index_map(expr, dimension_map)) # Add the user-supplied substitutions mapper.update(subs) # Apply mapper to expression @@ -184,14 +186,39 @@ def _make_dimension_map(expr): # Some Relationals may be pure SymPy objects, thus lacking the SubDomain dimension_map = {} else: - functions = [f for f in retrieve_functions(expr) if f._is_on_subdomain] + functions = retrieve_functions(expr) for f in functions: - dimension_map.update({d: expr.subdomain.dimension_map[d.root] - for d in f.space_dimensions if d.is_Sub}) + if f._is_on_subdomain: + dimension_map.update({d: expr.subdomain.dimension_map[d.root] + for d in f.space_dimensions if d.is_Sub}) + + # A subsampled Function, e.g. on a SubGrid, is subsampled from the + # SubDomain's Dimensions rather than from the whole Grid's + for f in functions: + dimension_map.update({ + d: d._rebuild(parent=expr.subdomain.dimension_map[d.parent]) + for d in f.dimensions + if d.is_Conditional and d.parent in expr.subdomain.dimension_map + }) return frozendict(dimension_map) +def _make_global_index_map(expr, dimension_map): + """ + Map the Dimensions of a distributed Grid to the global indices they stand for + when used as an expression, i.e. Eq(u, x + y) -> Eq(u[x, y], x + io_x + y + io_y). + Dimensions indexing an Indexed are unaffected as the Indexed is replaced as a whole. + """ + if not isinstance(expr, Eq): + return {} + + grids = {f.grid.root for f in retrieve_functions(expr) if f.grid is not None} + return {d: dimension_map.get(d, d) + o + for grid in grids if grid.distributor.is_parallel + for d, o in zip(grid.dimensions, grid.origin_isymbols, strict=True)} + + def concretize_subdims(exprs, **kwargs): """ Given a list of expressions, return a new list where all user-defined @@ -390,7 +417,7 @@ def generate_conditionals(expr, input_expr, ordering): # Replace the ConditionalDimensions in `expr` for d, cond in conditionals.items(): - index = d.index + index = d.index - d.offset if d.condition is not None and expr.has(d): index = index - relational_min(cond, d.parent) shift = relational_shift(cond, d.parent, d.symbolic_factor) diff --git a/devito/ir/support/guards.py b/devito/ir/support/guards.py index 3c4b573c03..1315484634 100644 --- a/devito/ir/support/guards.py +++ b/devito/ir/support/guards.py @@ -61,7 +61,7 @@ class GuardFactor(Guard, CondEq, Pickable): A guard for factor-based ConditionalDimensions. Given the ConditionalDimension `d` with factor `k`, create the - symbolic relational `d.parent % k == 0`. + symbolic relational `(d.parent - d.offset) % k == 0`. """ __rargs__ = ('d',) @@ -76,7 +76,7 @@ def __new__(cls, *args, **kwargs): d, = args assert d.is_Conditional - obj = super().__new__(cls, d.parent % d.symbolic_factor, 0) + obj = super().__new__(cls, (d.parent - d.offset) % d.symbolic_factor, 0) obj.d = d return obj diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 35596c39d6..fa5ea2e050 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -68,6 +68,7 @@ def __getattr__(self, name): 'Distributor', 'SparseDistributor', 'SubDistributor', + 'SubGridDistributor', 'devito_mpi_finalize', 'devito_mpi_init', ] @@ -375,10 +376,15 @@ def __init__(self, shape, dimensions, input_comm=None, topology=None): self._comm = MPI.COMM_NULL self._topology = tuple(1 for _ in range(len(shape))) - # The domain decomposition + self.__decomposition_setup__() + + def __decomposition_setup__(self): + """ + Set up the domain decomposition. + """ self._decomposition = [ Decomposition(np.array_split(range(i), j), c) - for i, j, c in zip(shape, self.topology, self.mycoords, strict=True) + for i, j, c in zip(self.glb_shape, self.topology, self.mycoords, strict=True) ] @cached_property @@ -485,6 +491,41 @@ def _rebuild(self, shape=None, dimensions=None, comm=None): comm or self.comm) +class SubGridDistributor(Distributor): + + """ + Decompose a SubGrid over a set of MPI processes in a manner consistent with + the parent distributor: a rank owns the SubGrid points on its parent points. + + Parameters + ---------- + subgrid : SubGrid + The SubGrid to be decomposed. + """ + + def __init__(self, subgrid): + self._parent = subgrid.parent.distributor + self._factors = subgrid.factor + + super().__init__(subgrid.shape, subgrid.dimensions, self.parent.comm, + tuple(self.parent.topology)) + + def __decomposition_setup__(self): + """ + Set up the decomposition, aligned with that of the parent Distributor. + """ + self._decomposition = [ + Decomposition([d[d % f == 0] // f for d in dec], c) + for dec, f, c in zip(self.parent.decomposition, self._factors, + self.mycoords, strict=True) + ] + + @property + def parent(self): + """The parent distributor of this SubGridDistributor.""" + return self._parent + + class SubDistributor(DenseDistributor): """ diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 38b2d4dc28..2fe0ebff16 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -43,7 +43,7 @@ split, timed_pass, timed_region ) from devito.types import Buffer, Evaluable, device_layer, disk_layer, host_layer -from devito.types.dimension import SubDimension, Thickness +from devito.types.dimension import SubDimension, SubsamplingOffset, Thickness from devito.warnings import warn __all__ = ['Operator'] @@ -415,6 +415,11 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Operation count before specialization init_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense) + # Strengthen factor-based ConditionalDimension guards into strided + # iteration. Before specialization, so that blocking tiles the swapped + # Dimensions + clusters = subsampling(clusters) + clusters = cls._specialize_clusters(clusters, **kwargs) # Operation count after specialization @@ -425,11 +430,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Generate implicit Clusters from higher level abstractions clusters = generate_implicit(clusters) - # Strengthen factor-based ConditionalDimension guards into strided - # iteration. Before `lower_index_derivatives` so that the Dimension - # swap is seen by everything downstream - clusters = subsampling(clusters) - # Lower all remaining high order symbolic objects clusters = lower_local_sums(clusters, **kwargs) clusters = lower_index_derivatives(clusters, **kwargs) @@ -695,15 +695,19 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs): # No need to do anything more if only checking the memory return args + # Process the symbols localized by MPI. A SubDimension's thicknesses are + # needed even when they only bound another Dimension, e.g. a subsampled one + localized = OrderedSet(*[p for p in self.parameters + if isinstance(p, (SubsamplingOffset, Thickness))]) + localized.update(t for d in toposort if isinstance(d, SubDimension) + for t in d.thickness) + for p in localized: + args.update(p._arg_values(grid=grid, **kwargs)) + # Process Dimensions for d in reversed(toposort): args.update(d._arg_values(self._dspace[d], grid, **kwargs)) - # Process Thicknesses - for p in self.parameters: - if isinstance(p, Thickness): - args.update(p._arg_values(grid=grid, **kwargs)) - # Process Objects for o in self.objects: args.update(o._arg_values(grid=grid, **kwargs)) diff --git a/devito/passes/clusters/subsampling.py b/devito/passes/clusters/subsampling.py index 7ef3e7f35e..6afd696c32 100644 --- a/devito/passes/clusters/subsampling.py +++ b/devito/passes/clusters/subsampling.py @@ -20,8 +20,8 @@ def subsampling(clusters): if (x%xf == 0) -> g[x/xf] = ... g[x/xf] = ... -The loop is taken over the subsampled Dimension `cd` itself, with a *unit* - stride, and every access to the parent `d` is rescaled to `cd*factor`: + The loop is taken over the subsampled Dimension `cd` itself, with a *unit* + stride, and every access to the parent `d` is mapped to `cd*factor + offset`: for (x = x_m; x <= x_M; x += 1) for (xs = xs_m; xs <= xs_M; xs++) if (x%xf == 0) -> g[xs] = u[xs*xf] @@ -30,8 +30,8 @@ def subsampling(clusters): A non-unitary increment is deliberately avoided: it would escape the Lamport-style dependence analysis the compiler relies on. The loop bounds are `cd`'s own runtime values, which `ConditionalDimension._arg_values` - derives as `ceil(x_m/factor) .. floor(x_M/factor)`, so no anchoring of the - loop minimum is needed either. + derives by inverting that map, so no anchoring of the loop minimum is needed + either. Parameters ---------- @@ -73,17 +73,18 @@ def callback(self, clusters, prefix): cd = guard.d factor = cd.symbolic_factor + offset = cd.offset processed = [] for c in clusters: # Iterate the subsampled Dimension itself, with unit stride, and - # rescale every parent access `d` to `cd*factor` + # map every parent access `d` to `cd*factor + offset` ispace = c.ispace.switch(d, cd) - # `IntDiv(d, factor)` is how the subsampled access was lowered; - # with `d -> cd*factor` it collapses to `cd`, so rewrite it - # directly rather than leaving `(cd*factor)/factor` behind - subs = {IntDiv(d, factor): cd, d: cd*factor} + # `IntDiv(d - offset, factor)` is how the subsampled access was + # lowered; with `d -> cd*factor + offset` it collapses to `cd`, so + # rewrite it directly rather than leaving the division behind + subs = {IntDiv(d - offset, factor): cd, d: cd*factor + offset} exprs = [uxreplace(e, subs) for e in c.exprs] # `cd` inherits `d`'s computational properties: it covers the very @@ -112,6 +113,10 @@ def _is_strideable(cluster, d, guard): if d not in cluster.ispace.itdims: return False + # Time keeps its guard, which later passes (e.g. buffering) rely upon + if d.is_Time: + return False + # Skipping points is only sound if no dependence is carried along `d`; # otherwise the discarded iterations may be doing work that `d`'s # remaining iterations rely upon diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 6b21493c14..259dc60950 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -942,6 +942,56 @@ class SubsamplingFactor(Scalar): pass +class SubsamplingOffset(DataSymbol): + + """ + The local index of the first subsampled point on an MPI rank. + + Subsampling is defined on global indices, but a rank need not start on a + subsampled point. E.g., with `factor=2` and two ranks: + + global 0 1 2 ... 10 | 11 12 13 14 ... 20 + local 0 1 2 ... 10 | 0 1 2 3 ... 9 + subsampled * . * ... * | . * . * ... * + offset 0 | 1 + + In local indices, `parent = factor*subsampled + offset`. + """ + + __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'factor') + + def __new__(cls, *args, root=None, factor=None, **kwargs): + newobj = super().__new__(cls, *args, **kwargs) + newobj._root = root + newobj._factor = factor + + return newobj + + def __init_finalize__(self, *args, **kwargs): + kwargs.setdefault('dtype', np.int32) + kwargs.setdefault('is_const', True) + kwargs.setdefault('nonnegative', True) + super().__init_finalize__(*args, **kwargs) + + @property + def root(self): + return self._root + + @property + def factor(self): + return self._factor + + def _arg_values(self, grid=None, **kwargs): + if grid is None: + return {self.name: 0} + + # Runtime arguments bind by name, even when different spacing dtypes + # make the compiled and runtime Grid Dimensions distinct symbols + root = next(d for d in grid.dimensions if d.name == self.root.name) + + return {self.name: -grid.origin_ioffset[root] % self.factor} + + class ConditionalDimension(DerivedDimension): """ @@ -971,6 +1021,9 @@ class ConditionalDimension(DerivedDimension): lowering for example combining Function's ConditionalDimension with an Equation's implicit_dim. All Dimensions within an equation must have `Or` relation for the final combined condition to be Or. + offset : SubsamplingOffset, optional, default=0 + The intercept of the map onto the parent Dimension, that is + ``parent = factor*self + offset``. Examples -------- @@ -1065,10 +1118,10 @@ class ConditionalDimension(DerivedDimension): STRICT = 'strict' __rkwargs__ = DerivedDimension.__rkwargs__ + \ - ('factor', 'condition', 'indirect', 'relation') + ('factor', 'condition', 'indirect', 'relation', 'offset') def __init_finalize__(self, name, parent=None, factor=None, condition=None, - indirect=False, relation=sympy.And, **kwargs): + indirect=False, relation=sympy.And, offset=0, **kwargs): # `parent=None` degenerates to a ConditionalDimension outside of # any iteration space if parent is None: @@ -1087,6 +1140,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None, else: raise ValueError("factor must be an integer") + self._offset = offset self._condition = condition self._indirect = indirect self._relation = relation @@ -1136,6 +1190,10 @@ def symbolic_factor(self): name=f'{self.name}f', dtype=np.int32, is_const=True ) + @property + def offset(self): + return self._offset + @property def condition(self): return self._condition @@ -1168,17 +1226,25 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs): fact = fact.data vals = {} + offset = args.get(self.offset.name, 0) if self.offset else 0 + + # The parent's rank-local bounds, e.g. including a SubDimension's thickness + bounds = (self.parent.symbolic_min, self.parent.symbolic_max) + values = {s: args[s.name] for b in bounds for s in b.free_symbols + if args.get(s.name) is not None} + parent_min, parent_max = (b.subs(values) for b in bounds) # The first subsampled point at or after the parent's min... with suppress(KeyError, TypeError): - vals[self.min_name] = math.ceil(kwargs.get(self.parent.min_name) / fact) + vals[self.min_name] = math.ceil((int(parent_min) - offset) / fact) # ...and the last one at or before its max. Note `floor` rather than # `ceil`: with e.g. `factor=3` and `x_M=11` the last subsampled point # is `x=9`, hence index 3, while `ceil(11/3) == 4` would overshoot the # Function by one with suppress(KeyError, TypeError): - vals[self.max_name] = kwargs.get(self.parent.max_name) // fact + vals[self.max_name] = min((int(parent_max) - offset) // fact, + args[self.max_name]) vals[self.symbolic_factor.name] = fact @@ -1207,8 +1273,10 @@ def _arg_defaults(self, _min=None, size=None, alias=None): f"{self.name}: {size} < {size0}") else: # Given a factor the last time index is factor*(size - 1) - # The maximum allowed value is then factor*size - 1 - defaults[dim.parent.max_name] = range(d0, d0 + factor*size) + # The maximum allowed value is then factor*size - 1, plus at + # most factor - 1 for an offset + stop = d0 + factor*size + (factor - 1 if dim.offset else 0) + defaults[dim.parent.max_name] = range(d0, stop) return defaults diff --git a/devito/types/grid.py b/devito/types/grid.py index 4568b852bf..060c3a8b9d 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -11,14 +11,14 @@ from devito.deprecations import deprecations from devito.finite_differences.elementary import Max, Min from devito.logger import warning -from devito.mpi import MPI, Distributor, SubDistributor +from devito.mpi import MPI, Distributor, SubDistributor, SubGridDistributor from devito.tools import ReducerMap, as_tuple, frozendict, is_integer from devito.types.args import ArgProvider from devito.types.basic import Scalar from devito.types.dense import Function from devito.types.dimension import ( ConditionalDimension, DefaultDimension, Dimension, MultiSubDimension, SpaceDimension, - Spacing, SteppingDimension, SubDimension, TimeDimension + Spacing, SteppingDimension, SubDimension, SubsamplingOffset, TimeDimension ) from devito.types.utils import DimensionTuple @@ -193,7 +193,7 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None, self._topology = None else: self._topology = None - self._distributor = Distributor(shape, dimensions, comm, self._topology) + self._distributor = self.__distributor_setup__(shape, dimensions, comm) # The physical extent and grid spacing if spacing is not None: @@ -212,6 +212,9 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None, self._origin_symbols = tuple(Scalar(name=f'o_{d.name}', dtype=dtype, is_const=True) for d in self.dimensions) + self._origin_isymbols = tuple(Scalar(name=f'io_{d.name}', dtype=np.int32, + is_const=True) + for d in self.dimensions) # Sanity check assert (self.dim == len(self.origin) == len(self.extent) == len(self.spacing)) @@ -262,10 +265,20 @@ def origin_map(self): """Map between origin symbols and their values.""" return dict(zip(self.origin_symbols, self.origin, strict=True)) + @property + def origin_isymbols(self): + """Symbols representing the global index of the rank's first grid point.""" + return self._origin_isymbols + + @property + def origin_imap(self): + """Map between index origin symbols and their values.""" + return dict(zip(self.origin_isymbols, self.origin_ioffset, strict=True)) + @property def origin_ioffset(self): """Offset index of the local (per-process) origin from the domain origin.""" - grid_origin = [min(i) for i in self.distributor.glb_numb] + grid_origin = [i.start for i in self.distributor.glb_slices.values()] assert len(grid_origin) == len(self.spacing) return DimensionTuple(*grid_origin, getters=self.dimensions) @@ -379,6 +392,7 @@ def _arg_names(self): ret = [] ret.append(self.time_dim.spacing.name) ret.extend([i.name for i in self.origin_map]) + ret.extend([i.name for i in self.origin_imap]) for i in self.spacing_map: try: ret.append(i.name) @@ -400,6 +414,7 @@ def _arg_defaults(self): # Grid origin args.update({k.name: v for k, v in self.origin_map.items()}) + args.update({k.name: v for k, v in self.origin_imap.items()}) # MPI-related objects if self.distributor.is_parallel: @@ -417,6 +432,9 @@ def _arg_values(self, **kwargs): return values + def __distributor_setup__(self, shape, dimensions, comm): + return Distributor(shape, dimensions, comm, self._topology) + def __getstate__(self): state = self.__dict__.copy() # A Distributor wraps an MPI communicator, which can't and shouldn't be pickled @@ -509,8 +527,10 @@ def __init__(self, parent, factor, name='s', **kwargs): ) dimensions = tuple( - d if f == 1 else ConditionalDimension(f'{d.name}{name}', parent=d, - factor=f) + d if f == 1 else ConditionalDimension( + f'{d.name}{name}', parent=d, factor=f, + offset=self.__offset_setup__(parent, d, f, name) + ) for d, f in zip(parent.dimensions, factors, strict=True) ) @@ -537,6 +557,15 @@ def __init__(self, parent, factor, name='s', **kwargs): # the very same buffer indices as one on the parent self._stepping_dim = parent.stepping_dim + def __distributor_setup__(self, shape, dimensions, comm): + return SubGridDistributor(self) + + @classmethod + def __offset_setup__(cls, parent, d, factor, name): + if not parent.distributor.is_parallel: + return 0 + return SubsamplingOffset(name=f'{d.name}{name}o', root=d, factor=factor) + @classmethod def __factors_setup__(cls, parent, factor): factors = as_tuple(factor) diff --git a/examples/userapi/05_conditional_dimension.ipynb b/examples/userapi/05_conditional_dimension.ipynb index a8e628545b..61a2965557 100644 --- a/examples/userapi/05_conditional_dimension.ipynb +++ b/examples/userapi/05_conditional_dimension.ipynb @@ -184,6 +184,9 @@ " lowering for example combining Function's ConditionalDimension with\n", " an Equation's implicit_dim. All Dimensions within an equation\n", " must have `Or` relation for the final combined condition to be Or.\n", + " offset : SubsamplingOffset, optional, default=0\n", + " The intercept of the map onto the parent Dimension, that is\n", + " ``parent = factor*self + offset``.\n", "\n", " Examples\n", " --------\n", diff --git a/tests/test_caching.py b/tests/test_caching.py index 3f17080312..2700df7c06 100644 --- a/tests/test_caching.py +++ b/tests/test_caching.py @@ -942,7 +942,7 @@ def test_solve(self, operate_on_empty_cache): # created by the finite difference (u.dt, u.dx2). We would have had # three extra references to u(t + dt), u(x - h_x) and u(x + h_x). # But this is not the case anymore! - assert len(_SymbolCache) == 8 + assert len(_SymbolCache) == 9 clear_cache() assert len(_SymbolCache) == 4 clear_cache() diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 7081be3beb..b05061150e 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -1137,6 +1137,60 @@ def check_halo_exchanges(op, exp0, exp1): return calls, tloop +class TestGlobalIndex: + + @pytest.mark.parallel(mode=[2, 4]) + def test_dimension_as_expression(self, mode): + """Dimensions in an expression are global indices (issue #2214).""" + grid = Grid(shape=(8, 8)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + + Operator(Eq(f, 10*x + y))() + + check = 10*np.arange(8)[:, None] + np.arange(8) + assert np.all(f.data[:] == check[f.local_indices]) + + @pytest.mark.parallel(mode=[2, 4]) + def test_dimension_as_expression_subdomain(self, mode): + """Same as above, but restricted to a SubDomain.""" + grid = Grid(shape=(8, 8)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + + Operator(Eq(f, 10*x + y, subdomain=grid.interior))() + + check = 10*np.arange(8)[:, None] + np.arange(8) + check[[0, -1], :] = check[:, [0, -1]] = 0 + assert np.all(f.data[:] == check[f.local_indices]) + + @pytest.mark.parallel(mode=[2]) + def test_dimension_indexing_unaffected(self, mode): + """Dimensions indexing a Function remain rank-local.""" + grid = Grid(shape=(8,)) + x = grid.dimensions[0] + f = Function(name='f', grid=grid, space_order=1) + g = Function(name='g', grid=grid) + + f.data[:] = np.arange(8) + Operator(Eq(g, f.subs(x, x + 1) + x))() + + # The last point reads a zero from the physical boundary halo + check = 2*np.arange(8) + 1 + check[-1] = 7 + assert np.all(g.data[:] == check[g.local_indices]) + + def test_serial_code_unchanged(self): + """No index offset is generated without MPI.""" + grid = Grid(shape=(4, 4)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + + op = Operator(Eq(f, x + y)) + + assert not any(p.name.startswith('io_') for p in op.parameters) + + class TestHaloWrites: @pytest.mark.parallel(mode=[1, 2]) diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index 3d235ca949..2ee0303e57 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -1,3 +1,4 @@ +import pickle import re import numpy as np @@ -6,11 +7,12 @@ from devito import ( CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Inc, Operator, - SpaceDimension, SubGrid, TimeFunction, switchconfig + SpaceDimension, SubDomain, SubGrid, TimeFunction, switchconfig ) from devito.ir.iet import Conditional, Expression, FindNodes, Iteration from devito.symbolics import retrieve_indexed from devito.tools import as_tuple +from devito.types.dimension import SubsamplingOffset def iteration(op, name): @@ -335,6 +337,11 @@ def test_time_subsampling(self): Operator([Eq(u.forward, u + 1.), Eq(usave, u)]).apply(time_M=nt-2) + # Time keeps its guard even where it is parallel + op = Operator(Eq(usave, 1.)) + assert FindNodes(Conditional).visit(op) + assert all(i.dim is not ts for i in FindNodes(Iteration).visit(op)) + # time runs 0..nt-2, so only the slots whose time index is reached # are written; `u` holds the step count at that point nwritten = (nt - 2)//factor + 1 @@ -533,11 +540,7 @@ def test_sequential_dimension(self): class TestInteractions: def test_with_blocking(self): - """ - Blocking rewrites `x` into BlockDimensions, so the subsampled - Dimension's parent leaves the iteration space and the guard is kept. - Correct, if not strided. - """ + """Striding comes first, so blocking tiles the subsampled Dimension.""" grid = Grid(shape=(32, 32)) x, y = grid.dimensions cd = ConditionalDimension(name='xs', parent=x, factor=4) @@ -546,12 +549,19 @@ def test_with_blocking(self): space_order=0) u.data[:] = np.arange(32*32).reshape(32, 32).astype(u.dtype) - op = Operator(Eq(g, u.dx2 + u.dy2), opt=('advanced', - {'blockrelax': True})) + eq = Eq(g, u.dx2 + u.dy2) + Operator(eq, opt='noop').apply() + expected = np.array(g.data) + g.data[:] = 0. + + op = Operator(eq, opt=('advanced', {'blockrelax': True})) op.apply() - assert len(FindNodes(Conditional).visit(op)) == 1 - assert np.all(np.isfinite(g.data)) + assert not FindNodes(Conditional).visit(op) + assert any(i.dim.is_Block and i.dim.parent is cd + for i in FindNodes(Iteration).visit(op)) + # `h^-2*eps*u` of float32 roundoff, as the two builds reorder flops + assert np.allclose(g.data, expected, atol=0.5) @switchconfig(language='openmp') def test_with_openmp(self): @@ -563,6 +573,22 @@ def test_with_openmp(self): assert np.array_equal(g.data, data[::2]) +class Middle(SubDomain): + + name = 'middle' + + def define(self, dimensions): + # Thicknesses deliberately not multiples of the SubGrid factors + return {d: ('middle', 2, 3) for d in dimensions} + + +def middle_of(data): + """`data` zeroed outside the `Middle` SubDomain.""" + ret = np.zeros_like(data) + ret[2:-3, 2:-3] = data[2:-3, 2:-3] + return ret + + class TestSubGrid: """The `SubGrid` API.""" @@ -625,6 +651,22 @@ def test_copy_from_parent(self, factor): fx, fy = sub.factor assert np.array_equal(f.data, data[::fx, ::fy]) + @pytest.mark.parametrize('factor', [2, 3, (2, 1), (3, 2)]) + def test_copy_on_subdomain(self, factor): + """Within a SubDomain, only the SubGrid points inside it are written.""" + grid = Grid(shape=(12, 12)) + sub = SubGrid(grid, factor=factor) + + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + data = np.arange(1, 145, dtype=u.dtype).reshape(12, 12) + u.data[:] = data + + Operator(Eq(f, u, subdomain=Middle(grid=grid))).apply() + + fx, fy = sub.factor + assert np.array_equal(f.data, middle_of(data)[::fx, ::fy]) + def test_copy_is_strided(self): """ No guard is emitted, every loop keeps a unit stride, and the Operator @@ -852,6 +894,92 @@ def test_custom_dimension_names(self): assert [d.name for d in sub.dimensions] == ['xc', 'yc'] +class TestBlocking: + + """ + Striding happens first, so that blocking tiles the subsampled Dimensions. + """ + + @staticmethod + def build(factor, shape=(21, 21, 21)): + grid = Grid(shape=shape) + sub = SubGrid(grid, factor=factor) + x = grid.dimensions[0] + + u = Function(name='u', grid=grid, space_order=2) + g = Function(name='g', grid=sub, space_order=2) + u.data[:] = np.arange(np.prod(shape), dtype=np.float64).reshape(shape) % 1009 + + eq = Eq(g, g + u.subs({x: x + 1})*u) + return eq, g, sub + + def test_blocked_and_strided(self): + """The block loops span the subsampled Dimensions, with no guards.""" + eq, _, sub = self.build(2) + + op = Operator(eq, opt=('advanced', {'blockrelax': True})) + + assert not FindNodes(Conditional).visit(op) + blocks = [i for i in FindNodes(Iteration).visit(op) if i.dim.is_Block] + assert {i.dim.parent for i in blocks if i.dim._depth == 1} == \ + set(sub.dimensions) + for d in sub.dimensions: + assert iteration(op, d.name).limits[2] == 1 + + @pytest.mark.parametrize('factor', [2, (2, 3, 4), (3, 2, 1), (1, 2, 2)]) + @pytest.mark.parametrize('block_size', [None, 3, 5]) + def test_numerics(self, factor, block_size): + """Any tiling of a strided nest visits each point exactly once.""" + eq, g, _ = self.build(factor) + Operator(eq, opt='noop').apply() + expected = np.array(g.data) + + g.data[:] = 0. + + op = Operator(eq, opt=('advanced', {'blockrelax': True})) + sizes = {p.name: block_size for p in op.parameters + if block_size and p.name.endswith('_size')} + assert sizes or block_size is None + op.apply(**sizes) + + assert np.array_equal(np.asarray(g.data), expected) + + def test_manual_conditional_dimensions(self): + """Hand-built ConditionalDimensions are tiled like a SubGrid's.""" + grid = Grid(shape=(21, 21, 21)) + x, y, z = grid.dimensions + xd, yd, zd = [ConditionalDimension(f'{d.name}d', parent=d, factor=2) + for d in grid.dimensions] + + u = Function(name='u', grid=grid, space_order=2) + g = Function(name='g', dimensions=(xd, yd, zd), shape=(11, 11, 11), + grid=grid, space_order=2) + + op = Operator(Eq(g, g + u[x + 1, y, z]*u[x - 1, y, z]), + opt=('advanced', {'blockrelax': True})) + + assert not FindNodes(Conditional).visit(op) + assert any(i.dim.is_Block and i.dim.parent is xd + for i in FindNodes(Iteration).visit(op)) + + def test_shared_nest_keeps_guards(self): + """A nest shared with unguarded work is blocked, not strided.""" + grid = Grid(shape=(21, 21, 21)) + sub = SubGrid(grid, factor=2) + x, y, z = grid.dimensions + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + g = Function(name='g', grid=sub, space_order=2) + u.data[:] = 2. + + op = Operator([Eq(v, u[x + 1, y, z] + u[x - 1, y, z]), Eq(g, g + v)], + opt=('advanced', {'blockrelax': True})) + op.apply() + + assert np.all(np.asarray(g.data[1:-1, :, :]) == 4.) + + class TestMPI: """ @@ -879,6 +1007,162 @@ def test_copy_from_parent(self, mode, factor): expected = glob[::factor, ::factor][g.local_indices] assert np.array_equal(np.asarray(g.data), expected) + @pytest.mark.parallel(mode=[2, 4]) + @pytest.mark.parametrize('factor', [2, 3]) + def test_copy_on_subdomain(self, mode, factor): + """The SubDomain bounds the SubGrid on every rank.""" + n = 24 + grid = Grid(shape=(n, n)) + sub = SubGrid(grid, factor=factor) + + u = Function(name='u', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + + glob = np.arange(1, n*n + 1, dtype=u.dtype).reshape(n, n) + u.data[:] = glob + + Operator(Eq(g, u, subdomain=Middle(grid=grid))).apply() + + expected = middle_of(glob)[::factor, ::factor][g.local_indices] + assert np.array_equal(np.asarray(g.data), expected) + + @pytest.mark.parallel(mode=[2, 4]) + @pytest.mark.parametrize('factor', [2, 3, 4]) + def test_rank_boundary_off_the_subgrid(self, mode, factor): + """The factor applies to global indices, wherever the ranks split.""" + n = 21 + grid = Grid(shape=(n, n)) + sub = SubGrid(grid, factor=factor) + + u = Function(name='u', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + + glob = np.arange(n*n, dtype=u.dtype).reshape(n, n) + u.data[:] = glob + + Operator(Eq(g, u)).apply() + + expected = glob[::factor, ::factor][g.local_indices] + assert np.array_equal(np.asarray(g.data), expected) + + @pytest.mark.parallel(mode=[2]) + def test_guarded_rank_boundary_off_the_subgrid(self, mode): + """A guard surviving in a shared nest also uses global indices.""" + n, factor = 21, 2 + grid = Grid(shape=(n, n)) + sub = SubGrid(grid, factor=factor) + + u = Function(name='u', grid=grid) + v = Function(name='v', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + + glob = np.arange(n*n, dtype=u.dtype).reshape(n, n) + u.data[:] = glob + + op = Operator([Eq(v, u + 1.), Eq(g, u)]) + + assert FindNodes(Conditional).visit(op) + op.apply() + + expected = glob[::factor, ::factor][g.local_indices] + assert np.array_equal(np.asarray(g.data), expected) + + @pytest.mark.parallel(mode=[2]) + def test_subgrid_only(self, mode): + """With no parent Function, the SubGrid alone sets the bounds.""" + grid = Grid(shape=(21, 21)) + sub = SubGrid(grid, factor=3) + + g = Function(name='g', grid=sub, space_order=0) + + Operator(Eq(g, 1.)).apply() + + assert np.all(np.asarray(g.data) == 1.) + + @pytest.mark.parallel(mode=[2]) + @pytest.mark.parametrize('factor', [2, 3]) + def test_blocked_rank_boundary_off_the_subgrid(self, mode, factor): + """Blocking a strided nest preserves the global sampling.""" + n = 21 + grid = Grid(shape=(n, n, n)) + sub = SubGrid(grid, factor=factor) + x = grid.dimensions[0] + + u = Function(name='u', grid=grid, space_order=2) + g = Function(name='g', grid=sub, space_order=2) + + glob = np.arange(n**3, dtype=np.float64).reshape(n, n, n) % 1009 + u.data[:] = glob + + op = Operator(Eq(g, g + u.subs({x: x + 1})*u), + opt=('advanced', {'blockrelax': True})) + + assert not FindNodes(Conditional).visit(op) + assert any(i.dim.is_Block for i in FindNodes(Iteration).visit(op)) + op.apply() + + padded = np.zeros((n + 1, n, n)) + padded[:n] = glob + expected = (padded[1:]*glob)[::factor, ::factor, ::factor] + assert np.array_equal(np.asarray(g.data), expected[g.local_indices]) + + @pytest.mark.parallel(mode=[4]) + def test_decomposition_follows_parent(self, mode): + """A rank owns exactly the subgrid points on its parent points.""" + n, factor = 13, 6 + grid = Grid(shape=(n,)) + sub = SubGrid(grid, factor=factor) + + for parent, child in zip(grid.distributor.decomposition[0], + sub.distributor.decomposition[0], strict=True): + assert np.array_equal(child*factor, parent[parent % factor == 0]) + + @pytest.mark.parallel(mode=[2]) + def test_manual_function_stays_in_rank(self, mode): + """A strided nest over a non-decomposed Function stops at the rank.""" + grid = Grid(shape=(21, 21)) + x, y = grid.dimensions + factor = 2 + + xd = ConditionalDimension('xd', parent=x, factor=factor) + yd = ConditionalDimension('yd', parent=y, factor=factor) + + u = Function(name='u', grid=grid) + g = Function(name='g', dimensions=(xd, yd), shape=(11, 11), grid=grid, + space_order=0) + + args = Operator(Eq(g, u)).arguments() + + assert args['xd_M'] == args['x_M'] // factor + + @pytest.mark.parallel(mode=[2]) + def test_offset_is_rank_data(self, mode): + """Each rank's offset lands its loop on the global subgrid.""" + grid = Grid(shape=(21,)) + sub = SubGrid(grid, factor=2) + offset = sub.dimensions[0].offset + + expected = {0: 0, 1: 1}[grid.distributor.myrank] + assert offset._arg_values(grid=grid) == {offset.name: expected} + + def test_offset_pickles(self): + """An offset is rebuilt from its factor, so it holds no rank data.""" + grid = Grid(shape=(21,)) + x = grid.dimensions[0] + offset = SubsamplingOffset(name='xso', root=x, factor=3) + + new = pickle.loads(pickle.dumps(offset)) + + assert new.name == offset.name + assert new.factor == 3 + assert new.root.name == x.name + + def test_serial_has_no_offset(self): + """Without a decomposition there's nothing to align.""" + sub = SubGrid(Grid(shape=(21,)), factor=2) + + assert sub.dimensions[0].offset == 0 + @pytest.mark.parallel(mode=[2, 4]) def test_strided_codegen(self, mode): """The guard still becomes a unit-stride subgrid loop under MPI."""