compiler: precompute interp pos/weights

Author: mloubout

.github/workflows/pytest-core-mpi.yaml

@@ -22,7 +22,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: ['3.10', '3.11']
+        python-version: ['3.11', '3.14']
 
     env:
       DEVITO_LANGUAGE: "openmp"

.github/workflows/pytest-core-nompi.yaml

@@ -39,7 +39,7 @@ jobs:
            pytest-ubuntu-py311-gcc10-noomp,
            pytest-ubuntu-py311-gcc11-cxxnoomp,
            pytest-ubuntu-py312-gcc12-cxxomp,
-           pytest-ubuntu-py312-gcc13-omp,
+           pytest-ubuntu-py314-gcc13-omp,
            pytest-ubuntu-py313-gcc14-omp
         ]
         set: [base, adjoint]
@@ -73,8 +73,8 @@ jobs:
           language: "C"
           sympy: "1.14"
 
-        - name: pytest-ubuntu-py312-gcc13-omp
-          python-version: '3.12'
+        - name: pytest-ubuntu-py314-gcc13-omp
+          python-version: '3.14'
           os: ubuntu-24.04
           arch: "gcc-13"
           language: "openmp"
@@ -134,7 +134,7 @@ jobs:
           fi
       id: set-tests
 
-    - name: Set pip flags for latest python (3.12)
+    - name: Set pip flags for Python 3.12+
       run: |
           ver="${{ matrix.python-version }}"
           major=${ver%%.*}

devito/operations/interpolators.py

@@ -15,8 +15,10 @@
 from devito.finite_differences.elementary import floor
 from devito.logger import warning
 from devito.symbolics import INT, retrieve_function_carriers, retrieve_functions
-from devito.tools import Pickable, as_tuple, filter_ordered, flatten, memoized_meth
-from devito.types import Eq, Evaluable, Inc, SubFunction, Symbol
+from devito.tools import (
+    Pickable, as_fp64_decimal, as_tuple, filter_ordered, flatten, memoized_meth
+)
+from devito.types import CustomDimension, Eq, Evaluable, Inc, SubFunction, Symbol
 from devito.types.utils import DimensionTuple
 
 __all__ = ['LinearInterpolator', 'PrecomputedInterpolator', 'SincInterpolator']
@@ -510,46 +512,195 @@ def _inject(self, field, expr, implicit_dims=None):
         return filter_ordered(temps) + eqns
 
 
