From 7144885220a9f54310dc123bbb9a051467599b4c Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Fri, 6 Nov 2026 14:31:38 -0500
Subject: [PATCH 1/3] compiler: Block strided subsampling and align SubGrids to
 global indices

Blocking replaced the iteration Dimensions, which degraded a GuardFactor to a
plain CondEq and left blocked Dimensions guarded. Run `subsampling` before
`_specialize_clusters` instead, so that tiles span the subsampled Dimensions.

Under MPI a SubGrid applied its factor to rank-local indices, so a rank not
starting on a subsampled point read the wrong parent points. Give
ConditionalDimension an `offset`, the intercept of its map onto the parent,
through a SubsamplingOffset symbol localized by MPI like a Thickness.

Decompose a SubGrid through SubGridDistributor, so that a rank owns exactly the
points lying on its own parent points, and derive the subsampled bounds from
the parent's rank-local ones, which stops strided loops overrunning a rank.
---
 devito/ir/equations/algorithms.py             |   2 +-
 devito/ir/support/guards.py                   |   4 +-
 devito/mpi/distributed.py                     |  45 +++-
 devito/operator/operator.py                   |  22 +-
 devito/passes/clusters/subsampling.py         |  19 +-
 devito/types/dimension.py                     |  76 +++++-
 devito/types/grid.py                          |  24 +-
 .../userapi/05_conditional_dimension.ipynb    |   3 +
 tests/test_subsampling.py                     | 246 +++++++++++++++++-
 9 files changed, 395 insertions(+), 46 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index de477b02922..44ad31a4f31 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -390,7 +390,7 @@ def generate_conditionals(expr, input_expr, ordering):
 
     # Replace the ConditionalDimensions in `expr`
     for d, cond in conditionals.items():
-        index = d.index
+        index = d.index - d.offset
         if d.condition is not None and expr.has(d):
             index = index - relational_min(cond, d.parent)
         shift = relational_shift(cond, d.parent, d.symbolic_factor)
diff --git a/devito/ir/support/guards.py b/devito/ir/support/guards.py
index 3c4b573c034..13154846347 100644
--- a/devito/ir/support/guards.py
+++ b/devito/ir/support/guards.py
@@ -61,7 +61,7 @@ class GuardFactor(Guard, CondEq, Pickable):
     A guard for factor-based ConditionalDimensions.
 
     Given the ConditionalDimension `d` with factor `k`, create the
-    symbolic relational `d.parent % k == 0`.
+    symbolic relational `(d.parent - d.offset) % k == 0`.
     """
 
     __rargs__ = ('d',)
@@ -76,7 +76,7 @@ def __new__(cls, *args, **kwargs):
         d, = args
         assert d.is_Conditional
 
-        obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
+        obj = super().__new__(cls, (d.parent - d.offset) % d.symbolic_factor, 0)
         obj.d = d
 
         return obj
diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py
index 35596c39d64..fa5ea2e0509 100644
--- a/devito/mpi/distributed.py
+++ b/devito/mpi/distributed.py
@@ -68,6 +68,7 @@ def __getattr__(self, name):
     'Distributor',
     'SparseDistributor',
     'SubDistributor',
+    'SubGridDistributor',
     'devito_mpi_finalize',
     'devito_mpi_init',
 ]
@@ -375,10 +376,15 @@ def __init__(self, shape, dimensions, input_comm=None, topology=None):
             self._comm = MPI.COMM_NULL
             self._topology = tuple(1 for _ in range(len(shape)))
 
-        # The domain decomposition
+        self.__decomposition_setup__()
+
+    def __decomposition_setup__(self):
+        """
+        Set up the domain decomposition.
+        """
         self._decomposition = [
             Decomposition(np.array_split(range(i), j), c)
-            for i, j, c in zip(shape, self.topology, self.mycoords, strict=True)
+            for i, j, c in zip(self.glb_shape, self.topology, self.mycoords, strict=True)
         ]
 
     @cached_property
@@ -485,6 +491,41 @@ def _rebuild(self, shape=None, dimensions=None, comm=None):
                            comm or self.comm)
 
 
+class SubGridDistributor(Distributor):
+
+    """
+    Decompose a SubGrid over a set of MPI processes in a manner consistent with
+    the parent distributor: a rank owns the SubGrid points on its parent points.
+
+    Parameters
+    ----------
+    subgrid : SubGrid
+        The SubGrid to be decomposed.
+    """
+
+    def __init__(self, subgrid):
+        self._parent = subgrid.parent.distributor
+        self._factors = subgrid.factor
+
+        super().__init__(subgrid.shape, subgrid.dimensions, self.parent.comm,
+                         tuple(self.parent.topology))
+
+    def __decomposition_setup__(self):
+        """
+        Set up the decomposition, aligned with that of the parent Distributor.
+        """
+        self._decomposition = [
+            Decomposition([d[d % f == 0] // f for d in dec], c)
+            for dec, f, c in zip(self.parent.decomposition, self._factors,
+                                 self.mycoords, strict=True)
+        ]
+
+    @property
+    def parent(self):
+        """The parent distributor of this SubGridDistributor."""
+        return self._parent
+
+
 class SubDistributor(DenseDistributor):
 
     """
diff --git a/devito/operator/operator.py b/devito/operator/operator.py
index 38b2d4dc28f..52bf4905982 100644
--- a/devito/operator/operator.py
+++ b/devito/operator/operator.py
@@ -43,7 +43,7 @@
     split, timed_pass, timed_region
 )
 from devito.types import Buffer, Evaluable, device_layer, disk_layer, host_layer
-from devito.types.dimension import SubDimension, Thickness
+from devito.types.dimension import SubDimension, SubsamplingOffset, Thickness
 from devito.warnings import warn
 
 __all__ = ['Operator']
@@ -415,6 +415,11 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         # Operation count before specialization
         init_ops = sum(estimate_cost(c.exprs) for c in clusters if c.is_dense)
 
