From 200302cf051a2fddd148baa15beacce512e1699b Mon Sep 17 00:00:00 2001 From: mloubout Date: Mon, 10 Aug 2026 14:48:22 -0400 Subject: [PATCH] api: Precompute sparse gridpoints and interpolation weights in fp64 on host --- devito/operations/interpolators.py | 225 +++++++++++++++++--- devito/tools/dtypes_lowering.py | 12 +- devito/types/dense.py | 8 +- devito/types/sparse.py | 27 ++- examples/userapi/06_sparse_operations.ipynb | 8 +- tests/test_dse.py | 6 +- tests/test_gpu_openacc.py | 5 +- tests/test_interpolation.py | 39 +++- tests/test_operator.py | 53 +++-- tests/test_pickle.py | 2 +- 10 files changed, 311 insertions(+), 74 deletions(-) diff --git a/devito/operations/interpolators.py b/devito/operations/interpolators.py index 6467e65df47..1bba5466355 100644 --- a/devito/operations/interpolators.py +++ b/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) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 2a8f5d78009..f3033aa6e90 100644 --- a/devito/tools/dtypes_lowering.py +++ 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 diff --git a/devito/types/dense.py b/devito/types/dense.py index a6eb2f358e9..98c6ed1c81b 100644 --- a/devito/types/dense.py +++ b/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: diff --git a/devito/types/sparse.py b/devito/types/sparse.py index e0f88e610c0..ccb330b2d6d 100644 --- a/devito/types/sparse.py +++ b/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): """ diff --git a/examples/userapi/06_sparse_operations.ipynb b/examples/userapi/06_sparse_operations.ipynb index 19141550603..4e2d08417d2 100644 --- a/examples/userapi/06_sparse_operations.ipynb +++ b/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" ] } diff --git a/tests/test_dse.py b/tests/test_dse.py index 7d25fbba534..d9f210bd9e3 100644 --- a/tests/test_dse.py +++ b/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) diff --git a/tests/test_gpu_openacc.py b/tests/test_gpu_openacc.py index ada536197e7..a335d55289b 100644 --- a/tests/test_gpu_openacc.py +++ b/tests/test_gpu_openacc.py @@ -111,8 +111,9 @@ def test_tile_insteadof_collapse(self, par_tile): assert trees[1][1].pragmas[0].ccode.value ==\ 'acc parallel loop tile(32,4) present(u)' strtile = ','.join([str(i) for i in stile]) + pres = 'src,src_gp,src_wx,src_wy,src_wz,u' assert trees[3][1].pragmas[0].ccode.value ==\ - f'acc parallel loop tile({strtile}) present(src,src_coords,u)' + f'acc parallel loop tile({strtile}) present({pres})' @pytest.mark.parametrize('par_tile', [((32, 4, 4), (8, 8)), ((32, 4), (8, 8)), ((32, 4, 4), (8, 8, 8)), @@ -141,7 +142,7 @@ def test_multiple_tile_sizes(self, par_tile): 'acc parallel loop tile(8,8) present(u)' sclause = 'collapse(4)' if par_tile[-1] is None else 'tile(8,8,8,8)' assert trees[3][1].pragmas[0].ccode.value ==\ - f'acc parallel loop {sclause} present(src,src_coords,u)' + f'acc parallel loop {sclause} present(src,src_gp,src_wx,src_wy,src_wz,u)' def test_multi_tile_blocking_structure(self): grid = Grid(shape=(8, 8, 8)) diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py index 75b125d83d0..5d8142a412c 100644 --- a/tests/test_interpolation.py +++ b/tests/test_interpolation.py @@ -517,13 +517,8 @@ def test_inject_staggered_mixed(self): eq = p.inject(v, expr=b * p).evaluate - # We should have - # - 3 injection equations v_x, v_y, v_z - # The standard 6 on node temps posx, posy, posz, px, py, pz - # 2 temps for the staggered in x vx posz_s1, px_s1 - # 2 temps for the staggered in y vy posz_s1, py_s1 - # 2 temps for the staggered in z vz posz_s1, pz_s1 - assert len(eq) == 3 + 6 + 2 + 2 + 2 + # 3 injection eqs + 3 position lookups per staggered field: 3 + 3*3 + assert len(eq) == 3 + 3 * 3 op = Operator(eq) # Should be a single loop nest with 3 injections @@ -848,13 +843,11 @@ def test_interp_default(self): @pytest.mark.parametrize('dtype, expected', [(np.complex64, np.float32), (np.complex128, np.float64)]) def test_point_symbol_types(self, dtype, expected): - """Test that positions are always real""" + """Interpolation weights must be real even for complex sparse fields.""" grid = Grid(shape=(11,)) s = SparseFunction(name='src', npoint=1, grid=grid, dtype=dtype) - point_symbol = s.interpolator._point_symbols()[0] - - assert point_symbol.dtype is expected + assert s.interpolator._coeffs()[0].dtype is expected def test_wrong_coords(self): grid = Grid(shape=(11, 11)) @@ -870,6 +863,30 @@ def test_wrong_coords(self): s.interpolate(u + s2) assert "Interpolation/injection with" in str(vinfo.value) + def test_position_map_fp64(self): + """ + Coord-based sparse interpolation must not do `(c-o)/h` and `floor(...)` + in the generated kernel: on fp32 that can push points across a cell + boundary. Gridpoints and coefficients are precomputed in fp64 on the + host and the kernel only indexes those tables. + """ + grid = Grid(shape=(30, 30), extent=(2.9, 2.9)) + assert grid.dtype is np.float32 + + sf = SparseTimeFunction(name='sf', grid=grid, npoint=1, nt=2) + u = TimeFunction(name='u', grid=grid, space_order=2, time_order=1) + sf.coordinates.data[0, :] = 0.6999999990000001 + sf.data[:] = 1.0 + + op = Operator(sf.inject(field=u.forward, expr=sf)) + code = str(op) + assert 'floor' not in code + assert 'o_x' not in code and 'o_y' not in code + + op.apply(time_M=0) + assert u.data[1, 7, 7] > 0.99 + assert u.data[1, 8, 8] == 0.0 + # --------------------------------------------------------------------------- # Complex-valued interpolation diff --git a/tests/test_operator.py b/tests/test_operator.py index b3f7ea295e8..f781819f7c5 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -34,7 +34,7 @@ from devito.ir.support import Any, Backward, Forward from devito.passes.iet.languages.C import CDataManager from devito.symbolics import ListInitializer, indexify, retrieve_indexed -from devito.tools import flatten, powerset, timed_region +from devito.tools import as_fp64_decimal, flatten, powerset, timed_region from devito.types import ( Array, Barrier, ConditionalDimension, CustomDimension, Indirection, Scalar, Symbol ) @@ -819,11 +819,13 @@ def test_default_sparse_functions(self): s.coordinates.data[:, 2] = np.arange(2., 5.) op = Operator(s.interpolate(f)) + # Note: coord-based SparseFunctions no longer emit `s_coords` in the + # compiled kernel; positions and coefficients are precomputed on the + # host and passed as `gps`/`ws{d}` tables. expected = { - 's': s, 's_coords': s.coordinates, + 's': s, # Default dimensions of the sparse data 'p_s_size': 3, 'p_s_m': 0, 'p_s_M': 2, - 'd_size': 3, 'd_m': 0, 'd_M': 2, 'time_size': 4, 'time_m': 0, 'time_M': 3, } self.verify_arguments(op.arguments(), expected) @@ -1049,9 +1051,15 @@ def test_override_sparse_data_fix_dim(self): # whether the override picks up the original coordinates or the changed ones args = op.arguments(src1=src2, time=0) - arg_name = src1.coordinates._arg_names[0] - assert(np.array_equal(src2.coordinates._C_as_ndarray(args[arg_name]), - np.asarray((new_coords,)))) + # Coord-based sparse functions no longer emit `s_coords` in the + # kernel; check the precomputed cell indices (`gp*`) instead. + gp = src1.interpolator._gridpoints() + spacing = np.array([as_fp64_decimal(h) for h in grid.spacing]) + origin = np.array([as_fp64_decimal(o) for o in grid.origin]) + expected = np.floor( + (np.asarray(new_coords, dtype=np.float64) - origin) / spacing + ).astype(np.int32) + assert np.array_equal(gp._C_as_ndarray(args[gp.name]).ravel(), expected) def test_override_sparse_data_default_dim(self): """ @@ -1073,9 +1081,13 @@ def test_override_sparse_data_default_dim(self): # whether the override picks up the original coordinates or the changed ones args = op.arguments(src1=src2, t=0) - arg_name = src1.coordinates._arg_names[0] - assert(np.array_equal(src2.coordinates._C_as_ndarray(args[arg_name]), - np.asarray((new_coords,)))) + gp = src1.interpolator._gridpoints() + spacing = np.array([as_fp64_decimal(h) for h in grid.spacing]) + origin = np.array([as_fp64_decimal(o) for o in grid.origin]) + expected = np.floor( + (np.asarray(new_coords, dtype=np.float64) - origin) / spacing + ).astype(np.int32) + assert np.array_equal(gp._C_as_ndarray(args[gp.name]).ravel(), expected) def test_argument_derivation_order(self, nt=100): """ Ensure the precedence order of arguments is respected @@ -2244,7 +2256,11 @@ def test_sparse(self, caplog, time): summary = op.estimate_memory() assert "Allocating" not in caplog.text - check = self.sum_sizes((f, src, src.coordinates)) + # Coord-based sparse functions ship precomputed gridpoints/coeffs + # tables as operator inputs (built on the host from `coordinates`). + tables = (src.interpolator._gridpoints(),) + \ + src.interpolator._coeffs() + check = self.sum_sizes((f, src, *tables)) self.parse_output(summary, check) @pytest.mark.parametrize('save', [None, Buffer(3), 10]) @@ -2278,7 +2294,11 @@ def test_mashup(self, caplog): summary = op.estimate_memory() assert "Allocating" not in caplog.text - check = self.sum_sizes((f, g, src0, src0.coordinates, src1, src1.coordinates)) + tables = ((src0.interpolator._gridpoints(),) + + src0.interpolator._coeffs() + + (src1.interpolator._gridpoints(),) + + src1.interpolator._coeffs()) + check = self.sum_sizes((f, g, src0, src1, *tables)) self.parse_output(summary, check) @pytest.mark.parametrize('override', [True, False]) @@ -2357,17 +2377,24 @@ def setup(size, npoint, nt, counter): with switchconfig(log_level='DEBUG'), caplog.at_level(logging.DEBUG): op = Operator([eq0, eq1] + s0_term + st0_term) + # Interpolator-owned tables are baked in from the *original* + # sparse functions and don't get re-sized by overrides. + tables = ((s0.interpolator._gridpoints(),) + + s0.interpolator._coeffs() + + (st0.interpolator._gridpoints(),) + + st0.interpolator._coeffs()) + # Apply overrides for the check summary0 = op.estimate_memory(f0=f1, tf0=tf1, s0=s1, st0=st1) - check0 = self.sum_sizes((f1, tf1, s1, s1.coordinates, st1, st1.coordinates)) + check0 = self.sum_sizes((f1, tf1, s1, st1, *tables)) self.parse_output(summary0, check0) # Check with a second set of overrides summary1 = op.estimate_memory(f0=f2, tf0=tf2, s0=s2, st0=st2) assert "Allocating" not in caplog.text - check1 = self.sum_sizes((f2, tf2, s2, s2.coordinates, st2, st2.coordinates)) + check1 = self.sum_sizes((f2, tf2, s2, st2, *tables)) self.parse_output(summary1, check1) def test_device(self, caplog): diff --git a/tests/test_pickle.py b/tests/test_pickle.py index f29a96649bc..8e77e36829d 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -841,7 +841,7 @@ def test_elemental(self, pickle): u = TimeFunction(name="u", grid=grid, time_order=2, space_order=2) rec_term = rec.interpolate(expr=u) - eq = rec_term.evaluate[2] + eq = rec_term.evaluate[3] eq = eq.func(eq.lhs, eq.rhs.args[0]) op = Operator(eq)