diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index 0bce1b9387..326562b5c9 100644 --- a/devito/ir/support/basic.py +++ b/devito/ir/support/basic.py @@ -869,8 +869,15 @@ def is_storage_related(self, dims=None): """ for d in self.findices: if d._defines & set(as_tuple(dims)): # noqa: SIM102 - if any(i.is_NonlinearDerived for i in d._defines) or \ - self.is_const(d): + # An injective Dimension maps distinct iterations to distinct + # locations by construction -- e.g. a factor-based + # ConditionalDimension, whose guard admits only + # `parent % factor == 0` and whose `i -> i//factor` is + # one-to-one on those. Unlike, say, a SteppingDimension, where + # `time % 2` deliberately revisits the same slot + nonlinear = [i for i in d._defines + if i.is_NonlinearDerived and not i.is_Injective] + if nonlinear or self.is_const(d): return True return False diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 8f3c02f983..13927d5761 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -34,7 +34,7 @@ from devito.passes import ( Graph, error_mapper, finalize_args, generate_implicit, generate_macros, is_on_device, lower_dtypes, lower_index_derivatives, lower_local_sums, minimize_symbols, - optimize_pows, unevaluate + optimize_pows, subsampling, unevaluate ) from devito.symbolics import estimate_cost, subs_op_args from devito.tools import ( @@ -423,6 +423,11 @@ 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) diff --git a/devito/passes/clusters/__init__.py b/devito/passes/clusters/__init__.py index bcdb8a4bd6..3a6cb0d060 100644 --- a/devito/passes/clusters/__init__.py +++ b/devito/passes/clusters/__init__.py @@ -5,6 +5,7 @@ from .factorization import * # noqa from .fusion import * # noqa from .blocking import * # noqa +from .subsampling import * # noqa from .asynchrony import * # noqa from .implicit import * # noqa from .misc import * # noqa diff --git a/devito/passes/clusters/subsampling.py b/devito/passes/clusters/subsampling.py new file mode 100644 index 0000000000..7ef3e7f35e --- /dev/null +++ b/devito/passes/clusters/subsampling.py @@ -0,0 +1,122 @@ +from devito.ir.clusters import Queue +from devito.ir.support import GuardFactor +from devito.symbolics import IntDiv, uxreplace +from devito.tools import timed_pass + +__all__ = ['subsampling'] + + +@timed_pass() +def subsampling(clusters): + """ + Strengthen factor-based ConditionalDimension guards into strided iteration. + + A Function subsampled by `factor` along `d` is written under a + `d % factor == 0` guard, so the enclosing loop visits every point of `d` + and discards all but one in `factor` of them. Where legal, drop the guard + and step the loop `factor` points at a time instead + + for (x = x_m; x <= x_M; x += 1) for (x = ...; x <= x_M; x += xf) + 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`: + + 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] + g[x/xf] = u[x] + + 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. + + Parameters + ---------- + clusters : tuple of Clusters + Input Clusters, subject of the optimization pass. + """ + return Subsampling().process(clusters) + + +class Subsampling(Queue): + + """ + Replace uniform factor-based guards with strided iteration. + + Clusters are grouped by iteration space prefix alone, deliberately *not* + by guard: a guard may only become a stride if every Cluster sharing the + prefix carries the very same one. Striding a loop nest that also hosts an + unguarded expression would run that expression on a subset of its + iterations -- and dropping the guard is what makes the nest shared in the + first place, since `stree_build` reuses a subtree across Clusters with + matching guards (see `reuse_whole_subtree`). + """ + + def callback(self, clusters, prefix): + if not prefix: + return clusters + + d = prefix[-1].dim + + guards = {c.guards.get(d) for c in clusters} + if len(guards) > 1: + # Non-uniformly guarded, hence `d`'s loop nest is shared with + # expressions that must keep running at every point + return clusters + + guard = guards.pop() + if not all(_is_strideable(c, d, guard) for c in clusters): + return clusters + + cd = guard.d + factor = cd.symbolic_factor + + processed = [] + for c in clusters: + # Iterate the subsampled Dimension itself, with unit stride, and + # rescale every parent access `d` to `cd*factor` + 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} + exprs = [uxreplace(e, subs) for e in c.exprs] + + # `cd` inherits `d`'s computational properties: it covers the very + # same points, just addressed in subsampled coordinates. Without + # this the loop would lose PARALLEL and so all of its OpenMP + properties = c.properties.drop(d) + properties = properties.add(cd, c.properties[d]) + + processed.append(c.rebuild(exprs=exprs, ispace=ispace, + guards=c.guards.popany(d), + properties=properties)) + + return processed + + +def _is_strideable(cluster, d, guard): + """ + True if iteration over `d` may step `guard`'s subsampling factor points at + a time in place of evaluating `guard` at every point. + """ + # The guard must be exactly a subsampling guard. Anything `And`/`Or`-ed in + # (e.g. a user-supplied `condition`) must still be evaluated pointwise + if not isinstance(guard, GuardFactor): + return False + + if d not in cluster.ispace.itdims: + 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 + if not cluster.properties.is_parallel_relaxed(d): + return False + + # Non-affine accesses may read the points that would be skipped + return cluster.properties.is_affine(d) diff --git a/devito/types/dimension.py b/devito/types/dimension.py index c81eb90d37..6b21493c14 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -116,6 +116,7 @@ class Dimension(ArgProvider): is_Custom = False is_Derived = False is_NonlinearDerived = False + is_Injective = False is_AbstractSub = False is_Sub = False is_MultiSub = False @@ -1090,6 +1091,19 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None, self._indirect = indirect self._relation = relation + @property + def is_Injective(self): + """ + True if distinct iterations of the parent that pass the guard map to + distinct indices of `self`, so that no two iterations can access the + same memory location through it. + + A pure factor subsampling is injective: the guard admits only + `parent % factor == 0`, and `i -> i//factor` is one-to-one on those. + An arbitrary `condition` carries no such promise. + """ + return self._factor is not None and self._condition is None + @property def uses_symbolic_factor(self): return self._factor is not None @@ -1153,13 +1167,18 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs): if isinstance(fact, Constant): fact = fact.data - toint = lambda x: math.ceil(x / fact) vals = {} + + # The first subsampled point at or after the parent's min... with suppress(KeyError, TypeError): - vals[self.min_name] = toint(kwargs.get(self.parent.min_name)) + vals[self.min_name] = math.ceil(kwargs.get(self.parent.min_name) / 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] = toint(kwargs.get(self.parent.max_name)) + vals[self.max_name] = kwargs.get(self.parent.max_name) // fact vals[self.symbolic_factor.name] = fact diff --git a/devito/types/grid.py b/devito/types/grid.py index 378d2543b4..4568b852bf 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -1,3 +1,4 @@ +import math from collections import namedtuple from functools import cached_property from itertools import product @@ -11,17 +12,17 @@ from devito.finite_differences.elementary import Max, Min from devito.logger import warning from devito.mpi import MPI, Distributor, SubDistributor -from devito.tools import ReducerMap, as_tuple, frozendict +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 ( - DefaultDimension, Dimension, MultiSubDimension, SpaceDimension, Spacing, - SteppingDimension, SubDimension, TimeDimension + ConditionalDimension, DefaultDimension, Dimension, MultiSubDimension, SpaceDimension, + Spacing, SteppingDimension, SubDimension, TimeDimension ) from devito.types.utils import DimensionTuple -__all__ = ['Border', 'Grid', 'SubDomain', 'SubDomainSet'] +__all__ = ['Border', 'Grid', 'SubDomain', 'SubDomainSet', 'SubGrid'] GlobalLocal = namedtuple('GlobalLocal', 'glb loc') @@ -428,6 +429,169 @@ def __setstate__(self, state): self._distributor = Distributor(self.shape, self.dimensions, MPI.COMM_SELF) +class SubGrid(Grid): + + """ + A coarser Grid whose points are a subset of a parent Grid's points. + + Along each Dimension subsampled by `factor`, the SubGrid retains every + `factor`-th point of the parent, so point `i` of the SubGrid sits at parent + index `i*factor`. The physical spacing is `factor` times the parent's, and + the origin is shared. + + Parameters + ---------- + parent : Grid + The Grid being subsampled. + factor : int or tuple of int + Subsampling factor, either one value per space Dimension of `parent`, + or a single value applied to all of them. A factor of 1 leaves the + corresponding Dimension untouched. + name : str, optional + Prefix for the names of the subsampled Dimensions; each is named + `f'{parent_dimension}{name}'`. Defaults to `'s'`, giving `xs`, `ys`. + + Examples + -------- + >>> from devito import Grid, SubGrid, Function + >>> grid = Grid(shape=(9, 9), extent=(8., 8.)) + >>> sub = SubGrid(grid, factor=2) + >>> sub.shape + (5, 5) + >>> sub.dimensions + (xs, ys) + + A Function defined on a SubGrid needs neither `shape` nor `dimensions`; + both are derived. + + >>> f = Function(name='f', grid=sub, space_order=0) + >>> f.shape + (5, 5) + + The spacing is scaled by the factor, so derivatives are taken with respect + to the coarser grid. + + >>> sub.spacing + (2.0, 2.0) + + A per-Dimension factor subsamples each axis independently; `1` leaves an + axis alone. + + >>> sub = SubGrid(grid, factor=(2, 1)) + >>> sub.shape + (5, 9) + + Notes + ----- + Because the subsampled points coincide with parent points, no interpolation + is involved: an `Eq` writing a SubGrid Function from a parent Function is a + strided copy, and the compiler emits a strided loop for it. + """ + + def __init__(self, parent, factor, name='s', **kwargs): + if not parent.is_Grid: + raise ValueError("`parent` must be a Grid, got " + f"`{type(parent).__name__}`") + + factors = self.__factors_setup__(parent, factor) + + if isinstance(parent, SubGrid): + # Nesting would need a ConditionalDimension to parent another one. + # Spell out the equivalent single-level call, since the composed + # factor is just the product of the two + combined = tuple(f0*f1 for f0, f1 in + zip(parent._factors, factors, strict=True)) + raise ValueError( + f"Cannot subsample a SubGrid. `{factor}` applied to a SubGrid " + f"of factor `{tuple(parent._factors)}` is the same as factor " + f"`{combined}` applied to its parent Grid directly, so use " + f"`SubGrid(parent.parent, factor={combined})`" + ) + + dimensions = tuple( + d if f == 1 else ConditionalDimension(f'{d.name}{name}', parent=d, + factor=f) + for d, f in zip(parent.dimensions, factors, strict=True) + ) + + # A subsampled axis keeps the parent's endpoints, so the number of + # retained points is `ceil(n/factor)` and the extent is unchanged up to + # the last retained point + shape = tuple((n + f - 1)//f + for n, f in zip(parent.shape, factors, strict=True)) + extent = tuple(s*(n - 1) for s, n in + zip(self.__spacing_setup__(parent, factors), shape, + strict=True)) + + self._parent = parent + self._factors = factors + + super().__init__(shape=shape, extent=extent, origin=parent.origin, + dimensions=dimensions, dtype=parent.dtype, + time_dimension=parent.time_dim, + comm=parent.distributor.comm, **kwargs) + + # Share the parent's stepping Dimension rather than the fresh + # `{time}_s` that Grid derives for a supplied `time_dimension`: a + # SubGrid coarsens space only, so a TimeFunction on it must step over + # the very same buffer indices as one on the parent + self._stepping_dim = parent.stepping_dim + + @classmethod + def __factors_setup__(cls, parent, factor): + factors = as_tuple(factor) + if len(factors) == 1: + factors = factors*len(parent.dimensions) + elif len(factors) != len(parent.dimensions): + raise ValueError(f"Expected 1 or {len(parent.dimensions)} factors " + f"for a {len(parent.dimensions)}D parent Grid, " + f"got {len(factors)}") + + for f in factors: + if not is_integer(f): + # Non-integer factors place SubGrid points between parent + # points, which requires interpolation rather than striding + raise ValueError( + f"Non-integer subsampling factor `{f}`. A SubGrid's points " + "must coincide with its parent's, so the factor must be a " + f"whole number of parent cells -- use " + f"`{max(1, math.floor(f))}` or `{max(1, math.ceil(f))}`. " + "Subsampling onto points that fall " + "between parent points requires interpolation, which " + "SubGrid does not perform" + ) + if f < 1: + raise ValueError(f"Subsampling factor must be >= 1, got `{f}`") + + return tuple(int(f) for f in factors) + + @classmethod + def __spacing_setup__(cls, parent, factors): + return tuple(s*f for s, f in zip(parent.spacing, factors, strict=True)) + + @property + def parent(self): + """The Grid being subsampled.""" + return self._parent + + @property + def factor(self): + """The subsampling factor along each Dimension.""" + return DimensionTuple(*self._factors, getters=self.dimensions) + + @property + def root(self): + # A SubGrid is a view of its parent, not a Grid in its own right, just + # as a SubDomain is. `Operator._prepare_arguments` collapses the + # discretizations it finds by `root` to tell "several views of one + # Grid" apart from "several unrelated Grids", and rejects the latter + # under MPI + return self._parent.root + + def __repr__(self): + return f"SubGrid[shape={self.shape}, factor={tuple(self._factors)}]" + + class AbstractSubDomain(CartesianDiscretization): """ diff --git a/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb new file mode 100644 index 0000000000..5a2b9b92df --- /dev/null +++ b/examples/userapi/10_subgrids.ipynb @@ -0,0 +1,1817 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Subgrids\n", + "\n", + "A *subgrid* is a coarser grid whose points are a subset of a parent `Grid`'s points.\n", + "`SubGrid` is the API for building one.\n", + "\n", + "This tutorial covers\n", + "\n", + "1. [Creating a SubGrid](#Creating-a-SubGrid) -- the constructor and what it derives\n", + "2. [Defining Functions on it](#Defining-Functions-on-a-SubGrid)\n", + "3. [Writing Operators](#Writing-Operators) -- moving data between the two grids\n", + "4. [Finite differences](#Finite-differences-on-a-SubGrid)\n", + "5. [What SubGrid does not cover](#What-SubGrid-does-not-cover)\n", + "6. [The generated code](#The-generated-code) -- how the compiler lowers it\n", + "\n", + "It covers **integer subsampling**, where subgrid points coincide exactly with parent\n", + "points, so no interpolation is involved." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Creating a SubGrid\n", + "\n", + "`SubGrid(parent, factor)` keeps every `factor`-th point of `parent` along each\n", + "`Dimension`, so subgrid point `i` sits at parent index `i*factor`." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:09.864846Z", + "iopub.status.busy": "2026-10-21T00:00:09.864277Z", + "iopub.status.idle": "2026-10-21T00:00:10.707068Z", + "shell.execute_reply": "2026-10-21T00:00:10.705504Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "SubGrid[shape=(4, 4), factor=(3, 3)]" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import numpy as np\n", + "from devito import (CondEq, ConditionalDimension, Eq, Function, Grid, Operator,\n", + " SubGrid, TimeFunction)\n", + "from devito.ir.iet import FindNodes, Iteration\n", + "\n", + "grid = Grid(shape=(12, 12), extent=(11., 11.))\n", + "x, y = grid.dimensions\n", + "\n", + "sub = SubGrid(grid, factor=3)\n", + "sub" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Everything else is derived from the parent. `shape` is `ceil(n/factor)` per axis:" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.761033Z", + "iopub.status.busy": "2026-10-21T00:00:10.760169Z", + "iopub.status.idle": "2026-10-21T00:00:10.769237Z", + "shell.execute_reply": "2026-10-21T00:00:10.767660Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "parent shape (12, 12) subgrid shape (4, 4)\n" + ] + } + ], + "source": [ + "print('parent shape', grid.shape, ' subgrid shape', sub.shape)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`spacing` is scaled by the factor, while `extent` and `origin` are inherited --\n", + "the subgrid spans the same physical domain with fewer points." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.775239Z", + "iopub.status.busy": "2026-10-21T00:00:10.774599Z", + "iopub.status.idle": "2026-10-21T00:00:10.787096Z", + "shell.execute_reply": "2026-10-21T00:00:10.785379Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "spacing parent (1.0, 1.0) subgrid (3.0, 3.0)\n", + "extent parent (11.0, 11.0) subgrid (9.0, 9.0)\n", + "origin parent (0.0, 0.0) subgrid (0.0, 0.0)\n" + ] + } + ], + "source": [ + "print('spacing parent', tuple(float(v) for v in grid.spacing),\n", + " ' subgrid', tuple(float(v) for v in sub.spacing))\n", + "print('extent parent', tuple(float(v) for v in grid.extent),\n", + " ' subgrid', tuple(float(v) for v in sub.extent))\n", + "print('origin parent', tuple(float(v) for v in grid.origin),\n", + " ' subgrid', tuple(float(v) for v in sub.origin))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`parent` and `factor` are the two things `SubGrid` adds over `Grid`:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.792743Z", + "iopub.status.busy": "2026-10-21T00:00:10.792135Z", + "iopub.status.idle": "2026-10-21T00:00:10.802173Z", + "shell.execute_reply": "2026-10-21T00:00:10.800578Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(True, (3, 3))" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sub.parent is grid, sub.factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Its `dimensions` are `ConditionalDimension`s derived from the parent's. These are\n", + "what carry the subsampling factor into the generated code." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.807643Z", + "iopub.status.busy": "2026-10-21T00:00:10.807100Z", + "iopub.status.idle": "2026-10-21T00:00:10.817272Z", + "shell.execute_reply": "2026-10-21T00:00:10.815647Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((xs, ys), [x, y])" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sub.dimensions, [d.parent for d in sub.dimensions]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Pass `name` to control how they are named -- useful when building several\n", + "SubGrids off one parent." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.822983Z", + "iopub.status.busy": "2026-10-21T00:00:10.822389Z", + "iopub.status.idle": "2026-10-21T00:00:10.836328Z", + "shell.execute_reply": "2026-10-21T00:00:10.834740Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(xc, yc)" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "SubGrid(grid, factor=2, name='c').dimensions" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A factor per `Dimension` subsamples each axis independently, and `1` leaves an\n", + "axis untouched:" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.841993Z", + "iopub.status.busy": "2026-10-21T00:00:10.841397Z", + "iopub.status.idle": "2026-10-21T00:00:10.852733Z", + "shell.execute_reply": "2026-10-21T00:00:10.851021Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "shape (4, 12)\n", + "dimensions (xs, y)\n" + ] + } + ], + "source": [ + "mixed = SubGrid(grid, factor=(3, 1))\n", + "print('shape ', mixed.shape)\n", + "print('dimensions', mixed.dimensions) # `y` passed through unchanged" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Being a `Grid` subclass, a SubGrid also carries the usual `Grid` attributes:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.858751Z", + "iopub.status.busy": "2026-10-21T00:00:10.858250Z", + "iopub.status.idle": "2026-10-21T00:00:10.867649Z", + "shell.execute_reply": "2026-10-21T00:00:10.866250Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(2, numpy.float32, time, t)" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sub.dim, sub.dtype, sub.time_dim, sub.stepping_dim" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### What is rejected\n", + "\n", + "Only factors that keep subgrid points on parent points are allowed. A non-integer\n", + "factor would place them in between, which needs interpolation rather than\n", + "subsampling." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.872862Z", + "iopub.status.busy": "2026-10-21T00:00:10.872263Z", + "iopub.status.idle": "2026-10-21T00:00:10.882012Z", + "shell.execute_reply": "2026-10-21T00:00:10.880300Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "factor=2.5 -> Non-integer subsampling factor `2.5`. A SubGrid's points must coincide with its parent's, so the factor must be a whole number of parent cells -- use `2` or `3`. Subsampling onto points that fall between parent points requires interpolation, which SubGrid does not perform\n", + "factor=0 -> Subsampling factor must be >= 1, got `0`\n", + "factor=(2, 2, 2) -> Expected 1 or 2 factors for a 2D parent Grid, got 3\n" + ] + } + ], + "source": [ + "for bad in (2.5, 0, (2, 2, 2)):\n", + " try:\n", + " SubGrid(grid, factor=bad)\n", + " except ValueError as e:\n", + " print(f'factor={bad!r:10} -> {e}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Nor can a SubGrid be subsampled again; combine the factors off the original\n", + "parent instead." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.887762Z", + "iopub.status.busy": "2026-10-21T00:00:10.887135Z", + "iopub.status.idle": "2026-10-21T00:00:10.895555Z", + "shell.execute_reply": "2026-10-21T00:00:10.893836Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Cannot subsample a SubGrid. `2` applied to a SubGrid of factor `(3, 3)` is the same as factor `(6, 6)` applied to its parent Grid directly, so use `SubGrid(parent.parent, factor=(6, 6))`\n" + ] + } + ], + "source": [ + "try:\n", + " SubGrid(sub, factor=2)\n", + "except ValueError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Defining Functions on a SubGrid\n", + "\n", + "This is the main convenience: a `Function` on a SubGrid needs neither `shape` nor\n", + "`dimensions`, both are taken from the grid -- exactly as for a normal `Grid`." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.901337Z", + "iopub.status.busy": "2026-10-21T00:00:10.900718Z", + "iopub.status.idle": "2026-10-21T00:00:10.919067Z", + "shell.execute_reply": "2026-10-21T00:00:10.917441Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "u (12, 12) (x, y)\n", + "usub (4, 4) (xs, ys)\n" + ] + } + ], + "source": [ + "u = Function(name='u', grid=grid) # on the parent\n", + "usub = Function(name='usub', grid=sub, space_order=0) # on the subgrid\n", + "\n", + "print('u ', u.shape, u.dimensions)\n", + "print('usub ', usub.shape, usub.dimensions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`TimeFunction` works too, and steps over the *parent's* time Dimension -- a\n", + "SubGrid coarsens space only." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.924379Z", + "iopub.status.busy": "2026-10-21T00:00:10.923847Z", + "iopub.status.idle": "2026-10-21T00:00:10.935092Z", + "shell.execute_reply": "2026-10-21T00:00:10.933488Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((t, xs, ys), True)" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "usave = TimeFunction(name='usave', grid=sub, space_order=0)\n", + "usave.dimensions, usave.dimensions[0] is grid.stepping_dim" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Equivalent, by hand\n", + "\n", + "`SubGrid` is a convenience over `ConditionalDimension`, which can still be used\n", + "directly. The manual form is equivalent but requires naming each Dimension and\n", + "computing the shape yourself." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.940307Z", + "iopub.status.busy": "2026-10-21T00:00:10.939772Z", + "iopub.status.idle": "2026-10-21T00:00:10.954034Z", + "shell.execute_reply": "2026-10-21T00:00:10.952426Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "True" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "xs_manual = ConditionalDimension(name='xsm', parent=x, factor=3)\n", + "ys_manual = ConditionalDimension(name='ysm', parent=y, factor=3)\n", + "f_manual = Function(name='f_manual', grid=grid,\n", + " dimensions=(xs_manual, ys_manual),\n", + " shape=(4, 4), space_order=0) # shape by hand\n", + "\n", + "f_manual.shape == usub.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Writing Operators\n", + "\n", + "An `Eq` between a parent Function and a SubGrid Function is a strided copy: Devito\n", + "matches the subgrid points to their parent counterparts." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:10.959450Z", + "iopub.status.busy": "2026-10-21T00:00:10.958997Z", + "iopub.status.idle": "2026-10-21T00:00:11.166399Z", + "shell.execute_reply": "2026-10-21T00:00:11.165173Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([[ 0., 3., 6., 9.],\n", + " [ 36., 39., 42., 45.],\n", + " [ 72., 75., 78., 81.],\n", + " [108., 111., 114., 117.]], dtype=float32)" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "usub.data[:] = 0.\n", + "\n", + "Operator(Eq(usub, u)).apply()\n", + "\n", + "assert np.array_equal(usub.data, u.data[::3, ::3])\n", + "usub.data" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Each row and column of `usub` comes from every third row and column of `u`.\n", + "\n", + "The reverse direction works the same way -- writing the parent from the subgrid\n", + "scatters onto the subsampled points, leaving the others untouched." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.171003Z", + "iopub.status.busy": "2026-10-21T00:00:11.170781Z", + "iopub.status.idle": "2026-10-21T00:00:11.301506Z", + "shell.execute_reply": "2026-10-21T00:00:11.300355Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([[ 7., -1., -1., 7.],\n", + " [-1., -1., -1., -1.],\n", + " [-1., -1., -1., -1.],\n", + " [ 7., -1., -1., 7.]], dtype=float32)" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = -1.\n", + "usub.data[:] = 7.\n", + "\n", + "Operator(Eq(u, usub)).apply()\n", + "\n", + "assert np.all(u.data[::3, ::3] == 7.)\n", + "assert np.all(u.data[1::3, :] == -1.) # untouched rows\n", + "u.data[:4, :4]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Restricting the iteration range maps onto the subgrid. With `x_m=2` and\n", + "`factor=3`, the first subgrid point at or after `x=2` is `x=3`, i.e. subgrid index 1,\n", + "so index 0 is never written." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.305798Z", + "iopub.status.busy": "2026-10-21T00:00:11.305576Z", + "iopub.status.idle": "2026-10-21T00:00:11.420274Z", + "shell.execute_reply": "2026-10-21T00:00:11.419116Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([-1., 1., 1., 1.], dtype=float32)" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = 1.\n", + "usub.data[:] = -1. # sentinel: never written\n", + "\n", + "Operator(Eq(usub, u)).apply(x_m=2)\n", + "\n", + "usub.data[:, 0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Several SubGrids can appear in one Operator, each with its own factor." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.424454Z", + "iopub.status.busy": "2026-10-21T00:00:11.424226Z", + "iopub.status.idle": "2026-10-21T00:00:11.614968Z", + "shell.execute_reply": "2026-10-21T00:00:11.606198Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "((6, 6), (3, 3))" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "sub_a = SubGrid(grid, factor=2, name='a')\n", + "sub_b = SubGrid(grid, factor=4, name='b')\n", + "fa = Function(name='fa', grid=sub_a, space_order=0)\n", + "fb = Function(name='fb', grid=sub_b, space_order=0)\n", + "\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "Operator([Eq(fa, u), Eq(fb, u)]).apply()\n", + "\n", + "assert np.array_equal(fa.data, u.data[::2, ::2])\n", + "assert np.array_equal(fb.data, u.data[::4, ::4])\n", + "fa.shape, fb.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Finite differences on a SubGrid\n", + "\n", + "Derivatives of a SubGrid Function are taken with respect to the *coarse* spacing,\n", + "since the Dimension's `spacing` is `factor*h`." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.621540Z", + "iopub.status.busy": "2026-10-21T00:00:11.620971Z", + "iopub.status.idle": "2026-10-21T00:00:11.634243Z", + "shell.execute_reply": "2026-10-21T00:00:11.632756Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(3*h_x, h_x)" + ] + }, + "execution_count": 18, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sub.dimensions[0].spacing, x.spacing" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "So the stencil steps `factor*h_x` and the weights scale accordingly:" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.639205Z", + "iopub.status.busy": "2026-10-21T00:00:11.638777Z", + "iopub.status.idle": "2026-10-21T00:00:11.669891Z", + "shell.execute_reply": "2026-10-21T00:00:11.668737Z" + } + }, + "outputs": [ + { + "data": { + "text/latex": [ + "$\\displaystyle - \\frac{0.333333333 fsub(xs, ys)}{h_{x}} + \\frac{0.333333333 fsub(xs + 3*h_x, ys)}{h_{x}}$" + ], + "text/plain": [ + "-0.333333333*fsub(xs, ys)/h_x + 0.333333333*fsub(xs + 3*h_x, ys)/h_x" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "fsub = Function(name='fsub', grid=sub, space_order=2)\n", + "fsub.dx.evaluate" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Checking it numerically on a linear field, whose derivative is exact:" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.674050Z", + "iopub.status.busy": "2026-10-21T00:00:11.673852Z", + "iopub.status.idle": "2026-10-21T00:00:11.853934Z", + "shell.execute_reply": "2026-10-21T00:00:11.852728Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "3.0" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "grid_l = Grid(shape=(21, 21), extent=(20., 20.))\n", + "sub_l = SubGrid(grid_l, factor=2)\n", + "\n", + "fl = Function(name='fl', grid=sub_l, space_order=2)\n", + "dfl = Function(name='dfl', grid=sub_l, space_order=2)\n", + "\n", + "# fl = 3*x, sampled on the coarse grid\n", + "for i in range(fl.shape[0]):\n", + " fl.data[i, :] = 3.*i*float(sub_l.spacing[0])\n", + "\n", + "Operator(Eq(dfl, fl.dx)).apply()\n", + "\n", + "assert np.allclose(dfl.data[1:-1, 1:-1], 3., atol=1e-4)\n", + "float(dfl.data[5, 5])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## What SubGrid does not cover\n", + "\n", + "Two related things are deliberately outside `SubGrid` and still want a\n", + "`ConditionalDimension` built by hand.\n", + "\n", + "**Subsampling in time.** A SubGrid coarsens space, so saving a wavefield every\n", + "`factor` timesteps uses a `ConditionalDimension` on the time Dimension." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:11.858435Z", + "iopub.status.busy": "2026-10-21T00:00:11.858182Z", + "iopub.status.idle": "2026-10-21T00:00:12.058665Z", + "shell.execute_reply": "2026-10-21T00:00:12.057133Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "[0.0, 4.0, 0.0]" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "nt, factor = 9, 4\n", + "grid_t = Grid(shape=(4, 4))\n", + "ts = ConditionalDimension(name='ts', parent=grid_t.time_dim, factor=factor)\n", + "\n", + "v = TimeFunction(name='v', grid=grid_t)\n", + "vsave = TimeFunction(name='vsave', grid=grid_t, time_dim=ts,\n", + " save=(nt + factor - 1)//factor)\n", + "\n", + "Operator([Eq(v.forward, v + 1.), Eq(vsave, v)]).apply(time_M=nt-2)\n", + "\n", + "[float(vsave.data[i, 0, 0]) for i in range(vsave.shape[0])]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Extra conditions.** `SubGrid` takes a factor, not an arbitrary predicate; for\n", + "`factor` *and* a `condition`, construct the `ConditionalDimension` directly." + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.064134Z", + "iopub.status.busy": "2026-10-21T00:00:12.063653Z", + "iopub.status.idle": "2026-10-21T00:00:12.082047Z", + "shell.execute_reply": "2026-10-21T00:00:12.080721Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(3, Eq(Mod(x, 2), 0))" + ] + }, + "execution_count": 22, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "xs_cond = ConditionalDimension(name='xsc', parent=x, factor=3,\n", + " condition=CondEq(x % 2, 0))\n", + "xs_cond.factor, xs_cond.condition" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Generic subsampling -- subgrid points that fall *between* parent points, requiring\n", + "interpolation -- is a separate topic and is not covered here." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## The generated code\n", + "\n", + "The rest of this tutorial shows how Devito lowers the above, which is useful when\n", + "reading or profiling the generated C.\n", + "\n", + "Devito iterates the **subgrid** directly: one unit-stride loop per subsampled\n", + "`Dimension`, with parent accesses rescaled -- rather than walking every parent point\n", + "under an `if (x%factor == 0)` guard." + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.087228Z", + "iopub.status.busy": "2026-10-21T00:00:12.086738Z", + "iopub.status.idle": "2026-10-21T00:00:12.095143Z", + "shell.execute_reply": "2026-10-21T00:00:12.093635Z" + } + }, + "outputs": [], + "source": [ + "def show(op):\n", + " \"\"\"Print the Operator's loop nest.\"\"\"\n", + " print(op.body.body[-1])\n", + "\n", + "\n", + "def steps(op):\n", + " \"\"\"The loop increment of each Iteration, by Dimension name.\"\"\"\n", + " return {i.dim.name: i.limits[2] for i in FindNodes(Iteration).visit(op)}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### One loop per subsampled Dimension" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.100120Z", + "iopub.status.busy": "2026-10-21T00:00:12.099707Z", + "iopub.status.idle": "2026-10-21T00:00:12.239134Z", + "shell.execute_reply": "2026-10-21T00:00:12.237862Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int xs = xs_m; xs <= xs_M; xs += 1)\n", + "{\n", + " #pragma omp simd aligned(u,usub:64)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", + " {\n", + " usub[xs][ys] = u[xs*xsf + 1][ys*ysf + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "op = Operator(Eq(usub, u))\n", + "show(op)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note what the loop iterates: `xs`, the subsampled `Dimension` itself, with a\n", + "**unit** increment, over its own bounds `xs_m .. xs_M`. There is no `if` in the nest.\n", + "\n", + "- the subgrid access is simply `usub[xs]`;\n", + "- the parent access is rescaled to `u[xs*xsf + 1]`.\n", + "\n", + "The unit stride is deliberate. A `x += factor` loop would also remove the guard, but a\n", + "non-unitary increment escapes the Lamport-style dependence analysis the compiler\n", + "relies on, so it is avoided.\n", + "\n", + "The bounds need no adjusting either: `xs_m` and `xs_M` are `ConditionalDimension`'s own\n", + "runtime arguments, derived as `ceil(x_m/factor)` and `floor(x_M/factor)`." + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.244318Z", + "iopub.status.busy": "2026-10-21T00:00:12.243861Z", + "iopub.status.idle": "2026-10-21T00:00:12.252593Z", + "shell.execute_reply": "2026-10-21T00:00:12.251182Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'xs': 1, 'ys': 1}" + ] + }, + "execution_count": 25, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "steps(op)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### The innermost Dimension\n", + "\n", + "Subsampling the innermost (vectorized) `Dimension` is handled the same way. The\n", + "alternative -- a guard inside the SIMD loop -- would mean a predicated store on every\n", + "iteration." + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.258305Z", + "iopub.status.busy": "2026-10-21T00:00:12.257700Z", + "iopub.status.idle": "2026-10-21T00:00:12.384155Z", + "shell.execute_reply": "2026-10-21T00:00:12.382996Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " #pragma omp simd aligned(u,usub_y:64)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", + " {\n", + " usub_y[x][ys] = u[x + 1][ys*ysf + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "sub_y = SubGrid(grid, factor=(1, 2)) # coarsen y only\n", + "usub_y = Function(name='usub_y', grid=sub_y, space_order=0)\n", + "\n", + "op_inner = Operator(Eq(usub_y, u))\n", + "show(op_inner)" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.388152Z", + "iopub.status.busy": "2026-10-21T00:00:12.387972Z", + "iopub.status.idle": "2026-10-21T00:00:12.393955Z", + "shell.execute_reply": "2026-10-21T00:00:12.392797Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'x': 1, 'ys': 1}" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "steps(op_inner)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The `y` loop is strided and the `#pragma omp simd` is retained." + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.398049Z", + "iopub.status.busy": "2026-10-21T00:00:12.397832Z", + "iopub.status.idle": "2026-10-21T00:00:12.520522Z", + "shell.execute_reply": "2026-10-21T00:00:12.518923Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([[ 0., 2., 4., 6., 8., 10.],\n", + " [12., 14., 16., 18., 20., 22.],\n", + " [24., 26., 28., 30., 32., 34.]], dtype=float32)" + ] + }, + "execution_count": 28, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "Operator(Eq(usub_y, u)).apply()\n", + "\n", + "assert np.array_equal(usub_y.data, u.data[:, ::2])\n", + "usub_y.data[:3]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Several Dimensions at once\n", + "\n", + "Each subsampled `Dimension` gets its own loop, independently." + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.525875Z", + "iopub.status.busy": "2026-10-21T00:00:12.525357Z", + "iopub.status.idle": "2026-10-21T00:00:12.755021Z", + "shell.execute_reply": "2026-10-21T00:00:12.753816Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int xs = xs_m; xs <= xs_M; xs += 1)\n", + "{\n", + " #pragma omp simd aligned(u,usub_xy:64)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", + " {\n", + " usub_xy[xs][ys] = u[xs*xsf + 1][ys*ysf + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "sub_xy = SubGrid(grid, factor=(3, 2))\n", + "usub_xy = Function(name='usub_xy', grid=sub_xy, space_order=0)\n", + "\n", + "op_both = Operator(Eq(usub_xy, u))\n", + "show(op_both)" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.758836Z", + "iopub.status.busy": "2026-10-21T00:00:12.758655Z", + "iopub.status.idle": "2026-10-21T00:00:12.889574Z", + "shell.execute_reply": "2026-10-21T00:00:12.887973Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "{'xs': 1, 'ys': 1}" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "Operator(Eq(usub_xy, u)).apply()\n", + "\n", + "assert np.array_equal(usub_xy.data, u.data[::3, ::2])\n", + "steps(op_both)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## When a guard is used instead\n", + "\n", + "Striding is applied only where it provably preserves semantics. Where it does not,\n", + "Devito visits every point under an `if (x%factor == 0)` guard. Results are the same\n", + "either way; only the generated loop differs. Four cases, each with its code." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1. The loop nest is shared with an unguarded equation\n", + "\n", + "`Eq(f, 1)` must run at **every** `x`, while `Eq(g, f)` runs every fourth. Both land in\n", + "one loop nest, so striding it would also stride `Eq(f, 1)`. The guard is kept." + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:12.894890Z", + "iopub.status.busy": "2026-10-21T00:00:12.894468Z", + "iopub.status.idle": "2026-10-21T00:00:13.027969Z", + "shell.execute_reply": "2026-10-21T00:00:13.026750Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " f[x + 1] = 1;\n", + " if ((x)%(xsf) == 0)\n", + " {\n", + " g[x / xsf] = f[x + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "grid1 = Grid(shape=(12,))\n", + "sub1 = SubGrid(grid1, factor=4)\n", + "\n", + "f = Function(name='f', grid=grid1)\n", + "g = Function(name='g', grid=sub1, space_order=0)\n", + "\n", + "op_shared = Operator([Eq(f, 1), Eq(g, f)])\n", + "show(op_shared)" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.031759Z", + "iopub.status.busy": "2026-10-21T00:00:13.031587Z", + "iopub.status.idle": "2026-10-21T00:00:13.051619Z", + "shell.execute_reply": "2026-10-21T00:00:13.050490Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "(Data([1., 1., 1.], dtype=float32), Data([1., 1., 1.], dtype=float32))" + ] + }, + "execution_count": 32, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "f.data[:] = 0.\n", + "g.data[:] = -1.\n", + "op_shared.apply()\n", + "\n", + "# `f` written at every point, `g` only on the subgrid\n", + "assert np.all(f.data == 1.)\n", + "assert np.all(g.data == 1.)\n", + "f.data[:3], g.data[:]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`f[x + 1] = 1` sits outside the `if`, `g[x / xs1f]` inside it -- which is exactly\n", + "what is required." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. The subsampled Dimension is sequential\n", + "\n", + "Striding skips points, which is only sound when no dependence is carried along the\n", + "`Dimension`. Here `g` reads its own previous subgrid value, so `x` is sequential." + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.055748Z", + "iopub.status.busy": "2026-10-21T00:00:13.055580Z", + "iopub.status.idle": "2026-10-21T00:00:13.189811Z", + "shell.execute_reply": "2026-10-21T00:00:13.188606Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " if ((x)%(xsf) == 0)\n", + " {\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " g2[x / xsf][y] = f2[x + 1][y + 1] + g2[(x / xsf) - 1][y];\n", + " }\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "grid2 = Grid(shape=(12, 12))\n", + "sub2 = SubGrid(grid2, factor=(3, 1))\n", + "xs3 = sub2.dimensions[0]\n", + "\n", + "f2 = Function(name='f2', grid=grid2)\n", + "g2 = Function(name='g2', grid=sub2, space_order=0)\n", + "\n", + "op_seq = Operator(Eq(g2, g2.subs(xs3, xs3-1) + f2))\n", + "show(op_seq)" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.193859Z", + "iopub.status.busy": "2026-10-21T00:00:13.193626Z", + "iopub.status.idle": "2026-10-21T00:00:13.218664Z", + "shell.execute_reply": "2026-10-21T00:00:13.217527Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Shrinking bounds (`x_m`, `x_M`); some `x` points will not be computed. Likely insufficient space_order for the derivatives.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/mloubout/DevitoCodes/devito/devito/operator/operator.py:718: DevitoWarning: Shrinking bounds (`x_m`, `x_M`); some `x` points will not be computed. Likely insufficient space_order for the derivatives.\n", + " warn(f\"Shrinking bounds (`{d.min_name}`, `{d.max_name}`); \"\n", + "Shrinking bounds (`xs_m`, `xs_M`); some `xs` points will not be computed. Likely insufficient space_order for the derivatives.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/mloubout/DevitoCodes/devito/devito/operator/operator.py:718: DevitoWarning: Shrinking bounds (`xs_m`, `xs_M`); some `xs` points will not be computed. Likely insufficient space_order for the derivatives.\n", + " warn(f\"Shrinking bounds (`{d.min_name}`, `{d.max_name}`); \"\n", + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([0., 1., 2., 3.], dtype=float32)" + ] + }, + "execution_count": 34, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "f2.data[:] = 1.\n", + "g2.data[:] = 0.\n", + "op_seq.apply()\n", + "\n", + "# Each subgrid row accumulates the one before it; row 0 reads the halo (zero)\n", + "assert np.allclose(g2.data[:, 0], [0., 1., 2., 3.])\n", + "g2.data[:, 0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Time subsampling is the everyday instance of this: the time loop is sequential, so a\n", + "saved-wavefield `Eq` keeps its guard. Note this one uses a `ConditionalDimension`\n", + "directly -- a `SubGrid` coarsens *space*, so it has no say over the time axis." + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.222737Z", + "iopub.status.busy": "2026-10-21T00:00:13.222568Z", + "iopub.status.idle": "2026-10-21T00:00:13.400826Z", + "shell.execute_reply": "2026-10-21T00:00:13.399620Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n", + "{\n", + " START(section0)\n", + " for (int x = x_m; x <= x_M; x += 1)\n", + " {\n", + " #pragma omp simd aligned(v:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " v[t1][x + 1][y + 1] = v[t0][x + 1][y + 1] + 1.0F;\n", + " }\n", + " }\n", + " STOP(section0,timers)\n", + "\n", + " if ((time)%(tsf) == 0)\n", + " {\n", + " START(section1)\n", + " for (int x = x_m; x <= x_M; x += 1)\n", + " {\n", + " #pragma omp simd aligned(v,vsave:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " vsave[time / tsf][x + 1][y + 1] = v[t0][x + 1][y + 1];\n", + " }\n", + " }\n", + " STOP(section1,timers)\n", + " }\n", + "}\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "nt, factor = 9, 4\n", + "grid3 = Grid(shape=(4, 4))\n", + "ts = ConditionalDimension(name='ts', parent=grid3.time_dim, factor=factor)\n", + "\n", + "v = TimeFunction(name='v', grid=grid3)\n", + "vsave = TimeFunction(name='vsave', grid=grid3, time_dim=ts,\n", + " save=(nt + factor - 1)//factor)\n", + "\n", + "op_time = Operator([Eq(v.forward, v + 1.), Eq(vsave, v)])\n", + "show(op_time)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The saved slices are correct -- time subsampling behaves exactly as it always has,\n", + "it simply does not take the strided form." + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.404763Z", + "iopub.status.busy": "2026-10-21T00:00:13.404595Z", + "iopub.status.idle": "2026-10-21T00:00:13.428563Z", + "shell.execute_reply": "2026-10-21T00:00:13.427431Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "[0.0, 4.0, 0.0]" + ] + }, + "execution_count": 36, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "op_time.apply(time_M=nt-2)\n", + "\n", + "nwritten = (nt - 2)//factor + 1\n", + "for i in range(nwritten):\n", + " assert np.allclose(vsave.data[i], i*factor)\n", + "\n", + "[float(vsave.data[i, 0, 0]) for i in range(vsave.shape[0])]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. A `condition` is attached as well\n", + "\n", + "A `ConditionalDimension` may carry an arbitrary `condition` alongside a `factor`. The\n", + "combined test is no longer \"every `factor`-th point\", so it is evaluated pointwise." + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.432786Z", + "iopub.status.busy": "2026-10-21T00:00:13.432609Z", + "iopub.status.idle": "2026-10-21T00:00:13.568839Z", + "shell.execute_reply": "2026-10-21T00:00:13.567306Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " if ((x & 1) == 0 && (x)%(xs4f) == 0)\n", + " {\n", + " #pragma omp simd aligned(u,w:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " w[x / xs4f][y] = u[x + 1][y + 1];\n", + " }\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "# `SubGrid` does not expose `condition`, so this one is built by hand\n", + "xs4 = ConditionalDimension(name='xs4', parent=x, factor=3,\n", + " condition=CondEq(x % 2, 0))\n", + "w = Function(name='w', grid=grid, dimensions=(xs4, y), shape=(4, 12),\n", + " space_order=0)\n", + "\n", + "u.data[:] = 1.\n", + "w.data[:] = -1.\n", + "op_cond = Operator(Eq(w, u))\n", + "op_cond.apply()\n", + "show(op_cond)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. Accesses are not affine in the subsampled Dimension\n", + "\n", + "An indirectly indexed read (`u[idx[x, y], y]`) is not affine in `x`, so the points a\n", + "stride would skip cannot be shown to be unneeded." + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-21T00:00:13.574150Z", + "iopub.status.busy": "2026-10-21T00:00:13.573710Z", + "iopub.status.idle": "2026-10-21T00:00:13.750574Z", + "shell.execute_reply": "2026-10-21T00:00:13.749383Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " if ((x)%(xsf) == 0)\n", + " {\n", + " #pragma omp simd aligned(gi,idx,u:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " if ((y)%(ysf) == 0)\n", + " {\n", + " int r0 = idx[x + 1][y + 1];\n", + " gi[x / xsf][y / ysf] = u[r0 + 1][y + 1];\n", + " }\n", + " }\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "idx = Function(name='idx', grid=grid, dtype=np.int32)\n", + "gi = Function(name='gi', grid=sub, space_order=0)\n", + "\n", + "idx.data[:] = 0 # read row 0 of `u` everywhere\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "op_indirect = Operator(Eq(gi, u[idx[x, y], y]))\n", + "op_indirect.apply()\n", + "\n", + "assert np.array_equal(gi.data, u.data[0, ::3] * np.ones((4, 1)))\n", + "show(op_indirect)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "`SubGrid(parent, factor)` builds a coarser grid from a parent `Grid`, keeping every\n", + "`factor`-th point. It is a `Grid` subclass, so Functions, TimeFunctions, Operators and\n", + "finite differences all work as usual:\n", + "\n", + "```python\n", + "sub = SubGrid(grid, factor=2) # or factor=(2, 3); 1 = passthrough\n", + "usub = Function(name='usub', grid=sub) # shape and dimensions derived\n", + "Operator(Eq(usub, u)).apply() # strided copy from the parent\n", + "```\n", + "\n", + "| | derived from the parent |\n", + "|---|---|\n", + "| `shape` | `ceil(n/factor)` per axis |\n", + "| `spacing` | `factor * parent.spacing` |\n", + "| `extent`, `origin`, `dtype`, `time_dim` | inherited unchanged |\n", + "| `dimensions` | `ConditionalDimension`s on the parent's |\n", + "| `parent`, `factor` | the two attributes `SubGrid` adds |\n", + "\n", + "Use a `ConditionalDimension` directly for subsampling in **time**, or to combine a\n", + "factor with an extra `condition` -- `SubGrid` covers neither.\n", + "\n", + "Under the hood Devito iterates the subgrid with a unit stride and rescales parent\n", + "accesses, falling back to a pointwise `if (x%factor == 0)` guard in the four cases\n", + "above. Both compute the same thing; only the generated loop differs.\n", + "\n", + "Generic subsampling -- subgrid points that fall *between* parent grid points and so\n", + "require interpolation -- is a separate topic, not covered here." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "devitoenv (3.11.0)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0rc1" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py new file mode 100644 index 0000000000..07ff1a5a1f --- /dev/null +++ b/tests/test_subsampling.py @@ -0,0 +1,1046 @@ +import re + +import numpy as np +import pytest +from sympy import Or + +from devito import ( + CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Inc, Operator, SubGrid, + TimeFunction, switchconfig +) +from devito.ir.iet import Conditional, Expression, FindNodes, Iteration +from devito.symbolics import retrieve_indexed +from devito.tools import as_tuple + + +def iteration(op, name): + """The Iteration over the Dimension named `name` in `op`.""" + for i in FindNodes(Iteration).visit(op): + if i.dim.name == name: + return i + raise KeyError(name) + + +def subsampled(shape, factor, ndim=2): + """ + A grid, a Function `f` over it, a Function `g` subsampled by `factor` along + the leading Dimension, and the ConditionalDimension tying them together. + """ + grid = Grid(shape=shape[:ndim]) + dims = grid.dimensions + cd = ConditionalDimension(name='xs', parent=dims[0], factor=factor) + + f = Function(name='f', grid=grid) + gshape = ((shape[0] + factor - 1)//factor,) + tuple(shape[1:ndim]) + g = Function(name='g', shape=gshape, dimensions=(cd,) + dims[1:], + grid=grid, space_order=0) + + return grid, f, g, cd + + +class TestStridedCodegen: + + """ + Subsampling iterates the subsampled Dimension directly, with unit stride, + rescaling parent accesses -- rather than striding the parent's loop. A + non-unitary increment would escape the Lamport-style dependence analysis + the compiler relies on. + """ + + def test_codegen(self): + _, f, g, cd = subsampled((11, 11), 2) + op = Operator(Eq(g, f)) + + # No guard... + assert not FindNodes(Conditional).visit(op) + # ...instead the loop is over `cd` itself, with unit stride + assert iteration(op, cd.name).limits[2] == 1 + with pytest.raises(KeyError): + iteration(op, 'x') + + @pytest.mark.parametrize('factor', [2, 3, 4, 5]) + def test_parent_access_rescaled(self, factor): + """ + Parent accesses become `cd*factor`; the subgrid access stays `cd`. + + Asserted on the expressions rather than the emitted C, which renders + accesses as `gL0(xs, y)` macros when linearization is on. + """ + _, f, g, cd = subsampled((16, 8), factor) + op = Operator(Eq(g, f)) + + # Pick the Expression writing `g`; linearization adds Size/Stride + # setup Expressions alongside it + expr, = [i.expr for i in FindNodes(Expression).visit(op) + if i.write is g] + + # The subgrid Function is indexed by `cd` itself... + lhs, = retrieve_indexed(expr.lhs) + assert lhs.indices[0] == cd + # ...while the parent read is rescaled to `cd*factor` + rhs, = retrieve_indexed(expr.rhs) + assert rhs.indices[0].has(cd.symbolic_factor) + + assert iteration(op, cd.name).limits[2] == 1 # no strided loop + assert not FindNodes(Conditional).visit(op) + + def test_unit_stride_everywhere(self): + """Every Iteration keeps a unit increment, so DDA stays applicable.""" + grid = Grid(shape=(12, 12)) + sub = SubGrid(grid, factor=(3, 2)) + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + + op = Operator(Eq(f, u)) + + for i in FindNodes(Iteration).visit(op): + assert i.limits[2] == 1 + + def test_innermost_strided(self): + """ + The innermost Dimension is strided too. The alternative is a guard + *inside* the SIMD loop -- a predicated store on every iteration -- + which is worse than striding on both instruction count and branching. + """ + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + ys = ConditionalDimension(name='ys', parent=y, factor=3) + u = Function(name='u', grid=grid) + usub = Function(name='usub', grid=grid, dimensions=(x, ys), + shape=(12, 4), space_order=0) + + op = Operator(Eq(usub, u)) + + assert iteration(op, ys.name).limits[2] == 1 + assert f'{ys.name}*{ys.symbolic_factor.name}' in str(op.ccode) + assert not FindNodes(Conditional).visit(op) + + def test_innermost_numerics(self): + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + ys = ConditionalDimension(name='ys', parent=y, factor=3) + u = Function(name='u', grid=grid) + usub = Function(name='usub', grid=grid, dimensions=(x, ys), + shape=(12, 4), space_order=0) + data = np.arange(144, dtype=u.dtype).reshape(12, 12) + u.data[:] = data + + Operator(Eq(usub, u)).apply() + + assert np.array_equal(usub.data, data[:, ::3]) + + def test_multiple_dimensions_strided(self): + """Each subsampled Dimension gets its own independent stride.""" + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + xs = ConditionalDimension(name='xs', parent=x, factor=3) + ys = ConditionalDimension(name='ys', parent=y, factor=2) + u = Function(name='u', grid=grid) + usub = Function(name='usub', grid=grid, dimensions=(xs, ys), + shape=(4, 6), space_order=0) + + op = Operator(Eq(usub, u)) + + assert not FindNodes(Conditional).visit(op) + for d in (xs, ys): + assert iteration(op, d.name).limits[2] == 1 + assert f'{d.name}*{d.symbolic_factor.name}' in str(op.ccode) + + def test_multiple_dimensions_numerics(self): + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + xs = ConditionalDimension(name='xs', parent=x, factor=3) + ys = ConditionalDimension(name='ys', parent=y, factor=2) + u = Function(name='u', grid=grid) + usub = Function(name='usub', grid=grid, dimensions=(xs, ys), + shape=(4, 6), space_order=0) + data = np.arange(144, dtype=u.dtype).reshape(12, 12) + u.data[:] = data + + Operator(Eq(usub, u)).apply() + + assert np.array_equal(usub.data, data[::3, ::2]) + + def test_properties_inherited(self): + """ + The subsampled Dimension must inherit the parent's computational + properties: it covers the same points, just in subsampled coordinates. + Losing them costs the loop its PARALLEL/VECTORIZED status, and with it + all of its OpenMP. + """ + grid = Grid(shape=(32, 32)) + x, y = grid.dimensions + sub = SubGrid(grid, factor=2) + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + + # What the same copy gets on the parent Grid... + v = Function(name='v', grid=grid) + ref = Operator(Eq(v, u)) + expected = {i.dim.name: i.properties + for i in FindNodes(Iteration).visit(ref)} + + # ...the subsampled version must get too + op = Operator(Eq(f, u)) + for i in FindNodes(Iteration).visit(op): + parent = i.dim.parent if i.dim.is_Conditional else i.dim + assert i.properties == expected[parent.name], i.dim.name + + @switchconfig(language='openmp') + def test_openmp_emitted(self): + """The subsampled loop nest is still parallelized and vectorized.""" + grid = Grid(shape=(32, 32)) + sub = SubGrid(grid, factor=2) + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + + ccode = str(Operator(Eq(f, u)).ccode) + + assert '#pragma omp for' in ccode + assert '#pragma omp simd' in ccode + + def test_bounds_are_subsampled(self): + """ + The loop runs over the subsampled Dimension's own runtime bounds, + which `ConditionalDimension._arg_values` already derives as + `ceil(x_m/factor) .. ceil(x_M/factor)`. No anchoring is needed. + """ + _, f, g, cd = subsampled((12, 12), 3) + op = Operator(Eq(g, f)) + + assert iteration(op, cd.name).limits[0] == cd.symbolic_min + args = op.arguments() + assert args[cd.min_name] == 0 + assert args[cd.max_name] == 3 + + +class TestNumerics: + + """Results must match an independent reference.""" + + @pytest.mark.parametrize('shape,factor,ndim', [ + ((12,), 2, 1), ((12,), 3, 1), ((13,), 4, 1), + ((12, 12), 2, 2), ((12, 12), 3, 2), ((11, 9), 4, 2), + ((9, 9, 9), 3, 3), + ]) + def test_against_numpy(self, shape, factor, ndim): + _, f, g, _ = subsampled(shape, factor, ndim=ndim) + data = np.arange(f.data.size, dtype=f.dtype).reshape(f.data.shape) + f.data[:] = data + Operator(Eq(g, f)).apply() + + assert np.array_equal(g.data, data[::factor]) + + @pytest.mark.parametrize('x_m,x_M', [ + (0, 11), (1, 11), (2, 11), (3, 11), (0, 9), (4, 10), + ]) + def test_anchoring(self, x_m, x_M): + """ + Exactly the points of `[x_m, x_M]` that lie on the subsampling grid + must be written; `x_m` need not be a multiple of `factor`. + """ + factor = 3 + _, f, g, _ = subsampled((12, 12), factor) + f.data[:] = 1. + g.data[:] = -1. + Operator(Eq(g, f)).apply(x_m=x_m, x_M=x_M) + + # The loop runs over the subsampled Dimension's own bounds, + # `ceil(x_m/factor) .. floor(x_M/factor)` + written = set(range(-(-x_m//factor), x_M//factor + 1)) + for i in range(g.data.shape[0]): + expected = 1. if i in written else -1. + assert np.all(g.data[i] == expected), (i, x_m, x_M) + + @pytest.mark.parametrize('x_M,expected', [ + (8, 2), (9, 3), (10, 3), (11, 3), + ]) + def test_max_bound_floors(self, x_M, expected): + """ + The subsampled max is `floor(x_M/factor)`, the last subsampled point at + or before `x_M` -- not `ceil`, which would run one point past the end + of the Function and corrupt the heap. + """ + _, f, g, cd = subsampled((12, 12), 3) + op = Operator(Eq(g, f)) + + args = op.arguments(x_M=x_M) + + assert args[cd.max_name] == expected + assert args[cd.max_name] < g.shape[0] + + def test_constant_factor(self): + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + factor = Constant(name='factor', dtype=np.int32, value=3) + cd = ConditionalDimension(name='xs', parent=x, factor=factor) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(4, 12), dimensions=(cd, y), grid=grid, + space_order=0) + data = np.arange(144, dtype=f.dtype).reshape(12, 12) + f.data[:] = data + Operator(Eq(g, f)).apply(factor=3) + + assert np.array_equal(g.data, data[::3]) + + def test_time_subsampling(self): + """ + Time is SEQUENTIAL, so this keeps its guard -- but must still save the + right slices. + """ + nt, factor = 9, 4 + grid = Grid(shape=(5, 5)) + ts = ConditionalDimension(name='ts', parent=grid.time_dim, + factor=factor) + u = TimeFunction(name='u', grid=grid) + usave = TimeFunction(name='usave', grid=grid, time_dim=ts, + save=(nt + factor - 1)//factor) + + Operator([Eq(u.forward, u + 1.), Eq(usave, u)]).apply(time_M=nt-2) + + # 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 + for i in range(nwritten): + assert np.allclose(usave.data[i], i*factor) + + +class TestEquivalence: + + """The strided loop must visit exactly the guard's point set.""" + + @pytest.mark.parametrize('factor', [1, 2, 3, 4, 5, 6]) + def test_point_set_identical(self, factor): + """ + Exhaustive check of the rewrite's arithmetic: iterating + `range(anchor, x_M + 1, factor)` with + `anchor = x_m + (factor - x_m%factor)%factor` enumerates exactly the + points of `[x_m, x_M]` that satisfy `x%factor == 0`. + """ + for x_m in range(0, 14): + for x_M in range(x_m, 20): + guarded = [x for x in range(x_m, x_M + 1) if x % factor == 0] + anchor = x_m + (factor - x_m % factor) % factor + strided = list(range(anchor, x_M + 1, factor)) + + assert guarded == strided, (factor, x_m, x_M) + + +class TestLegality: + + """Cases that must keep the pointwise guard.""" + + def test_mixed_guarded_and_unguarded(self): + """ + An unguarded expression sharing the loop nest must keep running at + every point, so the guard cannot become a stride. + """ + grid = Grid(shape=(12,)) + x, = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=4) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(3,), dimensions=(cd,), grid=grid, + space_order=0) + + op = Operator([Eq(f, 1), Eq(g, f)]) + + assert len(FindNodes(Conditional).visit(op)) == 1 + assert iteration(op, 'x').limits[2] == 1 + + def test_mixed_guards_reach_one_callback(self): + """ + The guarded and unguarded Clusters must be handed to the pass + *together*, as one group: that is what lets it see the nest is shared + and refuse the rewrite. Grouping by guard instead would hide the + unguarded Cluster and silently drop the guard. + """ + import importlib + + module = importlib.import_module('devito.passes.clusters.subsampling') + subsampling_queue = vars(module)['Subsampling'] + + seen = [] + original = subsampling_queue.callback + + def callback(self, clusters, prefix): + if prefix: + d = prefix[-1].dim + seen.append({type(c.guards.get(d)).__name__ for c in clusters}) + return original(self, clusters, prefix) + + grid = Grid(shape=(12,)) + x, = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=4) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(3,), dimensions=(cd,), grid=grid, + space_order=0) + + subsampling_queue.callback = callback + try: + Operator([Eq(f, 1), Eq(g, f)]) + finally: + subsampling_queue.callback = original + + assert {'BooleanTrue', 'GuardFactorEq'} in seen + + def test_mixed_numerics(self): + grid = Grid(shape=(12,)) + x, = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=4) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(3,), dimensions=(cd,), grid=grid, + space_order=0) + g.data[:] = -1 + + Operator([Eq(f, x), Eq(g, f)]).apply() + + assert np.array_equal(f.data, np.arange(12)) + assert np.array_equal(g.data, np.arange(12)[::4]) + + def test_non_affine_access(self): + """ + An indirectly indexed read is not affine in the subsampled Dimension, + so the points the stride would skip cannot be shown to be unneeded. + """ + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=3) + idx = Function(name='idx', grid=grid, dtype=np.int32) + u = Function(name='u', grid=grid) + g = Function(name='g', shape=(4, 12), dimensions=(cd, y), grid=grid, + space_order=0) + + op = Operator(Eq(g, u[idx[x, y], y])) + + assert iteration(op, 'x').limits[2] == 1 + assert FindNodes(Conditional).visit(op) + + def test_user_condition(self): + """A `condition=` must still be evaluated pointwise.""" + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + cd = ConditionalDimension(name='xs', parent=x, factor=3, + condition=CondEq(x % 2, 0)) + g = Function(name='g', shape=(4, 12), dimensions=(cd, y), grid=grid, + space_order=0) + + f.data[:] = 1. + op = Operator(Eq(g, f)) + op.apply() + + assert FindNodes(Conditional).visit(op) + assert iteration(op, 'x').limits[2] == 1 + + def test_or_relation(self): + """ + A factor guard `Or`-ed with another condition fires on more points + than the subsampling grid, so it cannot become a stride. + """ + grid = Grid(shape=(10, 10)) + x, y = grid.dimensions + c_even = ConditionalDimension(name='ceven', parent=x, factor=2, + relation=Or) + c_edge = ConditionalDimension(name='cedge', parent=x, + condition=CondEq(x, x.symbolic_max), + relation=Or) + f = Function(name='f', grid=grid, dimensions=(c_even, y), + shape=(5, 10), space_order=0) + + op = Operator(Eq(f, 1, implicit_dims=c_edge)) + + assert FindNodes(Conditional).visit(op) + assert iteration(op, 'x').limits[2] == 1 + + def test_uniformly_guarded_sequential(self): + """ + The case that isolates the parallelism gate: the subsampled Dimension + `x` is *uniformly* guarded (so the shared-nest check does not apply) + and is *not* innermost (so the SIMD check does not apply), but it + carries a flow dependence and is therefore SEQUENTIAL. Only the + parallelism gate can reject it; striding here would skip the + iterations that feed the retained ones. + """ + grid = Grid(shape=(12, 12)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=3) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(4, 12), dimensions=(cd, y), grid=grid, + space_order=0) + + op = Operator(Eq(g, g.subs(cd, cd-1) + f)) + + assert iteration(op, 'x').limits[2] == 1 + assert FindNodes(Conditional).visit(op) + + def test_sequential_dimension(self): + """ + A dependence carried along the subsampled Dimension forbids skipping + points: the discarded iterations feed the retained ones. + """ + grid = Grid(shape=(12,)) + x, = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=3) + f = Function(name='f', grid=grid) + g = Function(name='g', shape=(4,), dimensions=(cd,), grid=grid, + space_order=0) + + # `f[x] = f[x-1] + 1` carries a flow dependence along `x` + op = Operator([Eq(f, f.subs(x, x-1) + 1), Eq(g, f)]) + op.apply() + + assert iteration(op, 'x').limits[2] == 1 + assert np.array_equal(g.data, f.data[::3]) + + +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. + """ + grid = Grid(shape=(32, 32)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=4) + u = Function(name='u', grid=grid, space_order=4) + g = Function(name='g', shape=(8, 32), dimensions=(cd, y), grid=grid, + 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})) + op.apply() + + assert len(FindNodes(Conditional).visit(op)) == 1 + assert np.all(np.isfinite(g.data)) + + @switchconfig(language='openmp') + def test_with_openmp(self): + _, f, g, _ = subsampled((16, 16), 2) + data = np.arange(256, dtype=f.dtype).reshape(16, 16) + f.data[:] = data + Operator(Eq(g, f)).apply() + + assert np.array_equal(g.data, data[::2]) + + +class TestSubGrid: + + """The `SubGrid` API.""" + + def test_derived_properties(self): + grid = Grid(shape=(9, 9), extent=(8., 8.)) + sub = SubGrid(grid, factor=2) + + assert sub.parent is grid + assert sub.factor == (2, 2) + assert sub.shape == (5, 5) + assert all(d.is_Conditional for d in sub.dimensions) + assert all(d.parent is p + for d, p in zip(sub.dimensions, grid.dimensions, strict=True)) + + def test_spacing_scaled(self): + """Physical spacing is `factor` times the parent's; extent is kept.""" + grid = Grid(shape=(9, 9), extent=(8., 8.)) + sub = SubGrid(grid, factor=2) + + assert np.allclose(sub.spacing, [2*s for s in grid.spacing]) + assert np.allclose(sub.extent, grid.extent) + assert np.allclose(sub.origin, grid.origin) + + def test_function_needs_no_shape(self): + """`Function(grid=sub)` derives both dimensions and shape.""" + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=2) + + f = Function(name='f', grid=sub, space_order=0) + + assert f.dimensions == sub.dimensions + assert f.shape == (5, 5) + + @pytest.mark.parametrize('shape,factor,expected', [ + ((9, 9), 2, (5, 5)), + ((9, 9), (2, 1), (5, 9)), + ((10,), 3, (4,)), + ((10,), 5, (2,)), + ((8, 8, 8), (2, 4, 1), (4, 2, 8)), + ]) + def test_shape_derivation(self, shape, factor, expected): + sub = SubGrid(Grid(shape=shape), factor=factor) + + assert sub.shape == expected + + @pytest.mark.parametrize('factor', [2, 3, (2, 1), (1, 3), (3, 2)]) + def test_copy_from_parent(self, factor): + """Writing a SubGrid Function from a parent Function is a strided copy.""" + 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(144, dtype=u.dtype).reshape(12, 12) + u.data[:] = data + + Operator(Eq(f, u)).apply() + + fx, fy = sub.factor + assert np.array_equal(f.data, data[::fx, ::fy]) + + def test_copy_is_strided(self): + """ + No guard is emitted, every loop keeps a unit stride, and the Operator + runs to the right answer. + """ + grid = Grid(shape=(12, 12)) + sub = SubGrid(grid, factor=3) + + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + data = np.arange(144, dtype=u.dtype).reshape(12, 12) + u.data[:] = data + + op = Operator(Eq(f, u)) + op.apply() + + assert not FindNodes(Conditional).visit(op) + for d in sub.dimensions: + assert iteration(op, d.name).limits[2] == 1 + assert np.array_equal(f.data, data[::3, ::3]) + + @switchconfig(language='openmp') + def test_openmp_and_numerics(self): + """A SubGrid Operator is parallelized, vectorized, and correct.""" + grid = Grid(shape=(32, 32)) + sub = SubGrid(grid, factor=(4, 2)) + + u = Function(name='u', grid=grid) + f = Function(name='f', grid=sub, space_order=0) + data = np.arange(32*32, dtype=u.dtype).reshape(32, 32) + u.data[:] = data + + op = Operator(Eq(f, u)) + op.apply() + + ccode = str(op.ccode) + assert '#pragma omp for' in ccode + assert '#pragma omp simd' in ccode + assert np.array_equal(f.data, data[::4, ::2]) + + def test_shared_nest_numerics(self): + """ + A SubGrid Function sharing a nest with an unguarded Eq: the guard is + kept and both results are right. + """ + grid = Grid(shape=(12,)) + sub = SubGrid(grid, factor=4) + + f = Function(name='f', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + g.data[:] = -1. + + op = Operator([Eq(f, 1), Eq(g, f)]) + op.apply() + + assert FindNodes(Conditional).visit(op) + assert np.all(f.data == 1.) + assert np.all(g.data == 1.) + + def test_sequential_numerics(self): + """A SubGrid Function accumulating along the subsampled Dimension.""" + grid = Grid(shape=(12, 12)) + sub = SubGrid(grid, factor=(3, 1)) + xs = sub.dimensions[0] + + f = Function(name='f', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + f.data[:] = 1. + g.data[:] = 0. + + Operator(Eq(g, g.subs(xs, xs-1) + f)).apply() + + # Row 0 reads the halo (zero), then each row adds one + assert np.allclose(g.data[:, 0], [0., 1., 2., 3.]) + + def test_factor_one_is_passthrough(self): + """A factor of 1 leaves the Dimension untouched.""" + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=(2, 1)) + + assert sub.dimensions[0].is_Conditional + assert sub.dimensions[1] is grid.dimensions[1] + + def test_shares_parent_time_dimensions(self): + """ + A SubGrid coarsens space only, so a TimeFunction on it must step over + the same buffer indices as one on the parent. + """ + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=2) + + assert sub.time_dim is grid.time_dim + assert sub.stepping_dim is grid.stepping_dim + + u = TimeFunction(name='u', grid=grid) + us = TimeFunction(name='us', grid=sub, space_order=0) + + assert u.dimensions[0] is us.dimensions[0] + + def test_time_stepping_operator(self): + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=2) + + u = TimeFunction(name='u', grid=grid) + us = TimeFunction(name='us', grid=sub, space_order=0) + + Operator([Eq(u.forward, u + 1.), Eq(us.forward, u.forward)]).apply(time_M=4) + + assert np.allclose(us.data[-1], u.data[-1][::2, ::2]) + + def test_derivative_uses_coarse_spacing(self): + """ + `f.dx` on a SubGrid Function differentiates over the coarse spacing: + `ConditionalDimension.spacing` is `factor*h_x`, so the stencil steps + `factor*h_x` and the weights scale by `1/(factor*h_x)`. + """ + grid = Grid(shape=(11, 11), extent=(10., 10.)) + sub = SubGrid(grid, factor=2) + xs = sub.dimensions[0] + + assert xs.spacing == 2*grid.dimensions[0].spacing + + f = Function(name='f', grid=sub, space_order=2) + expr = f.dx.evaluate + + # Steps by the coarse spacing, with correspondingly halved weights + assert str(expr) == '-0.5*f(xs, ys)/h_x + 0.5*f(xs + 2*h_x, ys)/h_x' + + def test_derivative_numerics(self): + """A linear field differentiates exactly on the SubGrid.""" + grid = Grid(shape=(21, 21), extent=(20., 20.)) + sub = SubGrid(grid, factor=2) + xs, ys = sub.dimensions + + f = Function(name='f', grid=sub, space_order=2) + d = Function(name='d', grid=sub, space_order=2) + + # f = 3*x on the coarse grid, so df/dx == 3 everywhere inside + for i in range(f.shape[0]): + f.data[i, :] = 3.*i*float(sub.spacing[0]) + + Operator(Eq(d, f.dx)).apply() + + assert np.allclose(d.data[1:-1, 1:-1], 3., atol=1e-4) + + @pytest.mark.parametrize('factor,err', [ + (0, 'must be >= 1'), + (-1, 'must be >= 1'), + (2.5, 'Non-integer'), + ((2, 2, 2), 'Expected 1 or 2 factors'), + ]) + def test_invalid_factor(self, factor, err): + grid = Grid(shape=(9, 9)) + + with pytest.raises(ValueError, match=err): + SubGrid(grid, factor=factor) + + @pytest.mark.parametrize('f0,f1,combined', [ + (2, 3, (6, 6)), + (2, (3, 1), (6, 2)), + ((3, 2), (2, 4), (6, 8)), + ]) + def test_cannot_nest(self, f0, f1, combined): + """ + Nesting is rejected, and the message names the equivalent single-level + call rather than leaving the user to work out the arithmetic. + """ + grid = Grid(shape=(24, 24)) + sub = SubGrid(grid, factor=f0) + + with pytest.raises(ValueError) as e: + SubGrid(sub, factor=f1) + + assert f'factor={combined}' in str(e.value) + + @pytest.mark.parametrize('f0,f1,combined', [ + (2, 3, (6, 6)), + (2, (3, 1), (6, 2)), + ((3, 2), (2, 4), (6, 8)), + ]) + def test_nesting_suggestion_is_equivalent(self, f0, f1, combined): + """ + The combined factor the error suggests must really select the points a + two-step subsampling would have. + """ + grid = Grid(shape=(24, 24)) + sub = SubGrid(grid, factor=f0) + suggested = SubGrid(grid, factor=combined, name='c') + + u = Function(name='u', grid=grid) + f = Function(name='f', grid=suggested, space_order=0) + data = np.arange(24*24, dtype=u.dtype).reshape(24, 24) + u.data[:] = data + + Operator(Eq(f, u)).apply() + + # Two-step reference: coarsen by `f0`, then by `f1` + g0, g1 = sub.factor + h0, h1 = as_tuple(f1)*2 if len(as_tuple(f1)) == 1 else as_tuple(f1) + two_step = data[::g0, ::g1][::h0, ::h1] + + assert np.array_equal(f.data, two_step) + + def test_non_integer_factor_suggests_valid_factors(self): + """The suggested fallbacks must themselves be legal factors.""" + grid = Grid(shape=(9, 9)) + + for bad in (0.5, 1.2, 2.5): + with pytest.raises(ValueError) as e: + SubGrid(grid, factor=bad) + suggested = re.findall(r'use `(\d+)` or `(\d+)`', str(e.value)) + assert suggested, str(e.value) + for v in suggested[0]: + assert int(v) >= 1 + SubGrid(grid, factor=int(v)) # must not raise + + def test_parent_must_be_grid(self): + with pytest.raises(ValueError, match='must be a Grid'): + SubGrid(Grid(shape=(9, 9)).interior, factor=2) + + def test_custom_dimension_names(self): + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=2, name='c') + + assert [d.name for d in sub.dimensions] == ['xc', 'yc'] + + +class TestMPI: + + """ + Subsampling under domain decomposition. The subsampled Dimension's bounds + are derived per rank from the parent's, so each rank writes exactly its own + share of the subgrid. + """ + + @pytest.mark.parallel(mode=[2, 4]) + @pytest.mark.parametrize('factor', [2, 3, 4]) + def test_copy_from_parent(self, mode, factor): + """A strided copy must agree with numpy slicing 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(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, 4]) + def test_strided_codegen(self, mode): + """The guard still becomes a unit-stride subgrid loop under MPI.""" + n = 24 + grid = Grid(shape=(n, n)) + sub = SubGrid(grid, factor=3) + + u = Function(name='u', grid=grid) + g = Function(name='g', grid=sub, space_order=0) + + op = Operator(Eq(g, u)) + + assert not FindNodes(Conditional).visit(op) + for d in sub.dimensions: + assert iteration(op, d.name).limits[2] == 1 + + @pytest.mark.parallel(mode=[2, 4]) + def test_bounds_partition_the_subgrid(self, mode): + """ + Every subgrid point is written exactly once across all ranks: the + per-rank bounds must tile the global subgrid with no gaps or overlap. + Checked with a sentinel, since a reduction into a subsampled Function + is not supported under MPI. + """ + n, factor = 24, 3 + 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) + u.data[:] = 1. + g.data[:] = -7. # sentinel: never written + + Operator(Eq(g, u)).apply() + + # No gaps: this rank's whole slice was written + assert not np.any(np.asarray(g.data) == -7.) + assert np.all(np.asarray(g.data) == 1.) + + # No overlap: the local sizes sum to the global count + local = int(np.asarray(g.data).size) + total = grid.distributor.comm.allreduce(local) + assert total == sub.shape[0]*sub.shape[1] + + @pytest.mark.parallel(mode=[2, 4]) + def test_subgrid_shares_parent_decomposition(self, mode): + """ + The SubGrid must decompose along the same axes as its parent, so that + a rank's subgrid points come from its own parent points. + """ + grid = Grid(shape=(24, 24)) + sub = SubGrid(grid, factor=3) + + assert sub.root is grid.root + assert sub.distributor.nprocs == grid.distributor.nprocs + assert sub.distributor.topology == grid.distributor.topology + + @pytest.mark.parallel(mode=[2, 4]) + def test_time_subsampling(self, mode): + """Saving a wavefield every `factor` steps, decomposed in space.""" + nt, factor = 9, 4 + grid = Grid(shape=(16, 16)) + ts = ConditionalDimension(name='ts', parent=grid.time_dim, + factor=factor) + + v = TimeFunction(name='v', grid=grid) + vsave = TimeFunction(name='vsave', grid=grid, time_dim=ts, + save=(nt + factor - 1)//factor) + + Operator([Eq(v.forward, v + 1.), Eq(vsave, v)]).apply(time_M=nt-2) + + for i in range((nt - 2)//factor + 1): + assert np.allclose(vsave.data[i], i*factor) + + +class TestIncReduction: + + """ + `Inc` into a factor-subsampled Function must not be treated as a + reduction. The write index `i -> i//factor` is one-to-one on the guarded + points, so no two iterations touch the same location and the loop is + plainly parallel. + + Treating it as one emitted an OpenMP array reduction over the whole + Function; with several such Functions the private copies overflowed the + stack and the kernel segfaulted. + """ + + @pytest.mark.parametrize('factor', [1, 2, 3]) + def test_no_reduction_clause(self, factor): + n = 12 + grid = Grid(shape=(n, n)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=factor) + + f = Function(name='f', shape=((n + factor - 1)//factor, n), + dimensions=(cd, y), grid=grid, space_order=0) + + with switchconfig(language='openmp'): + op = Operator(Inc(f, 1)) + + assert 'reduction' not in str(op.ccode) + + @pytest.mark.parametrize('factor', [1, 2, 3]) + def test_dimensions_are_parallel(self, factor): + """ + The properties must match what the same `Inc` gets on a plain + Function, i.e. PARALLEL rather than PARALLEL_IF_ATOMIC. + """ + n = 12 + grid = Grid(shape=(n, n)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=factor) + + f = Function(name='f', shape=((n + factor - 1)//factor, n), + dimensions=(cd, y), grid=grid, space_order=0) + + op = Operator(Inc(f, 1)) + + for i in FindNodes(Iteration).visit(op): + assert i.is_Parallel + + @pytest.mark.parametrize('factor', [1, 2, 3]) + def test_accumulates(self, factor): + """Dropping the reduction must not drop the accumulation.""" + n = 12 + grid = Grid(shape=(n, n)) + x, y = grid.dimensions + cd = ConditionalDimension(name='xs', parent=x, factor=factor) + + f = Function(name='f', shape=((n + factor - 1)//factor, n), + dimensions=(cd, y), grid=grid, space_order=0) + f.data[:] = 0. + + op = Operator(Inc(f, 1)) + op.apply() + op.apply() + + assert np.all(f.data == 2.) + + def test_many_functions(self): + """ + The original failure: nine subsampled Functions incremented together. + The spurious array reductions needed ~13 MiB of thread-private copies + and overflowed the stack. + """ + n = 40 + grid = Grid(shape=(n, n, n)) + dims = tuple(ConditionalDimension(name=f'{d.name}s', parent=d, factor=1) + for d in grid.dimensions) + + fs = [Function(name=f'og_{i}', grid=grid, shape=(n, n, n), + dimensions=dims) for i in range(9)] + for f in fs: + f.data[:] = 0. + + with switchconfig(language='openmp'): + op = Operator([Inc(f, 1) for f in fs]) + assert 'reduction' not in str(op.ccode) + op.apply() + + for f in fs: + assert np.all(f.data == 1.) + + +class TestInjectivity: + + """`is_Injective` must be claimed only where it holds.""" + + def test_factor_is_injective(self): + grid = Grid(shape=(8, 8)) + x, _ = grid.dimensions + + assert ConditionalDimension(name='a', parent=x, factor=2).is_Injective + + def test_stepping_is_not(self): + """`time % 2` deliberately revisits the same slot.""" + grid = Grid(shape=(8, 8)) + + assert not grid.stepping_dim.is_Injective + + def test_condition_is_not(self): + """An arbitrary condition carries no one-to-one promise.""" + grid = Grid(shape=(8, 8)) + x, _ = grid.dimensions + cd = ConditionalDimension(name='b', parent=x, condition=CondEq(x, 0)) + + assert not cd.is_Injective + + def test_factor_with_condition_is_not(self): + grid = Grid(shape=(8, 8)) + x, _ = grid.dimensions + cd = ConditionalDimension(name='c', parent=x, factor=2, + condition=CondEq(x % 4, 0)) + + assert not cd.is_Injective + + def test_plain_dimension_is_not(self): + grid = Grid(shape=(8, 8)) + x, _ = grid.dimensions + + assert not x.is_Injective