Skip to content

Commit 307d39c

Browse files
committed
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.
1 parent e401545 commit 307d39c

12 files changed

Lines changed: 515 additions & 127 deletions

File tree

‎devito/operations/interpolators.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,9 @@ def _coeff_temps(self, implicit_dims, shifts=None):
315315
return []
316316

317317
def _positions(self, implicit_dims, shifts=None):
318+
# `pos = (c - o - s)/h` is computed in fp64 (see
319+
# `_position_map`); `posx = (int)floor(pos)` derives the integer
320+
# cell directly, no fp32 round-trip.
318321
return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims)
319322
for k, v in self.sfunction._position_map(shifts=shifts).items()]
320323

@@ -543,12 +546,19 @@ def _point_symbols(self, shifts=None):
543546
return DimensionTuple(*symbols, getters=self.grid.dimensions)
544547

545548
def _coeff_temps(self, implicit_dims, shifts=None):
546-
# Positions
547-
pmap = self.sfunction._position_map(shifts=shifts)
549+
# The fractional part is a weight in [0, 1] and can stay in fp32.
550+
# We compute it directly from the fp32 coord/origin/spacing and the
551+
# (already-tabulated) integer position -- no fp64 arithmetic, so
552+
# `pos` doesn't need to be tabulated at fp64 alongside `int_pos`.
553+
shifts = shifts or (0,) * len(self.grid.dimensions)
548554
psyms = self._point_symbols(shifts)
549-
poseq = [Eq(psyms[d], pos - floor(pos),
555+
int_syms = self.sfunction._pos_symbols(shifts=shifts)
556+
coords = self.sfunction._coordinate_symbols
557+
origins = self.grid.origin_symbols
558+
poseq = [Eq(psyms[d], (c - o - s)/d.spacing - ipos,
550559
implicit_dims=implicit_dims)
551-
for (d, pos) in zip(self._gdims, pmap.keys(), strict=True)]
560+
for d, c, o, s, ipos in zip(self._gdims, coords, origins,
561+
shifts, int_syms, strict=True)]
552562
return poseq
553563

554564

‎devito/passes/clusters/aliases.py‎

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
from devito.passes.clusters.cse import _cse
1616
from devito.passes.clusters.utils import expose_tuning_knobs
1717
from devito.symbolics import (
18-
Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search, sympy_dtype,
19-
uxreplace
18+
INT, Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search,
19+
sympy_dtype, uxreplace
2020
)
2121
from devito.tools import (
2222
Reconstructable, Stamp, as_mapper, as_tuple, flatten, frozendict, generator,
@@ -295,6 +295,10 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
295295

296296
class CireInvariants(CireTransformerLegacy, Queue):
297297

298+
# Predicate on Cluster used to pick which ones this pass fires on.
299+
# Subclasses override to target a different kind of cluster.
300+
_cluster_filter = staticmethod(lambda c: c.is_dense)
301+
298302
def __init__(self, sregistry, options, platform):
299303
super().__init__(sregistry, options, platform)
300304

@@ -324,7 +328,8 @@ def callback(self, clusters, prefix, xtracted=None):
324328
key = lambda c: self._lookup_key(c, d)
325329
processed = list(clusters)
326330
for ak, group in as_mapper(clusters, key=key).items():
327-
g = [c for c in group if c.is_dense and c not in xtracted]
331+
g = [c for c in group
332+
if self._cluster_filter(c) and c not in xtracted]
328333
if not g:
329334
continue
330335

@@ -387,6 +392,52 @@ def _generate(self, cgroup, exclude):
387392
yield self._do_generate(exprs, exclude, cbk_search)
388393

389394

395+
def _is_floor(e):
396+
return getattr(e, 'is_Function', False) and e.func.__name__ == 'floor'
397+
398+
399+
class CireInvariantsFloor(CireInvariants):
400+
401+
"""
402+
Hoist `INT(floor(x))` expressions -- typically the integer cell of a
403+
sparse position `INT(floor((c - o)/h))` -- into an int32 preamble Array.
404+
`CireInvariantsElementary`'s commutativity gate rejects both `floor(...)`
405+
(Cast in the argument makes `is_commutative == False`) and
406+
`INT(floor(...))` (a `Cast`, `is_commutative is None`), so we handle
407+
them here. Fires on every cluster: hoisting an invariant integer floor
408+
out of an inner loop is a strict win regardless of the cluster kind.
409+
"""
410+
411+
_cluster_filter = staticmethod(lambda c: True)
412+
413+
def _generate(self, cgroup, exclude):
414+
counter = generator()
415+
make_i32 = lambda: Symbol(name=f'dummy{counter()}', dtype=np.int32)
416+
417+
mapper = Uxmapper()
418+
for e in cgroup.exprs:
419+
for f in search(e, _is_floor, 'all', 'bfs'):
420+
cand = INT(f)
421+
if not {a.function for a in cand.free_symbols} & exclude:
422+
mapper.add(cand, make_i32, None)
423+
424+
yield mapper
425+
426+
def _choose(self, aliases, cgroup, mapper):
427+
# Skip score-based filtering: hoisting a floor out of an inner
428+
# loop is a strict flops win regardless of working-set size.
429+
exprs = cgroup.exprs
430+
431+
aliases = AliasList(aliases)
432+
if not aliases:
433+
return exprs, aliases
434+
435+
subs = {k: v for k, v in mapper.items()
436+
if v.free_symbols & set(aliases.aliaseds)}
437+
exprs = [uxreplace(e, subs) for e in exprs]
438+
return exprs, aliases
439+
440+
390441
class CireDerivatives(CireTransformerLegacy):
391442

392443
def __init__(self, sregistry, options, platform):
@@ -519,7 +570,8 @@ def _cbk_search2(self, expr, rank):
519570
# Subpass mapper
520571
modes = {
521572
'invariants': [CireInvariantsElementary,
522-
CireInvariantsDivs],
573+
CireInvariantsDivs,
574+
CireInvariantsFloor],
523575
'eval-derivs': [CireEvalDerivatives], # NOTE: legacy pass
524576
'index-derivs': [CireIndexDerivatives],
525577
}
@@ -707,8 +759,14 @@ def lower_aliases(aliases, meta, opt_maxpar, opt_block_temps):
707759
# not as an Indexed index. Then, it needs to be added to
708760
# the `writeto` region too
709761
interval = i
762+
elif writeto:
763+
# `d` is inner to the writeto region and unused by the
764+
# alias -- skip it, otherwise the preamble would iterate
765+
# over it and store the same value repeatedly.
766+
continue
710767
else:
711-
# E.g., `x0_blk0` or (`a[y_m+1]` => `y not in imapper`)
768+
# E.g., `x0_blk0` (outer to writeto, must stay in ispace
769+
# so later passes -- `Lift`, `Fuse` -- see it)
712770
intervals[d] = i
713771
continue
714772

@@ -1565,10 +1623,10 @@ def make_rotations_table(d, v):
15651623

15661624
def cit(ispace0, ispace1):
15671625
"""
1568-
The Common IterationIntervals of two IterationSpaces.
1626+
The Common IterationIntervals of two IterationSpaces (the shared prefix).
15691627
"""
15701628
found = []
1571-
for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=True):
1629+
for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=False):
15721630
if it0 == it1:
15731631
found.append(it0)
15741632
else:

‎devito/passes/clusters/cse.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@
1313

1414
from devito.finite_differences.differentiable import IndexDerivative
1515
from devito.ir import Cluster, Scope, cluster_pass
16-
from devito.symbolics import Reserved, estimate_cost, q_leaf, q_terminal, search
16+
from devito.symbolics import (
17+
Reserved, estimate_cost, q_leaf, q_terminal, search, sympy_dtype
18+
)
1719
from devito.symbolics.manipulation import _uxreplace
1820
from devito.tools import DAG, as_list, as_tuple, extract_dtype, frozendict
1921
from devito.types import Eq, Symbol, Temp
@@ -104,7 +106,10 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
104106
return cluster
105107

106108
def make(e):
107-
edtype = cse_dtype(e.dtype, dtype)
109+
# `sympy_dtype` respects Cast atoms (e.g. `DOUBLE(...)`); `e.dtype`
110+
# from Differentiable only inspects Function operands and would
111+
# silently downcast the temp, losing precision.
112+
edtype = cse_dtype(sympy_dtype(e.expr, default=e.dtype), dtype)
108113
return CTemp(name=sregistry.make_name(), dtype=edtype)
109114

110115
exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)

‎devito/symbolics/inspection.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,10 +316,24 @@ def sympy_dtype(expr, base=None, default=None, smin=None):
316316
if expr is None:
317317
return default
318318