+        # Strengthen factor-based ConditionalDimension guards into strided
+        # iteration. Before specialization, so that blocking tiles the swapped
+        # Dimensions
+        clusters = subsampling(clusters)
+
         clusters = cls._specialize_clusters(clusters, **kwargs)
 
         # Operation count after specialization
@@ -425,11 +430,6 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         # Generate implicit Clusters from higher level abstractions
         clusters = generate_implicit(clusters)
 
-        # Strengthen factor-based ConditionalDimension guards into strided
-        # iteration. Before `lower_index_derivatives` so that the Dimension
-        # swap is seen by everything downstream
-        clusters = subsampling(clusters)
-
         # Lower all remaining high order symbolic objects
         clusters = lower_local_sums(clusters, **kwargs)
         clusters = lower_index_derivatives(clusters, **kwargs)
@@ -695,15 +695,15 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
             # No need to do anything more if only checking the memory
             return args
 
+        # Process the symbols localized by MPI
+        for p in self.parameters:
+            if isinstance(p, (SubsamplingOffset, Thickness)):
+                args.update(p._arg_values(grid=grid, **kwargs))
+
         # Process Dimensions
         for d in reversed(toposort):
             args.update(d._arg_values(self._dspace[d], grid, **kwargs))
 
-        # Process Thicknesses
-        for p in self.parameters:
-            if isinstance(p, Thickness):
-                args.update(p._arg_values(grid=grid, **kwargs))
-
         # Process Objects
         for o in self.objects:
             args.update(o._arg_values(grid=grid, **kwargs))
diff --git a/devito/passes/clusters/subsampling.py b/devito/passes/clusters/subsampling.py
index 7ef3e7f35e6..c138d31968a 100644
--- a/devito/passes/clusters/subsampling.py
+++ b/devito/passes/clusters/subsampling.py
@@ -20,8 +20,8 @@ def subsampling(clusters):
           if (x%xf == 0)                 ->     g[x/xf] = ...
             g[x/xf] = ...
 
-The loop is taken over the subsampled Dimension `cd` itself, with a *unit*
-    stride, and every access to the parent `d` is rescaled to `cd*factor`:
+    The loop is taken over the subsampled Dimension `cd` itself, with a *unit*
+    stride, and every access to the parent `d` is mapped to `cd*factor + offset`:
 
         for (x = x_m; x <= x_M; x += 1)      for (xs = xs_m; xs <= xs_M; xs++)
           if (x%xf == 0)                 ->     g[xs] = u[xs*xf]
@@ -30,8 +30,8 @@ def subsampling(clusters):
     A non-unitary increment is deliberately avoided: it would escape the
     Lamport-style dependence analysis the compiler relies on. The loop bounds
     are `cd`'s own runtime values, which `ConditionalDimension._arg_values`
-    derives as `ceil(x_m/factor) .. floor(x_M/factor)`, so no anchoring of the
-    loop minimum is needed either.
+    derives by inverting that map, so no anchoring of the loop minimum is needed
+    either.
 
     Parameters
     ----------
@@ -73,17 +73,18 @@ def callback(self, clusters, prefix):
 
         cd = guard.d
         factor = cd.symbolic_factor
+        offset = cd.offset
 
         processed = []
         for c in clusters:
             # Iterate the subsampled Dimension itself, with unit stride, and
-            # rescale every parent access `d` to `cd*factor`
+            # map every parent access `d` to `cd*factor + offset`
             ispace = c.ispace.switch(d, cd)
 
-            # `IntDiv(d, factor)` is how the subsampled access was lowered;
-            # with `d -> cd*factor` it collapses to `cd`, so rewrite it
-            # directly rather than leaving `(cd*factor)/factor` behind
-            subs = {IntDiv(d, factor): cd, d: cd*factor}
+            # `IntDiv(d - offset, factor)` is how the subsampled access was
+            # lowered; with `d -> cd*factor + offset` it collapses to `cd`, so
+            # rewrite it directly rather than leaving the division behind
+            subs = {IntDiv(d - offset, factor): cd, d: cd*factor + offset}
             exprs = [uxreplace(e, subs) for e in c.exprs]
 
             # `cd` inherits `d`'s computational properties: it covers the very
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index 6b21493c146..f49927801bd 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -942,6 +942,56 @@ class SubsamplingFactor(Scalar):
     pass
 
 
