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/devito/arch/compiler.py b/devito/arch/compiler.py index d356765636..2bb9db57b3 100644 --- a/devito/arch/compiler.py +++ b/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): 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/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/passes/iet/mpi.py b/devito/passes/iet/mpi.py index db835400ae..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,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 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/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 6e913646a3..ca2e03382a 100644 --- a/devito/types/sparse.py +++ b/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)) diff --git a/tests/test_arch.py b/tests/test_arch.py index 305887d2e2..5106a80455 100644 --- a/tests/test_arch.py +++ b/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] 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 diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 7081be3beb..d4fcb09c40 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -934,6 +934,71 @@ 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_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)], + 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) + Operator(src.inject(field=(plain, shifted), expr=src), name='SparseSupport') + + # 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]) + 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) + + @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: @@ -2453,6 +2518,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: 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())