+def _shift_tag(shifts):
+    """Suffix used to distinguish per-staggering table names ("_s10", ...)."""
+    if not shifts or not any(shifts):
+        return ''
+    return '_s' + ''.join('1' if s else '0' for s in shifts)
+
+
+def _shift_values(shifts, grid, spacing):
+    """Physical half-cell offsets for each grid dim, as fp64."""
+    if not shifts:
+        return np.zeros(grid.dim, dtype=np.float64)
+    subs = {d.spacing: float(h)
+            for d, h in zip(grid.dimensions, spacing, strict=True)}
+    return np.array([float(sympy.sympify(s).xreplace(subs)) for s in shifts])
+
+
+class _HostTable(SubFunction):
+    """Base for a SubFunction populated on the host by the linear interpolator
+    at arg-prep time. Data honors runtime grid overrides (spacing/origin) and
+    the parent SparseFunction's scattered coordinates."""
+
+    def _arg_apply(self, *args, **kwargs):
+        return
+
+    def _emit(self, coords, spacing, origin):  # pragma: no cover - abstract
+        raise NotImplementedError
+
+    def _arg_defaults(self, alias=None, metadata=None, estimate_memory=False):
+        args = dict(self.dimensions[-1]._arg_defaults(_min=0,
+                                                      size=self.shape[-1]))
+        key = alias or self
+        if estimate_memory:
+            args[key.name] = self
+        else:
+            spacing, origin = _resolved_geometry(self._parent.grid, {})
+            args[key.name] = self._emit(self._parent.coordinates.data,
+                                        spacing, origin)
+        return args
+
+    def _arg_values(self, estimate_memory=False, **kwargs):
+        # Route through the parent so its dim defaults (and any override)
+        # are collected; then overlay this table's data honoring runtime
+        # coord / spacing / origin overrides.
+        values = super()._arg_values(estimate_memory=estimate_memory, **kwargs)
+        if estimate_memory:
+            return values
+        args = kwargs.get('args')
+        coords = (args.get(self._parent.coordinates.name)
+                  if args is not None else None)
+        if coords is None:
+            coords = self._parent.coordinates.data
+        spacing, origin = _resolved_geometry(self._parent.grid, kwargs)
+        values[self.name] = self._emit(coords, spacing, origin)
+        return values
+
+
+class Gridpoints(_HostTable):
+    """int32 cell indices per sparse point, shape ``(npoint, ndim)``."""
+
+    __rkwargs__ = _HostTable.__rkwargs__ + ('shifts',)
+
+    def __init_finalize__(self, *args, **kwargs):
+        self._shifts = kwargs.pop('shifts', None)
+        super().__init_finalize__(*args, **kwargs)
+
+    @property
+    def shifts(self):
+        return self._shifts
+
+    def _emit(self, coords, spacing, origin):
+        return _cell_indices(coords, self._parent, self._shifts,
+                             spacing, origin)
+
+
+class Coeffs(_HostTable):
+    """Per-dim `(1 - frac, frac)` interpolation weights, shape ``(npoint, 2)``."""
+
+    __rkwargs__ = _HostTable.__rkwargs__ + ('shifts', 'dim_index')
+
+    def __init_finalize__(self, *args, **kwargs):
+        self._shifts = kwargs.pop('shifts', None)
+        self._dim_index = kwargs.pop('dim_index')
+        super().__init_finalize__(*args, **kwargs)
+
+    @property
+    def shifts(self):
+        return self._shifts
+
+    @property
+    def dim_index(self):
+        return self._dim_index
+
+    def _emit(self, coords, spacing, origin):
+        return _linear_weights(coords, self._parent, self._shifts,
+                               self._dim_index, self.dtype, spacing, origin)
+
+
+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, sfunc, shifts, spacing, origin):
+    c64 = np.asarray(coords, dtype=np.float64)
+    return (c64 - origin - _shift_values(shifts, sfunc.grid, spacing)) / spacing
+
+
+def _cell_indices(coords, sfunc, shifts, spacing, origin):
+    return np.floor(_positions_fp64(coords, sfunc, shifts, spacing,
+                                    origin)).astype(np.int32)
+
+
+def _linear_weights(coords, sfunc, shifts, j, dtype, spacing, origin):
+    pos = _positions_fp64(coords, sfunc, shifts, spacing, origin)
+    frac = pos[:, j] - np.floor(pos[:, j])
+    data = np.empty((pos.shape[0], 2), dtype=dtype)
+    data[:, 0] = 1.0 - frac
+    data[:, 1] = frac
+    return data
+
+
 class LinearInterpolator(WeightedInterpolator):
     """
-    Concrete implementation of WeightedInterpolator implementing a Linear interpolation
-    scheme, i.e. Bilinear for 2D and Trilinear for 3D problems.
+    Linear (bilinear/trilinear) interpolator.
 
-    Parameters
-    ----------
-    sfunction: The SparseFunction that this Interpolator operates on.
+    Gridpoints and per-dim `(1-frac, frac)` weights are precomputed on the
+    host in fp64 (see `_arg_defaults`) and passed to the kernel as int32/fp
+    SubFunctions. The generated C only indexes those tables and never sees
+    `(c-o)/h` or `floor` on fp32.
     """
 
     _name = 'linear'
 
+    @cached_property
+    def _coeff_dtype(self):
+        # Weights are real even for complex-valued sparse fields.
+        dtype = np.dtype(self.sfunction.dtype)
+        if np.issubdtype(dtype, np.complexfloating):
+            return np.finfo(dtype).dtype.type
+        return dtype.type
+
     @memoized_meth
-    def _weights(self, subdomain=None, shifts=None):
-        rdim = self._rdim(subdomain=subdomain, shifts=shifts)
-        c = [(1 - p) * (1 - r) + p * r
-             for (p, d, r) in zip(self._point_symbols(shifts), self._gdims, rdim,
-                                  strict=True)]
-        return Mul(*c)
+    def _gridpoints(self, shifts=None):
+        name = f'gp{self.sfunction.name}{_shift_tag(shifts)}'
+        sfdim = self.sfunction._sparse_dim
+        ddim = CustomDimension(f'{name}d', 0, self.grid.dim - 1,
+                               self.grid.dim, sfdim)
+        return Gridpoints(name=name, dtype=np.int32,
+                          shape=(self.sfunction.npoint, self.grid.dim),
+                          dimensions=(sfdim, ddim), space_order=0,
+                          alias=self.sfunction.alias,
+                          parent=self.sfunction, shifts=shifts)
 
     @memoized_meth
-    def _point_symbols(self, shifts=None):
-        """Symbol for coordinate value in each Dimension of the point."""
-        dtype = self.sfunction.coordinates.dtype
-        symbols = []
-        for d in self.grid.dimensions:
-            if shifts and shifts[self.grid.dimensions.index(d)] != 0:
-                symbols.append(Symbol(name=f'p{d}_s1', dtype=dtype))
-            else:
-                symbols.append(Symbol(name=f'p{d}', dtype=dtype))
-        return DimensionTuple(*symbols, getters=self.grid.dimensions)
+    def _coeffs(self, shifts=None):
+        sfdim = self.sfunction._sparse_dim
+        tag = _shift_tag(shifts)
+        return tuple(
+            Coeffs(name=f'w{self.sfunction.name}{d.name}{tag}',
+                   dtype=self._coeff_dtype,
+                   shape=(self.sfunction.npoint, 2),
+                   dimensions=(sfdim, r), space_order=0,
+                   alias=self.sfunction.alias,
+                   parent=self.sfunction, shifts=shifts, dim_index=i)
+            for i, (d, r) in enumerate(zip(self._gdims, self._cdim, strict=True))
+        )
+
+    def _positions(self, implicit_dims, shifts=None):
+        gp = self._gridpoints(shifts=shifts)
+        ddim = gp.dimensions[-1]
+        return [Eq(p, gp._subs(ddim, di), implicit_dims=implicit_dims)
+                for (di, p) in enumerate(
+                    self.sfunction._pos_symbols(shifts=shifts))]
 
     def _coeff_temps(self, implicit_dims, shifts=None):
