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

devito/operations/interpolators.py

@@ -315,6 +315,9 @@ def _coeff_temps(self, implicit_dims, shifts=None):
         return []
 
     def _positions(self, implicit_dims, shifts=None):
+        # `pos = (c - o - s)/h` is computed in fp64 (see
+        # `_position_map`); `posx = (int)floor(pos)` derives the integer
+        # cell directly, no fp32 round-trip.
         return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims)
                 for k, v in self.sfunction._position_map(shifts=shifts).items()]
 
@@ -543,12 +546,19 @@ 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)
+        # The fractional part is a weight in [0, 1] and can stay in fp32.
+        # We compute it directly from the fp32 coord/origin/spacing and the
+        # (already-tabulated) integer position -- no fp64 arithmetic, so
+        # `pos` doesn't need to be tabulated at fp64 alongside `int_pos`.
+        shifts = shifts or (0,) * len(self.grid.dimensions)
         psyms = self._point_symbols(shifts)
-        poseq = [Eq(psyms[d], pos - floor(pos),
+        int_syms = self.sfunction._pos_symbols(shifts=shifts)
+        coords = self.sfunction._coordinate_symbols
+        origins = self.grid.origin_symbols
+        poseq = [Eq(psyms[d], (c - o - s)/d.spacing - ipos,
                     implicit_dims=implicit_dims)
-                 for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)]
+                 for d, c, o, s, ipos in zip(self._gdims, coords, origins,
+                                             shifts, int_syms, 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
+    INT, 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,52 @@ 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 CireInvariantsFloor(CireInvariants):
+
+    """
+    Hoist `INT(floor(x))` expressions -- typically the integer cell of a
+    sparse position `INT(floor((c - o)/h))` -- into an int32 preamble Array.
+    `CireInvariantsElementary`'s commutativity gate rejects both `floor(...)`
+    (Cast in the argument makes `is_commutative == False`) and
+    `INT(floor(...))` (a `Cast`, `is_commutative is None`), so we handle
+    them here. Fires on every cluster: hoisting an invariant integer floor
+    out of an inner loop is a strict win regardless of the cluster kind.
+    """
+
+    _cluster_filter = staticmethod(lambda c: True)
+
+    def _generate(self, cgroup, exclude):
+        counter = generator()
+        make_i32 = lambda: Symbol(name=f'dummy{counter()}', dtype=np.int32)
+
+        mapper = Uxmapper()
+        for e in cgroup.exprs:
+            for f in search(e, _is_floor, 'all', 'bfs'):
+                cand = INT(f)
+                if not {a.function for a in cand.free_symbols} & exclude:
+                    mapper.add(cand, make_i32, None)
+
+        yield mapper
+
+    def _choose(self, aliases, cgroup, mapper):
+        # Skip score-based filtering: hoisting a floor out of an inner
+        # loop is a strict flops win regardless of working-set size.
+        exprs = cgroup.exprs
+
+        aliases = AliasList(aliases)
+        if not aliases:
+            return exprs, aliases
+
+        subs = {k: v for k, v in mapper.items()
+                if v.free_symbols & set(aliases.aliaseds)}
+        exprs = [uxreplace(e, subs) for e in exprs]
+        return exprs, aliases
+
+
 class CireDerivatives(CireTransformerLegacy):
 
     def __init__(self, sregistry, options, platform):
@@ -519,7 +570,8 @@ def _cbk_search2(self, expr, rank):
 # Subpass mapper
 modes = {
     'invariants': [CireInvariantsElementary,
-                   CireInvariantsDivs],
+                   CireInvariantsDivs,
+                   CireInvariantsFloor],
     'eval-derivs': [CireEvalDerivatives],   # NOTE: legacy pass
     'index-derivs': [CireIndexDerivatives],
 }
@@ -707,8 +759,14 @@ def lower_aliases(aliases, meta, opt_maxpar, opt_block_temps):
                     # not as an Indexed index. Then, it needs to be added to
                     # the `writeto` region too
                     interval = i
+                elif writeto:
+                    # `d` is inner to the writeto region and unused by the
+                    # alias -- skip it, otherwise the preamble would iterate
+                    # over it and store the same value repeatedly.
+                    continue
                 else:
-                    # E.g., `x0_blk0` or (`a[y_m+1]` => `y not in imapper`)
+                    # E.g., `x0_blk0` (outer to writeto, must stay in ispace
+                    # so later passes -- `Lift`, `Fuse` -- see it)
                     intervals[d] = i
                     continue
 
@@ -1565,10 +1623,10 @@ def make_rotations_table(d, v):
 
 def cit(ispace0, ispace1):
     """
-    The Common IterationIntervals of two IterationSpaces.
+    The Common IterationIntervals of two IterationSpaces (the shared prefix).
     """
     found = []
-    for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=True):
+    for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=False):
         if it0 == it1:
             found.append(it0)
         else:

devito/passes/clusters/cse.py

@@ -13,7 +13,9 @@
 
 from devito.finite_differences.differentiable import IndexDerivative
 from devito.ir import Cluster, Scope, cluster_pass
-from devito.symbolics import Reserved, estimate_cost, q_leaf, q_terminal, search
+from devito.symbolics import (
+    Reserved, estimate_cost, q_leaf, q_terminal, search, sympy_dtype
+)
 from devito.symbolics.manipulation import _uxreplace
 from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict
 from devito.types import Eq, Symbol, Temp
@@ -104,7 +106,10 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
         return cluster
 
     def make(e):
-        edtype = cse_dtype(e.dtype, dtype)
+        # `sympy_dtype` respects Cast atoms (e.g. `DOUBLE(...)`); `e.dtype`
+        # from Differentiable only inspects Function operands and would
+        # silently downcast the temp, losing precision.
+        edtype = cse_dtype(sympy_dtype(e.expr, default=e.dtype), dtype)
         return CTemp(name=sregistry.make_name(), dtype=edtype)
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)

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,15 @@ 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.
+
+        Both origin and spacing are wrapped in `DOUBLE(...)` so the subtraction
+        and division happen in fp64 and pick the correct cell regardless of the
+        grid's fp32 rounding. Casting the spacing also prevents `1/h` from being
+        constant-folded to a fp32 reciprocal literal at compile time.
         """
         shifts = shifts or (0,) * len(self.grid.dimensions)
         return OrderedDict([
-            ((c - o - s)/d.spacing, p)
+            ((c - DOUBLE(o) - s)/DOUBLE(d.spacing), p)
             for p, c, d, o, s in zip(
                 self._pos_symbols(shifts=shifts),
                 self._coordinate_symbols,

tests/test_dle.py

@@ -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 ==\
+        # 1 preamble iteration (p_u) hoisted out of time by
+        # `CireInvariantsFloor`, then the time loop and the 4 injection iters.
+        assert len(iters) == 6
+        assert iters[1].is_Sequential
+        assert all(i.is_ParallelAtomic for i in iters[:1] + iters[2:])
+        assert iters[2].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[3:])
 
     @pytest.mark.parametrize('exprs,simd_level,expected', [
         (['Eq(y.symbolic_max, g[0, x], implicit_dims=(t, x))',