+class SubsamplingOffset(DataSymbol):
+
+    """
+    The local index of the first subsampled point on an MPI rank.
+
+    Subsampling is defined on global indices, but a rank need not start on a
+    subsampled point. E.g., with `factor=2` and two ranks:
+
+        global      0  1  2 ... 10 | 11 12 13 14 ... 20
+        local       0  1  2 ... 10 |  0  1  2  3 ...  9
+        subsampled  *  .  * ...  * |  .  *  .  * ...  *
+        offset      0              |     1
+
+    In local indices, `parent = factor*subsampled + offset`.
+    """
+
+    __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'factor')
+
+    def __new__(cls, *args, root=None, factor=None, **kwargs):
+        newobj = super().__new__(cls, *args, **kwargs)
+        newobj._root = root
+        newobj._factor = factor
+
+        return newobj
+
+    def __init_finalize__(self, *args, **kwargs):
+        kwargs.setdefault('dtype', np.int32)
+        kwargs.setdefault('is_const', True)
+        kwargs.setdefault('nonnegative', True)
+        super().__init_finalize__(*args, **kwargs)
+
+    @property
+    def root(self):
+        return self._root
+
+    @property
+    def factor(self):
+        return self._factor
+
+    def _arg_values(self, grid=None, **kwargs):
+        if grid is None:
+            return {self.name: 0}
+
+        # Runtime arguments bind by name, even when different spacing dtypes
+        # make the compiled and runtime Grid Dimensions distinct symbols
+        root = next(d for d in grid.dimensions if d.name == self.root.name)
+
+        return {self.name: -grid.origin_ioffset[root] % self.factor}
+
+
 class ConditionalDimension(DerivedDimension):
 
     """
@@ -971,6 +1021,9 @@ class ConditionalDimension(DerivedDimension):
         lowering for example combining Function's ConditionalDimension with
         an Equation's implicit_dim. All Dimensions within an equation
         must have `Or` relation for the final combined condition to be Or.
+    offset : SubsamplingOffset, optional, default=0
+        The intercept of the map onto the parent Dimension, that is
+        ``parent = factor*self + offset``.
 
     Examples
     --------
@@ -1065,10 +1118,10 @@ class ConditionalDimension(DerivedDimension):
     STRICT = 'strict'
 
     __rkwargs__ = DerivedDimension.__rkwargs__ + \
-        ('factor', 'condition', 'indirect', 'relation')
+        ('factor', 'condition', 'indirect', 'relation', 'offset')
 
     def __init_finalize__(self, name, parent=None, factor=None, condition=None,
-                          indirect=False, relation=sympy.And, **kwargs):
+                          indirect=False, relation=sympy.And, offset=0, **kwargs):
         # `parent=None` degenerates to a ConditionalDimension outside of
         # any iteration space
         if parent is None:
@@ -1087,6 +1140,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
         else:
             raise ValueError("factor must be an integer")
 
+        self._offset = offset
         self._condition = condition
         self._indirect = indirect
         self._relation = relation
@@ -1136,6 +1190,10 @@ def symbolic_factor(self):
                 name=f'{self.name}f', dtype=np.int32, is_const=True
             )
 
+    @property
+    def offset(self):
+        return self._offset
+
     @property
     def condition(self):
         return self._condition
@@ -1168,17 +1226,19 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs):
             fact = fact.data
 
         vals = {}
+        offset = args.get(self.offset.name, 0) if self.offset else 0
 
-        # The first subsampled point at or after the parent's min...
+        # The first subsampled point at or after the parent's rank-local min...
         with suppress(KeyError, TypeError):
-            vals[self.min_name] = math.ceil(kwargs.get(self.parent.min_name) / fact)
+            vals[self.min_name] = math.ceil((args[self.parent.min_name] - offset) / fact)
 
         # ...and the last one at or before its max. Note `floor` rather than
         # `ceil`: with e.g. `factor=3` and `x_M=11` the last subsampled point
         # is `x=9`, hence index 3, while `ceil(11/3) == 4` would overshoot the
         # Function by one
         with suppress(KeyError, TypeError):
-            vals[self.max_name] = kwargs.get(self.parent.max_name) // fact
+            vals[self.max_name] = min((args[self.parent.max_name] - offset) // fact,
+                                      args[self.max_name])
 
         vals[self.symbolic_factor.name] = fact
 
@@ -1207,8 +1267,10 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
                                      f"{self.name}: {size} < {size0}")
             else:
                 # Given a factor the last time index is factor*(size - 1)
-                # The maximum allowed value is then factor*size - 1
-                defaults[dim.parent.max_name] = range(d0, d0 + factor*size)
+                # The maximum allowed value is then factor*size - 1, plus at
+                # most factor - 1 for an offset
+                stop = d0 + factor*size + (factor - 1 if dim.offset else 0)
+                defaults[dim.parent.max_name] = range(d0, stop)
 
         return defaults
 
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 4568b852bf1..213241cc56e 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -11,14 +11,14 @@
 from devito.deprecations import deprecations
 from devito.finite_differences.elementary import Max, Min
 from devito.logger import warning
-from devito.mpi import MPI, Distributor, SubDistributor
+from devito.mpi import MPI, Distributor, SubDistributor, SubGridDistributor
 from devito.tools import ReducerMap, as_tuple, frozendict, is_integer
 from devito.types.args import ArgProvider
 from devito.types.basic import Scalar
 from devito.types.dense import Function
 from devito.types.dimension import (
     ConditionalDimension, DefaultDimension, Dimension, MultiSubDimension, SpaceDimension,
-    Spacing, SteppingDimension, SubDimension, TimeDimension
+    Spacing, SteppingDimension, SubDimension, SubsamplingOffset, TimeDimension
 )
 from devito.types.utils import DimensionTuple
 
@@ -193,7 +193,7 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
                 self._topology = None
         else:
             self._topology = None
-        self._distributor = Distributor(shape, dimensions, comm, self._topology)
+        self._distributor = self.__distributor_setup__(shape, dimensions, comm)
 
         # The physical extent and grid spacing
         if spacing is not None:
@@ -417,6 +417,9 @@ def _arg_values(self, **kwargs):
 
         return values
 
+    def __distributor_setup__(self, shape, dimensions, comm):
+        return Distributor(shape, dimensions, comm, self._topology)
+
     def __getstate__(self):
         state = self.__dict__.copy()
         # A Distributor wraps an MPI communicator, which can't and shouldn't be pickled
@@ -509,8 +512,10 @@ def __init__(self, parent, factor, name='s', **kwargs):
             )
 
         dimensions = tuple(
-            d if f == 1 else ConditionalDimension(f'{d.name}{name}', parent=d,
-                                                  factor=f)
+            d if f == 1 else ConditionalDimension(
+                f'{d.name}{name}', parent=d, factor=f,
+                offset=self.__offset_setup__(parent, d, f, name)
+            )
             for d, f in zip(parent.dimensions, factors, strict=True)
         )
 
@@ -537,6 +542,15 @@ def __init__(self, parent, factor, name='s', **kwargs):
         # the very same buffer indices as one on the parent
         self._stepping_dim = parent.stepping_dim
 
+    def __distributor_setup__(self, shape, dimensions, comm):
+        return SubGridDistributor(self)
+
+    @classmethod
+    def __offset_setup__(cls, parent, d, factor, name):
+        if not parent.distributor.is_parallel:
+            return 0
+        return SubsamplingOffset(name=f'{d.name}{name}o', root=d, factor=factor)
+
     @classmethod
     def __factors_setup__(cls, parent, factor):
         factors = as_tuple(factor)
diff --git a/examples/userapi/05_conditional_dimension.ipynb b/examples/userapi/05_conditional_dimension.ipynb
index a8e628545b1..61a29655573 100644
--- a/examples/userapi/05_conditional_dimension.ipynb
+++ b/examples/userapi/05_conditional_dimension.ipynb
@@ -184,6 +184,9 @@
       "        lowering for example combining Function's ConditionalDimension with\n",
       "        an Equation's implicit_dim. All Dimensions within an equation\n",
       "        must have `Or` relation for the final combined condition to be Or.\n",
+      "    offset : SubsamplingOffset, optional, default=0\n",
+      "        The intercept of the map onto the parent Dimension, that is\n",
+      "        ``parent = factor*self + offset``.\n",
       "\n",
       "    Examples\n",
       "    --------\n",
diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py
index 3d235ca9490..c47d62b7b6b 100644
--- a/tests/test_subsampling.py
+++ b/tests/test_subsampling.py
@@ -1,3 +1,4 @@
+import pickle
 import re
 
 import numpy as np
@@ -11,6 +12,7 @@
 from devito.ir.iet import Conditional, Expression, FindNodes, Iteration
 from devito.symbolics import retrieve_indexed
 from devito.tools import as_tuple
+from devito.types.dimension import SubsamplingOffset
 
 
 def iteration(op, name):
@@ -533,11 +535,7 @@ def test_sequential_dimension(self):
 class TestInteractions:
 
     def test_with_blocking(self):
-        """
-        Blocking rewrites `x` into BlockDimensions, so the subsampled
-        Dimension's parent leaves the iteration space and the guard is kept.
-        Correct, if not strided.
-        """
+        """Striding comes first, so blocking tiles the subsampled Dimension."""
         grid = Grid(shape=(32, 32))
         x, y = grid.dimensions
         cd = ConditionalDimension(name='xs', parent=x, factor=4)
@@ -546,12 +544,19 @@ def test_with_blocking(self):
                      space_order=0)
         u.data[:] = np.arange(32*32).reshape(32, 32).astype(u.dtype)
 
-        op = Operator(Eq(g, u.dx2 + u.dy2), opt=('advanced',
-                                                 {'blockrelax': True}))
+        eq = Eq(g, u.dx2 + u.dy2)
+        Operator(eq, opt='noop').apply()
+        expected = np.array(g.data)
+        g.data[:] = 0.
+
+        op = Operator(eq, opt=('advanced', {'blockrelax': True}))
         op.apply()
 
-        assert len(FindNodes(Conditional).visit(op)) == 1
-        assert np.all(np.isfinite(g.data))
+        assert not FindNodes(Conditional).visit(op)
+        assert any(i.dim.is_Block and i.dim.parent is cd
+                   for i in FindNodes(Iteration).visit(op))
+        # `h^-2*eps*u` of float32 roundoff, as the two builds reorder flops
+        assert np.allclose(g.data, expected, atol=0.5)
 
     @switchconfig(language='openmp')
     def test_with_openmp(self):
@@ -852,6 +857,92 @@ def test_custom_dimension_names(self):
         assert [d.name for d in sub.dimensions] == ['xc', 'yc']
 
 
+class TestBlocking:
+
+    """
+    Striding happens first, so that blocking tiles the subsampled Dimensions.
+    """
+
+    @staticmethod
+    def build(factor, shape=(21, 21, 21)):
+        grid = Grid(shape=shape)
+        sub = SubGrid(grid, factor=factor)
+        x = grid.dimensions[0]
+
+        u = Function(name='u', grid=grid, space_order=2)
+        g = Function(name='g', grid=sub, space_order=2)
+        u.data[:] = np.arange(np.prod(shape), dtype=np.float64).reshape(shape) % 1009
+
+        eq = Eq(g, g + u.subs({x: x + 1})*u)
+        return eq, g, sub
+
+    def test_blocked_and_strided(self):
+        """The block loops span the subsampled Dimensions, with no guards."""
+        eq, _, sub = self.build(2)
+
+        op = Operator(eq, opt=('advanced', {'blockrelax': True}))
+
+        assert not FindNodes(Conditional).visit(op)
+        blocks = [i for i in FindNodes(Iteration).visit(op) if i.dim.is_Block]
+        assert {i.dim.parent for i in blocks if i.dim._depth == 1} == \
+            set(sub.dimensions)
+        for d in sub.dimensions:
+            assert iteration(op, d.name).limits[2] == 1
+
+    @pytest.mark.parametrize('factor', [2, (2, 3, 4), (3, 2, 1), (1, 2, 2)])
+    @pytest.mark.parametrize('block_size', [None, 3, 5])
+    def test_numerics(self, factor, block_size):
+        """Any tiling of a strided nest visits each point exactly once."""
+        eq, g, _ = self.build(factor)
+        Operator(eq, opt='noop').apply()
+        expected = np.array(g.data)
+
+        g.data[:] = 0.
+
+        op = Operator(eq, opt=('advanced', {'blockrelax': True}))
+        sizes = {p.name: block_size for p in op.parameters
+                 if block_size and p.name.endswith('_size')}
+        assert sizes or block_size is None
+        op.apply(**sizes)
+
+        assert np.array_equal(np.asarray(g.data), expected)
+
+    def test_manual_conditional_dimensions(self):
+        """Hand-built ConditionalDimensions are tiled like a SubGrid's."""
+        grid = Grid(shape=(21, 21, 21))
+        x, y, z = grid.dimensions
+        xd, yd, zd = [ConditionalDimension(f'{d.name}d', parent=d, factor=2)
+                      for d in grid.dimensions]
+
+        u = Function(name='u', grid=grid, space_order=2)
+        g = Function(name='g', dimensions=(xd, yd, zd), shape=(11, 11, 11),
+                     grid=grid, space_order=2)
+
+        op = Operator(Eq(g, g + u[x + 1, y, z]*u[x - 1, y, z]),
+                      opt=('advanced', {'blockrelax': True}))
+
+        assert not FindNodes(Conditional).visit(op)
+        assert any(i.dim.is_Block and i.dim.parent is xd
+                   for i in FindNodes(Iteration).visit(op))
+
+    def test_shared_nest_keeps_guards(self):
+        """A nest shared with unguarded work is blocked, not strided."""
+        grid = Grid(shape=(21, 21, 21))
+        sub = SubGrid(grid, factor=2)
+        x, y, z = grid.dimensions
+
+        u = Function(name='u', grid=grid, space_order=2)
+        v = Function(name='v', grid=grid, space_order=2)
+        g = Function(name='g', grid=sub, space_order=2)
+        u.data[:] = 2.
+
+        op = Operator([Eq(v, u[x + 1, y, z] + u[x - 1, y, z]), Eq(g, g + v)],
+                      opt=('advanced', {'blockrelax': True}))
+        op.apply()
+
+        assert np.all(np.asarray(g.data[1:-1, :, :]) == 4.)
+
+
 class TestMPI:
 
     """
@@ -879,6 +970,143 @@ def test_copy_from_parent(self, mode, factor):
         expected = glob[::factor, ::factor][g.local_indices]
         assert np.array_equal(np.asarray(g.data), expected)
 
+    @pytest.mark.parallel(mode=[2, 4])
+    @pytest.mark.parametrize('factor', [2, 3, 4])
+    def test_rank_boundary_off_the_subgrid(self, mode, factor):
+        """The factor applies to global indices, wherever the ranks split."""
+        n = 21
+        grid = Grid(shape=(n, n))
+        sub = SubGrid(grid, factor=factor)
+
+        u = Function(name='u', grid=grid)
+        g = Function(name='g', grid=sub, space_order=0)
+
+        glob = np.arange(n*n, dtype=u.dtype).reshape(n, n)
+        u.data[:] = glob
+
+        Operator(Eq(g, u)).apply()
+
+        expected = glob[::factor, ::factor][g.local_indices]
+        assert np.array_equal(np.asarray(g.data), expected)
+
+    @pytest.mark.parallel(mode=[2])
+    def test_guarded_rank_boundary_off_the_subgrid(self, mode):
+        """A guard surviving in a shared nest also uses global indices."""
+        n, factor = 21, 2
+        grid = Grid(shape=(n, n))
+        sub = SubGrid(grid, factor=factor)
+
+        u = Function(name='u', grid=grid)
+        v = Function(name='v', grid=grid)
+        g = Function(name='g', grid=sub, space_order=0)
+
+        glob = np.arange(n*n, dtype=u.dtype).reshape(n, n)
+        u.data[:] = glob
+
+        op = Operator([Eq(v, u + 1.), Eq(g, u)])
+
+        assert FindNodes(Conditional).visit(op)
+        op.apply()
+
+        expected = glob[::factor, ::factor][g.local_indices]
+        assert np.array_equal(np.asarray(g.data), expected)
+
+    @pytest.mark.parallel(mode=[2])
+    def test_subgrid_only(self, mode):
+        """With no parent Function, the SubGrid alone sets the bounds."""
+        grid = Grid(shape=(21, 21))
+        sub = SubGrid(grid, factor=3)
+
+        g = Function(name='g', grid=sub, space_order=0)
+
+        Operator(Eq(g, 1.)).apply()
+
+        assert np.all(np.asarray(g.data) == 1.)
+
+    @pytest.mark.parallel(mode=[2])
+    @pytest.mark.parametrize('factor', [2, 3])
+    def test_blocked_rank_boundary_off_the_subgrid(self, mode, factor):
+        """Blocking a strided nest preserves the global sampling."""
+        n = 21
+        grid = Grid(shape=(n, n, n))
+        sub = SubGrid(grid, factor=factor)
+        x = grid.dimensions[0]
+
+        u = Function(name='u', grid=grid, space_order=2)
+        g = Function(name='g', grid=sub, space_order=2)
+
+        glob = np.arange(n**3, dtype=np.float64).reshape(n, n, n) % 1009
+        u.data[:] = glob
+
+        op = Operator(Eq(g, g + u.subs({x: x + 1})*u),
+                      opt=('advanced', {'blockrelax': True}))
+
+        assert not FindNodes(Conditional).visit(op)
+        assert any(i.dim.is_Block for i in FindNodes(Iteration).visit(op))
+        op.apply()
+
+        padded = np.zeros((n + 1, n, n))
+        padded[:n] = glob
+        expected = (padded[1:]*glob)[::factor, ::factor, ::factor]
+        assert np.array_equal(np.asarray(g.data), expected[g.local_indices])
+
+    @pytest.mark.parallel(mode=[4])
+    def test_decomposition_follows_parent(self, mode):
+        """A rank owns exactly the subgrid points on its parent points."""
+        n, factor = 13, 6
+        grid = Grid(shape=(n,))
+        sub = SubGrid(grid, factor=factor)
+
+        for parent, child in zip(grid.distributor.decomposition[0],
+                                 sub.distributor.decomposition[0], strict=True):
+            assert np.array_equal(child*factor, parent[parent % factor == 0])
+
+    @pytest.mark.parallel(mode=[2])
+    def test_manual_function_stays_in_rank(self, mode):
+        """A strided nest over a non-decomposed Function stops at the rank."""
+        grid = Grid(shape=(21, 21))
+        x, y = grid.dimensions
+        factor = 2
+
+        xd = ConditionalDimension('xd', parent=x, factor=factor)
+        yd = ConditionalDimension('yd', parent=y, factor=factor)
+
+        u = Function(name='u', grid=grid)
+        g = Function(name='g', dimensions=(xd, yd), shape=(11, 11), grid=grid,
+                     space_order=0)
+
+        args = Operator(Eq(g, u)).arguments()
+
+        assert args['xd_M'] == args['x_M'] // factor
+
+    @pytest.mark.parallel(mode=[2])
+    def test_offset_is_rank_data(self, mode):
+        """Each rank's offset lands its loop on the global subgrid."""
+        grid = Grid(shape=(21,))
+        sub = SubGrid(grid, factor=2)
+        offset = sub.dimensions[0].offset
+
+        expected = {0: 0, 1: 1}[grid.distributor.myrank]
+        assert offset._arg_values(grid=grid) == {offset.name: expected}
+
+    def test_offset_pickles(self):
+        """An offset is rebuilt from its factor, so it holds no rank data."""
+        grid = Grid(shape=(21,))
+        x = grid.dimensions[0]
+        offset = SubsamplingOffset(name='xso', root=x, factor=3)
+
+        new = pickle.loads(pickle.dumps(offset))
+
+        assert new.name == offset.name
+        assert new.factor == 3
+        assert new.root.name == x.name
+
+    def test_serial_has_no_offset(self):
+        """Without a decomposition there's nothing to align."""
+        sub = SubGrid(Grid(shape=(21,)), factor=2)
+
+        assert sub.dimensions[0].offset == 0
+
     @pytest.mark.parallel(mode=[2, 4])
     def test_strided_codegen(self, mode):
         """The guard still becomes a unit-stride subgrid loop under MPI."""

