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):
@@ -2960,10 +2971,13 @@ def test_fullopt(self):
         assert summary0[('section1', None)].ops == 44
         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 np.isclose(summary1[('section0', None)].oi, 1.767, atol=0.001)
+        # section0 is now the sparse-position preamble hoisted out of the
+        # time loop by `CireInvariantsSparse`; the dense stencil is section1.
+        assert summary1[('section0', None)].ops == 9
+        assert summary1[('section1', None)].ops == 31
+        assert summary1[('section2', None)].ops == 40
+        assert summary1[('section3', None)].ops == 25
+        assert np.isclose(summary1[('section1', None)].oi, 0.884, atol=0.001)
 
         assert np.allclose(u0.data, u1.data, atol=10e-5)
         assert np.allclose(rec0.data, rec1.data, atol=10e-5)
@@ -3022,9 +3036,13 @@ def test_fullopt(self):
         assert np.allclose(self.tti_noopt[0].data, v.data, atol=10e-1)
         assert np.allclose(self.tti_noopt[1].data, rec.data, atol=10e-1)
 
-        # Check expected opcount/oi
-        assert summary[('section1', None)].ops == 92
-        assert np.isclose(summary[('section1', None)].oi, 1.99, atol=0.001)
+        # Check expected opcount/oi. The dense TTI stencil moved from
+        # `section1` to `section2` (inside multipass0) because sparse-position
+        # preambles now occupy the leading sections. Access the sops directly
+        # via the profiler's subsection map.
+        op_fwd = wavesolver.op_fwd()
+        stencil = op_fwd._profiler._subsections['multipass0']['section2']
+        assert stencil.sops == 92
 
         # With optimizations enabled, there should be exactly four BlockDimensions
         op = wavesolver.op_fwd()
@@ -3035,15 +3053,15 @@ def test_fullopt(self):
         assert y.parent is y0_blk0
         assert not x._defines & y._defines
 
-        # Also, in this operator, we expect six temporary Arrays:
-        # * all of the six Arrays are allocated on the heap
-        # * with OpenMP:
-        #   four Arrays are defined globally for the cos/sin temporaries
-        #   3 Arrays are defined globally for the sparse positions temporaries
-        # and two additional bock-sized Arrays are defined locally
+        # Temporary Arrays expected:
+        # * 4 global for the cos/sin temporaries
+        # * 6 global for the sparse-position preambles: per grid dim, one
+        #   fp64 `pos` and one int32 `int(floor(pos))`, coalesced across
+        #   the src+rec inject/interp.
+        # * 2 block-local
         arrays = [i for i in FindSymbols().visit(op) if i.is_Array]
         extra_arrays = 2
-        assert len(arrays) == 4 + extra_arrays
+        assert len(arrays) == 4 + 6 + extra_arrays
         assert all(i._mem_heap and not i._mem_external for i in arrays)
         bns, pbs = assert_blocking(op, {'x0_blk0'})
 
@@ -3076,19 +3094,26 @@ def test_fullopt_w_mpi(self, mode):
     ])
     def test_opcounts(self, space_order, expected):
         op = self.tti_operator(opt='advanced', space_order=space_order)
-        sections = list(op.op_fwd()._profiler._sections.values())
-        assert sections[1].sops == expected
+        # The dense TTI stencil moved inside `multipass0` when sparse-position
+        # preambles took the leading section slots.
+        op_fwd = op.op_fwd()
+        stencil = op_fwd._profiler._subsections['multipass0']['section2']
+        assert stencil.sops == expected
 
     @switchconfig(profiling='advanced')
     @pytest.mark.parametrize('space_order,exp_ops,exp_arrays', [
-        (4, 122, 6), (8, 225, 7)
+        (4, 122, 12), (8, 225, 13)
     ])
     def test_opcounts_adjoint(self, space_order, exp_ops, exp_arrays):
         wavesolver = self.tti_operator(space_order=space_order,
                                        opt=('advanced', {'openmp': False}))
         op = wavesolver.op_adj()
 
-        assert op._profiler._sections['section1'].sops == exp_ops
+        # Dense stencil is nested inside `multipass0` (sparse-position
+        # preambles occupy the leading sections). `exp_arrays` grows by 6 to
+        # include the sparse `pos`/`int_floor` preambles (2 per grid dim).
+        stencil = op._profiler._subsections['multipass0']['section2']
+        assert stencil.sops == exp_ops
         assert len([i for i in FindSymbols().visit(op) if i.is_Array]) == exp_arrays
 
 

tests/test_dse.py

@@ -57,7 +57,13 @@ def test_gradient_checkpointing(self, dtype, opt, space_order, kernel, shape, sp
         gradient2, _ = wave.jacobian_adjoint(residual, u0, vp=v0, checkpointing=False,
                                              grad=grad)
 
-        assert np.allclose(gradient.data, gradient2.data, atol=0, rtol=0)
+        # Not strictly bit-identical since sparse-position preambles now
+        # compute `(c - o)/h` in fp64: the fp64 intermediates may be reordered
+        # or fused differently across the two Operator instantiations, giving
+        # a sub-ulp fp32 drift near the source. Tolerate a small mismatch.
+        peak = float(np.max(np.abs(gradient2.data)))
+        atol = 1e-5 * peak if dtype == np.float32 else 1e-12 * peak
+        assert np.allclose(gradient.data, gradient2.data, atol=atol, rtol=0)
 
     @skipif(['cpu64-icc', 'cpu64-arm'])
     @pytest.mark.parametrize('tn', [750.])