compiler: Fix MPI data distribution for free surface use cases

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

devito/arch/compiler.py

@@ -5,7 +5,7 @@
 from functools import partial
 from hashlib import sha1
 from itertools import filterfalse
-from os import environ, makedirs, path
+from os import environ, path
 from subprocess import DEVNULL, PIPE, CalledProcessError, check_call, check_output, run
 
 import numpy.ctypeslib as npct
@@ -346,9 +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)
-            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.
+            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):

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)

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

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,9 +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):
+    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

devito/tools/os_helper.py

@@ -37,18 +37,22 @@ 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
     reader never sees a partial file: it reads either the previous content or
     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

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."""

devito/types/sparse.py

@@ -307,15 +307,35 @@ 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)
+        return np.stack([minmax(index + s) for index in self.interpolator._support_indices
+                         for s in self._point_support], axis=2)
 
     @property
     def _dist_datamap(self):
         """
         Mapper ``M : MPI rank -> required sparse data``.
         """
-        return self.grid.distributor.glb_to_rank(self._support) or {}
+        # `_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.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(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()
+        return ret
 
     @property
     def npoint(self):
@@ -1004,10 +1024,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))

tests/test_arch.py

@@ -1,15 +1,19 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import suppress
+from io import BytesIO
+from pathlib import Path
+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
+    Compiler, CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
 )
+from devito.tools import os_helper
 
 
 @pytest.mark.parametrize("cc", [
@@ -127,3 +131,63 @@ 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
+
+
+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]

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