From eb60fc384704b819d36487176eb9c42fdecbdc63 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 02:03:39 +0000
Subject: [PATCH 01/10] mpi: Include staggered sparse stencil cells in
 redistribution

---
 devito/types/sparse.py | 16 ++++++++++++----
 tests/test_mpi.py      | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/devito/types/sparse.py b/devito/types/sparse.py
index 6e913646a3..5cbd592b38 100644
--- a/devito/types/sparse.py
+++ b/devito/types/sparse.py
@@ -11,9 +11,10 @@
 from devito.operations import (
     LinearInterpolator, NearestInterpolator, PrecomputedInterpolator, SincInterpolator
 )
+from devito.operations.interpolators import _cell_indices
 from devito.symbolics import indexify, retrieve_function_carriers
 from devito.tools import (
-    ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
+    ReducerMap, as_fp64_decimal, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
     memoized_meth, prod
 )
 from devito.types.basic import Symbol
@@ -307,8 +308,14 @@ def _support(self):
         """
         max_shape = np.array(self.grid.shape).reshape(1, self.grid.dim)
         minmax = lambda arr: np.minimum(max_shape, np.maximum(0, arr))
-        return np.stack([minmax(self._coords_indices + s) for s in self._point_support],
-                        axis=2)
+        # Staggered tables may use a different base cell across an MPI
+        # boundary. Match their fp64 geometry and actual stencil support.
+        spacing = np.array([as_fp64_decimal(h) for h in self.grid.spacing])
+        origin = np.array([as_fp64_decimal(o) for o in self.grid.origin])
+        indices = [_cell_indices(self.coordinates_data, self.grid, shift, spacing, origin)
+                   for shift in self.interpolator._shifts_used] or [self._coords_indices]
+        return np.stack([minmax(index + s) for index in indices
+                         for s in self._point_support], axis=2)
 
     @property
     def _dist_datamap(self):
@@ -1004,10 +1011,11 @@ def _decomposition(self):
         return tuple(mapper.get(d) for d in self.dimensions)
 
     def _arg_defaults(self, alias=None, estimate_memory=False):
+        key = alias or self
+        self.interpolator._shifts_used.update(key.interpolator._shifts_used)
         defaults = super()._arg_defaults(alias=alias, estimate_memory=estimate_memory)
         if estimate_memory:
             return defaults
-        key = alias or self
         coords = defaults.get(key.coordinates.name, self.coordinates.data)
         defaults.update(key.interpolator._arg_defaults(coords=coords,
                                                        sfunc=self))
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 7081be3beb..23e5f04501 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -934,6 +934,30 @@ class CoordSlowSparseFunction(SparseTimeFunction):
         op.apply(time_M=5)
         assert np.all(s.data == 1)
 
+    @pytest.mark.parallel(mode=[2, 4])
+    @pytest.mark.parametrize('override', [False, True])
+    def test_staggered_injection_support(self, mode, override):
+        """Half-cell shifts must redistribute across partition boundaries."""
+        grid = Grid(shape=(64,), extent=(63/32,))
+        coordinate = 2/grid.distributor.nprocs
+        x, = grid.dimensions
+        u = Function(name='u', grid=grid, space_order=2, staggered=x)
+        u.data[:] = 0
+        src = SparseFunction(name='src', grid=grid, npoint=1)
+        src.coordinates.data[:] = coordinate
+        src.data[:] = 1
+        op = Operator(src.inject(field=u, expr=src))
+        if override:
+            replacement = SparseFunction(name='replacement', grid=grid, npoint=1)
+            replacement.coordinates.data[:] = coordinate
+            replacement.data[:] = 1
+            op(src=replacement)
+        else:
+            op()
+
+        total = MPI.COMM_WORLD.allreduce(float(np.asarray(u.data).sum()), op=MPI.SUM)
+        assert np.isclose(total, 1, rtol=0, atol=1e-7)
+
 
 class TestOperatorSimple:
 

From 52b861247447fcd8afa4ddcc1ed7be59d28c12b7 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 02:44:38 +0000
Subject: [PATCH 02/10] mpi: Map tensor sparse support to ranks through compact
 intervals

---
 devito/types/sparse.py | 42 ++++++++++++++++++++++++++++++++++--------
 tests/test_mpi.py      | 21 +++++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/devito/types/sparse.py b/devito/types/sparse.py
index 5cbd592b38..a71f045278 100644
--- a/devito/types/sparse.py
+++ b/devito/types/sparse.py
@@ -300,6 +300,15 @@ def _coords_indices(self):
                 np.floor((self.coordinates_data - self.grid.origin) / self.grid.spacing)
             ).astype(int)
 
+    @property
+    def _support_indices(self):
+        # Staggered tables may use a different base cell across an MPI
+        # boundary. Match their fp64 geometry and actual stencil support.
+        spacing = np.array([as_fp64_decimal(h) for h in self.grid.spacing])
+        origin = np.array([as_fp64_decimal(o) for o in self.grid.origin])
+        return [_cell_indices(self.coordinates_data, self.grid, shift, spacing, origin)
+                for shift in self.interpolator._shifts_used] or [self._coords_indices]
+
     @property
     def _support(self):
         """
