compiler: precompute interp pos/weights

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_list, 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,205 @@ 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):
+    """SubFunction populated on the host by the parent interpolator's
+    `_arg_defaults` from the already-scattered coordinates, exactly the way
+    sinc's precomputed weights are populated.
+
+    `parent` links the table to its SparseFunction so the base
+    `SubFunction._arg_values` routing triggers `SparseFunction._arg_defaults`
+    -> `LinearInterpolator._arg_defaults`. The link is *not* preserved by
+    pickle (parent is not in `__rkwargs__`) so decoupled workers don't ship
+    the sfunction back through every table."""
+
+    def _arg_apply(self, *args, **kwargs):
+        return
+
+
+class Gridpoints(_HostTable):
+    """int32 cell indices per sparse point, shape ``(npoint, ndim)``."""
+
+
+class Coeffs(_HostTable):
+    """Per-dim `(1 - frac, frac)` interpolation weights, shape ``(npoint, 2)``."""
+
+
+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
+    that fp32 `floor((c - o)/h)` can produce for coordinates that sit on an
+    exact cell boundary."""
+    c64 = np.asarray(coords, dtype=np.float64)
+    return (c64 - origin - _shift_values(shifts, grid, spacing)) / spacing
+
+
+def _cell_indices(coords, grid, shifts, spacing, origin):
+    """Int32 base cell index per sparse point, one entry per grid Dimension."""
+    return np.floor(_positions_fp64(coords, grid, shifts, spacing,
+                                    origin)).astype(np.int32)
+
+
+def _linear_weights(coords, grid, shifts, j, dtype, spacing, origin):
+    """`(1 - frac, frac)` linear interpolation weights along dim `j`; `frac`
+    is the fractional cell offset from the base index returned by
+    `_cell_indices`."""
+    pos = _positions_fp64(coords, grid, 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'
 
-    @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 __init__(self, sfunction, shifts=()):
+        super().__init__(sfunction)
+        # Every shift set the interpolator has been asked to produce tables
+        # for. Persisted with the parent SparseFunction via `__rkwargs__` so
+        # a pickled/rebuilt interpolator (decoupled workers) knows which
+        # tables to emit in `_arg_defaults`.
+        self._shifts_used = set(tuple(s) if s else None for s in shifts)
+
+    @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 _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 _generate_coeffs(self, key):
+        """Create the ``(gridpoints, coeffs_per_dim)`` SubFunction tuple for
+        a given shift set. ``key`` is either ``None`` (plain, non-staggered)
+        or a tuple of per-dim shifts. Mirrors sinc's ``interpolation_coeffs``
+        cached_property but keyed on ``shifts``."""
+        # Record that the caller has emitted tables for this shift set, so
+        # `_arg_defaults` can regenerate them on the fly.
+        self._shifts_used.add(key)
+
+        shifts = as_list(key)
+        tag = _shift_tag(shifts)
+        sfname = self.sfunction.name
+        sfdim = self.sfunction._sparse_dim
+
+        # Gridpoints: `(npoint, ndim)` int32 base cell index per sparse point.
+        gp_name = f'{sfname}_gp{tag}'
+        ddim = CustomDimension(f'{gp_name}d', 0, self.grid.dim - 1,
+                               self.grid.dim, sfdim)
+        gp = Gridpoints(name=gp_name, dtype=np.int32,
+                        shape=(self.sfunction.npoint, self.grid.dim),
+                        dimensions=(sfdim, ddim), space_order=0,
+                        alias=self.sfunction.alias,
+                        parent=self.sfunction)
+
+        # Per-dim linear weights: `(npoint, 2)` holding `(1 - frac, frac)`.
+        coeffs = tuple(
+            Coeffs(name=f'{sfname}_w{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)
+            for d, r in zip(self._gdims, self._cdim, strict=True)
+        )
+
+        return gp, coeffs
+
+    def _gridpoints(self, shifts=None):
+        return self._generate_coeffs(tuple(shifts) if shifts else None)[0]
+
+    def _coeffs(self, shifts=None):
+        return self._generate_coeffs(tuple(shifts) if shifts else None)[1]
+
+    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)
+        ])
+
+    def _arg_defaults(self, coords=None, sfunc=None, origin=None):
+        """Fill the gridpoints/coeffs tables from the already-scattered
+        ``coords`` handed in by ``SparseFunction._arg_defaults``. Mirrors
+        sinc's `_arg_defaults`: regenerates the tables from the persisted
+        shift set, so a pickled/rebuilt interpolator (decoupled workers)
+        still emits data for every table the operator was compiled with."""
+        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)])
+
+        args = {}
+        for key in self._shifts_used or {None}:
+            shifts = as_list(key)
+            gp, coeffs = self._generate_coeffs(key)
+
+            # Tabulated data (int32 cell indices + fp linear weights per dim).
+            args[gp.name] = _cell_indices(coords, grid, shifts, spacing, origin)
+            for i, w in enumerate(coeffs):
+                args[w.name] = _linear_weights(
+                    coords, grid, shifts, i, w.dtype, spacing, origin
+                )
+
+            # Bounds for each table's dimensions, matching the computed data.
+            for f in (gp, *coeffs):
+                for d, s in zip(f.dimensions, args[f.name].shape, strict=True):
+                    args.update(d._arg_defaults(_min=0, size=s))
+
+        return args
 
 
 class PrecomputedInterpolator(WeightedInterpolator):
@@ -634,7 +795,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/dense.py

@@ -1627,14 +1627,16 @@ def __padding_setup__(self, **kwargs):
     def _halo_exchange(self):
         return
 
-    def _arg_values(self, **kwargs):
+    def _arg_values(self, estimate_memory=False, **kwargs):
         if self._parent is not None and self.parent.name not in kwargs:
-            return self._parent._arg_defaults(alias=self._parent).reduce_all()
+            return self._parent._arg_defaults(
+                alias=self._parent, estimate_memory=estimate_memory
+            ).reduce_all()
         elif self.name in kwargs:
             raise RuntimeError(f"`{self.name}` is a SubFunction, so it can't be assigned "
                                "a value dynamically")
         else:
-            return self._arg_defaults(alias=self)
+            return self._arg_defaults(alias=self, estimate_memory=estimate_memory)
 
     def _arg_apply(self, *args, **kwargs):
         if self._parent is not None:

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,13 +999,30 @@ 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,
                                                        sfunc=key))
         return defaults
 
+    def _arg_values(self, estimate_memory=False, **kwargs):
+        values = super()._arg_values(estimate_memory=estimate_memory, **kwargs)
+        if estimate_memory:
+            return values
+
+        # Resolve the runtime grid origin (honours `o_x`/`o_y`/... overrides)
+        # and hand it to the interpolator so tables reflect the actual frame
+        # of reference used by the kernel.
+        onames = [o.name for o in self.grid.origin_symbols]
+        origin = tuple(kwargs.get(n, o) for n, o in
+                       zip(onames, self.grid.origin, strict=True))
+        coords = values.get(self.coordinates.name, self.coordinates.data)
+        values.update(self.interpolator._arg_defaults(
+            coords=coords, sfunc=self, origin=origin
+        ))
+
+        return values
+
 
 class SparseTimeFunction(AbstractSparseTimeFunction, SparseFunction):
     """

