From 13bce359ff21d9a7f91b5375dbd31cec16569a64 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 12:20:39 -0400 Subject: [PATCH 01/11] compiler: Lower integer subsampling to a strided loop A Function subsampled by `factor` along `d` was written under a `d % factor == 0` guard, so the enclosing loop visited every point of `d` and discarded all but one in `factor` of them. Where legal, drop the guard and step the loop `factor` points at a time instead. The rewrite is recorded as a parameterized `Strided(factor)` property on `d`; `NodeIteration.limits` applies the stride and rounds the loop min up onto the subsampling grid. The iteration space stays keyed on `d`, so runtime bounds, domain decomposition and non-subsampled accesses are all unaffected. Clusters are grouped by iteration space prefix alone, not by guard: a guard may only become a stride if every Cluster sharing the prefix carries the same one, since dropping it is what makes a loop nest shared. --- devito/core/cpu.py | 7 +- devito/core/gpu.py | 7 +- devito/ir/stree/tree.py | 35 ++++++++-- devito/ir/support/properties.py | 38 +++++++++++ devito/passes/clusters/__init__.py | 1 + devito/passes/clusters/subsampling.py | 95 +++++++++++++++++++++++++++ 6 files changed, 177 insertions(+), 6 deletions(-) create mode 100644 devito/passes/clusters/subsampling.py diff --git a/devito/core/cpu.py b/devito/core/cpu.py index 3435a472e62..1d7b4217bc0 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -6,7 +6,7 @@ from devito.passes import stream_dimensions from devito.passes.clusters import ( Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse, - optimize_hyperplanes, optimize_pows + optimize_hyperplanes, optimize_pows, subsampling ) from devito.passes.equations import collect_derivatives from devito.passes.iet import ( @@ -195,6 +195,11 @@ def _specialize_clusters(cls, clusters, **kwargs): # Reduce flops clusters = cse(clusters, **kwargs) + # Strided iteration in place of subsampling guards. After the last + # `fuse`, so that no unguarded Cluster can be merged into a strided + # loop, and before blocking, which must see the final step + clusters = subsampling(clusters) + # Blocking to improve data locality if options['blocklazy']: clusters = blocking(clusters, sregistry, options) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 536a06422e2..76275994856 100644 --- a/devito/core/gpu.py +++ b/devito/core/gpu.py @@ -8,7 +8,7 @@ from devito.passes import is_on_device, stream_dimensions from devito.passes.clusters import ( Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse, - memcpy_prefetch, optimize_pows, tasking + memcpy_prefetch, optimize_pows, subsampling, tasking ) from devito.passes.equations import collect_derivatives from devito.passes.iet import ( @@ -242,6 +242,11 @@ def _specialize_clusters(cls, clusters, **kwargs): # Reduce flops clusters = cse(clusters, **kwargs) + # Strided iteration in place of subsampling guards. After the last + # `fuse`, so that no unguarded Cluster can be merged into a strided + # loop, and before blocking, which must see the final step + clusters = subsampling(clusters) + # Blocking to define thread blocks if options['blocklazy']: clusters = blocking(clusters, sregistry, options) diff --git a/devito/ir/stree/tree.py b/devito/ir/stree/tree.py index 96e498396d2..ec894073624 100644 --- a/devito/ir/stree/tree.py +++ b/devito/ir/stree/tree.py @@ -1,6 +1,7 @@ +import sympy from anytree import ContStyle, NodeMixin, PostOrderIter, RenderTree -from devito.ir.support import PrefetchUpdate, WithLock +from devito.ir.support import PrefetchUpdate, WithLock, strided_factor __all__ = [ "NodeConditional", @@ -64,9 +65,18 @@ def dim(self): @property def limits(self): - return (self.dim.symbolic_min + self.interval.lower, - self.dim.symbolic_max + self.interval.upper, - self.dim.symbolic_incr) + _min = self.dim.symbolic_min + self.interval.lower + _max = self.dim.symbolic_max + self.interval.upper + + # A Strided Dimension steps `factor` points at a time in place of a + # `dim % factor == 0` guard. The min must be advanced to the first + # point on the subsampling grid, as `_min` is not necessarily a + # multiple of `factor` + factor = strided_factor(self.properties or ()) + if factor is not None: + return (_round_up(_min, factor), _max, factor) + + return (_min, _max, self.dim.symbolic_incr) @property def direction(self): @@ -146,3 +156,20 @@ def __repr_render__(self): def render(stree): return RenderTree(stree, style=ContStyle()).by_attr('__repr_render__') + + +def _round_up(expr, factor): + """ + The smallest multiple of `factor` that is `>= expr`, for `expr >= 0`. + + Written as `expr + (factor - expr%factor)%factor` rather than the usual + `((expr + factor - 1)/factor)*factor`: the latter is rewritten by + `IntDiv.__mul__` into plain `expr + factor - 1`, and the otherwise + equivalent `expr + (-expr)%factor` would rely on a floored modulo, whereas + C truncates towards zero and would place the loop min *below* `expr`. + `evaluate=False` keeps SymPy from folding the double modulo back. + """ + rem = sympy.Mod(expr, factor) + shift = sympy.Mod(sympy.Add(factor, -rem, evaluate=False), factor, evaluate=False) + + return sympy.Add(expr, shift, evaluate=False) diff --git a/devito/ir/support/properties.py b/devito/ir/support/properties.py index f2675ae8b1f..952ab0be040 100644 --- a/devito/ir/support/properties.py +++ b/devito/ir/support/properties.py @@ -117,6 +117,37 @@ def __init__(self, name, val=None): # Bundles PARALLELS = {PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT} +STRIDED = 'strided' +"""Name of the parameterized Strided Property; see `Strided`.""" + + +def Strided(factor): + """ + Property marking a Dimension iterated with a non-unitary step `factor`, a + strengthened form of a factor-based ConditionalDimension guard. Instead of + visiting every point and testing `d % factor == 0`, the loop steps `factor` + points at a time, anchored so that the first iteration satisfies the guard. + + Only legal on a PARALLEL Dimension, where no dependence is carried along the + Dimension and hence skipping the guarded-out points cannot drop a required + ordering. + + The `factor` is carried as the Property value, so two Strided properties + with different factors compare unequal and therefore do not survive a + fusion that merges them. + """ + return Property(STRIDED, factor) + + +def strided_factor(properties): + """ + The step of the Strided Property in `properties`, or None if absent. + """ + for i in properties: + if i.name == STRIDED: + return i.val + return None + def normalize_properties(*args): if not args: @@ -147,6 +178,13 @@ def normalize_properties(*args): if not all(i in p for p in args): drop.add(i) + # A Strided property must be dropped unless every `args` entry carries the + # very same one; fusing an unguarded Cluster into a strided loop would + # otherwise silently skip iterations + strided = {frozenset(i for i in p if i.name == STRIDED) for p in args} + if len(strided) > 1: + drop.update(*strided) + properties = set() for p in args: properties.update(p - drop) diff --git a/devito/passes/clusters/__init__.py b/devito/passes/clusters/__init__.py index bcdb8a4bd64..3a6cb0d0600 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 00000000000..b613443d5f8 --- /dev/null +++ b/devito/passes/clusters/subsampling.py @@ -0,0 +1,95 @@ +from devito.ir.clusters import Queue +from devito.ir.support import GuardFactor, Strided +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 rewrite is recorded as a `Strided` property on `d`; the stride and the + anchoring of the loop min are applied by `NodeIteration.limits`. The + iteration space itself is left keyed on `d`, so the runtime `x_m`/`x_M` + values, the domain decomposition and any non-subsampled access `f[x + 1]` + are all unaffected. + + 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 + + processed = [] + for c in clusters: + properties = c.properties.add(d, Strided(guard.d.symbolic_factor)) + processed.append(c.rebuild(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) From 3cf66eb02b0fbc40bb8e78baf7959d965872ed8d Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 12:21:11 -0400 Subject: [PATCH 02/11] tests: Add coverage for strided subsampling Codegen, numerics against numpy slicing, loop-min anchoring across offset bounds, and the cases that must keep a pointwise guard: a loop nest shared with an unguarded equation, a sequential Dimension, an extra `condition`, and non-affine accesses. --- tests/test_subsampling.py | 422 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 422 insertions(+) create mode 100644 tests/test_subsampling.py diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py new file mode 100644 index 00000000000..b5cea667074 --- /dev/null +++ b/tests/test_subsampling.py @@ -0,0 +1,422 @@ +import numpy as np +import pytest +from sympy import Or + +from devito import ( + CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, TimeFunction, + switchconfig +) +from devito.ir.iet import Conditional, FindNodes, Iteration + + +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 lowers to a loop stride, not a guard.""" + + def test_codegen(self): + _, f, g, cd = subsampled((11, 11), 2) + op = Operator(Eq(g, f)) + + # No guard... + assert not FindNodes(Conditional).visit(op) + # ...the parent Dimension's loop is strided instead + assert iteration(op, 'x').limits[2] == cd.symbolic_factor + + @pytest.mark.parametrize('factor', [2, 3, 4, 5]) + def test_factor_in_loop_header(self, factor): + _, f, g, cd = subsampled((16, 8), factor) + op = Operator(Eq(g, f)) + + assert f'+= {cd.symbolic_factor.name}' in str(op.ccode) + assert not FindNodes(Conditional).visit(op) + + 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, 'y').limits[2] == ys.symbolic_factor + 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) + assert iteration(op, 'x').limits[2] == xs.symbolic_factor + assert iteration(op, 'y').limits[2] == ys.symbolic_factor + + 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_anchored_min(self): + """The loop min is rounded up onto the subsampling grid.""" + _, f, g, cd = subsampled((12, 12), 3) + op = Operator(Eq(g, f)) + + _min = iteration(op, 'x').limits[0] + # Not simply `x_m`: it carries the round-up to the next multiple + assert _min != f.grid.dimensions[0].symbolic_min + assert cd.symbolic_factor in _min.free_symbols + + +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) + + written = {i//factor for i in range(x_m, x_M + 1) if i % factor == 0} + 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) + + 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]) From dcddc17b9d82ab146f62a28fcfe3949613ba68e7 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 12:21:24 -0400 Subject: [PATCH 03/11] docs: Add a subgrids tutorial Walks through declaring a subsampled Function and the code generated for it: striding the outer Dimension, the innermost one, and several at once, the anchoring of the loop min, and the four cases that fall back to a pointwise guard. Every case prints its generated C. --- examples/userapi/10_subgrids.ipynb | 987 +++++++++++++++++++++++++++++ 1 file changed, 987 insertions(+) create mode 100644 examples/userapi/10_subgrids.ipynb diff --git a/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb new file mode 100644 index 00000000000..cf9b49c1951 --- /dev/null +++ b/examples/userapi/10_subgrids.ipynb @@ -0,0 +1,987 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Subgrids: integer subsampling\n", + "\n", + "A *subgrid* is a coarser grid whose points are a subset of a parent `Grid`'s points.\n", + "Devito expresses this with a `ConditionalDimension` carrying a `factor`: a `Function`\n", + "indexed by it holds every `factor`-th point of its parent `Dimension`.\n", + "\n", + "This tutorial covers **integer subsampling**, where subgrid points coincide exactly\n", + "with parent grid points, so no interpolation is involved. Devito compiles it to a\n", + "**strided loop** -- `x += factor` -- rather than visiting every point under an\n", + "`if (x%factor == 0)` guard.\n", + "\n", + "Every section below prints the generated C so you can see exactly what each case\n", + "produces.\n", + "\n", + "If you are new to `ConditionalDimension`, read\n", + "[05_conditional_dimension.ipynb](05_conditional_dimension.ipynb) first." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Declaring a subsampled Function\n", + "\n", + "We take a 12x12 parent grid and subsample `x` by a factor of 3. The\n", + "`ConditionalDimension` `xs` stands for \"every third `x`\"." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:07.909893Z", + "iopub.status.busy": "2026-10-20T16:33:07.909609Z", + "iopub.status.idle": "2026-10-20T16:33:09.113178Z", + "shell.execute_reply": "2026-10-20T16:33:09.111976Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "(x, 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", + " TimeFunction)\n", + "from devito.ir.iet import FindNodes, Iteration\n", + "\n", + "grid = Grid(shape=(12, 12))\n", + "x, y = grid.dimensions\n", + "\n", + "xs = ConditionalDimension(name='xs', parent=x, factor=3)\n", + "xs.parent, xs.factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Function` over `(xs, y)` is coarse along `x`, full along `y`. Its shape must be\n", + "given explicitly: 12 points subsampled by 3 gives 4." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.173126Z", + "iopub.status.busy": "2026-10-20T16:33:09.172143Z", + "iopub.status.idle": "2026-10-20T16:33:09.195648Z", + "shell.execute_reply": "2026-10-20T16:33:09.193868Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((12, 12), (4, 12))" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "u = Function(name='u', grid=grid) # parent grid\n", + "usub = Function(name='usub', grid=grid, dimensions=(xs, y),\n", + " shape=(4, 12), space_order=0) # subgrid\n", + "\n", + "u.shape, usub.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`xs.symbolic_factor` is the symbol the generated code uses for the factor; it is\n", + "passed to the kernel as a run-time argument." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.201437Z", + "iopub.status.busy": "2026-10-20T16:33:09.200777Z", + "iopub.status.idle": "2026-10-20T16:33:09.212124Z", + "shell.execute_reply": "2026-10-20T16:33:09.210475Z" + } + }, + "outputs": [ + { + "data": { + "text/latex": [ + "$\\displaystyle xsf$" + ], + "text/plain": [ + "xsf" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "xs.symbolic_factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Throughout, this helper prints the loop nest of an `Operator` -- the part we care\n", + "about -- rather than the whole kernel." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.217649Z", + "iopub.status.busy": "2026-10-20T16:33:09.217140Z", + "iopub.status.idle": "2026-10-20T16:33:09.226240Z", + "shell.execute_reply": "2026-10-20T16:33:09.224517Z" + } + }, + "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": [ + "## The generated loop\n", + "\n", + "`Eq(usub, u)` reads the parent grid and writes the subgrid." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.232008Z", + "iopub.status.busy": "2026-10-20T16:33:09.231407Z", + "iopub.status.idle": "2026-10-20T16:33:09.403808Z", + "shell.execute_reply": "2026-10-20T16:33:09.401929Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m + (xsf - (x_m)%(xsf))%(xsf); x <= x_M; x += xsf)\n", + "{\n", + " #pragma omp simd aligned(u,usub:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " usub[x / xsf][y] = u[x + 1][y + 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": [ + "Two features of the `x` loop header matter:\n", + "\n", + "- the increment is `x += xsf`, so only one point in three is visited -- there is no\n", + " `if` anywhere in the nest;\n", + "- the minimum is `x_m + (xsf - x_m%xsf)%xsf` rather than `x_m`.\n", + "\n", + "That second expression is the **anchor**. A subgrid point is one where\n", + "`x % factor == 0`, and `x_m` is not necessarily such a point, so iteration must start\n", + "at the first subgrid point at or after `x_m`.\n", + "\n", + "The body carries the index conversion `usub[x / xsf]`, mapping the parent coordinate\n", + "to a subgrid coordinate." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.409079Z", + "iopub.status.busy": "2026-10-20T16:33:09.408657Z", + "iopub.status.idle": "2026-10-20T16:33:09.417933Z", + "shell.execute_reply": "2026-10-20T16:33:09.416250Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'x': xsf, 'y': 1}" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "steps(op)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Checking the result\n", + "\n", + "`usub` should hold every third row of `u`." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.423552Z", + "iopub.status.busy": "2026-10-20T16:33:09.422941Z", + "iopub.status.idle": "2026-10-20T16:33:09.587287Z", + "shell.execute_reply": "2026-10-20T16:33:09.585623Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "Data([[ 0., 1., 2., 3.],\n", + " [ 36., 37., 38., 39.],\n", + " [ 72., 73., 74., 75.],\n", + " [108., 109., 110., 111.]], dtype=float32)" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = np.arange(144).reshape(12, 12)\n", + "usub.data[:] = 0.\n", + "Operator(Eq(usub, u)).apply()\n", + "\n", + "assert np.array_equal(usub.data, u.data[::3])\n", + "usub.data[:, :4]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### The anchor in action\n", + "\n", + "Run with `x_m=2`, so iteration no longer starts on a multiple of 3. The first subgrid\n", + "point is `x=3`, so the slot for `x=0` is never written." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.592039Z", + "iopub.status.busy": "2026-10-20T16:33:09.591836Z", + "iopub.status.idle": "2026-10-20T16:33:09.710261Z", + "shell.execute_reply": "2026-10-20T16:33:09.708734Z" + } + }, + "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": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = 1.\n", + "usub.data[:] = -1. # sentinel: never written\n", + "Operator(Eq(usub, u)).apply(x_m=2)\n", + "\n", + "usub.data[:, 0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The leading `-1.` is the untouched slot. Had iteration started at `x_m=2` instead\n", + "of `x=3`, the wrong points would have been written." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Subsampling the innermost Dimension\n", + "\n", + "Subsampling `y`, the innermost (vectorized) `Dimension`, is strided too. The\n", + "alternative -- a guard inside the SIMD loop -- would mean a predicated store on every\n", + "iteration, which is worse." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.715674Z", + "iopub.status.busy": "2026-10-20T16:33:09.715228Z", + "iopub.status.idle": "2026-10-20T16:33:09.840592Z", + "shell.execute_reply": "2026-10-20T16:33:09.839155Z" + } + }, + "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 y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", + " {\n", + " usub_y[x][y / ysf] = u[x + 1][y + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "ys = ConditionalDimension(name='ys', parent=y, factor=2)\n", + "usub_y = Function(name='usub_y', grid=grid, dimensions=(x, ys),\n", + " shape=(12, 6), space_order=0)\n", + "\n", + "op_inner = Operator(Eq(usub_y, u))\n", + "show(op_inner)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.845584Z", + "iopub.status.busy": "2026-10-20T16:33:09.844992Z", + "iopub.status.idle": "2026-10-20T16:33:09.855907Z", + "shell.execute_reply": "2026-10-20T16:33:09.854246Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'x': 1, 'y': ysf}" + ] + }, + "execution_count": 10, + "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": 11, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.861161Z", + "iopub.status.busy": "2026-10-20T16:33:09.860659Z", + "iopub.status.idle": "2026-10-20T16:33:09.992530Z", + "shell.execute_reply": "2026-10-20T16:33:09.991309Z" + } + }, + "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": 11, + "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": [ + "## Subsampling several Dimensions\n", + "\n", + "Each subsampled `Dimension` gets its own stride, independently." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:09.996633Z", + "iopub.status.busy": "2026-10-20T16:33:09.996417Z", + "iopub.status.idle": "2026-10-20T16:33:10.120737Z", + "shell.execute_reply": "2026-10-20T16:33:10.119405Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m + (xs2f - (x_m)%(xs2f))%(xs2f); x <= x_M; x += xs2f)\n", + "{\n", + " #pragma omp simd aligned(u,usub_xy:64)\n", + " for (int y = y_m + (ys2f - (y_m)%(ys2f))%(ys2f); y <= y_M; y += ys2f)\n", + " {\n", + " usub_xy[x / xs2f][y / ys2f] = u[x + 1][y + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "xs2 = ConditionalDimension(name='xs2', parent=x, factor=3)\n", + "ys2 = ConditionalDimension(name='ys2', parent=y, factor=2)\n", + "usub_xy = Function(name='usub_xy', grid=grid, dimensions=(xs2, ys2),\n", + " shape=(4, 6), space_order=0)\n", + "\n", + "op_both = Operator(Eq(usub_xy, u))\n", + "show(op_both)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.125137Z", + "iopub.status.busy": "2026-10-20T16:33:10.124963Z", + "iopub.status.idle": "2026-10-20T16:33:10.258398Z", + "shell.execute_reply": "2026-10-20T16:33:10.256939Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, + { + "data": { + "text/plain": [ + "{'x': xs2f, 'y': ys2f}" + ] + }, + "execution_count": 13, + "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": 14, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.263400Z", + "iopub.status.busy": "2026-10-20T16:33:10.263159Z", + "iopub.status.idle": "2026-10-20T16:33:10.391235Z", + "shell.execute_reply": "2026-10-20T16:33:10.389486Z" + } + }, + "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)%(xs1f) == 0)\n", + " {\n", + " g[x / xs1f] = f[x + 1];\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "grid1 = Grid(shape=(12,))\n", + "(x1,) = grid1.dimensions\n", + "xs1 = ConditionalDimension(name='xs1', parent=x1, factor=4)\n", + "\n", + "f = Function(name='f', grid=grid1)\n", + "g = Function(name='g', grid=grid1, dimensions=(xs1,), shape=(3,), space_order=0)\n", + "\n", + "op_shared = Operator([Eq(f, 1), Eq(g, f)])\n", + "show(op_shared)" + ] + }, + { + "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": 15, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.395695Z", + "iopub.status.busy": "2026-10-20T16:33:10.395520Z", + "iopub.status.idle": "2026-10-20T16:33:10.517016Z", + "shell.execute_reply": "2026-10-20T16:33:10.515669Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "START(section0)\n", + "for (int x = x_m; x <= x_M; x += 1)\n", + "{\n", + " if ((x)%(xs3f) == 0)\n", + " {\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " g2[x / xs3f][y] = f2[x + 1][y + 1] + g2[(x / xs3f) - 1][y];\n", + " }\n", + " }\n", + "}\n", + "STOP(section0,timers)\n" + ] + } + ], + "source": [ + "# NBVAL_IGNORE_OUTPUT\n", + "grid2 = Grid(shape=(12, 12))\n", + "x2, y2 = grid2.dimensions\n", + "xs3 = ConditionalDimension(name='xs3', parent=x2, factor=3)\n", + "\n", + "f2 = Function(name='f2', grid=grid2)\n", + "g2 = Function(name='g2', grid=grid2, dimensions=(xs3, y2), shape=(4, 12),\n", + " space_order=0)\n", + "\n", + "op_seq = Operator(Eq(g2, g2.subs(xs3, xs3-1) + f2))\n", + "show(op_seq)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Time subsampling is the everyday instance of this: the time loop is sequential, so\n", + "a saved-wavefield `Eq` keeps its guard." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.522400Z", + "iopub.status.busy": "2026-10-20T16:33:10.521906Z", + "iopub.status.idle": "2026-10-20T16:33:10.712059Z", + "shell.execute_reply": "2026-10-20T16:33:10.710256Z" + } + }, + "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": 17, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.717170Z", + "iopub.status.busy": "2026-10-20T16:33:10.716766Z", + "iopub.status.idle": "2026-10-20T16:33:10.757309Z", + "shell.execute_reply": "2026-10-20T16:33:10.755861Z" + } + }, + "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": 17, + "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": 18, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.762722Z", + "iopub.status.busy": "2026-10-20T16:33:10.762046Z", + "iopub.status.idle": "2026-10-20T16:33:10.986019Z", + "shell.execute_reply": "2026-10-20T16:33:10.984256Z" + } + }, + "outputs": [ + { + "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", + "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", + "op_cond = Operator(Eq(w, u))\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": 19, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:33:10.991639Z", + "iopub.status.busy": "2026-10-20T16:33:10.991028Z", + "iopub.status.idle": "2026-10-20T16:33:11.153763Z", + "shell.execute_reply": "2026-10-20T16:33:11.152805Z" + } + }, + "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", + " #pragma omp simd aligned(gi,idx,u:64)\n", + " for (int y = y_m; y <= y_M; y += 1)\n", + " {\n", + " int r0 = idx[x + 1][y + 1];\n", + " gi[x / xsf][y] = u[r0 + 1][y + 1];\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=grid, dimensions=(xs, y), shape=(4, 12),\n", + " space_order=0)\n", + "\n", + "op_indirect = Operator(Eq(gi, u[idx[x, y], y]))\n", + "show(op_indirect)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "Subsampling along a `Dimension` is declared with `ConditionalDimension(factor=n)` plus\n", + "a `Function` defined over it. Devito then generates, for that `Dimension`:\n", + "\n", + "| | strided loop | guarded loop |\n", + "|---|---|---|\n", + "| loop header | `x += factor`, anchored | `x += 1` |\n", + "| body | unconditional | wrapped in `if (x%factor == 0)` |\n", + "| iterations | one in `factor` | all |\n", + "| used when | the Dimension is parallel and accesses are affine | otherwise |\n", + "\n", + "The strided form is the default, including for the innermost (vectorized)\n", + "`Dimension` and for several subsampled `Dimension`s at once. The guarded form is the\n", + "fallback for the four cases above: a loop nest shared with an unguarded equation, a\n", + "sequential `Dimension` such as time, an extra `condition`, or non-affine accesses.\n", + "Both compute the same thing.\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": "devito", + "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 +} From a9ac791ca9abb484ccdcc261d9fec0a9860e28e0 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 13:07:26 -0400 Subject: [PATCH 04/11] api: Add SubGrid for integer subsampling A coarser Grid whose points are a subset of a parent Grid's points: along each Dimension subsampled by `factor`, point `i` sits at parent index `i*factor`. sub = SubGrid(grid, factor=2) f = Function(name='f', grid=sub) Dimensions, shape and spacing are all derived, so a Function on a SubGrid needs neither `shape` nor `dimensions`; previously both had to be worked out by hand. The Dimensions are ConditionalDimensions, so the strided lowering applies unchanged, and a factor of 1 reuses the parent Dimension untouched. The parent's stepping Dimension is shared rather than derived afresh: a SubGrid coarsens space only, so a TimeFunction on it must step over the same buffer indices as one on the parent. Non-integer factors are rejected for now -- they place SubGrid points between parent points, which requires interpolation. --- devito/types/grid.py | 147 +++++++++- examples/userapi/10_subgrids.ipynb | 421 ++++++++++++++++++----------- tests/test_subsampling.py | 180 +++++++++++- 3 files changed, 584 insertions(+), 164 deletions(-) diff --git a/devito/types/grid.py b/devito/types/grid.py index 378d2543b44..ea3aaf0df0c 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -11,17 +11,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 +428,145 @@ 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__}`") + if isinstance(parent, SubGrid): + # Would require a ConditionalDimension to parent another one + raise ValueError("Cannot subsample a SubGrid; derive from " + "`parent.parent` with a combined factor instead") + + factors = self.__factors_setup__(parent, factor) + + 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}` is not " + "supported; SubGrid points must coincide with " + "parent Grid points") + 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) + + 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 index cf9b49c1951..a90f27f9c92 100644 --- a/examples/userapi/10_subgrids.ipynb +++ b/examples/userapi/10_subgrids.ipynb @@ -7,29 +7,26 @@ "# Subgrids: integer subsampling\n", "\n", "A *subgrid* is a coarser grid whose points are a subset of a parent `Grid`'s points.\n", - "Devito expresses this with a `ConditionalDimension` carrying a `factor`: a `Function`\n", - "indexed by it holds every `factor`-th point of its parent `Dimension`.\n", + "`SubGrid` builds one: along each `Dimension` subsampled by `factor`, it keeps every\n", + "`factor`-th point of the parent, so subgrid point `i` sits at parent index\n", + "`i*factor`.\n", "\n", - "This tutorial covers **integer subsampling**, where subgrid points coincide exactly\n", - "with parent grid points, so no interpolation is involved. Devito compiles it to a\n", - "**strided loop** -- `x += factor` -- rather than visiting every point under an\n", - "`if (x%factor == 0)` guard.\n", + "Because the points coincide with parent points, no interpolation is involved. Devito\n", + "compiles a subgrid write to a **strided loop** -- `x += factor` -- rather than\n", + "visiting every point under an `if (x%factor == 0)` guard.\n", "\n", "Every section below prints the generated C so you can see exactly what each case\n", - "produces.\n", - "\n", - "If you are new to `ConditionalDimension`, read\n", - "[05_conditional_dimension.ipynb](05_conditional_dimension.ipynb) first." + "produces." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "## Declaring a subsampled Function\n", + "## Creating a SubGrid\n", "\n", - "We take a 12x12 parent grid and subsample `x` by a factor of 3. The\n", - "`ConditionalDimension` `xs` stands for \"every third `x`\"." + "`SubGrid(parent, factor)` derives everything from the parent: the dimensions, the\n", + "shape, and the physical spacing." ] }, { @@ -37,17 +34,17 @@ "execution_count": 1, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:07.909893Z", - "iopub.status.busy": "2026-10-20T16:33:07.909609Z", - "iopub.status.idle": "2026-10-20T16:33:09.113178Z", - "shell.execute_reply": "2026-10-20T16:33:09.111976Z" + "iopub.execute_input": "2026-10-20T16:54:53.254537Z", + "iopub.status.busy": "2026-10-20T16:54:53.253925Z", + "iopub.status.idle": "2026-10-20T16:54:54.490124Z", + "shell.execute_reply": "2026-10-20T16:54:54.488467Z" } }, "outputs": [ { "data": { "text/plain": [ - "(x, 3)" + "SubGrid[shape=(4, 4), factor=(3, 3)]" ] }, "execution_count": 1, @@ -58,22 +55,22 @@ "source": [ "import numpy as np\n", "from devito import (CondEq, ConditionalDimension, Eq, Function, Grid, Operator,\n", - " TimeFunction)\n", + " SubGrid, TimeFunction)\n", "from devito.ir.iet import FindNodes, Iteration\n", "\n", - "grid = Grid(shape=(12, 12))\n", + "grid = Grid(shape=(12, 12), extent=(11., 11.))\n", "x, y = grid.dimensions\n", + "sub = SubGrid(grid, factor=3)\n", "\n", - "xs = ConditionalDimension(name='xs', parent=x, factor=3)\n", - "xs.parent, xs.factor" + "sub" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "A `Function` over `(xs, y)` is coarse along `x`, full along `y`. Its shape must be\n", - "given explicitly: 12 points subsampled by 3 gives 4." + "The shape is `ceil(12/3) = 4` along each axis, and the spacing is three times\n", + "the parent's -- the subgrid spans the same physical extent with fewer points." ] }, { @@ -81,68 +78,177 @@ "execution_count": 2, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.173126Z", - "iopub.status.busy": "2026-10-20T16:33:09.172143Z", - "iopub.status.idle": "2026-10-20T16:33:09.195648Z", - "shell.execute_reply": "2026-10-20T16:33:09.193868Z" + "iopub.execute_input": "2026-10-20T16:54:54.545852Z", + "iopub.status.busy": "2026-10-20T16:54:54.544993Z", + "iopub.status.idle": "2026-10-20T16:54:54.555564Z", + "shell.execute_reply": "2026-10-20T16:54:54.554030Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "subgrid shape (4, 4) spacing (3.0, 3.0)\n", + "parent shape (12, 12) spacing (1.0, 1.0)\n" + ] + } + ], + "source": [ + "print('subgrid shape ', sub.shape, ' spacing', tuple(float(s) for s in sub.spacing))\n", + "print('parent shape ', grid.shape, ' spacing', tuple(float(s) for s in grid.spacing))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Its dimensions are `ConditionalDimension`s derived from the parent's, which is\n", + "what carries the subsampling factor into the generated code." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:54:54.561212Z", + "iopub.status.busy": "2026-10-20T16:54:54.560550Z", + "iopub.status.idle": "2026-10-20T16:54:54.570827Z", + "shell.execute_reply": "2026-10-20T16:54:54.569162Z" } }, "outputs": [ { "data": { "text/plain": [ - "((12, 12), (4, 12))" + "((xs, ys), (3, 3))" ] }, - "execution_count": 2, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "u = Function(name='u', grid=grid) # parent grid\n", - "usub = Function(name='usub', grid=grid, dimensions=(xs, y),\n", - " shape=(4, 12), space_order=0) # subgrid\n", + "sub.dimensions, sub.factor" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A `Function` on a `SubGrid` needs neither `shape` nor `dimensions` -- unlike the\n", + "hand-rolled equivalent, where both must be worked out by hand." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:54:54.576697Z", + "iopub.status.busy": "2026-10-20T16:54:54.576122Z", + "iopub.status.idle": "2026-10-20T16:54:54.595157Z", + "shell.execute_reply": "2026-10-20T16:54:54.593905Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((12, 12), (4, 4), (xs, ys))" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "u = Function(name='u', grid=grid) # parent grid\n", + "usub = Function(name='usub', grid=sub, space_order=0) # subgrid\n", "\n", - "u.shape, usub.shape" + "u.shape, usub.shape, usub.dimensions" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "`xs.symbolic_factor` is the symbol the generated code uses for the factor; it is\n", - "passed to the kernel as a run-time argument." + "A factor per `Dimension` subsamples each axis independently; `1` leaves an axis\n", + "untouched." ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 5, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.201437Z", - "iopub.status.busy": "2026-10-20T16:33:09.200777Z", - "iopub.status.idle": "2026-10-20T16:33:09.212124Z", - "shell.execute_reply": "2026-10-20T16:33:09.210475Z" + "iopub.execute_input": "2026-10-20T16:54:54.600146Z", + "iopub.status.busy": "2026-10-20T16:54:54.599726Z", + "iopub.status.idle": "2026-10-20T16:54:54.611723Z", + "shell.execute_reply": "2026-10-20T16:54:54.609903Z" } }, "outputs": [ { "data": { - "text/latex": [ - "$\\displaystyle xsf$" - ], "text/plain": [ - "xsf" + "((4, 12), (xs, y))" ] }, - "execution_count": 3, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "mixed = SubGrid(grid, factor=(3, 1))\n", + "mixed.shape, mixed.dimensions" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Equivalent, by hand\n", + "\n", + "`SubGrid` is a convenience over `ConditionalDimension`, which can still be used\n", + "directly when the dimensions need individual treatment. The two are equivalent; the\n", + "manual form just requires computing the shape yourself." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T16:54:54.617414Z", + "iopub.status.busy": "2026-10-20T16:54:54.616914Z", + "iopub.status.idle": "2026-10-20T16:54:54.629162Z", + "shell.execute_reply": "2026-10-20T16:54:54.627654Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "False" + ] + }, + "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "xs.symbolic_factor" + "xs_manual = ConditionalDimension(name='xs_m', parent=grid.dimensions[0], factor=3)\n", + "f_manual = Function(name='f_manual', grid=grid,\n", + " dimensions=(xs_manual, grid.dimensions[1]),\n", + " shape=(4, 12), space_order=0) # shape by hand\n", + "\n", + "f_manual.shape == usub.shape" ] }, { @@ -155,13 +261,13 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 7, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.217649Z", - "iopub.status.busy": "2026-10-20T16:33:09.217140Z", - "iopub.status.idle": "2026-10-20T16:33:09.226240Z", - "shell.execute_reply": "2026-10-20T16:33:09.224517Z" + "iopub.execute_input": "2026-10-20T16:54:54.634508Z", + "iopub.status.busy": "2026-10-20T16:54:54.634002Z", + "iopub.status.idle": "2026-10-20T16:54:54.642529Z", + "shell.execute_reply": "2026-10-20T16:54:54.640851Z" } }, "outputs": [], @@ -187,13 +293,13 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 8, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.232008Z", - "iopub.status.busy": "2026-10-20T16:33:09.231407Z", - "iopub.status.idle": "2026-10-20T16:33:09.403808Z", - "shell.execute_reply": "2026-10-20T16:33:09.401929Z" + "iopub.execute_input": "2026-10-20T16:54:54.648073Z", + "iopub.status.busy": "2026-10-20T16:54:54.647463Z", + "iopub.status.idle": "2026-10-20T16:54:54.825256Z", + "shell.execute_reply": "2026-10-20T16:54:54.823521Z" } }, "outputs": [ @@ -205,9 +311,9 @@ "for (int x = x_m + (xsf - (x_m)%(xsf))%(xsf); x <= x_M; x += xsf)\n", "{\n", " #pragma omp simd aligned(u,usub:64)\n", - " for (int y = y_m; y <= y_M; y += 1)\n", + " for (int y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", " {\n", - " usub[x / xsf][y] = u[x + 1][y + 1];\n", + " usub[x / xsf][y / ysf] = u[x + 1][y + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -240,23 +346,23 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 9, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.409079Z", - "iopub.status.busy": "2026-10-20T16:33:09.408657Z", - "iopub.status.idle": "2026-10-20T16:33:09.417933Z", - "shell.execute_reply": "2026-10-20T16:33:09.416250Z" + "iopub.execute_input": "2026-10-20T16:54:54.830512Z", + "iopub.status.busy": "2026-10-20T16:54:54.830098Z", + "iopub.status.idle": "2026-10-20T16:54:54.840469Z", + "shell.execute_reply": "2026-10-20T16:54:54.838656Z" } }, "outputs": [ { "data": { "text/plain": [ - "{'x': xsf, 'y': 1}" + "{'x': xsf, 'y': ysf}" ] }, - "execution_count": 6, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -276,13 +382,13 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 10, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.423552Z", - "iopub.status.busy": "2026-10-20T16:33:09.422941Z", - "iopub.status.idle": "2026-10-20T16:33:09.587287Z", - "shell.execute_reply": "2026-10-20T16:33:09.585623Z" + "iopub.execute_input": "2026-10-20T16:54:54.846210Z", + "iopub.status.busy": "2026-10-20T16:54:54.845610Z", + "iopub.status.idle": "2026-10-20T16:54:55.010474Z", + "shell.execute_reply": "2026-10-20T16:54:55.009286Z" } }, "outputs": [ @@ -296,13 +402,13 @@ { "data": { "text/plain": [ - "Data([[ 0., 1., 2., 3.],\n", - " [ 36., 37., 38., 39.],\n", - " [ 72., 73., 74., 75.],\n", - " [108., 109., 110., 111.]], dtype=float32)" + "Data([[ 0., 3., 6., 9.],\n", + " [ 36., 39., 42., 45.],\n", + " [ 72., 75., 78., 81.],\n", + " [108., 111., 114., 117.]], dtype=float32)" ] }, - "execution_count": 7, + "execution_count": 10, "metadata": {}, "output_type": "execute_result" } @@ -313,8 +419,8 @@ "usub.data[:] = 0.\n", "Operator(Eq(usub, u)).apply()\n", "\n", - "assert np.array_equal(usub.data, u.data[::3])\n", - "usub.data[:, :4]" + "assert np.array_equal(usub.data, u.data[::3, ::3])\n", + "usub.data" ] }, { @@ -329,13 +435,13 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 11, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.592039Z", - "iopub.status.busy": "2026-10-20T16:33:09.591836Z", - "iopub.status.idle": "2026-10-20T16:33:09.710261Z", - "shell.execute_reply": "2026-10-20T16:33:09.708734Z" + "iopub.execute_input": "2026-10-20T16:54:55.014690Z", + "iopub.status.busy": "2026-10-20T16:54:55.014442Z", + "iopub.status.idle": "2026-10-20T16:54:55.144762Z", + "shell.execute_reply": "2026-10-20T16:54:55.143244Z" } }, "outputs": [ @@ -352,7 +458,7 @@ "Data([-1., 1., 1., 1.], dtype=float32)" ] }, - "execution_count": 8, + "execution_count": 11, "metadata": {}, "output_type": "execute_result" } @@ -387,13 +493,13 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 12, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.715674Z", - "iopub.status.busy": "2026-10-20T16:33:09.715228Z", - "iopub.status.idle": "2026-10-20T16:33:09.840592Z", - "shell.execute_reply": "2026-10-20T16:33:09.839155Z" + "iopub.execute_input": "2026-10-20T16:54:55.150500Z", + "iopub.status.busy": "2026-10-20T16:54:55.149948Z", + "iopub.status.idle": "2026-10-20T16:54:55.270974Z", + "shell.execute_reply": "2026-10-20T16:54:55.269737Z" } }, "outputs": [ @@ -416,9 +522,8 @@ ], "source": [ "# NBVAL_IGNORE_OUTPUT\n", - "ys = ConditionalDimension(name='ys', parent=y, factor=2)\n", - "usub_y = Function(name='usub_y', grid=grid, dimensions=(x, ys),\n", - " shape=(12, 6), space_order=0)\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)" @@ -426,13 +531,13 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 13, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.845584Z", - "iopub.status.busy": "2026-10-20T16:33:09.844992Z", - "iopub.status.idle": "2026-10-20T16:33:09.855907Z", - "shell.execute_reply": "2026-10-20T16:33:09.854246Z" + "iopub.execute_input": "2026-10-20T16:54:55.276046Z", + "iopub.status.busy": "2026-10-20T16:54:55.275462Z", + "iopub.status.idle": "2026-10-20T16:54:55.287010Z", + "shell.execute_reply": "2026-10-20T16:54:55.285334Z" } }, "outputs": [ @@ -442,7 +547,7 @@ "{'x': 1, 'y': ysf}" ] }, - "execution_count": 10, + "execution_count": 13, "metadata": {}, "output_type": "execute_result" } @@ -460,13 +565,13 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 14, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.861161Z", - "iopub.status.busy": "2026-10-20T16:33:09.860659Z", - "iopub.status.idle": "2026-10-20T16:33:09.992530Z", - "shell.execute_reply": "2026-10-20T16:33:09.991309Z" + "iopub.execute_input": "2026-10-20T16:54:55.292500Z", + "iopub.status.busy": "2026-10-20T16:54:55.291964Z", + "iopub.status.idle": "2026-10-20T16:54:55.428150Z", + "shell.execute_reply": "2026-10-20T16:54:55.426871Z" } }, "outputs": [ @@ -485,7 +590,7 @@ " [24., 26., 28., 30., 32., 34.]], dtype=float32)" ] }, - "execution_count": 11, + "execution_count": 14, "metadata": {}, "output_type": "execute_result" } @@ -510,13 +615,13 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 15, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:09.996633Z", - "iopub.status.busy": "2026-10-20T16:33:09.996417Z", - "iopub.status.idle": "2026-10-20T16:33:10.120737Z", - "shell.execute_reply": "2026-10-20T16:33:10.119405Z" + "iopub.execute_input": "2026-10-20T16:54:55.433311Z", + "iopub.status.busy": "2026-10-20T16:54:55.432979Z", + "iopub.status.idle": "2026-10-20T16:54:55.548544Z", + "shell.execute_reply": "2026-10-20T16:54:55.547178Z" } }, "outputs": [ @@ -525,12 +630,12 @@ "output_type": "stream", "text": [ "START(section0)\n", - "for (int x = x_m + (xs2f - (x_m)%(xs2f))%(xs2f); x <= x_M; x += xs2f)\n", + "for (int x = x_m + (xsf - (x_m)%(xsf))%(xsf); x <= x_M; x += xsf)\n", "{\n", " #pragma omp simd aligned(u,usub_xy:64)\n", - " for (int y = y_m + (ys2f - (y_m)%(ys2f))%(ys2f); y <= y_M; y += ys2f)\n", + " for (int y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", " {\n", - " usub_xy[x / xs2f][y / ys2f] = u[x + 1][y + 1];\n", + " usub_xy[x / xsf][y / ysf] = u[x + 1][y + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -539,10 +644,8 @@ ], "source": [ "# NBVAL_IGNORE_OUTPUT\n", - "xs2 = ConditionalDimension(name='xs2', parent=x, factor=3)\n", - "ys2 = ConditionalDimension(name='ys2', parent=y, factor=2)\n", - "usub_xy = Function(name='usub_xy', grid=grid, dimensions=(xs2, ys2),\n", - " shape=(4, 6), space_order=0)\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)" @@ -550,13 +653,13 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 16, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.125137Z", - "iopub.status.busy": "2026-10-20T16:33:10.124963Z", - "iopub.status.idle": "2026-10-20T16:33:10.258398Z", - "shell.execute_reply": "2026-10-20T16:33:10.256939Z" + "iopub.execute_input": "2026-10-20T16:54:55.552823Z", + "iopub.status.busy": "2026-10-20T16:54:55.552658Z", + "iopub.status.idle": "2026-10-20T16:54:55.663895Z", + "shell.execute_reply": "2026-10-20T16:54:55.662572Z" } }, "outputs": [ @@ -570,10 +673,10 @@ { "data": { "text/plain": [ - "{'x': xs2f, 'y': ys2f}" + "{'x': xsf, 'y': ysf}" ] }, - "execution_count": 13, + "execution_count": 16, "metadata": {}, "output_type": "execute_result" } @@ -610,13 +713,13 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 17, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.263400Z", - "iopub.status.busy": "2026-10-20T16:33:10.263159Z", - "iopub.status.idle": "2026-10-20T16:33:10.391235Z", - "shell.execute_reply": "2026-10-20T16:33:10.389486Z" + "iopub.execute_input": "2026-10-20T16:54:55.669766Z", + "iopub.status.busy": "2026-10-20T16:54:55.669150Z", + "iopub.status.idle": "2026-10-20T16:54:55.802136Z", + "shell.execute_reply": "2026-10-20T16:54:55.800522Z" } }, "outputs": [ @@ -670,13 +773,13 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 18, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.395695Z", - "iopub.status.busy": "2026-10-20T16:33:10.395520Z", - "iopub.status.idle": "2026-10-20T16:33:10.517016Z", - "shell.execute_reply": "2026-10-20T16:33:10.515669Z" + "iopub.execute_input": "2026-10-20T16:54:55.806474Z", + "iopub.status.busy": "2026-10-20T16:54:55.806303Z", + "iopub.status.idle": "2026-10-20T16:54:55.926680Z", + "shell.execute_reply": "2026-10-20T16:54:55.925460Z" } }, "outputs": [ @@ -723,13 +826,13 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 19, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.522400Z", - "iopub.status.busy": "2026-10-20T16:33:10.521906Z", - "iopub.status.idle": "2026-10-20T16:33:10.712059Z", - "shell.execute_reply": "2026-10-20T16:33:10.710256Z" + "iopub.execute_input": "2026-10-20T16:54:55.931233Z", + "iopub.status.busy": "2026-10-20T16:54:55.931060Z", + "iopub.status.idle": "2026-10-20T16:54:56.100092Z", + "shell.execute_reply": "2026-10-20T16:54:56.098527Z" } }, "outputs": [ @@ -791,13 +894,13 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 20, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.717170Z", - "iopub.status.busy": "2026-10-20T16:33:10.716766Z", - "iopub.status.idle": "2026-10-20T16:33:10.757309Z", - "shell.execute_reply": "2026-10-20T16:33:10.755861Z" + "iopub.execute_input": "2026-10-20T16:54:56.104592Z", + "iopub.status.busy": "2026-10-20T16:54:56.104423Z", + "iopub.status.idle": "2026-10-20T16:54:56.136868Z", + "shell.execute_reply": "2026-10-20T16:54:56.135085Z" } }, "outputs": [ @@ -814,7 +917,7 @@ "[0.0, 4.0, 0.0]" ] }, - "execution_count": 17, + "execution_count": 20, "metadata": {}, "output_type": "execute_result" } @@ -842,13 +945,13 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 21, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.762722Z", - "iopub.status.busy": "2026-10-20T16:33:10.762046Z", - "iopub.status.idle": "2026-10-20T16:33:10.986019Z", - "shell.execute_reply": "2026-10-20T16:33:10.984256Z" + "iopub.execute_input": "2026-10-20T16:54:56.142381Z", + "iopub.status.busy": "2026-10-20T16:54:56.141852Z", + "iopub.status.idle": "2026-10-20T16:54:56.372159Z", + "shell.execute_reply": "2026-10-20T16:54:56.371300Z" } }, "outputs": [ @@ -895,13 +998,13 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 22, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:33:10.991639Z", - "iopub.status.busy": "2026-10-20T16:33:10.991028Z", - "iopub.status.idle": "2026-10-20T16:33:11.153763Z", - "shell.execute_reply": "2026-10-20T16:33:11.152805Z" + "iopub.execute_input": "2026-10-20T16:54:56.376170Z", + "iopub.status.busy": "2026-10-20T16:54:56.375965Z", + "iopub.status.idle": "2026-10-20T16:54:56.537965Z", + "shell.execute_reply": "2026-10-20T16:54:56.537300Z" } }, "outputs": [ @@ -917,8 +1020,11 @@ " #pragma omp simd aligned(gi,idx,u:64)\n", " for (int y = y_m; y <= y_M; y += 1)\n", " {\n", - " int r0 = idx[x + 1][y + 1];\n", - " gi[x / xsf][y] = u[r0 + 1][y + 1];\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", @@ -929,8 +1035,7 @@ "source": [ "# NBVAL_IGNORE_OUTPUT\n", "idx = Function(name='idx', grid=grid, dtype=np.int32)\n", - "gi = Function(name='gi', grid=grid, dimensions=(xs, y), shape=(4, 12),\n", - " space_order=0)\n", + "gi = Function(name='gi', grid=sub, space_order=0)\n", "\n", "op_indirect = Operator(Eq(gi, u[idx[x, y], y]))\n", "show(op_indirect)" diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index b5cea667074..7bb32942625 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -3,8 +3,8 @@ from sympy import Or from devito import ( - CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, TimeFunction, - switchconfig + CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, SubGrid, + TimeFunction, switchconfig ) from devito.ir.iet import Conditional, FindNodes, Iteration @@ -420,3 +420,179 @@ def test_with_openmp(self): 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: each subsampled Dimension's loop is strided.""" + 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) + + op = Operator(Eq(f, u)) + + assert not FindNodes(Conditional).visit(op) + for d in sub.dimensions: + assert iteration(op, d.parent.name).limits[2] == d.symbolic_factor + + 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) + + def test_cannot_nest(self): + grid = Grid(shape=(9, 9)) + sub = SubGrid(grid, factor=2) + + with pytest.raises(ValueError, match='Cannot subsample a SubGrid'): + SubGrid(sub, factor=2) + + 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'] From f10c537c8a67583fff0df8274d4ea284d87745f0 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 15:16:36 -0400 Subject: [PATCH 05/11] compiler: Iterate subsampled Dimensions with unit stride Address review on #3080. Rather than striding the parent's loop by `factor`, iterate the subsampled Dimension itself with a unit increment and rescale parent accesses: for (x = x_m; x <= x_M; x += 1) for (xs = xs_m; xs <= xs_M; xs += 1) if (x%xf == 0) -> g[xs] = u[xs*xf] g[x/xf] = u[x] A non-unitary increment escapes the Lamport-style dependence analysis, so it is avoided. The loop bounds are the ConditionalDimension's own runtime arguments, which removes the need to anchor the loop minimum, and with it the `Strided` property, the `NodeIteration.limits` hook and the hand-rolled round-up helper. The pass moves into `_lower_clusters` so every pipeline picks it up rather than just the CPU and GPU `advanced` ones; it now applies under `noop` too. Also fix `ConditionalDimension._arg_values`, which derived the subsampled max with `ceil` rather than `floor`. With `factor=3` and `x_M=11` that gave `xs_M=4` while the last subsampled point is `x=9`, i.e. index 3 -- harmless while a guard filtered the extra point, but a one-past-the-end write once the bound drives the loop directly. --- devito/core/cpu.py | 7 +- devito/core/gpu.py | 7 +- devito/ir/stree/tree.py | 35 +--- devito/ir/support/properties.py | 38 ---- devito/operator/operator.py | 7 +- devito/passes/clusters/subsampling.py | 36 +++- devito/types/dimension.py | 11 +- examples/userapi/10_subgrids.ipynb | 248 +++++++++++++------------- tests/test_subsampling.py | 81 +++++++-- 9 files changed, 236 insertions(+), 234 deletions(-) diff --git a/devito/core/cpu.py b/devito/core/cpu.py index 1d7b4217bc0..3435a472e62 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -6,7 +6,7 @@ from devito.passes import stream_dimensions from devito.passes.clusters import ( Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse, - optimize_hyperplanes, optimize_pows, subsampling + optimize_hyperplanes, optimize_pows ) from devito.passes.equations import collect_derivatives from devito.passes.iet import ( @@ -195,11 +195,6 @@ def _specialize_clusters(cls, clusters, **kwargs): # Reduce flops clusters = cse(clusters, **kwargs) - # Strided iteration in place of subsampling guards. After the last - # `fuse`, so that no unguarded Cluster can be merged into a strided - # loop, and before blocking, which must see the final step - clusters = subsampling(clusters) - # Blocking to improve data locality if options['blocklazy']: clusters = blocking(clusters, sregistry, options) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 76275994856..536a06422e2 100644 --- a/devito/core/gpu.py +++ b/devito/core/gpu.py @@ -8,7 +8,7 @@ from devito.passes import is_on_device, stream_dimensions from devito.passes.clusters import ( Lift, apply_par_tiles, blocking, buffering, cire, cse, factorize, fission, fuse, - memcpy_prefetch, optimize_pows, subsampling, tasking + memcpy_prefetch, optimize_pows, tasking ) from devito.passes.equations import collect_derivatives from devito.passes.iet import ( @@ -242,11 +242,6 @@ def _specialize_clusters(cls, clusters, **kwargs): # Reduce flops clusters = cse(clusters, **kwargs) - # Strided iteration in place of subsampling guards. After the last - # `fuse`, so that no unguarded Cluster can be merged into a strided - # loop, and before blocking, which must see the final step - clusters = subsampling(clusters) - # Blocking to define thread blocks if options['blocklazy']: clusters = blocking(clusters, sregistry, options) diff --git a/devito/ir/stree/tree.py b/devito/ir/stree/tree.py index ec894073624..96e498396d2 100644 --- a/devito/ir/stree/tree.py +++ b/devito/ir/stree/tree.py @@ -1,7 +1,6 @@ -import sympy from anytree import ContStyle, NodeMixin, PostOrderIter, RenderTree -from devito.ir.support import PrefetchUpdate, WithLock, strided_factor +from devito.ir.support import PrefetchUpdate, WithLock __all__ = [ "NodeConditional", @@ -65,18 +64,9 @@ def dim(self): @property def limits(self): - _min = self.dim.symbolic_min + self.interval.lower - _max = self.dim.symbolic_max + self.interval.upper - - # A Strided Dimension steps `factor` points at a time in place of a - # `dim % factor == 0` guard. The min must be advanced to the first - # point on the subsampling grid, as `_min` is not necessarily a - # multiple of `factor` - factor = strided_factor(self.properties or ()) - if factor is not None: - return (_round_up(_min, factor), _max, factor) - - return (_min, _max, self.dim.symbolic_incr) + return (self.dim.symbolic_min + self.interval.lower, + self.dim.symbolic_max + self.interval.upper, + self.dim.symbolic_incr) @property def direction(self): @@ -156,20 +146,3 @@ def __repr_render__(self): def render(stree): return RenderTree(stree, style=ContStyle()).by_attr('__repr_render__') - - -def _round_up(expr, factor): - """ - The smallest multiple of `factor` that is `>= expr`, for `expr >= 0`. - - Written as `expr + (factor - expr%factor)%factor` rather than the usual - `((expr + factor - 1)/factor)*factor`: the latter is rewritten by - `IntDiv.__mul__` into plain `expr + factor - 1`, and the otherwise - equivalent `expr + (-expr)%factor` would rely on a floored modulo, whereas - C truncates towards zero and would place the loop min *below* `expr`. - `evaluate=False` keeps SymPy from folding the double modulo back. - """ - rem = sympy.Mod(expr, factor) - shift = sympy.Mod(sympy.Add(factor, -rem, evaluate=False), factor, evaluate=False) - - return sympy.Add(expr, shift, evaluate=False) diff --git a/devito/ir/support/properties.py b/devito/ir/support/properties.py index 952ab0be040..f2675ae8b1f 100644 --- a/devito/ir/support/properties.py +++ b/devito/ir/support/properties.py @@ -117,37 +117,6 @@ def __init__(self, name, val=None): # Bundles PARALLELS = {PARALLEL, PARALLEL_INDEP, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT} -STRIDED = 'strided' -"""Name of the parameterized Strided Property; see `Strided`.""" - - -def Strided(factor): - """ - Property marking a Dimension iterated with a non-unitary step `factor`, a - strengthened form of a factor-based ConditionalDimension guard. Instead of - visiting every point and testing `d % factor == 0`, the loop steps `factor` - points at a time, anchored so that the first iteration satisfies the guard. - - Only legal on a PARALLEL Dimension, where no dependence is carried along the - Dimension and hence skipping the guarded-out points cannot drop a required - ordering. - - The `factor` is carried as the Property value, so two Strided properties - with different factors compare unequal and therefore do not survive a - fusion that merges them. - """ - return Property(STRIDED, factor) - - -def strided_factor(properties): - """ - The step of the Strided Property in `properties`, or None if absent. - """ - for i in properties: - if i.name == STRIDED: - return i.val - return None - def normalize_properties(*args): if not args: @@ -178,13 +147,6 @@ def normalize_properties(*args): if not all(i in p for p in args): drop.add(i) - # A Strided property must be dropped unless every `args` entry carries the - # very same one; fusing an unguarded Cluster into a strided loop would - # otherwise silently skip iterations - strided = {frozenset(i for i in p if i.name == STRIDED) for p in args} - if len(strided) > 1: - drop.update(*strided) - properties = set() for p in args: properties.update(p - drop) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 8f3c02f983b..13927d5761f 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/subsampling.py b/devito/passes/clusters/subsampling.py index b613443d5f8..46fc0cc3e3c 100644 --- a/devito/passes/clusters/subsampling.py +++ b/devito/passes/clusters/subsampling.py @@ -1,5 +1,6 @@ from devito.ir.clusters import Queue -from devito.ir.support import GuardFactor, Strided +from devito.ir.support import GuardFactor +from devito.symbolics import IntDiv, uxreplace from devito.tools import timed_pass __all__ = ['subsampling'] @@ -19,11 +20,18 @@ def subsampling(clusters): if (x%xf == 0) -> g[x/xf] = ... g[x/xf] = ... - The rewrite is recorded as a `Strided` property on `d`; the stride and the - anchoring of the loop min are applied by `NodeIteration.limits`. The - iteration space itself is left keyed on `d`, so the runtime `x_m`/`x_M` - values, the domain decomposition and any non-subsampled access `f[x + 1]` - are all unaffected. +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 ---------- @@ -63,11 +71,21 @@ def callback(self, clusters, prefix): 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: - properties = c.properties.add(d, Strided(guard.d.symbolic_factor)) - processed.append(c.rebuild(guards=c.guards.popany(d), - properties=properties)) + # 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 + exprs = [uxreplace(e, {IntDiv(d, factor): cd, d: cd*factor}) + for e in c.exprs] + processed.append(c.rebuild(exprs=exprs, ispace=ispace, + guards=c.guards.popany(d))) return processed diff --git a/devito/types/dimension.py b/devito/types/dimension.py index c81eb90d37d..a0b4d5fede8 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -1153,13 +1153,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/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb index a90f27f9c92..d803238efa2 100644 --- a/examples/userapi/10_subgrids.ipynb +++ b/examples/userapi/10_subgrids.ipynb @@ -12,8 +12,9 @@ "`i*factor`.\n", "\n", "Because the points coincide with parent points, no interpolation is involved. Devito\n", - "compiles a subgrid write to a **strided loop** -- `x += factor` -- rather than\n", - "visiting every point under an `if (x%factor == 0)` guard.\n", + "iterates the **subgrid** directly -- one unit-stride loop per subsampled `Dimension`,\n", + "with parent accesses rescaled -- rather than walking every parent point under an\n", + "`if (x%factor == 0)` guard.\n", "\n", "Every section below prints the generated C so you can see exactly what each case\n", "produces." @@ -34,10 +35,10 @@ "execution_count": 1, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:53.254537Z", - "iopub.status.busy": "2026-10-20T16:54:53.253925Z", - "iopub.status.idle": "2026-10-20T16:54:54.490124Z", - "shell.execute_reply": "2026-10-20T16:54:54.488467Z" + "iopub.execute_input": "2026-10-20T18:55:57.044352Z", + "iopub.status.busy": "2026-10-20T18:55:57.043772Z", + "iopub.status.idle": "2026-10-20T18:55:58.261963Z", + "shell.execute_reply": "2026-10-20T18:55:58.261215Z" } }, "outputs": [ @@ -78,10 +79,10 @@ "execution_count": 2, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.545852Z", - "iopub.status.busy": "2026-10-20T16:54:54.544993Z", - "iopub.status.idle": "2026-10-20T16:54:54.555564Z", - "shell.execute_reply": "2026-10-20T16:54:54.554030Z" + "iopub.execute_input": "2026-10-20T18:55:58.318377Z", + "iopub.status.busy": "2026-10-20T18:55:58.317767Z", + "iopub.status.idle": "2026-10-20T18:55:58.325639Z", + "shell.execute_reply": "2026-10-20T18:55:58.323989Z" } }, "outputs": [ @@ -112,10 +113,10 @@ "execution_count": 3, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.561212Z", - "iopub.status.busy": "2026-10-20T16:54:54.560550Z", - "iopub.status.idle": "2026-10-20T16:54:54.570827Z", - "shell.execute_reply": "2026-10-20T16:54:54.569162Z" + "iopub.execute_input": "2026-10-20T18:55:58.330762Z", + "iopub.status.busy": "2026-10-20T18:55:58.330349Z", + "iopub.status.idle": "2026-10-20T18:55:58.338422Z", + "shell.execute_reply": "2026-10-20T18:55:58.337065Z" } }, "outputs": [ @@ -147,10 +148,10 @@ "execution_count": 4, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.576697Z", - "iopub.status.busy": "2026-10-20T16:54:54.576122Z", - "iopub.status.idle": "2026-10-20T16:54:54.595157Z", - "shell.execute_reply": "2026-10-20T16:54:54.593905Z" + "iopub.execute_input": "2026-10-20T18:55:58.343306Z", + "iopub.status.busy": "2026-10-20T18:55:58.342858Z", + "iopub.status.idle": "2026-10-20T18:55:58.362448Z", + "shell.execute_reply": "2026-10-20T18:55:58.361023Z" } }, "outputs": [ @@ -185,10 +186,10 @@ "execution_count": 5, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.600146Z", - "iopub.status.busy": "2026-10-20T16:54:54.599726Z", - "iopub.status.idle": "2026-10-20T16:54:54.611723Z", - "shell.execute_reply": "2026-10-20T16:54:54.609903Z" + "iopub.execute_input": "2026-10-20T18:55:58.367555Z", + "iopub.status.busy": "2026-10-20T18:55:58.367076Z", + "iopub.status.idle": "2026-10-20T18:55:58.377548Z", + "shell.execute_reply": "2026-10-20T18:55:58.376239Z" } }, "outputs": [ @@ -224,10 +225,10 @@ "execution_count": 6, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.617414Z", - "iopub.status.busy": "2026-10-20T16:54:54.616914Z", - "iopub.status.idle": "2026-10-20T16:54:54.629162Z", - "shell.execute_reply": "2026-10-20T16:54:54.627654Z" + "iopub.execute_input": "2026-10-20T18:55:58.382603Z", + "iopub.status.busy": "2026-10-20T18:55:58.382146Z", + "iopub.status.idle": "2026-10-20T18:55:58.393704Z", + "shell.execute_reply": "2026-10-20T18:55:58.392240Z" } }, "outputs": [ @@ -264,10 +265,10 @@ "execution_count": 7, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.634508Z", - "iopub.status.busy": "2026-10-20T16:54:54.634002Z", - "iopub.status.idle": "2026-10-20T16:54:54.642529Z", - "shell.execute_reply": "2026-10-20T16:54:54.640851Z" + "iopub.execute_input": "2026-10-20T18:55:58.399039Z", + "iopub.status.busy": "2026-10-20T18:55:58.398520Z", + "iopub.status.idle": "2026-10-20T18:55:58.407099Z", + "shell.execute_reply": "2026-10-20T18:55:58.405629Z" } }, "outputs": [], @@ -296,10 +297,10 @@ "execution_count": 8, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.648073Z", - "iopub.status.busy": "2026-10-20T16:54:54.647463Z", - "iopub.status.idle": "2026-10-20T16:54:54.825256Z", - "shell.execute_reply": "2026-10-20T16:54:54.823521Z" + "iopub.execute_input": "2026-10-20T18:55:58.412262Z", + "iopub.status.busy": "2026-10-20T18:55:58.411719Z", + "iopub.status.idle": "2026-10-20T18:55:58.583987Z", + "shell.execute_reply": "2026-10-20T18:55:58.583103Z" } }, "outputs": [ @@ -308,12 +309,11 @@ "output_type": "stream", "text": [ "START(section0)\n", - "for (int x = x_m + (xsf - (x_m)%(xsf))%(xsf); x <= x_M; x += xsf)\n", + "for (int xs = xs_m; xs <= xs_M; xs += 1)\n", "{\n", - " #pragma omp simd aligned(u,usub:64)\n", - " for (int y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", " {\n", - " usub[x / xsf][y / ysf] = u[x + 1][y + 1];\n", + " usub[xs][ys] = u[xs*xsf + 1][ys*ysf + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -330,18 +330,18 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Two features of the `x` loop header matter:\n", + "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 increment is `x += xsf`, so only one point in three is visited -- there is no\n", - " `if` anywhere in the nest;\n", - "- the minimum is `x_m + (xsf - x_m%xsf)%xsf` rather than `x_m`.\n", + "- the subgrid access is simply `usub[xs]`;\n", + "- the parent access is rescaled to `u[xs*xsf + 1]`.\n", "\n", - "That second expression is the **anchor**. A subgrid point is one where\n", - "`x % factor == 0`, and `x_m` is not necessarily such a point, so iteration must start\n", - "at the first subgrid point at or after `x_m`.\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 body carries the index conversion `usub[x / xsf]`, mapping the parent coordinate\n", - "to a subgrid coordinate." + "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)`." ] }, { @@ -349,17 +349,17 @@ "execution_count": 9, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.830512Z", - "iopub.status.busy": "2026-10-20T16:54:54.830098Z", - "iopub.status.idle": "2026-10-20T16:54:54.840469Z", - "shell.execute_reply": "2026-10-20T16:54:54.838656Z" + "iopub.execute_input": "2026-10-20T18:55:58.587918Z", + "iopub.status.busy": "2026-10-20T18:55:58.587706Z", + "iopub.status.idle": "2026-10-20T18:55:58.593394Z", + "shell.execute_reply": "2026-10-20T18:55:58.592617Z" } }, "outputs": [ { "data": { "text/plain": [ - "{'x': xsf, 'y': ysf}" + "{'xs': 1, 'ys': 1}" ] }, "execution_count": 9, @@ -385,10 +385,10 @@ "execution_count": 10, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:54.846210Z", - "iopub.status.busy": "2026-10-20T16:54:54.845610Z", - "iopub.status.idle": "2026-10-20T16:54:55.010474Z", - "shell.execute_reply": "2026-10-20T16:54:55.009286Z" + "iopub.execute_input": "2026-10-20T18:55:58.597664Z", + "iopub.status.busy": "2026-10-20T18:55:58.597434Z", + "iopub.status.idle": "2026-10-20T18:55:58.749481Z", + "shell.execute_reply": "2026-10-20T18:55:58.748644Z" } }, "outputs": [ @@ -427,10 +427,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### The anchor in action\n", + "### Restricting the range\n", "\n", - "Run with `x_m=2`, so iteration no longer starts on a multiple of 3. The first subgrid\n", - "point is `x=3`, so the slot for `x=0` is never written." + "Overriding the parent's bounds maps onto the subgrid's. With `x_m=2` and `factor=3`,\n", + "the first subgrid point at or after `x=2` is `x=3`, i.e. subgrid index 1 -- so index 0\n", + "is never written." ] }, { @@ -438,10 +439,10 @@ "execution_count": 11, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.014690Z", - "iopub.status.busy": "2026-10-20T16:54:55.014442Z", - "iopub.status.idle": "2026-10-20T16:54:55.144762Z", - "shell.execute_reply": "2026-10-20T16:54:55.143244Z" + "iopub.execute_input": "2026-10-20T18:55:58.753674Z", + "iopub.status.busy": "2026-10-20T18:55:58.753460Z", + "iopub.status.idle": "2026-10-20T18:55:58.865492Z", + "shell.execute_reply": "2026-10-20T18:55:58.864704Z" } }, "outputs": [ @@ -476,8 +477,8 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The leading `-1.` is the untouched slot. Had iteration started at `x_m=2` instead\n", - "of `x=3`, the wrong points would have been written." + "The leading `-1.` is the untouched slot. Devito derived `xs_m = ceil(2/3) = 1`\n", + "from the supplied `x_m=2`." ] }, { @@ -496,10 +497,10 @@ "execution_count": 12, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.150500Z", - "iopub.status.busy": "2026-10-20T16:54:55.149948Z", - "iopub.status.idle": "2026-10-20T16:54:55.270974Z", - "shell.execute_reply": "2026-10-20T16:54:55.269737Z" + "iopub.execute_input": "2026-10-20T18:55:58.869707Z", + "iopub.status.busy": "2026-10-20T18:55:58.869504Z", + "iopub.status.idle": "2026-10-20T18:55:58.973102Z", + "shell.execute_reply": "2026-10-20T18:55:58.972276Z" } }, "outputs": [ @@ -510,10 +511,9 @@ "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 y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", " {\n", - " usub_y[x][y / ysf] = u[x + 1][y + 1];\n", + " usub_y[x][ys] = u[x + 1][ys*ysf + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -534,17 +534,17 @@ "execution_count": 13, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.276046Z", - "iopub.status.busy": "2026-10-20T16:54:55.275462Z", - "iopub.status.idle": "2026-10-20T16:54:55.287010Z", - "shell.execute_reply": "2026-10-20T16:54:55.285334Z" + "iopub.execute_input": "2026-10-20T18:55:58.976977Z", + "iopub.status.busy": "2026-10-20T18:55:58.976774Z", + "iopub.status.idle": "2026-10-20T18:55:58.982211Z", + "shell.execute_reply": "2026-10-20T18:55:58.981417Z" } }, "outputs": [ { "data": { "text/plain": [ - "{'x': 1, 'y': ysf}" + "{'x': 1, 'ys': 1}" ] }, "execution_count": 13, @@ -568,10 +568,10 @@ "execution_count": 14, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.292500Z", - "iopub.status.busy": "2026-10-20T16:54:55.291964Z", - "iopub.status.idle": "2026-10-20T16:54:55.428150Z", - "shell.execute_reply": "2026-10-20T16:54:55.426871Z" + "iopub.execute_input": "2026-10-20T18:55:58.986323Z", + "iopub.status.busy": "2026-10-20T18:55:58.986116Z", + "iopub.status.idle": "2026-10-20T18:55:59.305167Z", + "shell.execute_reply": "2026-10-20T18:55:59.304219Z" } }, "outputs": [ @@ -618,10 +618,10 @@ "execution_count": 15, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.433311Z", - "iopub.status.busy": "2026-10-20T16:54:55.432979Z", - "iopub.status.idle": "2026-10-20T16:54:55.548544Z", - "shell.execute_reply": "2026-10-20T16:54:55.547178Z" + "iopub.execute_input": "2026-10-20T18:55:59.309504Z", + "iopub.status.busy": "2026-10-20T18:55:59.309210Z", + "iopub.status.idle": "2026-10-20T18:55:59.430048Z", + "shell.execute_reply": "2026-10-20T18:55:59.429199Z" } }, "outputs": [ @@ -630,12 +630,11 @@ "output_type": "stream", "text": [ "START(section0)\n", - "for (int x = x_m + (xsf - (x_m)%(xsf))%(xsf); x <= x_M; x += xsf)\n", + "for (int xs = xs_m; xs <= xs_M; xs += 1)\n", "{\n", - " #pragma omp simd aligned(u,usub_xy:64)\n", - " for (int y = y_m + (ysf - (y_m)%(ysf))%(ysf); y <= y_M; y += ysf)\n", + " for (int ys = ys_m; ys <= ys_M; ys += 1)\n", " {\n", - " usub_xy[x / xsf][y / ysf] = u[x + 1][y + 1];\n", + " usub_xy[xs][ys] = u[xs*xsf + 1][ys*ysf + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -656,10 +655,10 @@ "execution_count": 16, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.552823Z", - "iopub.status.busy": "2026-10-20T16:54:55.552658Z", - "iopub.status.idle": "2026-10-20T16:54:55.663895Z", - "shell.execute_reply": "2026-10-20T16:54:55.662572Z" + "iopub.execute_input": "2026-10-20T18:55:59.433660Z", + "iopub.status.busy": "2026-10-20T18:55:59.433464Z", + "iopub.status.idle": "2026-10-20T18:55:59.735035Z", + "shell.execute_reply": "2026-10-20T18:55:59.734049Z" } }, "outputs": [ @@ -673,7 +672,7 @@ { "data": { "text/plain": [ - "{'x': xsf, 'y': ysf}" + "{'xs': 1, 'ys': 1}" ] }, "execution_count": 16, @@ -716,10 +715,10 @@ "execution_count": 17, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.669766Z", - "iopub.status.busy": "2026-10-20T16:54:55.669150Z", - "iopub.status.idle": "2026-10-20T16:54:55.802136Z", - "shell.execute_reply": "2026-10-20T16:54:55.800522Z" + "iopub.execute_input": "2026-10-20T18:55:59.738894Z", + "iopub.status.busy": "2026-10-20T18:55:59.738581Z", + "iopub.status.idle": "2026-10-20T18:55:59.860834Z", + "shell.execute_reply": "2026-10-20T18:55:59.859974Z" } }, "outputs": [ @@ -776,10 +775,10 @@ "execution_count": 18, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.806474Z", - "iopub.status.busy": "2026-10-20T16:54:55.806303Z", - "iopub.status.idle": "2026-10-20T16:54:55.926680Z", - "shell.execute_reply": "2026-10-20T16:54:55.925460Z" + "iopub.execute_input": "2026-10-20T18:55:59.864891Z", + "iopub.status.busy": "2026-10-20T18:55:59.864681Z", + "iopub.status.idle": "2026-10-20T18:55:59.984698Z", + "shell.execute_reply": "2026-10-20T18:55:59.983855Z" } }, "outputs": [ @@ -829,10 +828,10 @@ "execution_count": 19, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:55.931233Z", - "iopub.status.busy": "2026-10-20T16:54:55.931060Z", - "iopub.status.idle": "2026-10-20T16:54:56.100092Z", - "shell.execute_reply": "2026-10-20T16:54:56.098527Z" + "iopub.execute_input": "2026-10-20T18:55:59.988713Z", + "iopub.status.busy": "2026-10-20T18:55:59.988516Z", + "iopub.status.idle": "2026-10-20T18:56:00.157395Z", + "shell.execute_reply": "2026-10-20T18:56:00.156386Z" } }, "outputs": [ @@ -897,10 +896,10 @@ "execution_count": 20, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:56.104592Z", - "iopub.status.busy": "2026-10-20T16:54:56.104423Z", - "iopub.status.idle": "2026-10-20T16:54:56.136868Z", - "shell.execute_reply": "2026-10-20T16:54:56.135085Z" + "iopub.execute_input": "2026-10-20T18:56:00.161369Z", + "iopub.status.busy": "2026-10-20T18:56:00.161151Z", + "iopub.status.idle": "2026-10-20T18:56:00.182345Z", + "shell.execute_reply": "2026-10-20T18:56:00.181508Z" } }, "outputs": [ @@ -948,10 +947,10 @@ "execution_count": 21, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:56.142381Z", - "iopub.status.busy": "2026-10-20T16:54:56.141852Z", - "iopub.status.idle": "2026-10-20T16:54:56.372159Z", - "shell.execute_reply": "2026-10-20T16:54:56.371300Z" + "iopub.execute_input": "2026-10-20T18:56:00.186297Z", + "iopub.status.busy": "2026-10-20T18:56:00.186089Z", + "iopub.status.idle": "2026-10-20T18:56:00.410953Z", + "shell.execute_reply": "2026-10-20T18:56:00.409969Z" } }, "outputs": [ @@ -1001,10 +1000,10 @@ "execution_count": 22, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T16:54:56.376170Z", - "iopub.status.busy": "2026-10-20T16:54:56.375965Z", - "iopub.status.idle": "2026-10-20T16:54:56.537965Z", - "shell.execute_reply": "2026-10-20T16:54:56.537300Z" + "iopub.execute_input": "2026-10-20T18:56:00.415047Z", + "iopub.status.busy": "2026-10-20T18:56:00.414843Z", + "iopub.status.idle": "2026-10-20T18:56:00.558146Z", + "shell.execute_reply": "2026-10-20T18:56:00.557201Z" } }, "outputs": [ @@ -1047,17 +1046,18 @@ "source": [ "## Summary\n", "\n", - "Subsampling along a `Dimension` is declared with `ConditionalDimension(factor=n)` plus\n", - "a `Function` defined over it. Devito then generates, for that `Dimension`:\n", + "Subsampling along a `Dimension` is declared with `SubGrid` (or a\n", + "`ConditionalDimension` directly). Devito then generates, for that `Dimension`:\n", "\n", - "| | strided loop | guarded loop |\n", + "| | subsampled loop | guarded loop |\n", "|---|---|---|\n", - "| loop header | `x += factor`, anchored | `x += 1` |\n", - "| body | unconditional | wrapped in `if (x%factor == 0)` |\n", + "| iterates | `xs`, unit stride | `x`, unit stride |\n", + "| bounds | `ceil(x_m/f) .. floor(x_M/f)` | `x_m .. x_M` |\n", + "| body | unconditional, parent access `xs*f` | wrapped in `if (x%f == 0)` |\n", "| iterations | one in `factor` | all |\n", "| used when | the Dimension is parallel and accesses are affine | otherwise |\n", "\n", - "The strided form is the default, including for the innermost (vectorized)\n", + "The subsampled form is the default, including for the innermost (vectorized)\n", "`Dimension` and for several subsampled `Dimension`s at once. The guarded form is the\n", "fallback for the four cases above: a loop nest shared with an unguarded equation, a\n", "sequential `Dimension` such as time, an extra `condition`, or non-affine accesses.\n", diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index 7bb32942625..c3e6f63482b 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -36,7 +36,12 @@ def subsampled(shape, factor, ndim=2): class TestStridedCodegen: - """Subsampling lowers to a loop stride, not a guard.""" + """ + 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) @@ -44,17 +49,36 @@ def test_codegen(self): # No guard... assert not FindNodes(Conditional).visit(op) - # ...the parent Dimension's loop is strided instead - assert iteration(op, 'x').limits[2] == cd.symbolic_factor + # ...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_factor_in_loop_header(self, factor): + def test_parent_access_rescaled(self, factor): + """Parent accesses become `cd*factor`; subgrid accesses become `cd`.""" _, f, g, cd = subsampled((16, 8), factor) op = Operator(Eq(g, f)) + ccode = str(op.ccode) - assert f'+= {cd.symbolic_factor.name}' in str(op.ccode) + fname = cd.symbolic_factor.name + assert f'{cd.name}*{fname}' in ccode # parent, rescaled + assert f'g[{cd.name}]' in ccode # subgrid, direct + assert f'+= {fname}' not in ccode # 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 @@ -70,7 +94,8 @@ def test_innermost_strided(self): op = Operator(Eq(usub, u)) - assert iteration(op, 'y').limits[2] == ys.symbolic_factor + 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): @@ -100,8 +125,9 @@ def test_multiple_dimensions_strided(self): op = Operator(Eq(usub, u)) assert not FindNodes(Conditional).visit(op) - assert iteration(op, 'x').limits[2] == xs.symbolic_factor - assert iteration(op, 'y').limits[2] == ys.symbolic_factor + 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)) @@ -118,15 +144,19 @@ def test_multiple_dimensions_numerics(self): assert np.array_equal(usub.data, data[::3, ::2]) - def test_anchored_min(self): - """The loop min is rounded up onto the subsampling grid.""" + 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)) - _min = iteration(op, 'x').limits[0] - # Not simply `x_m`: it carries the round-up to the next multiple - assert _min != f.grid.dimensions[0].symbolic_min - assert cd.symbolic_factor in _min.free_symbols + 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: @@ -160,11 +190,30 @@ def test_anchoring(self, x_m, x_M): g.data[:] = -1. Operator(Eq(g, f)).apply(x_m=x_m, x_M=x_M) - written = {i//factor for i in range(x_m, x_M + 1) if i % factor == 0} + # 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 @@ -496,7 +545,7 @@ def test_copy_is_strided(self): assert not FindNodes(Conditional).visit(op) for d in sub.dimensions: - assert iteration(op, d.parent.name).limits[2] == d.symbolic_factor + assert iteration(op, d.name).limits[2] == 1 def test_factor_one_is_passthrough(self): """A factor of 1 leaves the Dimension untouched.""" From 35fb9ea7596c2e0f36ae3824623e2336869d5197 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 16:04:03 -0400 Subject: [PATCH 06/11] compiler: Carry properties onto the subsampled Dimension `IterationSpace.switch` remaps intervals, sub-iterators and directions, but not properties, so swapping the parent Dimension for the subsampled one left the latter with none: the loop lost PARALLEL and VECTORIZED, and with them every OpenMP pragma. Transfer them explicitly, as `blocking` does for its BlockDimensions. The subsampled Dimension covers the same points as its parent, just addressed in subsampled coordinates, so it inherits the parent's properties unchanged. Also broaden the tests and tutorial to exercise SubGrid end to end: the shared-nest, sequential and non-affine cases now build their Functions from a SubGrid and run the Operator with checked results, rather than only inspecting the generated code. --- devito/passes/clusters/subsampling.py | 15 +- examples/userapi/10_subgrids.ipynb | 340 ++++++++++++++++++-------- tests/test_subsampling.py | 101 +++++++- 3 files changed, 345 insertions(+), 111 deletions(-) diff --git a/devito/passes/clusters/subsampling.py b/devito/passes/clusters/subsampling.py index 46fc0cc3e3c..7ef3e7f35e6 100644 --- a/devito/passes/clusters/subsampling.py +++ b/devito/passes/clusters/subsampling.py @@ -79,13 +79,22 @@ def callback(self, clusters, prefix): # 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 - exprs = [uxreplace(e, {IntDiv(d, factor): cd, d: cd*factor}) - for e in c.exprs] + 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))) + guards=c.guards.popany(d), + properties=properties)) return processed diff --git a/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb index d803238efa2..3d66f1efdd3 100644 --- a/examples/userapi/10_subgrids.ipynb +++ b/examples/userapi/10_subgrids.ipynb @@ -35,10 +35,10 @@ "execution_count": 1, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:57.044352Z", - "iopub.status.busy": "2026-10-20T18:55:57.043772Z", - "iopub.status.idle": "2026-10-20T18:55:58.261963Z", - "shell.execute_reply": "2026-10-20T18:55:58.261215Z" + "iopub.execute_input": "2026-10-20T19:47:37.266954Z", + "iopub.status.busy": "2026-10-20T19:47:37.266342Z", + "iopub.status.idle": "2026-10-20T19:47:38.482933Z", + "shell.execute_reply": "2026-10-20T19:47:38.482018Z" } }, "outputs": [ @@ -79,10 +79,10 @@ "execution_count": 2, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.318377Z", - "iopub.status.busy": "2026-10-20T18:55:58.317767Z", - "iopub.status.idle": "2026-10-20T18:55:58.325639Z", - "shell.execute_reply": "2026-10-20T18:55:58.323989Z" + "iopub.execute_input": "2026-10-20T19:47:38.541091Z", + "iopub.status.busy": "2026-10-20T19:47:38.540134Z", + "iopub.status.idle": "2026-10-20T19:47:38.550952Z", + "shell.execute_reply": "2026-10-20T19:47:38.549260Z" } }, "outputs": [ @@ -113,10 +113,10 @@ "execution_count": 3, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.330762Z", - "iopub.status.busy": "2026-10-20T18:55:58.330349Z", - "iopub.status.idle": "2026-10-20T18:55:58.338422Z", - "shell.execute_reply": "2026-10-20T18:55:58.337065Z" + "iopub.execute_input": "2026-10-20T19:47:38.554978Z", + "iopub.status.busy": "2026-10-20T19:47:38.554374Z", + "iopub.status.idle": "2026-10-20T19:47:38.564543Z", + "shell.execute_reply": "2026-10-20T19:47:38.562985Z" } }, "outputs": [ @@ -148,10 +148,10 @@ "execution_count": 4, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.343306Z", - "iopub.status.busy": "2026-10-20T18:55:58.342858Z", - "iopub.status.idle": "2026-10-20T18:55:58.362448Z", - "shell.execute_reply": "2026-10-20T18:55:58.361023Z" + "iopub.execute_input": "2026-10-20T19:47:38.569006Z", + "iopub.status.busy": "2026-10-20T19:47:38.568430Z", + "iopub.status.idle": "2026-10-20T19:47:38.588366Z", + "shell.execute_reply": "2026-10-20T19:47:38.586951Z" } }, "outputs": [ @@ -186,10 +186,10 @@ "execution_count": 5, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.367555Z", - "iopub.status.busy": "2026-10-20T18:55:58.367076Z", - "iopub.status.idle": "2026-10-20T18:55:58.377548Z", - "shell.execute_reply": "2026-10-20T18:55:58.376239Z" + "iopub.execute_input": "2026-10-20T19:47:38.593261Z", + "iopub.status.busy": "2026-10-20T19:47:38.592855Z", + "iopub.status.idle": "2026-10-20T19:47:38.602233Z", + "shell.execute_reply": "2026-10-20T19:47:38.600963Z" } }, "outputs": [ @@ -225,10 +225,10 @@ "execution_count": 6, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.382603Z", - "iopub.status.busy": "2026-10-20T18:55:58.382146Z", - "iopub.status.idle": "2026-10-20T18:55:58.393704Z", - "shell.execute_reply": "2026-10-20T18:55:58.392240Z" + "iopub.execute_input": "2026-10-20T19:47:38.607357Z", + "iopub.status.busy": "2026-10-20T19:47:38.606926Z", + "iopub.status.idle": "2026-10-20T19:47:38.618469Z", + "shell.execute_reply": "2026-10-20T19:47:38.616998Z" } }, "outputs": [ @@ -265,10 +265,10 @@ "execution_count": 7, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.399039Z", - "iopub.status.busy": "2026-10-20T18:55:58.398520Z", - "iopub.status.idle": "2026-10-20T18:55:58.407099Z", - "shell.execute_reply": "2026-10-20T18:55:58.405629Z" + "iopub.execute_input": "2026-10-20T19:47:38.623332Z", + "iopub.status.busy": "2026-10-20T19:47:38.622862Z", + "iopub.status.idle": "2026-10-20T19:47:38.630816Z", + "shell.execute_reply": "2026-10-20T19:47:38.629412Z" } }, "outputs": [], @@ -297,10 +297,10 @@ "execution_count": 8, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.412262Z", - "iopub.status.busy": "2026-10-20T18:55:58.411719Z", - "iopub.status.idle": "2026-10-20T18:55:58.583987Z", - "shell.execute_reply": "2026-10-20T18:55:58.583103Z" + "iopub.execute_input": "2026-10-20T19:47:38.635815Z", + "iopub.status.busy": "2026-10-20T19:47:38.635346Z", + "iopub.status.idle": "2026-10-20T19:47:38.807419Z", + "shell.execute_reply": "2026-10-20T19:47:38.806525Z" } }, "outputs": [ @@ -311,6 +311,7 @@ "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", @@ -349,10 +350,10 @@ "execution_count": 9, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.587918Z", - "iopub.status.busy": "2026-10-20T18:55:58.587706Z", - "iopub.status.idle": "2026-10-20T18:55:58.593394Z", - "shell.execute_reply": "2026-10-20T18:55:58.592617Z" + "iopub.execute_input": "2026-10-20T19:47:38.811559Z", + "iopub.status.busy": "2026-10-20T19:47:38.811363Z", + "iopub.status.idle": "2026-10-20T19:47:38.816795Z", + "shell.execute_reply": "2026-10-20T19:47:38.816021Z" } }, "outputs": [ @@ -385,10 +386,10 @@ "execution_count": 10, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.597664Z", - "iopub.status.busy": "2026-10-20T18:55:58.597434Z", - "iopub.status.idle": "2026-10-20T18:55:58.749481Z", - "shell.execute_reply": "2026-10-20T18:55:58.748644Z" + "iopub.execute_input": "2026-10-20T19:47:38.821139Z", + "iopub.status.busy": "2026-10-20T19:47:38.820907Z", + "iopub.status.idle": "2026-10-20T19:47:38.964401Z", + "shell.execute_reply": "2026-10-20T19:47:38.963553Z" } }, "outputs": [ @@ -439,10 +440,10 @@ "execution_count": 11, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.753674Z", - "iopub.status.busy": "2026-10-20T18:55:58.753460Z", - "iopub.status.idle": "2026-10-20T18:55:58.865492Z", - "shell.execute_reply": "2026-10-20T18:55:58.864704Z" + "iopub.execute_input": "2026-10-20T19:47:38.968547Z", + "iopub.status.busy": "2026-10-20T19:47:38.968332Z", + "iopub.status.idle": "2026-10-20T19:47:39.084027Z", + "shell.execute_reply": "2026-10-20T19:47:39.083270Z" } }, "outputs": [ @@ -497,10 +498,10 @@ "execution_count": 12, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.869707Z", - "iopub.status.busy": "2026-10-20T18:55:58.869504Z", - "iopub.status.idle": "2026-10-20T18:55:58.973102Z", - "shell.execute_reply": "2026-10-20T18:55:58.972276Z" + "iopub.execute_input": "2026-10-20T19:47:39.088359Z", + "iopub.status.busy": "2026-10-20T19:47:39.088139Z", + "iopub.status.idle": "2026-10-20T19:47:39.194975Z", + "shell.execute_reply": "2026-10-20T19:47:39.194129Z" } }, "outputs": [ @@ -511,6 +512,7 @@ "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", @@ -534,10 +536,10 @@ "execution_count": 13, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.976977Z", - "iopub.status.busy": "2026-10-20T18:55:58.976774Z", - "iopub.status.idle": "2026-10-20T18:55:58.982211Z", - "shell.execute_reply": "2026-10-20T18:55:58.981417Z" + "iopub.execute_input": "2026-10-20T19:47:39.198698Z", + "iopub.status.busy": "2026-10-20T19:47:39.198494Z", + "iopub.status.idle": "2026-10-20T19:47:39.204332Z", + "shell.execute_reply": "2026-10-20T19:47:39.203208Z" } }, "outputs": [ @@ -568,10 +570,10 @@ "execution_count": 14, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:58.986323Z", - "iopub.status.busy": "2026-10-20T18:55:58.986116Z", - "iopub.status.idle": "2026-10-20T18:55:59.305167Z", - "shell.execute_reply": "2026-10-20T18:55:59.304219Z" + "iopub.execute_input": "2026-10-20T19:47:39.208617Z", + "iopub.status.busy": "2026-10-20T19:47:39.208349Z", + "iopub.status.idle": "2026-10-20T19:47:39.328256Z", + "shell.execute_reply": "2026-10-20T19:47:39.326573Z" } }, "outputs": [ @@ -618,10 +620,10 @@ "execution_count": 15, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:59.309504Z", - "iopub.status.busy": "2026-10-20T18:55:59.309210Z", - "iopub.status.idle": "2026-10-20T18:55:59.430048Z", - "shell.execute_reply": "2026-10-20T18:55:59.429199Z" + "iopub.execute_input": "2026-10-20T19:47:39.333768Z", + "iopub.status.busy": "2026-10-20T19:47:39.333152Z", + "iopub.status.idle": "2026-10-20T19:47:39.467859Z", + "shell.execute_reply": "2026-10-20T19:47:39.467015Z" } }, "outputs": [ @@ -632,6 +634,7 @@ "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", @@ -655,10 +658,10 @@ "execution_count": 16, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:59.433660Z", - "iopub.status.busy": "2026-10-20T18:55:59.433464Z", - "iopub.status.idle": "2026-10-20T18:55:59.735035Z", - "shell.execute_reply": "2026-10-20T18:55:59.734049Z" + "iopub.execute_input": "2026-10-20T19:47:39.471512Z", + "iopub.status.busy": "2026-10-20T19:47:39.471331Z", + "iopub.status.idle": "2026-10-20T19:47:39.586008Z", + "shell.execute_reply": "2026-10-20T19:47:39.585209Z" } }, "outputs": [ @@ -715,10 +718,10 @@ "execution_count": 17, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:59.738894Z", - "iopub.status.busy": "2026-10-20T18:55:59.738581Z", - "iopub.status.idle": "2026-10-20T18:55:59.860834Z", - "shell.execute_reply": "2026-10-20T18:55:59.859974Z" + "iopub.execute_input": "2026-10-20T19:47:39.590169Z", + "iopub.status.busy": "2026-10-20T19:47:39.589960Z", + "iopub.status.idle": "2026-10-20T19:47:39.705734Z", + "shell.execute_reply": "2026-10-20T19:47:39.704904Z" } }, "outputs": [ @@ -730,9 +733,9 @@ "for (int x = x_m; x <= x_M; x += 1)\n", "{\n", " f[x + 1] = 1;\n", - " if ((x)%(xs1f) == 0)\n", + " if ((x)%(xsf) == 0)\n", " {\n", - " g[x / xs1f] = f[x + 1];\n", + " g[x / xsf] = f[x + 1];\n", " }\n", "}\n", "STOP(section0,timers)\n" @@ -742,16 +745,57 @@ "source": [ "# NBVAL_IGNORE_OUTPUT\n", "grid1 = Grid(shape=(12,))\n", - "(x1,) = grid1.dimensions\n", - "xs1 = ConditionalDimension(name='xs1', parent=x1, factor=4)\n", + "sub1 = SubGrid(grid1, factor=4)\n", "\n", "f = Function(name='f', grid=grid1)\n", - "g = Function(name='g', grid=grid1, dimensions=(xs1,), shape=(3,), space_order=0)\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": 18, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T19:47:39.709325Z", + "iopub.status.busy": "2026-10-20T19:47:39.709138Z", + "iopub.status.idle": "2026-10-20T19:47:39.726057Z", + "shell.execute_reply": "2026-10-20T19:47:39.725324Z" + } + }, + "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": 18, + "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": {}, @@ -772,13 +816,13 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 19, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:59.864891Z", - "iopub.status.busy": "2026-10-20T18:55:59.864681Z", - "iopub.status.idle": "2026-10-20T18:55:59.984698Z", - "shell.execute_reply": "2026-10-20T18:55:59.983855Z" + "iopub.execute_input": "2026-10-20T19:47:39.729855Z", + "iopub.status.busy": "2026-10-20T19:47:39.729686Z", + "iopub.status.idle": "2026-10-20T19:47:39.858640Z", + "shell.execute_reply": "2026-10-20T19:47:39.857801Z" } }, "outputs": [ @@ -789,11 +833,11 @@ "START(section0)\n", "for (int x = x_m; x <= x_M; x += 1)\n", "{\n", - " if ((x)%(xs3f) == 0)\n", + " if ((x)%(xsf) == 0)\n", " {\n", " for (int y = y_m; y <= y_M; y += 1)\n", " {\n", - " g2[x / xs3f][y] = f2[x + 1][y + 1] + g2[(x / xs3f) - 1][y];\n", + " g2[x / xsf][y] = f2[x + 1][y + 1] + g2[(x / xsf) - 1][y];\n", " }\n", " }\n", "}\n", @@ -804,34 +848,93 @@ "source": [ "# NBVAL_IGNORE_OUTPUT\n", "grid2 = Grid(shape=(12, 12))\n", - "x2, y2 = grid2.dimensions\n", - "xs3 = ConditionalDimension(name='xs3', parent=x2, factor=3)\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=grid2, dimensions=(xs3, y2), shape=(4, 12),\n", - " space_order=0)\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": 20, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T19:47:39.862144Z", + "iopub.status.busy": "2026-10-20T19:47:39.861969Z", + "iopub.status.idle": "2026-10-20T19:47:39.884751Z", + "shell.execute_reply": "2026-10-20T19:47:39.884003Z" + } + }, + "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": 20, + "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\n", - "a saved-wavefield `Eq` keeps its guard." + "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": 19, + "execution_count": 21, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:55:59.988713Z", - "iopub.status.busy": "2026-10-20T18:55:59.988516Z", - "iopub.status.idle": "2026-10-20T18:56:00.157395Z", - "shell.execute_reply": "2026-10-20T18:56:00.156386Z" + "iopub.execute_input": "2026-10-20T19:47:39.888453Z", + "iopub.status.busy": "2026-10-20T19:47:39.888230Z", + "iopub.status.idle": "2026-10-20T19:47:40.067896Z", + "shell.execute_reply": "2026-10-20T19:47:40.067083Z" } }, "outputs": [ @@ -893,13 +996,13 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 22, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:56:00.161369Z", - "iopub.status.busy": "2026-10-20T18:56:00.161151Z", - "iopub.status.idle": "2026-10-20T18:56:00.182345Z", - "shell.execute_reply": "2026-10-20T18:56:00.181508Z" + "iopub.execute_input": "2026-10-20T19:47:40.071659Z", + "iopub.status.busy": "2026-10-20T19:47:40.071480Z", + "iopub.status.idle": "2026-10-20T19:47:40.092014Z", + "shell.execute_reply": "2026-10-20T19:47:40.091258Z" } }, "outputs": [ @@ -916,7 +1019,7 @@ "[0.0, 4.0, 0.0]" ] }, - "execution_count": 20, + "execution_count": 22, "metadata": {}, "output_type": "execute_result" } @@ -944,16 +1047,23 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 23, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:56:00.186297Z", - "iopub.status.busy": "2026-10-20T18:56:00.186089Z", - "iopub.status.idle": "2026-10-20T18:56:00.410953Z", - "shell.execute_reply": "2026-10-20T18:56:00.409969Z" + "iopub.execute_input": "2026-10-20T19:47:40.095918Z", + "iopub.status.busy": "2026-10-20T19:47:40.095708Z", + "iopub.status.idle": "2026-10-20T19:47:40.540497Z", + "shell.execute_reply": "2026-10-20T19:47:40.539593Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -976,12 +1086,16 @@ ], "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)" ] }, @@ -997,16 +1111,23 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 24, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T18:56:00.415047Z", - "iopub.status.busy": "2026-10-20T18:56:00.414843Z", - "iopub.status.idle": "2026-10-20T18:56:00.558146Z", - "shell.execute_reply": "2026-10-20T18:56:00.557201Z" + "iopub.execute_input": "2026-10-20T19:47:40.545872Z", + "iopub.status.busy": "2026-10-20T19:47:40.545436Z", + "iopub.status.idle": "2026-10-20T19:47:40.895126Z", + "shell.execute_reply": "2026-10-20T19:47:40.894398Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, { "name": "stdout", "output_type": "stream", @@ -1036,7 +1157,12 @@ "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)" ] }, diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index c3e6f63482b..63546f9306d 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -144,6 +144,44 @@ def test_multiple_dimensions_numerics(self): 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, @@ -534,18 +572,79 @@ def test_copy_from_parent(self, factor): assert np.array_equal(f.data, data[::fx, ::fy]) def test_copy_is_strided(self): - """No guard is emitted: each subsampled Dimension's loop is strided.""" + """ + 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.""" From 87248bc960f2efed1122f358d874b269db594f45 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 16:19:09 -0400 Subject: [PATCH 07/11] docs: Restructure the subgrids tutorial around the SubGrid API The tutorial was organised around the compiler -- "the generated loop", "when a guard is used instead" -- with SubGrid relegated to a setup step, so the API it is meant to document was never actually covered. Lead with the API instead: creating a SubGrid and what it derives from the parent, what it rejects and why, defining Functions and TimeFunctions on it, writing Operators in both directions, finite differences over the coarse spacing, and the two things SubGrid deliberately leaves to a bare ConditionalDimension (time subsampling and extra conditions). The codegen material follows as a closing section for readers profiling the generated C. Every Operator in the tutorial now runs, with the result asserted. --- examples/userapi/10_subgrids.ipynb | 1225 +++++++++++++++++++++------- 1 file changed, 912 insertions(+), 313 deletions(-) diff --git a/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb index 3d66f1efdd3..4b0f6d3c2e2 100644 --- a/examples/userapi/10_subgrids.ipynb +++ b/examples/userapi/10_subgrids.ipynb @@ -4,271 +4,974 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Subgrids: integer subsampling\n", + "# Subgrids\n", "\n", "A *subgrid* is a coarser grid whose points are a subset of a parent `Grid`'s points.\n", - "`SubGrid` builds one: along each `Dimension` subsampled by `factor`, it keeps every\n", - "`factor`-th point of the parent, so subgrid point `i` sits at parent index\n", - "`i*factor`.\n", + "`SubGrid` is the API for building one.\n", "\n", - "Because the points coincide with parent points, no interpolation is involved. Devito\n", - "iterates the **subgrid** directly -- one unit-stride loop per subsampled `Dimension`,\n", - "with parent accesses rescaled -- rather than walking every parent point under an\n", - "`if (x%factor == 0)` guard.\n", + "This tutorial covers\n", "\n", - "Every section below prints the generated C so you can see exactly what each case\n", - "produces." + "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-20T20:14:13.112164Z", + "iopub.status.busy": "2026-10-20T20:14:13.111560Z", + "iopub.status.idle": "2026-10-20T20:14:14.275586Z", + "shell.execute_reply": "2026-10-20T20:14:14.274253Z" + } + }, + "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-20T20:14:14.333353Z", + "iopub.status.busy": "2026-10-20T20:14:14.332459Z", + "iopub.status.idle": "2026-10-20T20:14:14.341238Z", + "shell.execute_reply": "2026-10-20T20:14:14.339576Z" + } + }, + "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-20T20:14:14.345306Z", + "iopub.status.busy": "2026-10-20T20:14:14.344702Z", + "iopub.status.idle": "2026-10-20T20:14:14.356836Z", + "shell.execute_reply": "2026-10-20T20:14:14.355369Z" + } + }, + "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-20T20:14:14.361979Z", + "iopub.status.busy": "2026-10-20T20:14:14.361535Z", + "iopub.status.idle": "2026-10-20T20:14:14.370227Z", + "shell.execute_reply": "2026-10-20T20:14:14.368724Z" + } + }, + "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-20T20:14:14.375589Z", + "iopub.status.busy": "2026-10-20T20:14:14.374997Z", + "iopub.status.idle": "2026-10-20T20:14:14.385326Z", + "shell.execute_reply": "2026-10-20T20:14:14.383772Z" + } + }, + "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-20T20:14:14.390898Z", + "iopub.status.busy": "2026-10-20T20:14:14.390362Z", + "iopub.status.idle": "2026-10-20T20:14:14.403118Z", + "shell.execute_reply": "2026-10-20T20:14:14.401567Z" + } + }, + "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-20T20:14:14.408340Z", + "iopub.status.busy": "2026-10-20T20:14:14.407861Z", + "iopub.status.idle": "2026-10-20T20:14:14.417364Z", + "shell.execute_reply": "2026-10-20T20:14:14.415772Z" + } + }, + "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-20T20:14:14.422587Z", + "iopub.status.busy": "2026-10-20T20:14:14.422076Z", + "iopub.status.idle": "2026-10-20T20:14:14.431670Z", + "shell.execute_reply": "2026-10-20T20:14:14.430081Z" + } + }, + "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-20T20:14:14.436916Z", + "iopub.status.busy": "2026-10-20T20:14:14.436362Z", + "iopub.status.idle": "2026-10-20T20:14:14.445765Z", + "shell.execute_reply": "2026-10-20T20:14:14.444113Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "factor=2.5 -> Non-integer subsampling factor `2.5` is not supported; SubGrid points must coincide with parent Grid points\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-20T20:14:14.451358Z", + "iopub.status.busy": "2026-10-20T20:14:14.450751Z", + "iopub.status.idle": "2026-10-20T20:14:14.458857Z", + "shell.execute_reply": "2026-10-20T20:14:14.457238Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Cannot subsample a SubGrid; derive from `parent.parent` with a combined factor instead\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-20T20:14:14.464489Z", + "iopub.status.busy": "2026-10-20T20:14:14.463882Z", + "iopub.status.idle": "2026-10-20T20:14:14.482669Z", + "shell.execute_reply": "2026-10-20T20:14:14.481119Z" + } + }, + "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-20T20:14:14.487890Z", + "iopub.status.busy": "2026-10-20T20:14:14.487330Z", + "iopub.status.idle": "2026-10-20T20:14:14.499570Z", + "shell.execute_reply": "2026-10-20T20:14:14.498050Z" + } + }, + "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-20T20:14:14.504630Z", + "iopub.status.busy": "2026-10-20T20:14:14.504153Z", + "iopub.status.idle": "2026-10-20T20:14:14.517643Z", + "shell.execute_reply": "2026-10-20T20:14:14.516143Z" + } + }, + "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-20T20:14:14.522843Z", + "iopub.status.busy": "2026-10-20T20:14:14.522326Z", + "iopub.status.idle": "2026-10-20T20:14:14.717725Z", + "shell.execute_reply": "2026-10-20T20:14:14.716632Z" + } + }, + "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-20T20:14:14.721662Z", + "iopub.status.busy": "2026-10-20T20:14:14.721484Z", + "iopub.status.idle": "2026-10-20T20:14:14.853508Z", + "shell.execute_reply": "2026-10-20T20:14:14.851990Z" + } + }, + "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": [ - "## Creating a SubGrid\n", - "\n", - "`SubGrid(parent, factor)` derives everything from the parent: the dimensions, the\n", - "shape, and the physical spacing." + "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": 1, + "execution_count": 16, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:37.266954Z", - "iopub.status.busy": "2026-10-20T19:47:37.266342Z", - "iopub.status.idle": "2026-10-20T19:47:38.482933Z", - "shell.execute_reply": "2026-10-20T19:47:38.482018Z" + "iopub.execute_input": "2026-10-20T20:14:14.858427Z", + "iopub.status.busy": "2026-10-20T20:14:14.857809Z", + "iopub.status.idle": "2026-10-20T20:14:14.986309Z", + "shell.execute_reply": "2026-10-20T20:14:14.985485Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, { "data": { "text/plain": [ - "SubGrid[shape=(4, 4), factor=(3, 3)]" + "Data([-1., 1., 1., 1.], dtype=float32)" ] }, - "execution_count": 1, + "execution_count": 16, "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", + "# NBVAL_IGNORE_OUTPUT\n", + "u.data[:] = 1.\n", + "usub.data[:] = -1. # sentinel: never written\n", "\n", - "grid = Grid(shape=(12, 12), extent=(11., 11.))\n", - "x, y = grid.dimensions\n", - "sub = SubGrid(grid, factor=3)\n", + "Operator(Eq(usub, u)).apply(x_m=2)\n", "\n", - "sub" + "usub.data[:, 0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "The shape is `ceil(12/3) = 4` along each axis, and the spacing is three times\n", - "the parent's -- the subgrid spans the same physical extent with fewer points." + "Several SubGrids can appear in one Operator, each with its own factor." ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 17, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.541091Z", - "iopub.status.busy": "2026-10-20T19:47:38.540134Z", - "iopub.status.idle": "2026-10-20T19:47:38.550952Z", - "shell.execute_reply": "2026-10-20T19:47:38.549260Z" + "iopub.execute_input": "2026-10-20T20:14:14.990434Z", + "iopub.status.busy": "2026-10-20T20:14:14.990205Z", + "iopub.status.idle": "2026-10-20T20:14:15.176291Z", + "shell.execute_reply": "2026-10-20T20:14:15.174774Z" } }, "outputs": [ { - "name": "stdout", + "name": "stderr", "output_type": "stream", "text": [ - "subgrid shape (4, 4) spacing (3.0, 3.0)\n", - "parent shape (12, 12) spacing (1.0, 1.0)\n" + "Operator `Kernel` ran in 0.01 s\n" ] + }, + { + "data": { + "text/plain": [ + "((6, 6), (3, 3))" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ - "print('subgrid shape ', sub.shape, ' spacing', tuple(float(s) for s in sub.spacing))\n", - "print('parent shape ', grid.shape, ' spacing', tuple(float(s) for s in grid.spacing))" + "# 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": [ - "Its dimensions are `ConditionalDimension`s derived from the parent's, which is\n", - "what carries the subsampling factor into the generated code." + "## 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": 3, + "execution_count": 18, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.554978Z", - "iopub.status.busy": "2026-10-20T19:47:38.554374Z", - "iopub.status.idle": "2026-10-20T19:47:38.564543Z", - "shell.execute_reply": "2026-10-20T19:47:38.562985Z" + "iopub.execute_input": "2026-10-20T20:14:15.181715Z", + "iopub.status.busy": "2026-10-20T20:14:15.180798Z", + "iopub.status.idle": "2026-10-20T20:14:15.192100Z", + "shell.execute_reply": "2026-10-20T20:14:15.190596Z" } }, "outputs": [ { "data": { "text/plain": [ - "((xs, ys), (3, 3))" + "(3*h_x, h_x)" ] }, - "execution_count": 3, + "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "sub.dimensions, sub.factor" + "sub.dimensions[0].spacing, x.spacing" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "A `Function` on a `SubGrid` needs neither `shape` nor `dimensions` -- unlike the\n", - "hand-rolled equivalent, where both must be worked out by hand." + "So the stencil steps `factor*h_x` and the weights scale accordingly:" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 19, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.569006Z", - "iopub.status.busy": "2026-10-20T19:47:38.568430Z", - "iopub.status.idle": "2026-10-20T19:47:38.588366Z", - "shell.execute_reply": "2026-10-20T19:47:38.586951Z" + "iopub.execute_input": "2026-10-20T20:14:15.197416Z", + "iopub.status.busy": "2026-10-20T20:14:15.196754Z", + "iopub.status.idle": "2026-10-20T20:14:15.227420Z", + "shell.execute_reply": "2026-10-20T20:14:15.226557Z" } }, "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": [ - "((12, 12), (4, 4), (xs, ys))" + "-0.333333333*fsub(xs, ys)/h_x + 0.333333333*fsub(xs + 3*h_x, ys)/h_x" ] }, - "execution_count": 4, + "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "u = Function(name='u', grid=grid) # parent grid\n", - "usub = Function(name='usub', grid=sub, space_order=0) # subgrid\n", - "\n", - "u.shape, usub.shape, usub.dimensions" + "fsub = Function(name='fsub', grid=sub, space_order=2)\n", + "fsub.dx.evaluate" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "A factor per `Dimension` subsamples each axis independently; `1` leaves an axis\n", - "untouched." + "Checking it numerically on a linear field, whose derivative is exact:" ] }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 20, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.593261Z", - "iopub.status.busy": "2026-10-20T19:47:38.592855Z", - "iopub.status.idle": "2026-10-20T19:47:38.602233Z", - "shell.execute_reply": "2026-10-20T19:47:38.600963Z" + "iopub.execute_input": "2026-10-20T20:14:15.231982Z", + "iopub.status.busy": "2026-10-20T20:14:15.231745Z", + "iopub.status.idle": "2026-10-20T20:14:15.402026Z", + "shell.execute_reply": "2026-10-20T20:14:15.400448Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, { "data": { "text/plain": [ - "((4, 12), (xs, y))" + "3.0" ] }, - "execution_count": 5, + "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "mixed = SubGrid(grid, factor=(3, 1))\n", - "mixed.shape, mixed.dimensions" + "# 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": [ - "### Equivalent, by hand\n", + "## What SubGrid does not cover\n", "\n", - "`SubGrid` is a convenience over `ConditionalDimension`, which can still be used\n", - "directly when the dimensions need individual treatment. The two are equivalent; the\n", - "manual form just requires computing the shape yourself." + "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": 6, + "execution_count": 21, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.607357Z", - "iopub.status.busy": "2026-10-20T19:47:38.606926Z", - "iopub.status.idle": "2026-10-20T19:47:38.618469Z", - "shell.execute_reply": "2026-10-20T19:47:38.616998Z" + "iopub.execute_input": "2026-10-20T20:14:15.407404Z", + "iopub.status.busy": "2026-10-20T20:14:15.406878Z", + "iopub.status.idle": "2026-10-20T20:14:15.616910Z", + "shell.execute_reply": "2026-10-20T20:14:15.615467Z" } }, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Operator `Kernel` ran in 0.01 s\n" + ] + }, { "data": { "text/plain": [ - "False" + "[0.0, 4.0, 0.0]" ] }, - "execution_count": 6, + "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "xs_manual = ConditionalDimension(name='xs_m', parent=grid.dimensions[0], factor=3)\n", - "f_manual = Function(name='f_manual', grid=grid,\n", - " dimensions=(xs_manual, grid.dimensions[1]),\n", - " shape=(4, 12), space_order=0) # shape by hand\n", + "# 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", - "f_manual.shape == usub.shape" + "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": [ - "Throughout, this helper prints the loop nest of an `Operator` -- the part we care\n", - "about -- rather than the whole kernel." + "**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": 7, + "execution_count": 22, + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-20T20:14:15.620712Z", + "iopub.status.busy": "2026-10-20T20:14:15.620540Z", + "iopub.status.idle": "2026-10-20T20:14:15.630216Z", + "shell.execute_reply": "2026-10-20T20:14:15.628859Z" + } + }, + "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-20T19:47:38.623332Z", - "iopub.status.busy": "2026-10-20T19:47:38.622862Z", - "iopub.status.idle": "2026-10-20T19:47:38.630816Z", - "shell.execute_reply": "2026-10-20T19:47:38.629412Z" + "iopub.execute_input": "2026-10-20T20:14:15.634450Z", + "iopub.status.busy": "2026-10-20T20:14:15.634210Z", + "iopub.status.idle": "2026-10-20T20:14:15.640015Z", + "shell.execute_reply": "2026-10-20T20:14:15.638571Z" } }, "outputs": [], @@ -287,20 +990,18 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## The generated loop\n", - "\n", - "`Eq(usub, u)` reads the parent grid and writes the subgrid." + "### One loop per subsampled Dimension" ] }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 24, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.635815Z", - "iopub.status.busy": "2026-10-20T19:47:38.635346Z", - "iopub.status.idle": "2026-10-20T19:47:38.807419Z", - "shell.execute_reply": "2026-10-20T19:47:38.806525Z" + "iopub.execute_input": "2026-10-20T20:14:15.644270Z", + "iopub.status.busy": "2026-10-20T20:14:15.643924Z", + "iopub.status.idle": "2026-10-20T20:14:15.769538Z", + "shell.execute_reply": "2026-10-20T20:14:15.768248Z" } }, "outputs": [ @@ -347,13 +1048,13 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 25, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.811559Z", - "iopub.status.busy": "2026-10-20T19:47:38.811363Z", - "iopub.status.idle": "2026-10-20T19:47:38.816795Z", - "shell.execute_reply": "2026-10-20T19:47:38.816021Z" + "iopub.execute_input": "2026-10-20T20:14:15.773741Z", + "iopub.status.busy": "2026-10-20T20:14:15.773548Z", + "iopub.status.idle": "2026-10-20T20:14:15.779411Z", + "shell.execute_reply": "2026-10-20T20:14:15.778348Z" } }, "outputs": [ @@ -363,7 +1064,7 @@ "{'xs': 1, 'ys': 1}" ] }, - "execution_count": 9, + "execution_count": 25, "metadata": {}, "output_type": "execute_result" } @@ -376,132 +1077,22 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### Checking the result\n", - "\n", - "`usub` should hold every third row of `u`." - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.821139Z", - "iopub.status.busy": "2026-10-20T19:47:38.820907Z", - "iopub.status.idle": "2026-10-20T19:47:38.964401Z", - "shell.execute_reply": "2026-10-20T19:47:38.963553Z" - } - }, - "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": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# NBVAL_IGNORE_OUTPUT\n", - "u.data[:] = np.arange(144).reshape(12, 12)\n", - "usub.data[:] = 0.\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": [ - "### Restricting the range\n", - "\n", - "Overriding the parent's bounds maps onto the subgrid's. With `x_m=2` and `factor=3`,\n", - "the first subgrid point at or after `x=2` is `x=3`, i.e. subgrid index 1 -- so index 0\n", - "is never written." - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "metadata": { - "execution": { - "iopub.execute_input": "2026-10-20T19:47:38.968547Z", - "iopub.status.busy": "2026-10-20T19:47:38.968332Z", - "iopub.status.idle": "2026-10-20T19:47:39.084027Z", - "shell.execute_reply": "2026-10-20T19:47:39.083270Z" - } - }, - "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": 11, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# NBVAL_IGNORE_OUTPUT\n", - "u.data[:] = 1.\n", - "usub.data[:] = -1. # sentinel: never written\n", - "Operator(Eq(usub, u)).apply(x_m=2)\n", - "\n", - "usub.data[:, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The leading `-1.` is the untouched slot. Devito derived `xs_m = ceil(2/3) = 1`\n", - "from the supplied `x_m=2`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Subsampling the innermost Dimension\n", + "### The innermost Dimension\n", "\n", - "Subsampling `y`, the innermost (vectorized) `Dimension`, is strided too. The\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, which is worse." + "iteration." ] }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 26, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.088359Z", - "iopub.status.busy": "2026-10-20T19:47:39.088139Z", - "iopub.status.idle": "2026-10-20T19:47:39.194975Z", - "shell.execute_reply": "2026-10-20T19:47:39.194129Z" + "iopub.execute_input": "2026-10-20T20:14:15.784450Z", + "iopub.status.busy": "2026-10-20T20:14:15.783862Z", + "iopub.status.idle": "2026-10-20T20:14:15.908451Z", + "shell.execute_reply": "2026-10-20T20:14:15.907115Z" } }, "outputs": [ @@ -533,13 +1124,13 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 27, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.198698Z", - "iopub.status.busy": "2026-10-20T19:47:39.198494Z", - "iopub.status.idle": "2026-10-20T19:47:39.204332Z", - "shell.execute_reply": "2026-10-20T19:47:39.203208Z" + "iopub.execute_input": "2026-10-20T20:14:15.912403Z", + "iopub.status.busy": "2026-10-20T20:14:15.912234Z", + "iopub.status.idle": "2026-10-20T20:14:15.917486Z", + "shell.execute_reply": "2026-10-20T20:14:15.916677Z" } }, "outputs": [ @@ -549,7 +1140,7 @@ "{'x': 1, 'ys': 1}" ] }, - "execution_count": 13, + "execution_count": 27, "metadata": {}, "output_type": "execute_result" } @@ -567,13 +1158,13 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 28, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.208617Z", - "iopub.status.busy": "2026-10-20T19:47:39.208349Z", - "iopub.status.idle": "2026-10-20T19:47:39.328256Z", - "shell.execute_reply": "2026-10-20T19:47:39.326573Z" + "iopub.execute_input": "2026-10-20T20:14:15.921644Z", + "iopub.status.busy": "2026-10-20T20:14:15.921447Z", + "iopub.status.idle": "2026-10-20T20:14:16.042292Z", + "shell.execute_reply": "2026-10-20T20:14:16.041261Z" } }, "outputs": [ @@ -592,7 +1183,7 @@ " [24., 26., 28., 30., 32., 34.]], dtype=float32)" ] }, - "execution_count": 14, + "execution_count": 28, "metadata": {}, "output_type": "execute_result" } @@ -610,20 +1201,20 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Subsampling several Dimensions\n", + "### Several Dimensions at once\n", "\n", - "Each subsampled `Dimension` gets its own stride, independently." + "Each subsampled `Dimension` gets its own loop, independently." ] }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 29, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.333768Z", - "iopub.status.busy": "2026-10-20T19:47:39.333152Z", - "iopub.status.idle": "2026-10-20T19:47:39.467859Z", - "shell.execute_reply": "2026-10-20T19:47:39.467015Z" + "iopub.execute_input": "2026-10-20T20:14:16.046166Z", + "iopub.status.busy": "2026-10-20T20:14:16.045954Z", + "iopub.status.idle": "2026-10-20T20:14:16.259465Z", + "shell.execute_reply": "2026-10-20T20:14:16.258428Z" } }, "outputs": [ @@ -655,13 +1246,13 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 30, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.471512Z", - "iopub.status.busy": "2026-10-20T19:47:39.471331Z", - "iopub.status.idle": "2026-10-20T19:47:39.586008Z", - "shell.execute_reply": "2026-10-20T19:47:39.585209Z" + "iopub.execute_input": "2026-10-20T20:14:16.263293Z", + "iopub.status.busy": "2026-10-20T20:14:16.263111Z", + "iopub.status.idle": "2026-10-20T20:14:16.377852Z", + "shell.execute_reply": "2026-10-20T20:14:16.376981Z" } }, "outputs": [ @@ -678,7 +1269,7 @@ "{'xs': 1, 'ys': 1}" ] }, - "execution_count": 16, + "execution_count": 30, "metadata": {}, "output_type": "execute_result" } @@ -715,13 +1306,13 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 31, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.590169Z", - "iopub.status.busy": "2026-10-20T19:47:39.589960Z", - "iopub.status.idle": "2026-10-20T19:47:39.705734Z", - "shell.execute_reply": "2026-10-20T19:47:39.704904Z" + "iopub.execute_input": "2026-10-20T20:14:16.382311Z", + "iopub.status.busy": "2026-10-20T20:14:16.382138Z", + "iopub.status.idle": "2026-10-20T20:14:16.495138Z", + "shell.execute_reply": "2026-10-20T20:14:16.494108Z" } }, "outputs": [ @@ -756,13 +1347,13 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 32, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.709325Z", - "iopub.status.busy": "2026-10-20T19:47:39.709138Z", - "iopub.status.idle": "2026-10-20T19:47:39.726057Z", - "shell.execute_reply": "2026-10-20T19:47:39.725324Z" + "iopub.execute_input": "2026-10-20T20:14:16.498895Z", + "iopub.status.busy": "2026-10-20T20:14:16.498725Z", + "iopub.status.idle": "2026-10-20T20:14:16.515895Z", + "shell.execute_reply": "2026-10-20T20:14:16.514457Z" } }, "outputs": [ @@ -779,7 +1370,7 @@ "(Data([1., 1., 1.], dtype=float32), Data([1., 1., 1.], dtype=float32))" ] }, - "execution_count": 18, + "execution_count": 32, "metadata": {}, "output_type": "execute_result" } @@ -816,13 +1407,13 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 33, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.729855Z", - "iopub.status.busy": "2026-10-20T19:47:39.729686Z", - "iopub.status.idle": "2026-10-20T19:47:39.858640Z", - "shell.execute_reply": "2026-10-20T19:47:39.857801Z" + "iopub.execute_input": "2026-10-20T20:14:16.520291Z", + "iopub.status.busy": "2026-10-20T20:14:16.520074Z", + "iopub.status.idle": "2026-10-20T20:14:16.650888Z", + "shell.execute_reply": "2026-10-20T20:14:16.650053Z" } }, "outputs": [ @@ -860,13 +1451,13 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 34, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.862144Z", - "iopub.status.busy": "2026-10-20T19:47:39.861969Z", - "iopub.status.idle": "2026-10-20T19:47:39.884751Z", - "shell.execute_reply": "2026-10-20T19:47:39.884003Z" + "iopub.execute_input": "2026-10-20T20:14:16.655049Z", + "iopub.status.busy": "2026-10-20T20:14:16.654876Z", + "iopub.status.idle": "2026-10-20T20:14:16.678551Z", + "shell.execute_reply": "2026-10-20T20:14:16.677075Z" } }, "outputs": [ @@ -901,7 +1492,7 @@ "Data([0., 1., 2., 3.], dtype=float32)" ] }, - "execution_count": 20, + "execution_count": 34, "metadata": {}, "output_type": "execute_result" } @@ -928,13 +1519,13 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 35, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:39.888453Z", - "iopub.status.busy": "2026-10-20T19:47:39.888230Z", - "iopub.status.idle": "2026-10-20T19:47:40.067896Z", - "shell.execute_reply": "2026-10-20T19:47:40.067083Z" + "iopub.execute_input": "2026-10-20T20:14:16.683928Z", + "iopub.status.busy": "2026-10-20T20:14:16.683367Z", + "iopub.status.idle": "2026-10-20T20:14:16.880620Z", + "shell.execute_reply": "2026-10-20T20:14:16.879741Z" } }, "outputs": [ @@ -996,13 +1587,13 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 36, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:40.071659Z", - "iopub.status.busy": "2026-10-20T19:47:40.071480Z", - "iopub.status.idle": "2026-10-20T19:47:40.092014Z", - "shell.execute_reply": "2026-10-20T19:47:40.091258Z" + "iopub.execute_input": "2026-10-20T20:14:16.884679Z", + "iopub.status.busy": "2026-10-20T20:14:16.884473Z", + "iopub.status.idle": "2026-10-20T20:14:16.906315Z", + "shell.execute_reply": "2026-10-20T20:14:16.904866Z" } }, "outputs": [ @@ -1019,7 +1610,7 @@ "[0.0, 4.0, 0.0]" ] }, - "execution_count": 22, + "execution_count": 36, "metadata": {}, "output_type": "execute_result" } @@ -1047,13 +1638,13 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": 37, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:40.095918Z", - "iopub.status.busy": "2026-10-20T19:47:40.095708Z", - "iopub.status.idle": "2026-10-20T19:47:40.540497Z", - "shell.execute_reply": "2026-10-20T19:47:40.539593Z" + "iopub.execute_input": "2026-10-20T20:14:16.911529Z", + "iopub.status.busy": "2026-10-20T20:14:16.910937Z", + "iopub.status.idle": "2026-10-20T20:14:17.060446Z", + "shell.execute_reply": "2026-10-20T20:14:17.059354Z" } }, "outputs": [ @@ -1111,13 +1702,13 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": 38, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T19:47:40.545872Z", - "iopub.status.busy": "2026-10-20T19:47:40.545436Z", - "iopub.status.idle": "2026-10-20T19:47:40.895126Z", - "shell.execute_reply": "2026-10-20T19:47:40.894398Z" + "iopub.execute_input": "2026-10-20T20:14:17.064287Z", + "iopub.status.busy": "2026-10-20T20:14:17.064083Z", + "iopub.status.idle": "2026-10-20T20:14:17.222603Z", + "shell.execute_reply": "2026-10-20T20:14:17.221456Z" } }, "outputs": [ @@ -1172,22 +1763,30 @@ "source": [ "## Summary\n", "\n", - "Subsampling along a `Dimension` is declared with `SubGrid` (or a\n", - "`ConditionalDimension` directly). Devito then generates, for that `Dimension`:\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", - "| | subsampled loop | guarded loop |\n", - "|---|---|---|\n", - "| iterates | `xs`, unit stride | `x`, unit stride |\n", - "| bounds | `ceil(x_m/f) .. floor(x_M/f)` | `x_m .. x_M` |\n", - "| body | unconditional, parent access `xs*f` | wrapped in `if (x%f == 0)` |\n", - "| iterations | one in `factor` | all |\n", - "| used when | the Dimension is parallel and accesses are affine | otherwise |\n", + "Use a `ConditionalDimension` directly for subsampling in **time**, or to combine a\n", + "factor with an extra `condition` -- `SubGrid` covers neither.\n", "\n", - "The subsampled form is the default, including for the innermost (vectorized)\n", - "`Dimension` and for several subsampled `Dimension`s at once. The guarded form is the\n", - "fallback for the four cases above: a loop nest shared with an unguarded equation, a\n", - "sequential `Dimension` such as time, an extra `condition`, or non-affine accesses.\n", - "Both compute the same thing.\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." @@ -1196,7 +1795,7 @@ ], "metadata": { "kernelspec": { - "display_name": "devito", + "display_name": "devitoenv (3.11.0)", "language": "python", "name": "python3" }, From d07af9168cab6725d8460ea4d974fbbd0501e01c Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 16:29:52 -0400 Subject: [PATCH 08/11] api: Make SubGrid's rejection messages actionable Both said what was disallowed without saying what to do instead. Nesting now names the equivalent single-level call, since the composed factor is just the product of the two: Cannot subsample a SubGrid. `3` applied to a SubGrid of factor `(2, 2)` is the same as factor `(6, 6)` applied to its parent Grid directly, so use `SubGrid(parent.parent, factor=(6, 6))` This needs the factors validated first, so the check moves after `__factors_setup__`. A non-integer factor now names the nearest legal values, clamped to >= 1 so the suggestion is never itself rejected. Tests assert the suggested factors are both legal and equivalent to the two-step subsampling they replace, so the messages cannot drift. --- devito/types/grid.py | 30 ++++++++++++++----- tests/test_subsampling.py | 63 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 82 insertions(+), 11 deletions(-) diff --git a/devito/types/grid.py b/devito/types/grid.py index ea3aaf0df0c..91170c905ce 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 @@ -491,13 +492,22 @@ 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__}`") - if isinstance(parent, SubGrid): - # Would require a ConditionalDimension to parent another one - raise ValueError("Cannot subsample a SubGrid; derive from " - "`parent.parent` with a combined factor instead") 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) @@ -541,9 +551,15 @@ def __factors_setup__(cls, parent, factor): 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}` is not " - "supported; SubGrid points must coincide with " - "parent Grid points") + 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}`") diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index 63546f9306d..22ce5ed177b 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -1,3 +1,5 @@ +import re + import numpy as np import pytest from sympy import Or @@ -7,6 +9,7 @@ TimeFunction, switchconfig ) from devito.ir.iet import Conditional, FindNodes, Iteration +from devito.tools import as_tuple def iteration(op, name): @@ -728,12 +731,64 @@ def test_invalid_factor(self, factor, err): with pytest.raises(ValueError, match=err): SubGrid(grid, factor=factor) - def test_cannot_nest(self): + @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)) - sub = SubGrid(grid, factor=2) - with pytest.raises(ValueError, match='Cannot subsample a SubGrid'): - SubGrid(sub, factor=2) + 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'): From aab92cb048d9b1106e158e91465d4352446b12e8 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 20:45:49 -0400 Subject: [PATCH 09/11] tests: Fix the subsampling checks under linearization Two CI failures on this branch, both in the new material rather than the compiler. `test_parent_access_rescaled` asserted on literal substrings of the emitted C (`g[xs]`), which only hold when linearization is off; the CXX jobs render the same access as the macro `gL0(xs, y)`. Assert on the lowered Expression instead, selecting the one writing `g` since linearization adds Size/Stride Expressions alongside it. Verified against all four CI language settings: C, CXX, openmp, CXXopenmp. The tutorial's stored outputs were stale: the error-message cells still showed the wording from before the previous commit reworded them, so nbval compared them against the new text and failed. Re-executed. --- examples/userapi/10_subgrids.ipynb | 308 ++++++++++++++--------------- tests/test_subsampling.py | 28 ++- 2 files changed, 175 insertions(+), 161 deletions(-) diff --git a/examples/userapi/10_subgrids.ipynb b/examples/userapi/10_subgrids.ipynb index 4b0f6d3c2e2..5a2b9b92df8 100644 --- a/examples/userapi/10_subgrids.ipynb +++ b/examples/userapi/10_subgrids.ipynb @@ -37,10 +37,10 @@ "execution_count": 1, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:13.112164Z", - "iopub.status.busy": "2026-10-20T20:14:13.111560Z", - "iopub.status.idle": "2026-10-20T20:14:14.275586Z", - "shell.execute_reply": "2026-10-20T20:14:14.274253Z" + "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": [ @@ -80,10 +80,10 @@ "execution_count": 2, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.333353Z", - "iopub.status.busy": "2026-10-20T20:14:14.332459Z", - "iopub.status.idle": "2026-10-20T20:14:14.341238Z", - "shell.execute_reply": "2026-10-20T20:14:14.339576Z" + "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": [ @@ -112,10 +112,10 @@ "execution_count": 3, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.345306Z", - "iopub.status.busy": "2026-10-20T20:14:14.344702Z", - "iopub.status.idle": "2026-10-20T20:14:14.356836Z", - "shell.execute_reply": "2026-10-20T20:14:14.355369Z" + "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": [ @@ -150,10 +150,10 @@ "execution_count": 4, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.361979Z", - "iopub.status.busy": "2026-10-20T20:14:14.361535Z", - "iopub.status.idle": "2026-10-20T20:14:14.370227Z", - "shell.execute_reply": "2026-10-20T20:14:14.368724Z" + "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": [ @@ -185,10 +185,10 @@ "execution_count": 5, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.375589Z", - "iopub.status.busy": "2026-10-20T20:14:14.374997Z", - "iopub.status.idle": "2026-10-20T20:14:14.385326Z", - "shell.execute_reply": "2026-10-20T20:14:14.383772Z" + "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": [ @@ -220,10 +220,10 @@ "execution_count": 6, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.390898Z", - "iopub.status.busy": "2026-10-20T20:14:14.390362Z", - "iopub.status.idle": "2026-10-20T20:14:14.403118Z", - "shell.execute_reply": "2026-10-20T20:14:14.401567Z" + "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": [ @@ -255,10 +255,10 @@ "execution_count": 7, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.408340Z", - "iopub.status.busy": "2026-10-20T20:14:14.407861Z", - "iopub.status.idle": "2026-10-20T20:14:14.417364Z", - "shell.execute_reply": "2026-10-20T20:14:14.415772Z" + "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": [ @@ -289,10 +289,10 @@ "execution_count": 8, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.422587Z", - "iopub.status.busy": "2026-10-20T20:14:14.422076Z", - "iopub.status.idle": "2026-10-20T20:14:14.431670Z", - "shell.execute_reply": "2026-10-20T20:14:14.430081Z" + "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": [ @@ -327,10 +327,10 @@ "execution_count": 9, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.436916Z", - "iopub.status.busy": "2026-10-20T20:14:14.436362Z", - "iopub.status.idle": "2026-10-20T20:14:14.445765Z", - "shell.execute_reply": "2026-10-20T20:14:14.444113Z" + "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": [ @@ -338,7 +338,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "factor=2.5 -> Non-integer subsampling factor `2.5` is not supported; SubGrid points must coincide with parent Grid points\n", + "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" ] @@ -365,10 +365,10 @@ "execution_count": 10, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.451358Z", - "iopub.status.busy": "2026-10-20T20:14:14.450751Z", - "iopub.status.idle": "2026-10-20T20:14:14.458857Z", - "shell.execute_reply": "2026-10-20T20:14:14.457238Z" + "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": [ @@ -376,7 +376,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Cannot subsample a SubGrid; derive from `parent.parent` with a combined factor instead\n" + "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" ] } ], @@ -402,10 +402,10 @@ "execution_count": 11, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.464489Z", - "iopub.status.busy": "2026-10-20T20:14:14.463882Z", - "iopub.status.idle": "2026-10-20T20:14:14.482669Z", - "shell.execute_reply": "2026-10-20T20:14:14.481119Z" + "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": [ @@ -439,10 +439,10 @@ "execution_count": 12, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.487890Z", - "iopub.status.busy": "2026-10-20T20:14:14.487330Z", - "iopub.status.idle": "2026-10-20T20:14:14.499570Z", - "shell.execute_reply": "2026-10-20T20:14:14.498050Z" + "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": [ @@ -478,10 +478,10 @@ "execution_count": 13, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.504630Z", - "iopub.status.busy": "2026-10-20T20:14:14.504153Z", - "iopub.status.idle": "2026-10-20T20:14:14.517643Z", - "shell.execute_reply": "2026-10-20T20:14:14.516143Z" + "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": [ @@ -521,10 +521,10 @@ "execution_count": 14, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.522843Z", - "iopub.status.busy": "2026-10-20T20:14:14.522326Z", - "iopub.status.idle": "2026-10-20T20:14:14.717725Z", - "shell.execute_reply": "2026-10-20T20:14:14.716632Z" + "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": [ @@ -575,10 +575,10 @@ "execution_count": 15, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.721662Z", - "iopub.status.busy": "2026-10-20T20:14:14.721484Z", - "iopub.status.idle": "2026-10-20T20:14:14.853508Z", - "shell.execute_reply": "2026-10-20T20:14:14.851990Z" + "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": [ @@ -629,10 +629,10 @@ "execution_count": 16, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.858427Z", - "iopub.status.busy": "2026-10-20T20:14:14.857809Z", - "iopub.status.idle": "2026-10-20T20:14:14.986309Z", - "shell.execute_reply": "2026-10-20T20:14:14.985485Z" + "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": [ @@ -676,10 +676,10 @@ "execution_count": 17, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:14.990434Z", - "iopub.status.busy": "2026-10-20T20:14:14.990205Z", - "iopub.status.idle": "2026-10-20T20:14:15.176291Z", - "shell.execute_reply": "2026-10-20T20:14:15.174774Z" + "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": [ @@ -731,10 +731,10 @@ "execution_count": 18, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.181715Z", - "iopub.status.busy": "2026-10-20T20:14:15.180798Z", - "iopub.status.idle": "2026-10-20T20:14:15.192100Z", - "shell.execute_reply": "2026-10-20T20:14:15.190596Z" + "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": [ @@ -765,10 +765,10 @@ "execution_count": 19, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.197416Z", - "iopub.status.busy": "2026-10-20T20:14:15.196754Z", - "iopub.status.idle": "2026-10-20T20:14:15.227420Z", - "shell.execute_reply": "2026-10-20T20:14:15.226557Z" + "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": [ @@ -803,10 +803,10 @@ "execution_count": 20, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.231982Z", - "iopub.status.busy": "2026-10-20T20:14:15.231745Z", - "iopub.status.idle": "2026-10-20T20:14:15.402026Z", - "shell.execute_reply": "2026-10-20T20:14:15.400448Z" + "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": [ @@ -864,10 +864,10 @@ "execution_count": 21, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.407404Z", - "iopub.status.busy": "2026-10-20T20:14:15.406878Z", - "iopub.status.idle": "2026-10-20T20:14:15.616910Z", - "shell.execute_reply": "2026-10-20T20:14:15.615467Z" + "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": [ @@ -917,10 +917,10 @@ "execution_count": 22, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.620712Z", - "iopub.status.busy": "2026-10-20T20:14:15.620540Z", - "iopub.status.idle": "2026-10-20T20:14:15.630216Z", - "shell.execute_reply": "2026-10-20T20:14:15.628859Z" + "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": [ @@ -968,10 +968,10 @@ "execution_count": 23, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.634450Z", - "iopub.status.busy": "2026-10-20T20:14:15.634210Z", - "iopub.status.idle": "2026-10-20T20:14:15.640015Z", - "shell.execute_reply": "2026-10-20T20:14:15.638571Z" + "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": [], @@ -998,10 +998,10 @@ "execution_count": 24, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.644270Z", - "iopub.status.busy": "2026-10-20T20:14:15.643924Z", - "iopub.status.idle": "2026-10-20T20:14:15.769538Z", - "shell.execute_reply": "2026-10-20T20:14:15.768248Z" + "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": [ @@ -1051,10 +1051,10 @@ "execution_count": 25, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.773741Z", - "iopub.status.busy": "2026-10-20T20:14:15.773548Z", - "iopub.status.idle": "2026-10-20T20:14:15.779411Z", - "shell.execute_reply": "2026-10-20T20:14:15.778348Z" + "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": [ @@ -1089,10 +1089,10 @@ "execution_count": 26, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.784450Z", - "iopub.status.busy": "2026-10-20T20:14:15.783862Z", - "iopub.status.idle": "2026-10-20T20:14:15.908451Z", - "shell.execute_reply": "2026-10-20T20:14:15.907115Z" + "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": [ @@ -1127,10 +1127,10 @@ "execution_count": 27, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.912403Z", - "iopub.status.busy": "2026-10-20T20:14:15.912234Z", - "iopub.status.idle": "2026-10-20T20:14:15.917486Z", - "shell.execute_reply": "2026-10-20T20:14:15.916677Z" + "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": [ @@ -1161,10 +1161,10 @@ "execution_count": 28, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:15.921644Z", - "iopub.status.busy": "2026-10-20T20:14:15.921447Z", - "iopub.status.idle": "2026-10-20T20:14:16.042292Z", - "shell.execute_reply": "2026-10-20T20:14:16.041261Z" + "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": [ @@ -1211,10 +1211,10 @@ "execution_count": 29, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.046166Z", - "iopub.status.busy": "2026-10-20T20:14:16.045954Z", - "iopub.status.idle": "2026-10-20T20:14:16.259465Z", - "shell.execute_reply": "2026-10-20T20:14:16.258428Z" + "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": [ @@ -1249,10 +1249,10 @@ "execution_count": 30, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.263293Z", - "iopub.status.busy": "2026-10-20T20:14:16.263111Z", - "iopub.status.idle": "2026-10-20T20:14:16.377852Z", - "shell.execute_reply": "2026-10-20T20:14:16.376981Z" + "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": [ @@ -1309,10 +1309,10 @@ "execution_count": 31, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.382311Z", - "iopub.status.busy": "2026-10-20T20:14:16.382138Z", - "iopub.status.idle": "2026-10-20T20:14:16.495138Z", - "shell.execute_reply": "2026-10-20T20:14:16.494108Z" + "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": [ @@ -1350,10 +1350,10 @@ "execution_count": 32, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.498895Z", - "iopub.status.busy": "2026-10-20T20:14:16.498725Z", - "iopub.status.idle": "2026-10-20T20:14:16.515895Z", - "shell.execute_reply": "2026-10-20T20:14:16.514457Z" + "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": [ @@ -1410,10 +1410,10 @@ "execution_count": 33, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.520291Z", - "iopub.status.busy": "2026-10-20T20:14:16.520074Z", - "iopub.status.idle": "2026-10-20T20:14:16.650888Z", - "shell.execute_reply": "2026-10-20T20:14:16.650053Z" + "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": [ @@ -1454,10 +1454,10 @@ "execution_count": 34, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.655049Z", - "iopub.status.busy": "2026-10-20T20:14:16.654876Z", - "iopub.status.idle": "2026-10-20T20:14:16.678551Z", - "shell.execute_reply": "2026-10-20T20:14:16.677075Z" + "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": [ @@ -1522,10 +1522,10 @@ "execution_count": 35, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.683928Z", - "iopub.status.busy": "2026-10-20T20:14:16.683367Z", - "iopub.status.idle": "2026-10-20T20:14:16.880620Z", - "shell.execute_reply": "2026-10-20T20:14:16.879741Z" + "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": [ @@ -1590,10 +1590,10 @@ "execution_count": 36, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.884679Z", - "iopub.status.busy": "2026-10-20T20:14:16.884473Z", - "iopub.status.idle": "2026-10-20T20:14:16.906315Z", - "shell.execute_reply": "2026-10-20T20:14:16.904866Z" + "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": [ @@ -1641,10 +1641,10 @@ "execution_count": 37, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:16.911529Z", - "iopub.status.busy": "2026-10-20T20:14:16.910937Z", - "iopub.status.idle": "2026-10-20T20:14:17.060446Z", - "shell.execute_reply": "2026-10-20T20:14:17.059354Z" + "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": [ @@ -1705,10 +1705,10 @@ "execution_count": 38, "metadata": { "execution": { - "iopub.execute_input": "2026-10-20T20:14:17.064287Z", - "iopub.status.busy": "2026-10-20T20:14:17.064083Z", - "iopub.status.idle": "2026-10-20T20:14:17.222603Z", - "shell.execute_reply": "2026-10-20T20:14:17.221456Z" + "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": [ diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index 22ce5ed177b..8e7643d3cea 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -8,7 +8,8 @@ CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, SubGrid, TimeFunction, switchconfig ) -from devito.ir.iet import Conditional, FindNodes, Iteration +from devito.ir.iet import Conditional, Expression, FindNodes, Iteration +from devito.symbolics import retrieve_indexed from devito.tools import as_tuple @@ -59,15 +60,28 @@ def test_codegen(self): @pytest.mark.parametrize('factor', [2, 3, 4, 5]) def test_parent_access_rescaled(self, factor): - """Parent accesses become `cd*factor`; subgrid accesses become `cd`.""" + """ + 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)) - ccode = str(op.ccode) - fname = cd.symbolic_factor.name - assert f'{cd.name}*{fname}' in ccode # parent, rescaled - assert f'g[{cd.name}]' in ccode # subgrid, direct - assert f'+= {fname}' not in ccode # no strided loop + # 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): From 1ffb50a5e2cf4fd255f8f44a71ad6da9a93d3647 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 21:25:13 -0400 Subject: [PATCH 10/11] tests: Add MPI coverage for subsampling Fix `SubGrid.root` to return the parent's root. Like a SubDomain, a SubGrid is a view of its parent, and `Operator._prepare_arguments` collapses the discretizations it finds by `root` to tell several views of one Grid apart from several unrelated Grids -- rejecting the latter under MPI. Without this, any Operator touching both a Grid and a SubGrid of it raised "Multiple Grids found", so subsampling could not be used with MPI at all. With that fixed, add the tests: a strided copy against numpy slicing for factors 2/3/4, the codegen still being a unit-stride subgrid loop, the per-rank bounds tiling the global subgrid with no gaps or overlap, the SubGrid sharing its parent's decomposition, and time subsampling. All on 2 and 4 ranks. Note a reduction into a subsampled Function (`Inc(g, u)`) is rejected under MPI with "Unsupported MPI for the given equations", so the partitioning check uses a sentinel plus an allreduce over the local sizes rather than Inc. --- devito/types/grid.py | 9 ++++ tests/test_subsampling.py | 102 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/devito/types/grid.py b/devito/types/grid.py index 91170c905ce..4568b852bf1 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -579,6 +579,15 @@ 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)}]" diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index 8e7643d3cea..daa9956ae3d 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -813,3 +813,105 @@ def test_custom_dimension_names(self): 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) From c5e56cdde4fa0947a897f1e97aa101543dc76ef5 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 20 Oct 2026 23:45:51 -0400 Subject: [PATCH 11/11] compiler: Don't treat a subsampled Inc as a reduction `Inc` into a factor-subsampled Function was classified PARALLEL_IF_ATOMIC, so an OpenMP array reduction was emitted over the whole Function. With nine such Functions the thread-private copies came to ~13 MiB and the kernel segfaulted; four arrays worked, six did not, which is the size of the default stack. The dependence is reported with an all-zero distance vector, so the write is plainly independent across iterations. What pushed it into the reduction branch was `is_storage_related`, which treats any NonlinearDerived Dimension as capable of mapping several iterations onto one location. That is true of a SteppingDimension, where `time % 2` deliberately revisits a slot, but not of a factor-based ConditionalDimension: the guard admits only `parent % factor == 0`, and `i -> i//factor` is one-to-one on those. Add `Dimension.is_Injective`, true exactly for that case, and exempt it. An arbitrary `condition`, or a factor combined with a condition, carries no such promise and stays non-injective. --- devito/ir/support/basic.py | 11 +++- devito/types/dimension.py | 14 ++++ tests/test_subsampling.py | 131 ++++++++++++++++++++++++++++++++++++- 3 files changed, 153 insertions(+), 3 deletions(-) diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py index 0bce1b9387e..326562b5c9f 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/types/dimension.py b/devito/types/dimension.py index a0b4d5fede8..6b21493c146 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 diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py index daa9956ae3d..07ff1a5a1f5 100644 --- a/tests/test_subsampling.py +++ b/tests/test_subsampling.py @@ -5,7 +5,7 @@ from sympy import Or from devito import ( - CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Operator, SubGrid, + CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Inc, Operator, SubGrid, TimeFunction, switchconfig ) from devito.ir.iet import Conditional, Expression, FindNodes, Iteration @@ -915,3 +915,132 @@ def test_time_subsampling(self, mode): 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