api: Precompute sparse gridpoints and interpolation weights in fp64 on host

Author: mloubout

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,220 @@ 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 filled on the host from the parent SparseFunction's
+    scattered coordinates by ``LinearInterpolator._arg_defaults`` (invoked
+    via ``SparseFunction._arg_defaults``). Gathering is a no-op — these
+    tables are input-only."""
+
+    def _arg_apply(self, *args, **kwargs):
+        return
+
+    def _arg_values(self, estimate_memory=False, **kwargs):
+        if self._parent.name in kwargs:
+            return {}
+        return self._parent._arg_defaults(
+            alias=self._parent, estimate_memory=estimate_memory
+        ).reduce_all()
+
+    def compute(self, coords, grid, spacing, origin):  # pragma: no cover
+        raise NotImplementedError
+
+    def _arg_defaults_data(self, coords, grid, spacing, origin):
+        """``(name -> data)`` plus this table's dim entries."""
+        data = self.compute(coords, grid, spacing, origin)
+        args = {self.name: data}
+        for d, s in zip(self.dimensions, data.shape, strict=True):
+            args.update(d._arg_defaults(_min=0, size=s))
+        return args
+
+
+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 compute(self, coords, grid, spacing, origin):
+        return _cell_indices(coords, grid, 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 compute(self, coords, grid, spacing, origin):
+        return _linear_weights(coords, grid, 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, grid, shifts, spacing, origin):
+    c64 = np.asarray(coords, dtype=np.float64)
+    return (c64 - origin - _shift_values(shifts, grid, spacing)) / spacing
+
+
+def _cell_indices(coords, grid, shifts, spacing, origin):
+    return np.floor(_positions_fp64(coords, grid, shifts, spacing,
+                                    origin)).astype(np.int32)
+
+
+def _linear_weights(coords, grid, shifts, j, dtype, spacing, origin):
+    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):
+        super().__init__(sfunction)
+        # Per-shift `(gridpoints, coeffs_per_dim)` produced by `_positions` /
+        # `_weights` during equation construction. Iterated at arg-prep time
+        # by `_arg_defaults` to populate the host-side tables.
+        self._tables = {}
 
-    @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)
+    @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
+
+    def _make_tables(self, shifts):
+        # Names must sort alphabetically *after* the parent SparseFunction:
+        # the operator's default-arg loop populates the parent (which fills
+        # these tables via `interpolator._arg_defaults`) before it would fall
+        # back to each SubFunction's own zero-buffer default. See how the
+        # sinc `wsinc*` tables rely on the same ordering.
+        tag = _shift_tag(shifts)
+        sfname = self.sfunction.name
+        sfdim = self.sfunction._sparse_dim
+        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, shifts=shifts)
+        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, shifts=shifts, dim_index=i)
+            for i, (d, r) in enumerate(zip(self._gdims, self._cdim, strict=True))
+        )
+        return gp, coeffs
+
+    def _gridpoints(self, shifts=None):
+        key = tuple(shifts) if shifts else None
+        if key not in self._tables:
+            self._tables[key] = self._make_tables(shifts)
+        return self._tables[key][0]
+
+    def _coeffs(self, shifts=None):
+        key = tuple(shifts) if shifts else None
+        if key not in self._tables:
+            self._tables[key] = self._make_tables(shifts)
+        return self._tables[key][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 referenced by the operator from
+        the already-scattered ``coords`` handed in by
+        ``SparseFunction._arg_defaults``."""
+        if coords is None or sfunc is None:
+            raise ValueError("No coordinates or sparse function provided")
+        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 gp, coeffs in self._tables.values():
+            args.update(gp._arg_defaults_data(coords, grid, spacing, origin))
+            for w in coeffs:
+                args.update(w._arg_defaults_data(coords, grid, spacing, origin))
+        return args
 
 
 class PrecomputedInterpolator(WeightedInterpolator):
@@ -634,7 +810,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,13 +999,25 @@ 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
+        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)