api: compute sparse position/floor in fp64 to fix off-by-one cell shift

The linear interpolator computes `(c - o)/h` in the grid's dtype, which
for fp32 grids rounds coord/origin/spacing to fp32. That rounding can
push the position across an integer boundary, so `floor((c - o)/h)`
picks a different cell than the fp64-truth cell -- and CPU vs GPU can
disagree on the fractional part while agreeing on the integer position,
producing inconsistent injection weights.

Compute `pos = (c - o)/h` and `floor(pos)` in fp64 by casting the fp32
free symbols and by substituting the spacing symbols with their fp64
decimal value (recovered via the fp32 short-form round-trip).

A new CIRE subpass `CireInvariantsSparse` hoists `pos` and `floor(pos)`
out of the per-stencil-point inner loop into per-source preamble Arrays.
The floor tab is stored as int32 (half the memory of an fp64 tab, no
precision loss), and bare `floor(pos)` uses inside `pos - floor(pos)`
are rewritten to `DOUBLE(int_tab)` so both consumers share the tab.

The `sympy_dtype` inference is taught to recognize `Cast` (outermost
and inner) so the printer emits `floor` (fp64) instead of `floorf`
(fp32) when a `DOUBLE(...)` cast is present in the expression.

Author: mloubout

devito/operations/interpolators.py

@@ -543,11 +543,9 @@ def _point_symbols(self, shifts=None):
         return DimensionTuple(*symbols, getters=self.grid.dimensions)
 
     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)
+        poseq = [Eq(psyms[d], pos - floor(pos), implicit_dims=implicit_dims)
                  for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)]
         return poseq
 

devito/passes/clusters/aliases.py

@@ -15,8 +15,8 @@
 from devito.passes.clusters.cse import _cse
 from devito.passes.clusters.utils import expose_tuning_knobs
 from devito.symbolics import (
-    Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search, sympy_dtype,
-    uxreplace
+    DOUBLE, INT, Cast, Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched,
+    search, sympy_dtype, uxreplace
 )
 from devito.tools import (
     Reconstructable, Stamp, as_mapper, as_tuple, flatten, frozendict, generator,
@@ -295,6 +295,10 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
 
 class CireInvariants(CireTransformerLegacy, Queue):
 
+    # Predicate on Cluster used to pick which ones this pass fires on.
+    # Subclasses override to target a different kind of cluster.
+    _cluster_filter = staticmethod(lambda c: c.is_dense)
+
     def __init__(self, sregistry, options, platform):
         super().__init__(sregistry, options, platform)
 
@@ -324,7 +328,8 @@ def callback(self, clusters, prefix, xtracted=None):
         key = lambda c: self._lookup_key(c, d)
         processed = list(clusters)
         for ak, group in as_mapper(clusters, key=key).items():
-            g = [c for c in group if c.is_dense and c not in xtracted]
+            g = [c for c in group
+                 if self._cluster_filter(c) and c not in xtracted]
             if not g:
                 continue
 
@@ -387,6 +392,73 @@ def _generate(self, cgroup, exclude):
         yield self._do_generate(exprs, exclude, cbk_search)
 
 
+def _is_floor(e):
+    return getattr(e, 'is_Function', False) and e.func.__name__ == 'floor'
+
+
+class CireInvariantsSparse(CireInvariants):
+
+    """
+    Hoist sparse-point position temps `pos = (c - o)/h` and their
+    `floor(pos)` out of the per-stencil-point inner loop into preamble
+    Arrays computed once per source. The inner loop then reads
+    tabulated values instead of recomputing `floor((c - o)/h)` for each
+    `(rp_srcx, rp_srcy, rp_srcz)` combination. `Lift` then moves the
+    preamble out of the time loop.
+    """
+
+    _cluster_filter = staticmethod(lambda c: not c.is_dense)
+
+    def _generate(self, cgroup, exclude):
+        # Tabulate `pos` as fp64 and `INT(floor(pos))` as int32. The int32
+        # tab feeds both integer position lookups and (via DOUBLE cast) the
+        # bare `floor(pos)` uses in `pos - floor(pos)`.
+        counter = generator()
+        make_f64 = lambda: Symbol(name=f'dummy{counter()}', dtype=np.float64)
+        make_i32 = lambda: Symbol(name=f'dummy{counter()}', dtype=np.int32)
+
+        mapper = Uxmapper()
+
+        def _add(expr, make):
+            if expr.is_commutative is False:
+                return
+            if {a.function for a in expr.free_symbols} & exclude:
+                return
+            mapper.add(expr, make, None)
+
+        for e in cgroup.exprs:
+            for f in search(e, _is_floor, 'all', 'bfs'):
+                _add(f.args[0], make_f64)
+                _add(INT(f), make_i32)
+
+        yield mapper
+
+    def _choose(self, aliases, cgroup, mapper):
+        # Skip score-based filtering and fold bare `floor(pos)` onto the
+        # int32 alias built for `INT(floor(pos))` via a `DOUBLE(...)` cast.
+        exprs = cgroup.exprs
+
+        aliases = AliasList(aliases)
+        if not aliases:
+            return exprs, aliases
+
+        aliaseds = set(aliases.aliaseds)
+        subs = {k: v for k, v in mapper.items() if v.free_symbols & aliaseds}
+
+        for k, v in list(mapper.items()):
+            if not isinstance(k, Cast):
+                continue
+            if not (isinstance(k.dtype, type) and
+                    issubclass(k.dtype, np.integer)):
+                continue
+            inner = k.base
+            if _is_floor(inner) and (v.free_symbols & aliaseds):
+                subs[inner] = DOUBLE(v)
+
+        exprs = [uxreplace(e, subs) for e in exprs]
+        return exprs, aliases
+
+
 class CireDerivatives(CireTransformerLegacy):
 
     def __init__(self, sregistry, options, platform):
@@ -519,7 +591,8 @@ def _cbk_search2(self, expr, rank):
 # Subpass mapper
 modes = {
     'invariants': [CireInvariantsElementary,
-                   CireInvariantsDivs],
+                   CireInvariantsDivs,
+                   CireInvariantsSparse],
     'eval-derivs': [CireEvalDerivatives],   # NOTE: legacy pass
     'index-derivs': [CireIndexDerivatives],
 }