@@ -308,13 +317,7 @@ def _support(self):
         """
         max_shape = np.array(self.grid.shape).reshape(1, self.grid.dim)
         minmax = lambda arr: np.minimum(max_shape, np.maximum(0, arr))
-        # Staggered tables may use a different base cell across an MPI
-        # boundary. Match their fp64 geometry and actual stencil support.
-        spacing = np.array([as_fp64_decimal(h) for h in self.grid.spacing])
-        origin = np.array([as_fp64_decimal(o) for o in self.grid.origin])
-        indices = [_cell_indices(self.coordinates_data, self.grid, shift, spacing, origin)
-                   for shift in self.interpolator._shifts_used] or [self._coords_indices]
-        return np.stack([minmax(index + s) for index in indices
+        return np.stack([minmax(index + s) for index in self._support_indices
                          for s in self._point_support], axis=2)
 
     @property
@@ -322,7 +325,30 @@ def _dist_datamap(self):
         """
         Mapper ``M : MPI rank -> required sparse data``.
         """
-        return self.grid.distributor.glb_to_rank(self._support) or {}
+        offsets = self._point_support
+        lower, upper = offsets.min(axis=0), offsets.max(axis=0)
+        if (len(offsets) != prod(upper - lower + 1)
+                or len(np.unique(offsets, axis=0)) != len(offsets)):
+            return self.grid.distributor.glb_to_rank(self._support) or {}
+
+        # A tensor stencil intersects a Cartesian rank exactly when its
+        # intervals intersect in every dimension. Keep different shifts
+        # separate: their union need not be one rectangular stencil.
+        bounds = [(np.clip(index + lower, 0, self.grid.shape),
+                   np.clip(index + upper, 0, self.grid.shape))
+                  for index in self._support_indices]
+        ret = {}
+        for rank, ranges in enumerate(self.grid.distributor.all_ranges):
+            if any(not r for r in ranges):
+                continue
+            mins = np.array([r[0] for r in ranges])
+            maxs = np.array([r[-1] for r in ranges])
+            mask = np.zeros(bounds[0][0].shape[0], dtype=bool)
+            for low, high in bounds:
+                mask |= ((low <= maxs) & (high >= mins)).all(axis=1)
+            if mask.any():
+                ret[rank] = np.flatnonzero(mask).tolist()
+        return ret
 
     @property
     def npoint(self):
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 23e5f04501..8bec46d1fe 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -934,6 +934,27 @@ class CoordSlowSparseFunction(SparseTimeFunction):
         op.apply(time_M=5)
         assert np.all(s.data == 1)
 
+    @pytest.mark.parallel(mode=[2, 4])
+    @pytest.mark.parametrize('ndim', [1, 2, 3])
+    @pytest.mark.parametrize('interpolation,radius', [('nearest', 0), ('linear', 1), ('sinc', 4)])
+    def test_tensor_support_rank_intersections(self, mode, ndim, interpolation, radius):
+        """Compact intervals match every point of wide, shifted tensor stencils."""
+        grid = Grid(shape=(9,)*ndim, extent=(1.,)*ndim)
+        coordinates = np.array([[value + (axis % 2)*0.125 for axis in range(ndim)]
+                                for value in (-0.25, 0., 0.125, 0.49, 0.5, 0.875, 1., 1.25)],
+                               dtype=np.float32)
+        src = SparseFunction(name='support_src', grid=grid, npoint=len(coordinates),
+                             interpolation=interpolation, r=radius)
+        src.coordinates.data[:] = coordinates
+        plain = Function(name='support_plain', grid=grid, space_order=8)
+        shifted = Function(name='support_shifted', grid=grid, space_order=8,
+                           staggered=grid.dimensions)
+        src.inject(field=plain, expr=src)
+        src.inject(field=shifted, expr=src)
+
+        exhaustive = grid.distributor.glb_to_rank(src._support) or {}
+        assert src._dist_datamap == exhaustive
+
     @pytest.mark.parallel(mode=[2, 4])
     @pytest.mark.parametrize('override', [False, True])
     def test_staggered_injection_support(self, mode, override):

From 3e23854433678dc489f20128772dc10022efafc9 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 04:26:22 +0000
Subject: [PATCH 03/10] jit: Publish restored libraries atomically across MPI
 ranks

---
 devito/arch/compiler.py            | 19 +++++++++++--
 tests/test_compiler_publication.py | 44 ++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_compiler_publication.py

diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py
index d356765636..2c2fd37664 100644
--- a/devito/arch/compiler.py
+++ b/devito/arch/compiler.py
@@ -5,8 +5,9 @@
 from functools import partial
 from hashlib import sha1
 from itertools import filterfalse
-from os import environ, makedirs, path
+from os import environ, fsync, link, makedirs, path, unlink
 from subprocess import DEVNULL, PIPE, CalledProcessError, check_call, check_output, run
+from tempfile import NamedTemporaryFile
 
 import numpy.ctypeslib as npct
 from codepy.jit import compile_from_string
@@ -347,8 +348,20 @@ def save(self, soname, binary):
                   " as it already exists")
         else:
             makedirs(self.get_jit_dir(), exist_ok=True)
-            with open(str(sofile), 'wb') as f:
-                f.write(binary)
+            # MPI ranks may restore the same Operator concurrently. Publish
+            # only a complete binary: dlopen of a visible partial ELF can
+            # raise SIGBUS rather than a recoverable loader exception.
+            with NamedTemporaryFile(dir=self.get_jit_dir(), delete=False) as f:
+                temporary = f.name
+                try:
+                    f.write(binary)
+                    f.flush()
+                    fsync(f.fileno())
+                    # link is atomic and keeps an existing publication intact.
+                    with suppress(FileExistsError):
+                        link(temporary, sofile)
+                finally:
+                    unlink(temporary)
             debug(f"{self}: `{sofile.name}` successfully saved in `{self.get_jit_dir()}`")
 
     def make(self, loc, args):
diff --git a/tests/test_compiler_publication.py b/tests/test_compiler_publication.py
new file mode 100644
index 0000000000..58920533b6
--- /dev/null
+++ b/tests/test_compiler_publication.py
@@ -0,0 +1,44 @@
+"""Concurrent Operator restoration must never expose a partial library."""
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from types import SimpleNamespace
+
+from devito.arch.compiler import Compiler
+
+
+def test_complete_binary_publication(tmp_path, monkeypatch):
+    import devito.arch.compiler as module
+
+    compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
+    target = tmp_path/'shared.so'
+    binary = b'ELF-binary-content'*1024
+    original = module.link
+    publications = []
+
+    def observe(source, destination):
+        assert not target.exists()
+        assert Path(source).read_bytes() == binary
+        publications.append(source)
+        original(source, destination)
+
+    monkeypatch.setattr(module, 'link', observe)
+    Compiler.save(compiler, 'shared', binary)
+    assert publications
+    assert target.read_bytes() == binary
+    assert list(tmp_path.iterdir()) == [target]
+    Compiler.save(compiler, 'shared', b'different')
+    assert target.read_bytes() == binary
+
+
+def test_concurrent_binary_publication(tmp_path):
+    compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
+    target = tmp_path/'shared.so'
+    binary = b'ELF-binary-content'*(1 << 18)
+
+    def restore(_):
+        Compiler.save(compiler, 'shared', binary)
+        assert target.read_bytes() == binary
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        list(pool.map(restore, range(32)))
+    assert list(tmp_path.iterdir()) == [target]

From 2c21a345f21e58e305776dc1f16019ff7e93126d Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 04:26:22 +0000
Subject: [PATCH 04/10] mpi: Keep enclosing halo regions synchronous during
 overlap lowering

---
 devito/passes/iet/mpi.py |  5 +++++
 tests/test_mpi.py        | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py
index db835400ae..02f25aa935 100644
--- a/devito/passes/iet/mpi.py
+++ b/devito/passes/iet/mpi.py
@@ -278,6 +278,11 @@ def _mark_overlappable(iet):
     # Analysis
     found = []
     for hs in FindNodes(HaloSpot).visit(iet):
+        # Extracting this body into a compute callable would hide nested
+        # exchanges from the lowering traversal. Keep the enclosing exchange
+        # synchronous; its children can still overlap independently.
+        if FindNodes(HaloSpot).visit(hs.body):
+            continue
         # Heuristic: avoid comp/comm overlap for sparse Iteration nests
         iters = FindNodes(Iteration).visit(hs)
         if any(i.dim._defines & set(hs.halo_scheme.distributed_aindices) and
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 8bec46d1fe..4d0dff3508 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -31,6 +31,27 @@
 
 class TestMPIUtils:
 
+    @pytest.mark.parallel(mode=[(1, 'overlap2')])
+    def test_nested_halo_overlap(self, mode):
+        from devito.ir.iet import Expression, HaloSpot, List
+        from devito.ir import LoweredEq
+        from devito.mpi.halo_scheme import HaloScheme
+        from devito.passes.iet.mpi import OverlappableHaloSpot, _mark_overlappable
+
+        grid = Grid(shape=(16, 16, 16), topology=(1, '*', 1))
+        x, y, z = grid.dimensions
+        u = Function(name='nested_u', grid=grid, space_order=2)
+        v = Function(name='nested_v', grid=grid, space_order=2)
+        equation = LoweredEq(Eq(v, u[x, y+1, z]))
+        ispace = IterationSpace([Interval(d, 0, 0) for d in grid.dimensions])
+        scheme = HaloScheme([equation], ispace)
+        inner = HaloSpot(Expression(equation), scheme)
+        outer = HaloSpot(List(body=[inner]), scheme)
+        result = _mark_overlappable(outer)
+        spots = FindNodes(HaloSpot).visit(result)
+        assert type(spots[0]) is HaloSpot
+        assert isinstance(spots[1], OverlappableHaloSpot)
+
     @pytest.mark.parametrize('comm', [None, MPI.COMM_NULL], ids=['none', 'null'])
     @pytest.mark.parametrize('error', [None, '', 'invalid argument'])
     @pytest.mark.parametrize('exception', [ValueError, InvalidArgument])

From 123b55af2150f80f0f148f7e60ff9c74d2eb3578 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 10:26:02 +0000
Subject: [PATCH 05/10] arch: Convert NVIDIA memory measurements from MiB

---
 devito/arch/archinfo.py |  2 +-
 tests/test_arch.py      | 26 +++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/devito/arch/archinfo.py b/devito/arch/archinfo.py
index ccaed53cfe..923e83a103 100644
--- a/devito/arch/archinfo.py
+++ b/devito/arch/archinfo.py
@@ -247,7 +247,7 @@ def cbk(deviceid=0):
                         line = lines.pop(deviceid)
                         _, v, unit = re.split(r'([0-9]+)\s', line)
                         assert unit == 'MiB'
-                        return int(v)*10**6
+                        return int(v)*1024**2
                     except:
                         # We shouldn't really end up here, unless nvidia-smi changes
                         # the output format (though we still have tests in place that
diff --git a/tests/test_arch.py b/tests/test_arch.py
index 305887d2e2..19dbf6df2e 100644
--- a/tests/test_arch.py
+++ b/tests/test_arch.py
@@ -1,11 +1,13 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
+from io import BytesIO
+from types import SimpleNamespace
 
 import pytest
 from packaging.version import Version
 
 from devito import configuration, switchconfig
-from devito.arch import BLACKWELL, NVIDIAX, platform_registry
+from devito.arch import BLACKWELL, NVIDIAX, archinfo, platform_registry
 from devito.arch.archinfo import AppleArm, get_apple_platform
 from devito.arch.compiler import (
     CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
@@ -127,3 +129,25 @@ def test_fp32x2_support(platform, cc, expected, language, monkeypatch):
     monkeypatch.setattr('devito.arch.archinfo.get_nvidia_cc', lambda: cc)
     platform = platform_registry[platform]()
     assert platform.supports('fp32x2', language) is (expected and language == 'cuda')
+
+
+@pytest.mark.parametrize('memory', ['total', 'free', 'used'])
+def test_nvidia_memory_units(monkeypatch, memory):
+    responses = {
+        ('nvidia-smi', '-L'): (
+            b'GPU 0: NVIDIA A100-SXM4-80GB (UUID: GPU-0)\n'
+            b'GPU 1: NVIDIA A100-SXM4-80GB (UUID: GPU-1)\n'
+        ),
+        ('nvidia-smi', f'--query-gpu=memory.{memory}', '--format=csv'): (
+            f'memory.{memory} [MiB]\n1024 MiB\n81920 MiB\n'.encode()
+        )
+    }
+
+    def popen(command, **kwargs):
+        return SimpleNamespace(stdout=BytesIO(responses[tuple(command)]))
+
+    monkeypatch.setattr(archinfo, 'Popen', popen)
+    info = archinfo.get_gpu_info.func()
+
+    assert info[f'mem.{memory}'](0) == 1024**3
+    assert info[f'mem.{memory}'](1) == 80*1024**3

From 0ac74c5e4c322a4d28edc7376f7332e038af909c Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 12:11:22 +0000
Subject: [PATCH 06/10] data: Avoid quadratic membership scans when reshaping
 decompositions

---
 devito/data/decomposition.py |  5 ++++-
 tests/test_data.py           | 15 ++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/devito/data/decomposition.py b/devito/data/decomposition.py
index ce14dfb742..bc9114bac4 100644
--- a/devito/data/decomposition.py
+++ b/devito/data/decomposition.py
@@ -527,7 +527,10 @@ def reshape(self, *args):
                 else:
                     nright = arg.stop - self.glb_max - 1
             elif isinstance(arg, Iterable):
-                items = [np.array([j for j in i if j in arg]) for i in self]
+                # Sparse masks can be large; avoid scanning one per index.
+                if not isinstance(arg, np.ndarray):
+                    arg = tuple(arg)
+                items = [i[np.isin(i, arg)] for i in self]
                 for i, arr in enumerate(list(items)):
                     items[i] = np.arange(arr.size) + sum(j.size for j in items[:i])
                 return Decomposition(items, self.local)
diff --git a/tests/test_data.py b/tests/test_data.py
index 28698b9344..0438ea80ac 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -661,13 +661,18 @@ def test_reshape_slice(self):
         assert d.reshape(slice(3, -1)) == Decomposition([[], [0, 1], [2, 3, 4],
                                                          [5, 6, 7]], 2)
 
-    def test_reshape_iterable(self):
+    @pytest.mark.parametrize('container', [tuple, list, set, np.asarray])
+    def test_reshape_iterable(self, container):
         d = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 2)
 
-        assert d.reshape(()) == Decomposition([[], [], [], []], 2)
-        assert d.reshape((1, 3, 5)) == Decomposition([[0], [1], [2], []], 2)
-        assert d.reshape((1, 3, 10, 11)) == Decomposition([[0], [1], [], [2, 3]], 2)
-        assert d.reshape((1, 3, 10, 11, 14)) == Decomposition([[0], [1], [], [2, 3]], 2)
+        assert d.reshape(container(())) == Decomposition([[], [], [], []], 2)
+        assert d.reshape(container((1, 3, 5))) == Decomposition([[0], [1], [2], []], 2)
+        assert d.reshape(container((1, 3, 10, 11))) == \
+            Decomposition([[0], [1], [], [2, 3]], 2)
+        assert d.reshape(container((1, 3, 10, 11, 14))) == \
+            Decomposition([[0], [1], [], [2, 3]], 2)
+        assert d.reshape(container((11, 3, 3, 1, -1, 14))) == \
+            Decomposition([[0], [1], [], [2]], 2)
 
 
 # A reference global shape for the serial engine tests below; indices are

From f00d2f69754e2525c1d941a04c0af1ffbcb64e85 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 12:19:30 +0000
Subject: [PATCH 07/10] mpi: Restrict halo overlap to regions within loops

Select overlap candidates through their nearest enclosing Iteration, keeping
exchanges outside loops synchronous. Replace the hand-built HaloSpot test
with Operator-based structure and numerical checks for all three overlap
modes on one and two ranks.
---
 devito/passes/iet/mpi.py | 13 ++++-----
 tests/test_mpi.py        | 60 +++++++++++++++++++++++++---------------
 2 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py
index 02f25aa935..718008a280 100644
--- a/devito/passes/iet/mpi.py
+++ b/devito/passes/iet/mpi.py
@@ -12,7 +12,7 @@
 from devito.mpi.routines import HaloExchangeBuilder, ReductionBuilder
 from devito.passes.iet.engine import iet_pass
 from devito.symbolics import VectorAccess, search
-from devito.tools import generator
+from devito.tools import flatten, generator
 from devito.types import TensorMove
 
 __all__ = ['mpiize']
@@ -275,14 +275,13 @@ def _mark_overlappable(iet):
     Detect the HaloSpots allowing computation/communication overlap and turn
     them into OverlappableHaloSpots.
     """