From 9604a8230fdab15e95befd681e3d69aad20aac8d Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Fri, 6 Nov 2026 20:03:56 -0500
Subject: [PATCH 2/3] compiler: Lower Dimensions used as expressions to global
 indices

A bare Dimension in an expression compiled to the rank-local loop index, so
Eq(u, x + y) restarted from zero on each MPI rank. Offset it by the new
Grid.origin_isymbols, passed like origin_symbols. Dimensions indexing a
Function stay rank-local, and the code generated without MPI is unchanged.

Fixes #2214
---
 devito/ir/equations/algorithms.py | 17 ++++++++++
 devito/types/grid.py              | 17 +++++++++-
 tests/test_caching.py             |  2 +-
 tests/test_mpi.py                 | 54 +++++++++++++++++++++++++++++++
 4 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index 44ad31a4f31..38868c4ed94 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -156,6 +156,8 @@ def _lower_exprs(expressions, subs):
         # Add dimensions map to the mapper in case dimensions are used
         # as an expression, i.e. Eq(u, x, subdomain=xleft)
         mapper.update(dimension_map)
+        # Dimensions used as an expression are global indices, unlike loop indices
+        mapper.update(_make_global_index_map(expr, dimension_map))
         # Add the user-supplied substitutions
         mapper.update(subs)
         # Apply mapper to expression