devito/symbolics/inspection.py

@@ -316,10 +316,24 @@ def sympy_dtype(expr, base=None, default=None, smin=None):
     if expr is None:
         return default
 
+    # An outermost Cast (e.g. `DOUBLE(...)`, `INT(...)`) declares the dtype
+    # of the whole expression regardless of its operand's free symbols.
+    if isinstance(expr, Cast):
+        cast_dtype = expr.dtype
+        if isinstance(cast_dtype, type) and issubclass(cast_dtype, np.generic):
+            return cast_dtype
+
     dtypes = set()
+    # A Cast inside the tree bumps the inferred dtype: it forces the
+    # operation to occur at (at least) the Cast's precision.
+    for c in expr.atoms(Cast):
+        cd = c.dtype
+        if isinstance(cd, type) and issubclass(cd, np.generic):
+            dtypes.add(cd)
     for i in expr.free_symbols:
         with suppress(AttributeError):
             dtypes.add(i.dtype)
+    dtypes.discard(None)
 
     if not dtypes or not np.issubdtype(base, np.complexfloating):
         dtypes.update({base} - {None})

devito/types/sparse.py

@@ -11,7 +11,7 @@
 from devito.operations import (
     LinearInterpolator, PrecomputedInterpolator, SincInterpolator
 )