+    # Only overlap exchanges within an Iteration
+    mapper = MapNodes(Iteration, HaloSpot, 'immediate').visit(iet)
+    mapper.pop(None, None)
+
     # Analysis
     found = []
-    for hs in FindNodes(HaloSpot).visit(iet):
-        # Extracting this body into a compute callable would hide nested
-        # exchanges from the lowering traversal. Keep the enclosing exchange
-        # synchronous; its children can still overlap independently.
-        if FindNodes(HaloSpot).visit(hs.body):
-            continue
+    for hs in flatten(mapper.values()):
         # Heuristic: avoid comp/comm overlap for sparse Iteration nests
         iters = FindNodes(Iteration).visit(hs)
         if any(i.dim._defines & set(hs.halo_scheme.distributed_aindices) and
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 4d0dff3508..11e798f3f5 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -31,27 +31,6 @@
 
 class TestMPIUtils:
 
-    @pytest.mark.parallel(mode=[(1, 'overlap2')])
-    def test_nested_halo_overlap(self, mode):
-        from devito.ir.iet import Expression, HaloSpot, List
-        from devito.ir import LoweredEq
-        from devito.mpi.halo_scheme import HaloScheme
-        from devito.passes.iet.mpi import OverlappableHaloSpot, _mark_overlappable
-
-        grid = Grid(shape=(16, 16, 16), topology=(1, '*', 1))
-        x, y, z = grid.dimensions
-        u = Function(name='nested_u', grid=grid, space_order=2)
-        v = Function(name='nested_v', grid=grid, space_order=2)
-        equation = LoweredEq(Eq(v, u[x, y+1, z]))
-        ispace = IterationSpace([Interval(d, 0, 0) for d in grid.dimensions])
-        scheme = HaloScheme([equation], ispace)
-        inner = HaloSpot(Expression(equation), scheme)
-        outer = HaloSpot(List(body=[inner]), scheme)
-        result = _mark_overlappable(outer)
-        spots = FindNodes(HaloSpot).visit(result)
-        assert type(spots[0]) is HaloSpot
-        assert isinstance(spots[1], OverlappableHaloSpot)
-
     @pytest.mark.parametrize('comm', [None, MPI.COMM_NULL], ids=['none', 'null'])
     @pytest.mark.parametrize('error', [None, '', 'invalid argument'])
     @pytest.mark.parametrize('exception', [ValueError, InvalidArgument])
@@ -957,12 +936,14 @@ class CoordSlowSparseFunction(SparseTimeFunction):
 
     @pytest.mark.parallel(mode=[2, 4])
     @pytest.mark.parametrize('ndim', [1, 2, 3])
-    @pytest.mark.parametrize('interpolation,radius', [('nearest', 0), ('linear', 1), ('sinc', 4)])
+    @pytest.mark.parametrize('interpolation,radius',
+                             [('nearest', 0), ('linear', 1), ('sinc', 4)])
     def test_tensor_support_rank_intersections(self, mode, ndim, interpolation, radius):
         """Compact intervals match every point of wide, shifted tensor stencils."""
         grid = Grid(shape=(9,)*ndim, extent=(1.,)*ndim)
         coordinates = np.array([[value + (axis % 2)*0.125 for axis in range(ndim)]
-                                for value in (-0.25, 0., 0.125, 0.49, 0.5, 0.875, 1., 1.25)],
+                                for value in (-0.25, 0., 0.125, 0.49, 0.5, 0.875,
+                                              1., 1.25)],
                                dtype=np.float32)
         src = SparseFunction(name='support_src', grid=grid, npoint=len(coordinates),
                              interpolation=interpolation, r=radius)