@@ -192,6 +194,21 @@ def _make_dimension_map(expr):
     return frozendict(dimension_map)
 
 
+def _make_global_index_map(expr, dimension_map):
+    """
+    Map the Dimensions of a distributed Grid to the global indices they stand for
+    when used as an expression, i.e. Eq(u, x + y) -> Eq(u[x, y], x + io_x + y + io_y).
+    Dimensions indexing an Indexed are unaffected as the Indexed is replaced as a whole.
+    """
+    if not isinstance(expr, Eq):
+        return {}
+
+    grids = {f.grid.root for f in retrieve_functions(expr) if f.grid is not None}
+    return {d: dimension_map.get(d, d) + o
+            for grid in grids if grid.distributor.is_parallel
+            for d, o in zip(grid.dimensions, grid.origin_isymbols, strict=True)}
+
+
 def concretize_subdims(exprs, **kwargs):
     """
     Given a list of expressions, return a new list where all user-defined
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 213241cc56e..060c3a8b9d8 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -212,6 +212,9 @@ def __init__(self, shape, extent=None, origin=None, dimensions=None,
         self._origin_symbols = tuple(Scalar(name=f'o_{d.name}', dtype=dtype,
                                             is_const=True)
                                      for d in self.dimensions)
+        self._origin_isymbols = tuple(Scalar(name=f'io_{d.name}', dtype=np.int32,
+                                             is_const=True)
+                                      for d in self.dimensions)
 
         # Sanity check
         assert (self.dim == len(self.origin) == len(self.extent) == len(self.spacing))
@@ -262,10 +265,20 @@ def origin_map(self):
         """Map between origin symbols and their values."""
         return dict(zip(self.origin_symbols, self.origin, strict=True))
 
+    @property
+    def origin_isymbols(self):
+        """Symbols representing the global index of the rank's first grid point."""
+        return self._origin_isymbols
+
+    @property
+    def origin_imap(self):
+        """Map between index origin symbols and their values."""
+        return dict(zip(self.origin_isymbols, self.origin_ioffset, strict=True))
+
     @property
     def origin_ioffset(self):
         """Offset index of the local (per-process) origin from the domain origin."""