319+
# An outermost Cast (e.g. `DOUBLE(...)`, `INT(...)`) declares the dtype
320+
# of the whole expression regardless of its operand's free symbols.
321+
if isinstance(expr, Cast):
322+
cast_dtype = expr.dtype
323+
if isinstance(cast_dtype, type) and issubclass(cast_dtype, np.generic):
324+
return cast_dtype
325+
319326
dtypes = set()
327+
# A Cast inside the tree bumps the inferred dtype: it forces the
328+
# operation to occur at (at least) the Cast's precision.
329+
for c in expr.atoms(Cast):
330+
cd = c.dtype
331+
if isinstance(cd, type) and issubclass(cd, np.generic):
332+
dtypes.add(cd)
320333
for i in expr.free_symbols:
321334
with suppress(AttributeError):
322335
dtypes.add(i.dtype)
336+
dtypes.discard(None)
323337

324338
if not dtypes or not np.issubdtype(base, np.complexfloating):
325339
dtypes.update({base} - {None})

‎devito/types/sparse.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from devito.operations import (
1212
LinearInterpolator, PrecomputedInterpolator, SincInterpolator
1313
)
14-
from devito.symbolics import indexify, retrieve_function_carriers
14+
from devito.symbolics import DOUBLE, indexify, retrieve_function_carriers
1515
from devito.tools import (
1616
ReducerMap, as_tuple, dtype_to_mpidtype, filter_ordered, flatten, is_integer,
1717
memoized_meth, prod
@@ -394,10 +394,15 @@ def _position_map(self, shifts=None):
394394
Dimension, of additional physical offsets to subtract from the sparse
395395
coordinates (e.g. ``h_x/2`` for a field staggered in ``x``). If ``shifts``
396396
is None, only the grid origin is subtracted.
397+
398+
Both origin and spacing are wrapped in `DOUBLE(...)` so the subtraction
399+
and division happen in fp64 and pick the correct cell regardless of the
400+
grid's fp32 rounding. Casting the spacing also prevents `1/h` from being
401+
constant-folded to a fp32 reciprocal literal at compile time.
397402
"""
398403
shifts = shifts or (0,) * len(self.grid.dimensions)
399404
return OrderedDict([
400-
((c - o - s)/d.spacing, p)
405+
((c - DOUBLE(o) - s)/DOUBLE(d.spacing), p)
401406
for p, c, d, o, s in zip(
402407
self._pos_symbols(shifts=shifts),
403408
self._coordinate_symbols,

‎examples/userapi/06_sparse_operations.ipynb‎

Lines changed: 241 additions & 53 deletions
Large diffs are not rendered by default.

‎tests/test_dle.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -204,9 +204,9 @@ def test_basic(self):
204204

205205
op = Operator(eqns, opt=('advanced', {'blockrelax': True}))
206206

207-
bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0'})
207+
bns, _ = assert_blocking(op, {'x0_blk0', 'p_src0_blk0', 'p_src1_blk0'})
208208

209-
iters = FindNodes(Iteration).visit(bns['p_src0_blk0'])
209+
iters = FindNodes(Iteration).visit(bns['p_src1_blk0'])
210210
assert len(iters) == 5
211211
assert iters[0].dim.is_Block
212212
assert iters[1].dim.is_Block
@@ -1070,12 +1070,14 @@ def test_incr_perfect_sparse_outer(self):
10701070
'openmp': True}))
10711071

10721072
iters = FindNodes(Iteration).visit(op)
1073-
assert len(iters) == 5
1074-
assert iters[0].is_Sequential
1075-
assert all(i.is_ParallelAtomic for i in iters[1:])
1076-
assert iters[1].pragmas[0].ccode.value ==\
1073+
# 1 preamble iteration (p_u) hoisted out of time by
1074+
# `CireInvariantsFloor`, then the time loop and the 4 injection iters.
1075+
assert len(iters) == 6
1076+
assert iters[1].is_Sequential
1077+
assert all(i.is_ParallelAtomic for i in iters[:1] + iters[2:])
1078+
assert iters[2].pragmas[0].ccode.value ==\
10771079
'omp for schedule(dynamic,chunk_size)'
1078-
assert all(not i.pragmas for i in iters[2:])
1080+
assert all(not i.pragmas for i in iters[3:])
10791081

10801082
@pytest.mark.parametrize('exprs,simd_level,expected', [
10811083
(['Eq(y.symbolic_max, g[0, x], implicit_dims=(t, x))',

0 commit comments

Comments
 (0)