@@ -2519,6 +2500,39 @@ def define(self, dimensions):
         # to Function on SubDomain
         check_halo_exchanges(op, 0, 0)
 
+    @pytest.mark.parallel(mode=[(1, 'overlap'), (1, 'overlap2'), (1, 'full'),
+                                (2, 'overlap'), (2, 'overlap2'), (2, 'full')])
+    @pytest.mark.parametrize('time_dependent', [False, True])
+    def test_overlap_within_iteration(self, mode, time_dependent):
+        grid = Grid(shape=(16, 16), extent=(15., 15.), topology=('*', 1))
+        if time_dependent:
+            u = TimeFunction(name='u', grid=grid, space_order=2)
+            v = TimeFunction(name='v', grid=grid, space_order=2)
+            u.data_with_halo[:] = 0
+            equations = [Eq(u.forward, u + u.laplace + 1),
+                         Eq(v.forward, u.forward + u.forward.laplace)]
+        else:
+            u = Function(name='u', grid=grid, space_order=2)
+            v = Function(name='v', grid=grid, space_order=2)
+            equations = [Eq(u, 1), Eq(v, u + u.laplace)]
+
+        op = Operator(equations, name='OverlapWithinIteration')
+        calls = FindNodes(ComputeCall).visit(op)
+        assert len(calls) == int(time_dependent)
+        if time_dependent:
+            check_halo_exchanges(op, 2, 1)
+            assert calls == FindNodes(ComputeCall).visit(get_time_loop(op))
+            op.apply(time_M=0)
+            values = np.asarray(v.data[1])
+        else:
+            op.apply()
+            values = np.asarray(v.data)
+
+        x, y = np.meshgrid(*grid.distributor.glb_numb, indexing='ij')
+        # The zero exterior halo subtracts one at each physical boundary
+        expected = 1 - (x == 0) - (x == 15) - (y == 0) - (y == 15)
+        np.testing.assert_allclose(values, expected, rtol=0, atol=1e-7)
+
 
 class TestOperatorAdvanced:
 