-        grid_origin = [min(i) for i in self.distributor.glb_numb]
+        grid_origin = [i.start for i in self.distributor.glb_slices.values()]
         assert len(grid_origin) == len(self.spacing)
         return DimensionTuple(*grid_origin, getters=self.dimensions)
 
@@ -379,6 +392,7 @@ def _arg_names(self):
         ret = []
         ret.append(self.time_dim.spacing.name)
         ret.extend([i.name for i in self.origin_map])
+        ret.extend([i.name for i in self.origin_imap])
         for i in self.spacing_map:
             try:
                 ret.append(i.name)
@@ -400,6 +414,7 @@ def _arg_defaults(self):
 
         # Grid origin
         args.update({k.name: v for k, v in self.origin_map.items()})
+        args.update({k.name: v for k, v in self.origin_imap.items()})
 
         # MPI-related objects
         if self.distributor.is_parallel:
diff --git a/tests/test_caching.py b/tests/test_caching.py
index 3f170803123..2700df7c060 100644
--- a/tests/test_caching.py
+++ b/tests/test_caching.py
@@ -942,7 +942,7 @@ def test_solve(self, operate_on_empty_cache):
         # created by the finite difference (u.dt, u.dx2). We would have had
         # three extra references to u(t + dt), u(x - h_x) and u(x + h_x).
         # But this is not the case anymore!
-        assert len(_SymbolCache) == 8
+        assert len(_SymbolCache) == 9
         clear_cache()
         assert len(_SymbolCache) == 4
         clear_cache()
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 7081be3beb9..b05061150ec 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -1137,6 +1137,60 @@ def check_halo_exchanges(op, exp0, exp1):
     return calls, tloop
 
 