-from devito.symbolics import indexify, retrieve_function_carriers
+from devito.symbolics import DOUBLE, indexify, retrieve_function_carriers
 from devito.tools import (
     ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
     memoized_meth, prod
@@ -394,10 +394,14 @@ def _position_map(self, shifts=None):
         Dimension, of additional physical offsets to subtract from the sparse
         coordinates (e.g. ``h_x/2`` for a field staggered in ``x``). If ``shifts``
         is None, only the grid origin is subtracted.
+
+        The origin is cast to double so the subtraction and subsequent division
+        by the (fp32) spacing happen in fp64 and pick the correct cell
+        regardless of the grid's fp32 rounding.
         """
         shifts = shifts or (0,) * len(self.grid.dimensions)
         return OrderedDict([
-            ((c - o - s)/d.spacing, p)
+            ((c - DOUBLE(o) - s)/d.spacing, p)
             for p, c, d, o, s in zip(
                 self._pos_symbols(shifts=shifts),
                 self._coordinate_symbols,

examples/userapi/06_sparse_operations.ipynb

@@ -204,9 +204,9 @@ def test_basic(self):
 
         op = Operator(eqns, opt=('advanced', {'blockrelax': True}))
 
-        bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0'})
+        bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0', 'p_src1_blk0'})
 
-        iters = FindNodes(Iteration).visit(bns['p_src0_blk0'])
+        iters = FindNodes(Iteration).visit(bns['p_src1_blk0'])
         assert len(iters) == 5
         assert iters[0].dim.is_Block
         assert iters[1].dim.is_Block
@@ -1070,12 +1070,14 @@ def test_incr_perfect_sparse_outer(self):
                                               'openmp': True}))
 
         iters = FindNodes(Iteration).visit(op)
-        assert len(iters) == 5
-        assert iters[0].is_Sequential
-        assert all(i.is_ParallelAtomic for i in iters[1:])
-        assert iters[1].pragmas[0].ccode.value ==\
+        # 4 preamble iterations (p_u + rp_u{x,y,z}) hoisted out of time by
+        # `CireInvariantsSparse`, then the time loop and the 4 injection iters.
+        assert len(iters) == 9
+        assert iters[4].is_Sequential
+        assert all(i.is_ParallelAtomic for i in iters[:4] + iters[5:])
+        assert iters[5].pragmas[0].ccode.value ==\
             'omp for schedule(dynamic,chunk_size)'
-        assert all(not i.pragmas for i in iters[2:])
+        assert all(not i.pragmas for i in iters[6:])
 
     @pytest.mark.parametrize('exprs,simd_level,expected', [
         (['Eq(y.symbolic_max, g[0, x], implicit_dims=(t, x))',

tests/test_dle.py

@@ -50,12 +50,18 @@ def test_scheduling_after_rewrite():
     op = Operator([eqn1] + eqn2 + [eqn3])
     trees = retrieve_iteration_tree(op)
 
-    # Check loop nest structure
-    assert all(
-        i.dim is j for i, j in zip(trees[0], grid.dimensions, strict=True)
-    )  # time invariant
-    assert trees[1].root.dim is grid.time_dim
-    assert all(trees[1].root.dim is tree.root.dim for tree in trees[1:])
+    # Check loop nest structure. `CireInvariantsSparse` hoists the sparse
+    # position preamble as an extra time-invariant tree ahead of the time
+    # loops, alongside the dense sin(const) preamble.
+    invariants, in_time = [], []
+    for t in trees:
+        (in_time if t.root.dim is grid.time_dim else invariants).append(t)
+    assert invariants  # sin(const) + the sparse-position preamble(s)
+    assert in_time     # the actual time-loop clusters
+    # Invariants must all precede the time-loop trees
+    for inv in invariants:
+        for tt in in_time:
+            assert trees.index(inv) < trees.index(tt)
 
 
 @pytest.mark.parametrize('expr,expected', [
@@ -1399,7 +1405,10 @@ def test_catch_duplicate_from_different_clusters(self):
         op = Operator(eqns, opt=('advanced', {'cire-mingain': 100}))
 
         arrays = [i for i in FindSymbols().visit(op) if i.is_Array]
-        assert len(arrays) == 3
+        # 3 dense-derivative aliases + 4 sparse-position preambles (pos_x,
+        # pos_y and their int(floor(...)) tabulations) from
+        # `CireInvariantsSparse`.
+        assert len(arrays) == 7
         assert all(i._mem_heap and not i._mem_external for i in arrays)
 
     def test_discarded_compound(self):
@@ -1542,7 +1551,9 @@ def test_drop_redundants_after_fusion(self, rotate):
         op = Operator(eqns, opt=('advanced', {'cire-rotate': rotate}))
 
         arrays = [i for i in FindSymbols().visit(op) if i.is_Array]
-        assert len(arrays) == 2
+        # 2 dense aliases + 4 sparse-position preambles per sparse function
+        # (pos_x, pos_y, int_floor_x, int_floor_y) x 1 SparseTimeFunction.
+        assert len(arrays) == 6
         assert all(i._mem_heap and not i._mem_external for i in arrays)
 
     def test_full_shape_big_temporaries(self):

tests/test_dse.py

@@ -12,6 +12,7 @@
     switchconfig
 )
 from devito.operations.interpolators import LinearInterpolator, SincInterpolator
+from devito.symbolics.extended_sympy import Cast
 from devito.tools import as_tuple
 from examples.seismic import (
     AcquisitionGeometry, Receiver, RickerSource, TimeAxis, demo_model
@@ -489,7 +490,7 @@ def test_inject_staggered(self, stagg):
         if stagg == 'NODE':
             assert np.isclose(a.data[5, 5, 5], 2, rtol=1e-6)
             # all other should be zero
-            assert np.sum(a.data) == 2
+            assert np.isclose(np.sum(a.data), 2, rtol=1e-6)
         else:
             # Bottom corner since the source position is left of the staggered field
             corner = [5, 5, 5] - np.array(a.staggered)
@@ -507,7 +508,8 @@ def test_inject_staggered(self, stagg):
             assert np.allclose(a.data[slices], interp_val, rtol=1e-6)
             # All other should be zero so should sum to the interp_val * number of points.
             # Use abs to make sure there is no +- cancellations
-            assert np.sum(np.abs(a.data)) == interp_val * 2**(sum(np.array(a.staggered)))
+            expected = float(interp_val * 2**(sum(np.array(a.staggered))))
+            assert np.isclose(np.sum(np.abs(a.data)), expected, rtol=1e-6)
 
     def test_inject_staggered_mixed(self):
         grid = Grid((11, 11, 11))
@@ -870,6 +872,45 @@ def test_wrong_coords(self):
             s.interpolate(u + s2)
         assert "Interpolation/injection with" in str(vinfo.value)
 
+    def test_position_map_fp64(self):
+        """
+        The linear sparse position `(c - o)/h` must be computed in fp64 even
+        when the grid dtype is fp32. Otherwise fp32 rounding of the origin
+        can push the position across an integer boundary, so `floor((c-o)/h)`
+        picks a different cell than the fp64-truth cell.
+
+        Uses `h = 0.1` (not fp32-exact) and a coordinate whose fp64 position
+        is just below cell `7`; without the fix the fp32 position rounds up
+        to exactly `7.0`, shifting the injection by one cell.
+        """
+        # h = 0.1 in fp64, but not exactly representable in fp32
+        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)
+
+        # Symbolic check: the origin is wrapped in a `DOUBLE(...)` cast so
+        # the subtraction and subsequent division promote to fp64.
+        pmap = sf._position_map()
+        for expr in pmap:
+            assert expr.find(Cast), \
+                f"position map missing fp64 cast on origin: {expr}"
+
+        # End-to-end: fp64 position ~ 6.99999988 -> floor = 6, fraction ~ 1.
+        # Without the fix the fp32 position is exactly 7.0, so floor = 7 and
+        # the fraction is 0, misplacing the injection at cell 8 in 2D.
+        u = TimeFunction(name='u', grid=grid, space_order=2, time_order=1)
+        c = 0.6999999990000001
+        sf.coordinates.data[0, :] = c
+        sf.data[:] = 1.0
+
+        op = Operator(sf.inject(field=u.forward, expr=sf))
+        op.apply(time_M=0)
+
+        # fp64-truth: nearly all mass at (7, 7), essentially none at (8, 8).
+        assert np.isclose(u.data[1, 7, 7], 1.0, atol=1e-5)
+        assert np.isclose(u.data[1, 8, 8], 0.0, atol=1e-10)
+
 
 # ---------------------------------------------------------------------------
 # Complex-valued interpolation

tests/test_interpolation.py

@@ -3350,11 +3350,15 @@ def test_adjoint_codegen(self, shape, kernel, space_order, save, mode):
         solver = acoustic_setup(shape=shape, spacing=[15. for _ in shape], kernel=kernel,
                                 tn=500, space_order=space_order, nrec=130,
                                 preset='layers-isotropic', dtype=np.float64)
+        # Count only meaningful kernel/comm calls; ignore memory-management
+        # (posix_memalign/free) that grows with sparse-position preambles.
+        skip = {'posix_memalign', 'free'}
+
         op_fwd = solver.op_fwd(save=save)
-        fwd_calls = FindNodes(Call).visit(op_fwd)
+        fwd_calls = [c for c in FindNodes(Call).visit(op_fwd) if c.name not in skip]
 
         op_adj = solver.op_adj()
-        adj_calls = FindNodes(Call).visit(op_adj)
+        adj_calls = [c for c in FindNodes(Call).visit(op_adj) if c.name not in skip]
 
         assert len(fwd_calls) == 1
         assert len(adj_calls) == 1