From fdb2d644382a1e731880855ddeafc4e3256fa98f Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 12:19:38 +0000
Subject: [PATCH 08/10] arch: Share atomic library publication through OS
 helpers

Extend the existing atomic_write helper with exclusive publication and use
binary mode when restoring shared libraries. Keep text and header writes
replacing by default, and cover cleanup and existing destinations.
---
 devito/arch/compiler.py            | 16 ++--------------
 devito/tools/os_helper.py          | 17 +++++++++++++----
 devito/types/sparse.py             |  4 ++--
 tests/test_compiler_publication.py |  7 +++----
 tests/test_tools.py                | 22 ++++++++++++++++++++++
 5 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/devito/arch/compiler.py b/devito/arch/compiler.py
index 2c2fd37664..2bb9db57b3 100644
--- a/devito/arch/compiler.py
+++ b/devito/arch/compiler.py
@@ -5,9 +5,8 @@
 from functools import partial
 from hashlib import sha1
 from itertools import filterfalse
-from os import environ, fsync, link, makedirs, path, unlink
+from os import environ, path
 from subprocess import DEVNULL, PIPE, CalledProcessError, check_call, check_output, run
-from tempfile import NamedTemporaryFile
 
 import numpy.ctypeslib as npct
 from codepy.jit import compile_from_string
@@ -347,21 +346,10 @@ def save(self, soname, binary):
             debug(f"{self}: `{sofile.name}` was not saved in `{self.get_jit_dir()}`"
                   " as it already exists")
         else:
-            makedirs(self.get_jit_dir(), exist_ok=True)
             # MPI ranks may restore the same Operator concurrently. Publish
             # only a complete binary: dlopen of a visible partial ELF can
             # raise SIGBUS rather than a recoverable loader exception.
-            with NamedTemporaryFile(dir=self.get_jit_dir(), delete=False) as f:
-                temporary = f.name
-                try:
-                    f.write(binary)
-                    f.flush()
-                    fsync(f.fileno())
-                    # link is atomic and keeps an existing publication intact.
-                    with suppress(FileExistsError):
-                        link(temporary, sofile)
-                finally:
-                    unlink(temporary)
+            atomic_write(sofile, binary, mode='wb', overwrite=False)
             debug(f"{self}: `{sofile.name}` successfully saved in `{self.get_jit_dir()}`")
 
     def make(self, loc, args):
diff --git a/devito/tools/os_helper.py b/devito/tools/os_helper.py
index 136ab34316..bc8ef635a4 100644
--- a/devito/tools/os_helper.py
+++ b/devito/tools/os_helper.py
@@ -37,7 +37,7 @@ def make_tempdir(prefix=None):
     return tmpdir
 
 
-def atomic_write(path, data, mode='w'):
+def atomic_write(path, data, mode='w', *, overwrite=True):
     """
     Write `data` to `path` through a private temporary file in the same
     directory, then rename it into place. The rename is atomic, so a concurrent
@@ -45,10 +45,14 @@ def atomic_write(path, data, mode='w'):
     the new one, and a reader that already opened the previous file keeps
     reading that one in full. The temporary file is removed if the write fails.
 
+    With `overwrite=False`, publish through an exclusive hard link instead,
+    preserving an existing destination. Missing parent directories are created.
+
     Note: the resulting file has the restrictive mode of the temporary file
     (0600, further restricted by the process umask).
     """
     path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
     fd, tmpfile = mkstemp(dir=path.parent, prefix=f'{path.name}.', suffix='.tmp')
     try:
         try:
@@ -58,9 +62,14 @@ def atomic_write(path, data, mode='w'):
             raise
         with f:
             f.write(data)
-        os.replace(tmpfile, path)
-    except BaseException:
+            f.flush()
+            os.fsync(f.fileno())
+        if overwrite:
+            os.replace(tmpfile, path)
+        else:
+            with suppress(FileExistsError):
+                os.link(tmpfile, path)
+    finally:
         with suppress(FileNotFoundError):
             os.unlink(tmpfile)
-        raise
     return path