+class TestGlobalIndex:
+
+    @pytest.mark.parallel(mode=[2, 4])
+    def test_dimension_as_expression(self, mode):
+        """Dimensions in an expression are global indices (issue #2214)."""
+        grid = Grid(shape=(8, 8))
+        x, y = grid.dimensions
+        f = Function(name='f', grid=grid)
+
+        Operator(Eq(f, 10*x + y))()
+
+        check = 10*np.arange(8)[:, None] + np.arange(8)
+        assert np.all(f.data[:] == check[f.local_indices])
+
+    @pytest.mark.parallel(mode=[2, 4])
+    def test_dimension_as_expression_subdomain(self, mode):
+        """Same as above, but restricted to a SubDomain."""
+        grid = Grid(shape=(8, 8))
+        x, y = grid.dimensions
+        f = Function(name='f', grid=grid)
+
+        Operator(Eq(f, 10*x + y, subdomain=grid.interior))()
+
+        check = 10*np.arange(8)[:, None] + np.arange(8)
+        check[[0, -1], :] = check[:, [0, -1]] = 0
+        assert np.all(f.data[:] == check[f.local_indices])
+
+    @pytest.mark.parallel(mode=[2])
+    def test_dimension_indexing_unaffected(self, mode):
+        """Dimensions indexing a Function remain rank-local."""
+        grid = Grid(shape=(8,))
+        x = grid.dimensions[0]
+        f = Function(name='f', grid=grid, space_order=1)
+        g = Function(name='g', grid=grid)
+
+        f.data[:] = np.arange(8)
+        Operator(Eq(g, f.subs(x, x + 1) + x))()
+
+        # The last point reads a zero from the physical boundary halo
+        check = 2*np.arange(8) + 1
+        check[-1] = 7
+        assert np.all(g.data[:] == check[g.local_indices])
+
+    def test_serial_code_unchanged(self):
+        """No index offset is generated without MPI."""
+        grid = Grid(shape=(4, 4))
+        x, y = grid.dimensions
+        f = Function(name='f', grid=grid)
+
+        op = Operator(Eq(f, x + y))
+
+        assert not any(p.name.startswith('io_') for p in op.parameters)
+
+
 class TestHaloWrites:
 
     @pytest.mark.parallel(mode=[1, 2])

From fb0ea5d59112aa500636309b7e270ae5c62ba0df Mon Sep 17 00:00:00 2001
From: mloubout <mathias.louboutin@gmail.com>
Date: Sat, 7 Nov 2026 22:05:20 -0500
Subject: [PATCH 3/3] compiler: Restrict SubGrid Functions to the SubDomain of
 their Eq

With subdomain=, a Function on a SubGrid kept iterating its whole subsampled
Dimension around the SubDomain loop, so every point accumulated over the
entire SubDomain. Subsample the SubDomain's Dimension instead, bound it by the
parent's actual extent, thickness included, and localize the thicknesses that
only bound such a Dimension.

Also keep time guarded in the subsampling pass: it now runs before cluster
specialization, where buffering and other time-based lowering rely upon the
guard.
---
 devito/ir/equations/algorithms.py     | 16 ++++++--
 devito/operator/operator.py           | 12 ++++--
 devito/passes/clusters/subsampling.py |  4 ++
 devito/types/dimension.py             | 12 ++++--
 tests/test_subsampling.py             | 58 ++++++++++++++++++++++++++-
 5 files changed, 91 insertions(+), 11 deletions(-)

diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py
index 38868c4ed94..4224b63134d 100644
--- a/devito/ir/equations/algorithms.py
+++ b/devito/ir/equations/algorithms.py
@@ -186,10 +186,20 @@ def _make_dimension_map(expr):
         # Some Relationals may be pure SymPy objects, thus lacking the SubDomain
         dimension_map = {}
     else:
-        functions = [f for f in retrieve_functions(expr) if f._is_on_subdomain]
+        functions = retrieve_functions(expr)
         for f in functions:
-            dimension_map.update({d: expr.subdomain.dimension_map[d.root]
-                                  for d in f.space_dimensions if d.is_Sub})
+            if f._is_on_subdomain:
+                dimension_map.update({d: expr.subdomain.dimension_map[d.root]
+                                      for d in f.space_dimensions if d.is_Sub})
+
+        # A subsampled Function, e.g. on a SubGrid, is subsampled from the
+        # SubDomain's Dimensions rather than from the whole Grid's
+        for f in functions:
+            dimension_map.update({
+                d: d._rebuild(parent=expr.subdomain.dimension_map[d.parent])
+                for d in f.dimensions
+                if d.is_Conditional and d.parent in expr.subdomain.dimension_map
+            })
 
     return frozendict(dimension_map)
 
diff --git a/devito/operator/operator.py b/devito/operator/operator.py
index 52bf4905982..2fe0ebff16a 100644
--- a/devito/operator/operator.py
+++ b/devito/operator/operator.py
@@ -695,10 +695,14 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
             # No need to do anything more if only checking the memory
             return args
 
-        # Process the symbols localized by MPI
-        for p in self.parameters:
-            if isinstance(p, (SubsamplingOffset, Thickness)):
-                args.update(p._arg_values(grid=grid, **kwargs))
+        # Process the symbols localized by MPI. A SubDimension's thicknesses are
+        # needed even when they only bound another Dimension, e.g. a subsampled one
+        localized = OrderedSet(*[p for p in self.parameters
+                                 if isinstance(p, (SubsamplingOffset, Thickness))])
+        localized.update(t for d in toposort if isinstance(d, SubDimension)
+                         for t in d.thickness)
+        for p in localized:
+            args.update(p._arg_values(grid=grid, **kwargs))
 
         # Process Dimensions
         for d in reversed(toposort):
diff --git a/devito/passes/clusters/subsampling.py b/devito/passes/clusters/subsampling.py
index c138d31968a..6afd696c326 100644
--- a/devito/passes/clusters/subsampling.py
+++ b/devito/passes/clusters/subsampling.py
@@ -113,6 +113,10 @@ def _is_strideable(cluster, d, guard):
     if d not in cluster.ispace.itdims:
         return False
 