examples/userapi/06_sparse_operations.ipynb

@@ -277,12 +277,10 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Eq(posx, (int)floor((-o_x + s_coords(p_s, 0))/h_x))\n",
-      "Eq(posy, (int)floor((-o_y + s_coords(p_s, 1))/h_y))\n",
-      "Eq(px, -floor((-o_x + s_coords(p_s, 0))/h_x) + (-o_x + s_coords(p_s, 0))/h_x)\n",
-      "Eq(py, -floor((-o_y + s_coords(p_s, 1))/h_y) + (-o_y + s_coords(p_s, 1))/h_y)\n",
+      "Eq(posx, s_gp(p_s, 0))\n",
+      "Eq(posy, s_gp(p_s, 1))\n",
       "Eq(sums, 0.0)\n",
-      "Inc(sums, (rp_sx*px + (1 - rp_sx)*(1 - px))*(rp_sy*py + (1 - rp_sy)*(1 - py))*f(t, rp_sx + posx, rp_sy + posy))\n",
+      "Inc(sums, s_wx(p_s, rp_sx)*s_wy(p_s, rp_sy)*f(t, rp_sx + posx, rp_sy + posy))\n",
       "Eq(s(time, p_s), sums)\n"
      ]
     }

tests/test_dse.py

@@ -2957,12 +2957,12 @@ def test_fullopt(self):
         bns, _ = assert_blocking(op1, {'x0_blk0'})  # due to loop blocking
 
         assert summary0[('section0', None)].ops == 55
-        assert summary0[('section1', None)].ops == 44
+        assert summary0[('section1', None)].ops == 8
         assert np.isclose(summary0[('section0', None)].oi, 3.136, atol=0.001)
 
         assert summary1[('section0', None)].ops == 31
-        assert summary1[('section1', None)].ops == 88
-        assert summary1[('section2', None)].ops == 25
+        assert summary1[('section1', None)].ops == 16
+        assert summary1[('section2', None)].ops == 4
         assert np.isclose(summary1[('section0', None)].oi, 1.767, atol=0.001)
 
         assert np.allclose(u0.data, u1.data, atol=10e-5)