-        # Positions
-        pmap = self.sfunction._position_map(shifts=shifts)
-        psyms = self._point_symbols(shifts)
-        poseq = [Eq(psyms[d], pos - floor(pos),
-                    implicit_dims=implicit_dims)
-                 for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)]
-        return poseq
+        return []
+
+    @memoized_meth
+    def _weights(self, subdomain=None, shifts=None):
+        rdims = self._rdim(subdomain=subdomain, shifts=shifts)
+        coeffs = self._coeffs(shifts=shifts)
+        return Mul(*[
+            w._subs(rd, rd - rd.parent.symbolic_min)
+            for (rd, w) in zip(rdims, coeffs, strict=True)
+        ])
 
 
 class PrecomputedInterpolator(WeightedInterpolator):
@@ -634,7 +785,7 @@ def _weights(self, subdomain=None, shifts=None):
             for (rd, w) in zip(rdims, self.interpolation_coeffs, strict=True)
         ])
 
-    def _arg_defaults(self, coords=None, sfunc=None):
+    def _arg_defaults(self, coords=None, sfunc=None, origin=None):
         args = {}
         b = self._b_table[self.r]
         b0 = i0(b)

devito/tools/dtypes_lowering.py

@@ -16,7 +16,17 @@
            'dtype_to_cstr', 'dtype_to_ctype', 'infer_datasize', 'dtype_to_mpitype',
            'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper',
            'is_external_ctype', 'infer_dtype', 'extract_dtype', 'CustomDtype',
-           'mpi4py_mapper']
+           'mpi4py_mapper', 'as_fp64_decimal']
+
+
+def as_fp64_decimal(v):
+    """
+    fp64 value of ``v`` matching its shortest round-tripping decimal.
+    For an `np.float32` this recovers the decimal the user wrote (e.g.
+    ``np.float32(0.1)`` -> ``0.1`` exact in fp64) rather than the widened
+    fp32 bit pattern (``0.10000000149...``).
+    """
+    return np.float64(np.format_float_positional(v, unique=True, trim='0'))
 
 
 # *** Custom np.dtypes

devito/types/sparse.py

@@ -760,8 +760,12 @@ def _arg_values(self, estimate_memory=False, **kwargs):
                 values = new._arg_defaults(alias=self,
                                            estimate_memory=estimate_memory).reduce_all()
             else:
-                # We've been provided a pure-data replacement (array)
-                values = {}
+                # Pure-data replacement (ndarray). Re-derive full defaults so
+                # any interpolator-owned SubFunctions get rebuilt alongside
+                # the scattered data.
+                values = self._arg_defaults(
+                    alias=self, estimate_memory=estimate_memory
+                ).reduce_all()
                 for k, v in self._dist_scatter(data=new).items():
                     values[k.name] = v
                     for i, s in zip(k.indices, v.shape, strict=True):
@@ -995,7 +999,6 @@ def _arg_defaults(self, alias=None, estimate_memory=False):
         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, key.coordinates.data)
         defaults.update(key.interpolator._arg_defaults(coords=coords,

pyproject.toml

@@ -12,7 +12,7 @@ description = "Finite Difference DSL for symbolic computation."
 license = { file = "LICENSE.md" }
 readme = "README.md"
 keywords = ["finite-difference", "DSL", "symbolic", "jit", "devito"]
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.10,<3.15"
 authors = [
   { name = "Imperial College London", email = "g.gorman@imperial.ac.uk" },
   { name = "Fabio Luporini",  email = "fabio@devitocodes.com" },
@@ -32,6 +32,8 @@ classifiers = [
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
   "Programming Language :: Python :: 3 :: Only",
   "Topic :: Scientific/Engineering",
   "Topic :: Scientific/Engineering :: Mathematics",

requirements-optional.txt

@@ -1,7 +1,7 @@
 matplotlib<3.10.10
 pillow>11,<12.2.1
-pyrevolve==2.2.7
-scipy<1.15.4
+pyrevolve==2.2.8
+scipy>=1.8.0,<1.18.1
 distributed<2026.7.2
 click<9.0
 cloudpickle<3.1.3

requirements-testing.txt

@@ -3,7 +3,7 @@ pytest-runner<6.0.2
 pytest-cov<7.1.1
 flake8-pyproject>=1.2.3,<1.2.5
 nbval<0.11.1
-scipy<1.15.4
+scipy>=1.8.0,<1.18.1
 pooch<1.9.1
 click<9.0
 cloudpickle<3.1.3