diff --git a/devito/types/sparse.py b/devito/types/sparse.py
index a71f045278..801e67fb34 100644
--- a/devito/types/sparse.py
+++ b/devito/types/sparse.py
@@ -14,8 +14,8 @@
 from devito.operations.interpolators import _cell_indices
 from devito.symbolics import indexify, retrieve_function_carriers
 from devito.tools import (
-    ReducerMap, as_fp64_decimal, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
-    memoized_meth, prod
+    ReducerMap, as_fp64_decimal, as_tuple, dtype_to_mpidtype, filter_ordered, flatten,
+    is_integer, memoized_meth, prod
 )
 from devito.types.basic import Symbol
 from devito.types.dense import DiscreteFunction, SubFunction
diff --git a/tests/test_compiler_publication.py b/tests/test_compiler_publication.py
index 58920533b6..b9f44095c0 100644
--- a/tests/test_compiler_publication.py
+++ b/tests/test_compiler_publication.py
@@ -4,15 +4,14 @@
 from types import SimpleNamespace
 
 from devito.arch.compiler import Compiler
+from devito.tools import os_helper
 
 
 def test_complete_binary_publication(tmp_path, monkeypatch):
-    import devito.arch.compiler as module
-
     compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
     target = tmp_path/'shared.so'
     binary = b'ELF-binary-content'*1024
-    original = module.link
+    original = os_helper.os.link
     publications = []
 
     def observe(source, destination):
@@ -21,7 +20,7 @@ def observe(source, destination):
         publications.append(source)
         original(source, destination)
 
-    monkeypatch.setattr(module, 'link', observe)
+    monkeypatch.setattr(os_helper.os, 'link', observe)
     Compiler.save(compiler, 'shared', binary)
     assert publications
     assert target.read_bytes() == binary
diff --git a/tests/test_tools.py b/tests/test_tools.py
index 2947866a85..63434e078c 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -351,3 +351,25 @@ def test_atomic_write(tmp_path):
         atomic_write(path, 1)
     assert path.read_text() == 'foo'
     assert list(tmp_path.iterdir()) == [path]
+
+    assert atomic_write(path, 'bar') == path
+    assert path.read_text() == 'bar'
+
+
+def test_atomic_write_exclusive(tmp_path):
+    target = tmp_path/'nested'/'file'
+    assert atomic_write(target, b'first', mode='wb', overwrite=False) == target
+    assert atomic_write(str(target), b'second', mode='wb', overwrite=False) == target
+    assert target.read_bytes() == b'first'
+    assert list(target.parent.iterdir()) == [target]
+
+
+@pytest.mark.parametrize('operation', ['fsync', 'link'])
+def test_atomic_write_failure(tmp_path, monkeypatch, operation):
+    def fail(*args):
+        raise OSError('Publication failed')
+
+    monkeypatch.setattr(os, operation, fail)
+    with pytest.raises(OSError, match='Publication failed'):
+        atomic_write(tmp_path/'file', b'contents', mode='wb', overwrite=False)
+    assert not list(tmp_path.iterdir())

From db6ffd9d9b83cea3d00e7b21e02eb2b54aea240b Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 15:29:41 +0000
Subject: [PATCH 09/10] tests: Group library publication checks with
 architecture tests

Move both restoration regressions into TestLibraryPublication and remove the standalone module, preserving the existing assertions and fixtures.
---
 tests/test_arch.py                 | 42 ++++++++++++++++++++++++++++-
 tests/test_compiler_publication.py | 43 ------------------------------
 2 files changed, 41 insertions(+), 44 deletions(-)
 delete mode 100644 tests/test_compiler_publication.py

diff --git a/tests/test_arch.py b/tests/test_arch.py
index 19dbf6df2e..5106a80455 100644
--- a/tests/test_arch.py
+++ b/tests/test_arch.py
@@ -1,6 +1,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
 from io import BytesIO
+from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
@@ -10,8 +11,9 @@
 from devito.arch import BLACKWELL, NVIDIAX, archinfo, platform_registry
 from devito.arch.archinfo import AppleArm, get_apple_platform
 from devito.arch.compiler import (
-    CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
+    Compiler, CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
 )