+    # Time keeps its guard, which later passes (e.g. buffering) rely upon
+    if d.is_Time:
+        return False
+
     # Skipping points is only sound if no dependence is carried along `d`;
     # otherwise the discarded iterations may be doing work that `d`'s
     # remaining iterations rely upon
diff --git a/devito/types/dimension.py b/devito/types/dimension.py
index f49927801bd..259dc60950e 100644
--- a/devito/types/dimension.py
+++ b/devito/types/dimension.py
@@ -1228,16 +1228,22 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs):
         vals = {}
         offset = args.get(self.offset.name, 0) if self.offset else 0
 
-        # The first subsampled point at or after the parent's rank-local min...
+        # The parent's rank-local bounds, e.g. including a SubDimension's thickness
+        bounds = (self.parent.symbolic_min, self.parent.symbolic_max)
+        values = {s: args[s.name] for b in bounds for s in b.free_symbols
+                  if args.get(s.name) is not None}
+        parent_min, parent_max = (b.subs(values) for b in bounds)
+
+        # The first subsampled point at or after the parent's min...
         with suppress(KeyError, TypeError):
-            vals[self.min_name] = math.ceil((args[self.parent.min_name] - offset) / fact)
+            vals[self.min_name] = math.ceil((int(parent_min) - offset) / fact)
 
         # ...and the last one at or before its max. Note `floor` rather than
         # `ceil`: with e.g. `factor=3` and `x_M=11` the last subsampled point
         # is `x=9`, hence index 3, while `ceil(11/3) == 4` would overshoot the
         # Function by one
         with suppress(KeyError, TypeError):
-            vals[self.max_name] = min((args[self.parent.max_name] - offset) // fact,
+            vals[self.max_name] = min((int(parent_max) - offset) // fact,
                                       args[self.max_name])
 
         vals[self.symbolic_factor.name] = fact
diff --git a/tests/test_subsampling.py b/tests/test_subsampling.py
index c47d62b7b6b..2ee0303e572 100644
--- a/tests/test_subsampling.py
+++ b/tests/test_subsampling.py
@@ -7,7 +7,7 @@
 
 from devito import (
     CondEq, ConditionalDimension, Constant, Eq, Function, Grid, Inc, Operator,
-    SpaceDimension, SubGrid, TimeFunction, switchconfig
+    SpaceDimension, SubDomain, SubGrid, TimeFunction, switchconfig
 )
 from devito.ir.iet import Conditional, Expression, FindNodes, Iteration
 from devito.symbolics import retrieve_indexed
@@ -337,6 +337,11 @@ def test_time_subsampling(self):
 
         Operator([Eq(u.forward, u + 1.), Eq(usave, u)]).apply(time_M=nt-2)
 
+        # Time keeps its guard even where it is parallel
+        op = Operator(Eq(usave, 1.))
+        assert FindNodes(Conditional).visit(op)
+        assert all(i.dim is not ts for i in FindNodes(Iteration).visit(op))
+
         # time runs 0..nt-2, so only the slots whose time index is reached
         # are written; `u` holds the step count at that point
         nwritten = (nt - 2)//factor + 1
@@ -568,6 +573,22 @@ def test_with_openmp(self):
         assert np.array_equal(g.data, data[::2])
 
 
+class Middle(SubDomain):
+
+    name = 'middle'
+
+    def define(self, dimensions):
+        # Thicknesses deliberately not multiples of the SubGrid factors
+        return {d: ('middle', 2, 3) for d in dimensions}
+
+
+def middle_of(data):
+    """`data` zeroed outside the `Middle` SubDomain."""
+    ret = np.zeros_like(data)
+    ret[2:-3, 2:-3] = data[2:-3, 2:-3]
+    return ret
+
+
 class TestSubGrid:
 
     """The `SubGrid` API."""
@@ -630,6 +651,22 @@ def test_copy_from_parent(self, factor):
         fx, fy = sub.factor
         assert np.array_equal(f.data, data[::fx, ::fy])
 
+    @pytest.mark.parametrize('factor', [2, 3, (2, 1), (3, 2)])
+    def test_copy_on_subdomain(self, factor):
+        """Within a SubDomain, only the SubGrid points inside it are written."""
+        grid = Grid(shape=(12, 12))
+        sub = SubGrid(grid, factor=factor)
+
+        u = Function(name='u', grid=grid)
+        f = Function(name='f', grid=sub, space_order=0)
+        data = np.arange(1, 145, dtype=u.dtype).reshape(12, 12)
+        u.data[:] = data
+
+        Operator(Eq(f, u, subdomain=Middle(grid=grid))).apply()
+
+        fx, fy = sub.factor
+        assert np.array_equal(f.data, middle_of(data)[::fx, ::fy])
+
     def test_copy_is_strided(self):
         """
         No guard is emitted, every loop keeps a unit stride, and the Operator
@@ -970,6 +1007,25 @@ def test_copy_from_parent(self, mode, factor):
         expected = glob[::factor, ::factor][g.local_indices]
         assert np.array_equal(np.asarray(g.data), expected)
 
+    @pytest.mark.parallel(mode=[2, 4])
+    @pytest.mark.parametrize('factor', [2, 3])
+    def test_copy_on_subdomain(self, mode, factor):
+        """The SubDomain bounds the SubGrid on every rank."""
+        n = 24
+        grid = Grid(shape=(n, n))
+        sub = SubGrid(grid, factor=factor)
+
+        u = Function(name='u', grid=grid)
+        g = Function(name='g', grid=sub, space_order=0)
+
+        glob = np.arange(1, n*n + 1, dtype=u.dtype).reshape(n, n)
+        u.data[:] = glob
+
+        Operator(Eq(g, u, subdomain=Middle(grid=grid))).apply()
+
+        expected = middle_of(glob)[::factor, ::factor][g.local_indices]
+        assert np.array_equal(np.asarray(g.data), expected)
+
     @pytest.mark.parallel(mode=[2, 4])
     @pytest.mark.parametrize('factor', [2, 3, 4])
     def test_rank_boundary_off_the_subgrid(self, mode, factor):