+from devito.tools import os_helper
 
 
 @pytest.mark.parametrize("cc", [
@@ -151,3 +153,41 @@ def popen(command, **kwargs):
 
     assert info[f'mem.{memory}'](0) == 1024**3
     assert info[f'mem.{memory}'](1) == 80*1024**3
+
+
+class TestLibraryPublication:
+    """Concurrent Operator restoration must never expose a partial library."""
+
+    def test_complete_binary(self, tmp_path, monkeypatch):
+        compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
+        target = tmp_path/'shared.so'
+        binary = b'ELF-binary-content'*1024
+        original = os_helper.os.link
+        publications = []
+
+        def observe(source, destination):
+            assert not target.exists()
+            assert Path(source).read_bytes() == binary
+            publications.append(source)
+            original(source, destination)
+
+        monkeypatch.setattr(os_helper.os, 'link', observe)
+        Compiler.save(compiler, 'shared', binary)
+        assert publications
+        assert target.read_bytes() == binary
+        assert list(tmp_path.iterdir()) == [target]
+        Compiler.save(compiler, 'shared', b'different')
+        assert target.read_bytes() == binary
+
+    def test_concurrent_binary(self, tmp_path):
+        compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
+        target = tmp_path/'shared.so'
+        binary = b'ELF-binary-content'*(1 << 18)
+
+        def restore(_):
+            Compiler.save(compiler, 'shared', binary)
+            assert target.read_bytes() == binary
+
+        with ThreadPoolExecutor(max_workers=8) as pool:
+            list(pool.map(restore, range(32)))
+        assert list(tmp_path.iterdir()) == [target]
diff --git a/tests/test_compiler_publication.py b/tests/test_compiler_publication.py
deleted file mode 100644
index b9f44095c0..0000000000
--- a/tests/test_compiler_publication.py
+++ /dev/null
@@ -1,43 +0,0 @@
-"""Concurrent Operator restoration must never expose a partial library."""
-from concurrent.futures import ThreadPoolExecutor
-from pathlib import Path
-from types import SimpleNamespace
-
-from devito.arch.compiler import Compiler
-from devito.tools import os_helper
-
-
-def test_complete_binary_publication(tmp_path, monkeypatch):
-    compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
-    target = tmp_path/'shared.so'
-    binary = b'ELF-binary-content'*1024
-    original = os_helper.os.link
-    publications = []
-
-    def observe(source, destination):
-        assert not target.exists()
-        assert Path(source).read_bytes() == binary
-        publications.append(source)
-        original(source, destination)
-
-    monkeypatch.setattr(os_helper.os, 'link', observe)
-    Compiler.save(compiler, 'shared', binary)
-    assert publications
-    assert target.read_bytes() == binary
-    assert list(tmp_path.iterdir()) == [target]
-    Compiler.save(compiler, 'shared', b'different')
-    assert target.read_bytes() == binary
-
-
-def test_concurrent_binary_publication(tmp_path):
-    compiler = SimpleNamespace(get_jit_dir=lambda: tmp_path, so_ext='.so')
-    target = tmp_path/'shared.so'
-    binary = b'ELF-binary-content'*(1 << 18)
-
-    def restore(_):
-        Compiler.save(compiler, 'shared', binary)
-        assert target.read_bytes() == binary
-
-    with ThreadPoolExecutor(max_workers=8) as pool:
-        list(pool.map(restore, range(32)))
-    assert list(tmp_path.iterdir()) == [target]

From 9b1c076a89ee3c5ce2be6c68f201bbe33ee7c311 Mon Sep 17 00:00:00 2001
From: Fabio Luporini <fabio@devitocodes.com>
Date: Thu, 5 Nov 2026 17:01:17 +0000
Subject: [PATCH 10/10] mpi: Simplify sparse support rank mapping

Derive tensor support bounds directly from the radius and keep cell geometry in the interpolator. Preserve the interval optimization and test separate staggered support boxes.
---
 devito/operations/interpolators.py | 28 +++++++++++++-----------
 devito/types/grid.py               | 12 +++++++++-
 devito/types/sparse.py             | 35 ++++++++++--------------------
 tests/test_mpi.py                  | 30 ++++++++++++++++++++-----
 4 files changed, 61 insertions(+), 44 deletions(-)

diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py
index ff4e14b1de..050999aff4 100644
--- a/devito/operations/interpolators.py
+++ b/devito/operations/interpolators.py
@@ -244,6 +244,18 @@ def grid(self):
     def r(self):
         return self.sfunction.r
 
+    @property
+    def _support_indices(self):
+        """Base cells for every staggering used by the sparse operations."""
+        if not self._shifts_used:
+            return [self.sfunction._coords_indices]
+
+        # Redistribution must use the same fp64 cells as the interpolation tables
+        spacing, origin = self.grid._spacing_fp64, self.grid._origin_fp64
+        return [_cell_indices(self.sfunction.coordinates_data, self.grid,
+                              shift, spacing, origin)
+                for shift in self._shifts_used]
+
     @cached_property
     def _coeff_dtype(self):
         # Weights are real even for complex-valued sparse fields.
@@ -398,11 +410,10 @@ def _arg_defaults(self, coords=None, sfunc=None, origin=None):
         if coords is None or sfunc is None:
             raise ValueError("No coordinates or sparse function provided")
 
-        # Fp64 grid geometry -- avoids fp32 rounding on cell boundaries.
         grid = sfunc.grid
-        spacing = np.array([as_fp64_decimal(h) for h in grid.spacing])
-        origin = np.array([as_fp64_decimal(o)
-                           for o in (origin or grid.origin)])
+        spacing = grid._spacing_fp64
+        origin = (np.array([as_fp64_decimal(o) for o in origin])
+                  if origin is not None else grid._origin_fp64)
 
         args = {}
         for key in self._shifts_used or {None}:
@@ -654,15 +665,6 @@ class Coeffs(_HostTable):
     """Per-dim interpolation weights, shape ``(npoint, 2*r)``."""
 
 
-def _resolved_geometry(grid, kwargs):
-    """Fp64 (spacing, origin) tuple honoring runtime `h_x`/`o_x`/... overrides."""
-    spacing = np.array([as_fp64_decimal(kwargs.get(s.name, v)) for s, v
-                        in zip(grid.spacing_symbols, grid.spacing, strict=True)])
-    origin = np.array([as_fp64_decimal(kwargs.get(o.name, v)) for o, v
-                       in zip(grid.origin_symbols, grid.origin, strict=True)])
-    return spacing, origin
-
-
 def _positions_fp64(coords, grid, shifts, spacing, origin):
     """Fp64 fractional grid position `(c - o - shift)/h` for each sparse point,
     computed in double precision to avoid boundary-crossing rounding errors
diff --git a/devito/types/grid.py b/devito/types/grid.py
index 4568b852bf..35f2b0daae 100644
--- a/devito/types/grid.py
+++ b/devito/types/grid.py
@@ -12,7 +12,7 @@
 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, is_integer
+from devito.tools import ReducerMap, as_fp64_decimal, as_tuple, frozendict, is_integer
 from devito.types.args import ArgProvider
 from devito.types.basic import Scalar
 from devito.types.dense import Function
@@ -252,6 +252,11 @@ def origin(self):
         """Physical coordinates of the domain origin."""
         return self._origin
 
+    @cached_property
+    def _origin_fp64(self):
+        """Origin in fp64, preserving the grid's decimal values."""
+        return np.array([as_fp64_decimal(o) for o in self.origin])
+
     @property
     def origin_symbols(self):
         """Symbols representing the grid origin in each SpaceDimension."""
@@ -313,6 +318,11 @@ def spacing(self):
         spacing = (np.array(self.extent) / (np.array(self.shape) - 1)).astype(self.dtype)
         return as_tuple(spacing)
 
+    @cached_property
+    def _spacing_fp64(self):
+        """Spacing in fp64, preserving the grid's decimal values."""
+        return np.array([as_fp64_decimal(h) for h in self.spacing])
+
     @cached_property
     def spacing_symbols(self):
         """Symbols representing the grid spacing in each SpaceDimension."""
diff --git a/devito/types/sparse.py b/devito/types/sparse.py
index 801e67fb34..ca2e03382a 100644
--- a/devito/types/sparse.py
+++ b/devito/types/sparse.py
@@ -11,11 +11,10 @@
 from devito.operations import (
     LinearInterpolator, NearestInterpolator, PrecomputedInterpolator, SincInterpolator
 )
-from devito.operations.interpolators import _cell_indices
 from devito.symbolics import indexify, retrieve_function_carriers
 from devito.tools import (
-    ReducerMap, as_fp64_decimal, as_tuple, dtype_to_mpidtype, filter_ordered, flatten,
-    is_integer, memoized_meth, prod
+    ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
+    memoized_meth, prod
 )
 from devito.types.basic import Symbol
 from devito.types.dense import DiscreteFunction, SubFunction
@@ -300,15 +299,6 @@ def _coords_indices(self):
                 np.floor((self.coordinates_data - self.grid.origin) / self.grid.spacing)
             ).astype(int)
 
-    @property
-    def _support_indices(self):
-        # Staggered tables may use a different base cell across an MPI
-        # boundary. Match their fp64 geometry and actual stencil support.
-        spacing = np.array([as_fp64_decimal(h) for h in self.grid.spacing])
-        origin = np.array([as_fp64_decimal(o) for o in self.grid.origin])
-        return [_cell_indices(self.coordinates_data, self.grid, shift, spacing, origin)
-                for shift in self.interpolator._shifts_used] or [self._coords_indices]
-
     @property
     def _support(self):
         """
@@ -317,7 +307,7 @@ def _support(self):
         """
         max_shape = np.array(self.grid.shape).reshape(1, self.grid.dim)
         minmax = lambda arr: np.minimum(max_shape, np.maximum(0, arr))
-        return np.stack([minmax(index + s) for index in self._support_indices
+        return np.stack([minmax(index + s) for index in self.interpolator._support_indices
                          for s in self._point_support], axis=2)
 
     @property
@@ -325,26 +315,23 @@ def _dist_datamap(self):
         """
         Mapper ``M : MPI rank -> required sparse data``.
         """
-        offsets = self._point_support
-        lower, upper = offsets.min(axis=0), offsets.max(axis=0)
-        if (len(offsets) != prod(upper - lower + 1)
-                or len(np.unique(offsets, axis=0)) != len(offsets)):
-            return self.grid.distributor.glb_to_rank(self._support) or {}
-
-        # A tensor stencil intersects a Cartesian rank exactly when its
-        # intervals intersect in every dimension. Keep different shifts
-        # separate: their union need not be one rectangular stencil.
+        # `_point_increments` is a tensor product: [-r+1, r], or [0, 0] for nearest
+        lower, upper = min(1 - self.r, 0), self.r
+
+        # Optimization: avoid enumerating the stencil (e.g. 512 points for r=4 in 3D).
+        # Keep each staggering's box separate: their union may have gaps.
         bounds = [(np.clip(index + lower, 0, self.grid.shape),
                    np.clip(index + upper, 0, self.grid.shape))
-                  for index in self._support_indices]
+                  for index in self.interpolator._support_indices]
         ret = {}
         for rank, ranges in enumerate(self.grid.distributor.all_ranges):
             if any(not r for r in ranges):
                 continue
             mins = np.array([r[0] for r in ranges])
             maxs = np.array([r[-1] for r in ranges])
-            mask = np.zeros(bounds[0][0].shape[0], dtype=bool)
+            mask = np.zeros(self.npoint, dtype=bool)
             for low, high in bounds:
+                # Match every axis of a box; any used staggering can qualify
                 mask |= ((low <= maxs) & (high >= mins)).all(axis=1)
             if mask.any():
                 ret[rank] = np.flatnonzero(mask).tolist()
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
index 11e798f3f5..d4fcb09c40 100644
--- a/tests/test_mpi.py
+++ b/tests/test_mpi.py
@@ -938,9 +938,10 @@ class CoordSlowSparseFunction(SparseTimeFunction):
     @pytest.mark.parametrize('ndim', [1, 2, 3])
     @pytest.mark.parametrize('interpolation,radius',
                              [('nearest', 0), ('linear', 1), ('sinc', 4)])
-    def test_tensor_support_rank_intersections(self, mode, ndim, interpolation, radius):
-        """Compact intervals match every point of wide, shifted tensor stencils."""
+    def test_rank_map_equivalence(self, mode, ndim, interpolation, radius):
+        """Send each point to every rank owning a cell in its interpolation stencil."""
         grid = Grid(shape=(9,)*ndim, extent=(1.,)*ndim)
+        # Include partition boundaries and coordinates outside the global domain
         coordinates = np.array([[value + (axis % 2)*0.125 for axis in range(ndim)]
                                 for value in (-0.25, 0., 0.125, 0.49, 0.5, 0.875,
                                               1., 1.25)],
@@ -951,11 +952,11 @@ def test_tensor_support_rank_intersections(self, mode, ndim, interpolation, radi
         plain = Function(name='support_plain', grid=grid, space_order=8)
         shifted = Function(name='support_shifted', grid=grid, space_order=8,
                            staggered=grid.dimensions)
-        src.inject(field=plain, expr=src)
-        src.inject(field=shifted, expr=src)
+        Operator(src.inject(field=(plain, shifted), expr=src), name='SparseSupport')
 
-        exhaustive = grid.distributor.glb_to_rank(src._support) or {}
-        assert src._dist_datamap == exhaustive
+        # Reference: enumerate all stencil cells, including the half-cell shift
+        expected = grid.distributor.glb_to_rank(src._support) or {}
+        assert src._dist_datamap == expected
 
     @pytest.mark.parallel(mode=[2, 4])
     @pytest.mark.parametrize('override', [False, True])
@@ -981,6 +982,23 @@ def test_staggered_injection_support(self, mode, override):
         total = MPI.COMM_WORLD.allreduce(float(np.asarray(u.data).sum()), op=MPI.SUM)
         assert np.isclose(total, 1, rtol=0, atol=1e-7)
 
+    @pytest.mark.parallel(mode=4)
+    @pytest.mark.parametrize('interpolation', ['nearest', 'linear'])
+    def test_staggered_support_union(self, mode, interpolation):
+        grid = Grid(shape=(8, 8), extent=(7., 7.), topology=(2, 2))
+        x, y = grid.dimensions
+        u = Function(name='ux', grid=grid, space_order=2, staggered=x)
+        v = Function(name='uy', grid=grid, space_order=2, staggered=y)
+        src = SparseFunction(name='src', grid=grid, npoint=1,
+                             interpolation=interpolation)
+        src.coordinates.data[:] = 4.
+        Operator(src.inject(field=(u, v), expr=src), name='SparseSupportUnion')
+
+        # Boxing the union would also send the point to rank 0
+        ranks = (1, 2) if interpolation == 'nearest' else (1, 2, 3)
+        expected = {r: [0] for r in ranks} if src.npoint else {}
+        assert src._dist_datamap == expected
+
 
 class TestOperatorSimple: