diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 0f0354f6638..8f186df0f7a 100644 --- a/devito/core/gpu.py +++ b/devito/core/gpu.py @@ -162,7 +162,10 @@ def wrapper(expressions, mode='default', options=None, **kwargs1): # small kernels typically generated by recursive compilation par_tile0 = options0['par-tile'] par_tile = options.get('par-tile') - if par_tile0 and par_tile: + if par_tile is False: + # The caller explicitly opted out of tiling + options = {**options0, **options, 'par-tile': ParTile(None)} + elif par_tile0 and par_tile: options = {**options0, **options, 'par-tile': par_tile} elif par_tile0: par_tile = ParTile(par_tile0.default, default=par_tile0.default) diff --git a/devito/finite_differences/derivative.py b/devito/finite_differences/derivative.py index 08feb17839a..de28b508590 100644 --- a/devito/finite_differences/derivative.py +++ b/devito/finite_differences/derivative.py @@ -14,7 +14,7 @@ from devito.warnings import warn from .differentiable import Add, Differentiable, Mul, diffify, interp_for_fd -from .finite_difference import cross_derivative, generic_derivative +from .finite_difference import cross_derivative, generic_derivative, indices_at from .rsfd import d45 from .tools import direct, transpose @@ -575,6 +575,12 @@ def _eval_fd(self, expr, **kwargs): shited derivative. - 4: Apply substitutions. """ + # Differentiation is linear, and a sum of terms at different staggered + # locations must use it: `Add` reports its first argument's location, + # so `x0` would shift the other terms off the point they sat at. + if expr.is_Add and any(len(indices_at(expr, d)) > 1 for d in self.dims): + return expr.func(*[self._eval_fd(a, **kwargs) for a in expr.args]) + # Step 1: Evaluate non-derivative x0. We currently enforce a simple 2nd order # interpolation to avoid very expensive finite differences on top of it x0_deriv = self._filter_dims(self.x0) diff --git a/devito/finite_differences/finite_difference.py b/devito/finite_differences/finite_difference.py index 11c2946ed15..de2e92898df 100644 --- a/devito/finite_differences/finite_difference.py +++ b/devito/finite_differences/finite_difference.py @@ -100,6 +100,29 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, side=None, **kw return expr +def indices_at(expr, dim): + """ + The locations `expr`'s terms sit at along `dim`. + + Terms with no location of their own, a scalar say, contribute none. + """ + indices = set() + for i in (expr.args if expr.is_Add else (expr,)): + try: + indices.add(i.indices_ref[dim]) + except (AttributeError, KeyError, IndexError, TypeError): + continue + return indices + + +def index_at(expr, dim): + """ + Where `expr` sits along `dim`, or None if it does not say. + """ + indices = indices_at(expr, dim) + return indices.pop() if len(indices) == 1 else None + + @check_input def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, coefficients='taylor', expand=True, weights=None, side=None): @@ -139,8 +162,9 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None, if deriv_order == 1 and fd_order == 2 and side is None: fd_order = 1 - # Zeroth order derivative is just the expression itself if not shifted - if deriv_order == 0 and not x0: + # Zeroth order is the identity when `expr` already sits at `x0`, not a + # stencil centred there. + if deriv_order == 0 and (not x0 or index_at(expr, dim) == x0.get(dim)): return expr # Enforce stable time coefficients diff --git a/devito/ir/clusters/algorithms.py b/devito/ir/clusters/algorithms.py index 5daffb78b0c..86170eafa11 100644 --- a/devito/ir/clusters/algorithms.py +++ b/devito/ir/clusters/algorithms.py @@ -742,11 +742,10 @@ def _normalize_reductions_dense(cluster, mapper, sregistry, platform): elif rhs in mapper: # Seen this RHS already, so reuse the Array that was created for it processed.append(e.func(lhs, mapper[rhs].indexify())) - elif rf and rf.is_Array and sum(flatten(rf._size_nodomain)) == 0: - # Special case: the RHS is an Array with no halo/padding, meaning - # that the written data values are contiguous in memory, hence - # we can simply reuse the Array itself as we're already in the - # desired memory layout + elif rf and rf.is_Array and rf._is_reduction_ready: + # Special case: the RHS is an Array whose written data values are + # contiguous in memory, hence we can simply reuse the Array + # itself as we're already in the desired memory layout processed.append(e) else: name = sregistry.make_name() diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 81b12fa53b3..19a46044544 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -7,6 +7,7 @@ import ctypes from collections import OrderedDict from collections.abc import Callable, Generator, Iterable, Iterator, Sequence +from contextlib import suppress from itertools import chain, groupby from typing import Any, Generic, TypeVar @@ -17,8 +18,8 @@ from devito.exceptions import CompilationError from devito.ir.cgen.printer import get_printer from devito.ir.iet.nodes import ( - BlankLine, Call, Expression, ExpressionBundle, Iteration, Lambda, ListMajor, Node, - Section, _same_as_before + BlankLine, Call, Definition, Expression, ExpressionBundle, Iteration, Lambda, + ListMajor, Node, Section, _same_as_before ) from devito.ir.support.space import Backward from devito.symbolics import ( @@ -1275,6 +1276,23 @@ def visit_Call(self, o: Call, **kwargs) -> Iterator[ApplicationType]: except (AttributeError, TypeError): yield from self._visit(i) + def visit_Definition(self, o: Definition, **kwargs) -> Iterator[ApplicationType]: + # The defined object carries expressions in its constructor arguments + # and in its initializer, both of which end up in the generated code + f = o.function + if f.is_LocalObject: + candidates = (*f.cargs, f.initvalue) + elif f.is_Array: + candidates = as_tuple(f.initvalue) + else: + return + + for i in candidates: + # Not everything in there is a symbolic expression, e.g. a plain + # number, a string, or nothing at all + with suppress(AttributeError, TypeError): + yield from i.find(self.match) + class IsPerfectIteration(Visitor): diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 2f1cce8f108..ad07f2aa19a 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -20,7 +20,7 @@ VOID, Byref, DefFunction, FieldFromPointer, IndexedPointer, ListInitializer, SizeOf, as_long, pow_to_mul, unevaluate ) -from devito.tools import as_list, as_mapper, as_tuple, filter_sorted, flatten +from devito.tools import as_list, as_mapper, as_tuple, filter_sorted, flatten, is_integer from devito.types import ( Array, ComponentAccess, CustomDimension, DeviceMap, DeviceRM, Dimension, Eq, Symbol, size_t @@ -91,6 +91,27 @@ def __init__(self, rcompile=None, sregistry=None, platform=None, self.sregistry = sregistry self.platform = platform + # Off inside the recursive compilation of a zero-init itself, which + # would otherwise ask for a zero-init of its own, ad infinitum + self.zero_init = (options or {}).get('zero-init', True) + + def _zero_init(self, obj, storage): + """ + The nodes zeroing `obj` upfront, if it asks for it, plus the efuncs + they call, if any. + """ + if not (obj._is_zero_init and self.zero_init): + return (), () + + return self._make_zero_init(obj, storage) + + def _make_zero_init(self, obj, storage): + """How to zero `obj`'s whole allocation, padding included.""" + storage.include(self.langbb['header-memcpy']) + nbytes = SizeOf(obj._C_typedata)*as_long(obj.size) + + return (self.langbb['host-memset'](obj._C_symbol, 0, nbytes),), () + def _alloc_object_on_low_lat_mem(self, site, obj, storage): """ Allocate a LocalObject in the low latency memory. @@ -172,11 +193,13 @@ def _alloc_host_array_on_high_bw_mem(self, site, obj, storage, *args): memptr = VOID(Byref(obj._C_symbol), '**') alignment = obj._data_alignment nbytes = SizeOf(obj._C_typedata)*as_long(obj.size) - alloc = self.langbb['host-alloc'](memptr, alignment, nbytes) + zeroing, efuncs = self._zero_init(obj, storage) + allocs = [decl, self.langbb['host-alloc'](memptr, alignment, nbytes), + *zeroing] free = self.langbb['host-free'](obj._C_symbol) - storage.update(obj, site, allocs=(decl, alloc), frees=free) + storage.update(obj, site, allocs=tuple(allocs), frees=free, efuncs=efuncs) def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, *args): """ @@ -568,7 +591,7 @@ def __init__(self, options=None, **kwargs): self.gpu_create = options['gpu-create'] self.gpu_place_transfers = options.get('place-transfers') - super().__init__(**kwargs) + super().__init__(options=options, **kwargs) def _alloc_local_array_on_high_bw_mem(self, site, obj, storage): """ @@ -579,11 +602,22 @@ def _alloc_local_array_on_high_bw_mem(self, site, obj, storage): dofree = self.langbb['device-free'] nbytes = SizeOf(obj._C_typedata)*obj.size - init = doalloc(nbytes, deviceid, retobj=obj) + + zeroing, efuncs = self._zero_init(obj, storage) + allocs = [doalloc(nbytes, deviceid, retobj=obj), *zeroing] free = dofree(obj._C_name, deviceid) - storage.update(obj, site, allocs=init, frees=free) + storage.update(obj, site, allocs=tuple(allocs), frees=free, efuncs=efuncs) + + def _make_zero_init(self, obj, storage): + # No language here has a device-side memset, so use a kernel. It gains + # nothing from tiling, and nvc++ trips over the padded loop bounds + # when it is asked to tile them + efuncs, init = make_zero_init(obj, self.rcompile, self.sregistry, + options={'par-tile': False}) + + return (init,), efuncs def _map_array_on_high_bw_mem(self, site, obj, storage): """ @@ -702,18 +736,19 @@ def process(self, graph): self.place_casts(graph) -def make_zero_init(obj, rcompile, sregistry): +def make_zero_init(obj, rcompile, sregistry, options=None): cdims = [] - for d, (h0, h1), s in zip( - obj.dimensions, obj._size_halo, obj.symbolic_shape, strict=True + for d, (h0, h1), (p0, p1), s in zip( + obj.dimensions, obj._size_halo, obj._size_padding, obj.symbolic_shape, + strict=True ): if d.is_NonlinearDerived: - assert h0 == h1 == 0 + assert h0 == h1 == p0 == p1 == 0 m = 0 M = s - 1 else: m = d.symbolic_min - h0 - M = d.symbolic_max + h1 + M = d.symbolic_max + h1 + (0 if is_integer(p1) else p1) cdims.append(CustomDimension(name=d.name, parent=d, symbolic_min=m, symbolic_max=M)) @@ -722,7 +757,8 @@ def make_zero_init(obj, rcompile, sregistry): else: eqns = [Eq(obj[cdims], 0)] - irs, byproduct = rcompile(eqns) + irs, byproduct = rcompile(eqns, options={'zero-init': False, + **(options or {})}) init = irs.iet.body.body[0] diff --git a/devito/passes/iet/languages/C.py b/devito/passes/iet/languages/C.py index ddd61b325f4..c5ac413ac49 100644 --- a/devito/passes/iet/languages/C.py +++ b/devito/passes/iet/languages/C.py @@ -56,6 +56,8 @@ class CBB(LangBB): Call('free', (i,)), 'host-free-pin': lambda i: Call('free', (i,)), + 'host-memset': lambda i, j, k: + Call('memset', (i, j, k)), 'alloc-global-symbol': lambda i, j, k: Call('memcpy', (i, j, k)) } diff --git a/devito/passes/iet/languages/CXX.py b/devito/passes/iet/languages/CXX.py index c078f53de82..d7e392630cb 100644 --- a/devito/passes/iet/languages/CXX.py +++ b/devito/passes/iet/languages/CXX.py @@ -141,6 +141,8 @@ class CXXBB(LangBB): Call('free', (i,)), 'host-free-pin': lambda i: Call('free', (i,)), + 'host-memset': lambda i, j, k: + Call('memset', (i, j, k)), 'alloc-global-symbol': lambda i, j, k: Call('memcpy', (i, j, k)) } diff --git a/devito/types/basic.py b/devito/types/basic.py index 3f8e0f6ec0a..29028a27551 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -16,7 +16,7 @@ from devito.data import default_allocator from devito.parameters import configuration from devito.tools import ( - CustomDtype, Pickable, as_tuple, dtype_to_ctype, frozendict, memoized_meth, + CustomDtype, Pickable, as_tuple, dtype_to_ctype, flatten, frozendict, memoized_meth, sympy_mutex ) from devito.types.args import ArgProvider @@ -1351,6 +1351,20 @@ def _size_nodomain(self): return DimensionTuple(*sizes, getters=self.dimensions, left=left, right=right) + _is_zero_init = False + """ + Whether the entries outside `self`'s DOMAIN carry meaningful data rather + than scratch, in which case the whole allocation must be zeroed upfront. + """ + + @property + def _is_reduction_ready(self): + """ + True if a reduction over `self` may run over the whole allocated data + rather than over the DOMAIN alone, False otherwise. + """ + return not sum(flatten(self._size_nodomain)) + @cached_property def _size_ghost(self): """ diff --git a/devito/types/utils.py b/devito/types/utils.py index 8c5617f1258..ecb2b360f60 100644 --- a/devito/types/utils.py +++ b/devito/types/utils.py @@ -44,17 +44,21 @@ class Stagger(Tag): class DimensionTuple(EnrichedTuple): - def __getitem_hook__(self, dim): + def _getter(self, dim): + # Exact hit first: a derived Dimension carries its parent in + # `_defines`, so an overlap test alone matches the parent's entry. + if dim in self.getters: + return dim for d in self.getters: if d._defines & dim._defines: - return self.getters[d] + return d raise KeyError + def __getitem_hook__(self, dim): + return self.getters[self._getter(dim)] + def dindex(self, dim): - for d in self.getters: - if d._defines & dim._defines: - return list(self.getters).index(d) - raise KeyError + return list(self.getters).index(self._getter(dim)) class Staggering(DimensionTuple): diff --git a/tests/test_data.py b/tests/test_data.py index a263213094b..64069bf65d3 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -412,6 +412,18 @@ def test_temp_array_smart_padding_codegen_avoids_negative_mod(self): assert '(-z_size)' not in code assert 'z_size' in code + def test_is_reduction_ready(self): + grid = Grid(shape=(4, 4)) + + # No halo, no padding -- the DOMAIN spans the whole allocation, so a + # reduction may run straight off `r0` + r0 = TempArray(name='r0', dimensions=grid.dimensions, dtype=np.float32) + assert r0._is_reduction_ready + + r1 = TempArray(name='r1', dimensions=grid.dimensions, dtype=np.float32, + halo=((0, 0), (0, 1))) + assert not r1._is_reduction_ready + def test_w_halo_custom(self): grid = Grid(shape=(4, 4)) diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 7ed57b6c91f..362d3d1201f 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -1461,3 +1461,43 @@ def test_unevaluated(self): assert Derivative(self.x, self.t) assert Derivative(self.x, self.y, self.t) assert Derivative(self.x, (self.x, 0)) + + +@pytest.mark.parametrize('expand', [True, False]) +@pytest.mark.parametrize('deriv_order', [0, 1, 2]) +def test_deriv_sum_mixed_staggering(expand, deriv_order): + """ + A shifted derivative is linear: `D(a + b) == D(a) + D(b)`, at every order. + + Broke for terms at different staggered locations, `Add` reporting only its + first argument's. + """ + so = 8 + grid = Grid(shape=(41, 41), extent=(40., 40.)) + x, y = grid.dimensions + + vx = Function(name='vx', grid=grid, space_order=so, staggered=x) + vy = Function(name='vy', grid=grid, space_order=so, staggered=y) + out = Function(name='out', grid=grid, space_order=so, staggered=(x, y)) + + rng = np.random.default_rng(3) + for f in (vx, vy): + f.data[:] = rng.normal(size=f.shape) + + def shifted(expr): + return expr.diff(y, deriv_order=deriv_order, fd_order=2, + x0={y: y + y.spacing/2}) + + def run(expr): + out.data[:] = 0. + Operator(Eq(out, expr), opt=('advanced', {'expand': expand})).apply() + return np.array(out.data) + + s = slice(so + 3, -(so + 3)) + together = run(shifted(vx.dy + vy.dx))[s, s] + apart = (run(shifted(vx.dy)) + run(shifted(vy.dx)))[s, s] + + assert np.linalg.norm(apart) > 0 + # float32 reassociation only: the two forms sum the same terms in a + # different order + assert np.linalg.norm(together - apart) / np.linalg.norm(apart) < 1e-5 diff --git a/tests/test_linearize.py b/tests/test_linearize.py index c2b424d934f..c78880924d9 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -8,7 +8,7 @@ ) from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry from devito.passes import Graph, generate_macros, linearize -from devito.types import Array, Bundle, DefaultDimension +from devito.types import Array, Bundle, CustomDimension, DefaultDimension def test_basic(): @@ -716,3 +716,24 @@ def test_cire_n_strides(): # NOTE: not exact equality because `op2` slightly changes the order of # arithmetic operations, which in turn causes some rounding differences assert np.allclose(u.data, u1.data, rtol=1e-4) + + +def test_bundle_derived_dim_stride(): + """ + A Bundle's stride comes from the Dimension asked for, not from its parent. + + `rp._defines` contains `p`, so an overlap lookup gave the Bundle the number + of points as innermost stride instead of the number of weights. + """ + grid = Grid(shape=(4, 4)) + p = DefaultDimension(name='p', default_value=5) + rp = CustomDimension(name='rp', parent=p, symbolic_size=2) + + w0 = Function(name='w0', dimensions=(p, rp), shape=(5, 2)) + w1 = Function(name='w1', dimensions=(p, rp), shape=(5, 2)) + bundle = Bundle(name='w0w1', components=(w0, w1), grid=grid) + + assert rp._defines & p._defines # the overlap that used to mislead + assert bundle.symbolic_shape[p] is not bundle.symbolic_shape[rp] + assert bundle.symbolic_shape[rp] == 2 + assert bundle.symbolic_shape.dindex(rp) == 1 diff --git a/tests/test_operator.py b/tests/test_operator.py index de445cde016..a6090b54be4 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -1394,6 +1394,26 @@ def test_conditional_declarations(self): assert i[0].is_Expression assert i[0].expr.rhs is init_value + def test_zero_init_array(self): + """ + An Array whose entries outside the DOMAIN are data, rather than + scratch, is zeroed right after being allocated. + """ + grid = Grid(shape=(4, 4)) + + class ZeroInitArray(Array): + _is_zero_init = True + + a = ZeroInitArray(name='a', dimensions=grid.dimensions, + dtype=grid.dtype, space='local') + b = Array(name='b', dimensions=grid.dimensions, dtype=grid.dtype, + space='local') + + f = Function(name='f', grid=grid) + + assert 'memset(a' in str(Operator(Eq(f, a.indexify()))) + assert 'memset(b' not in str(Operator(Eq(f, b.indexify()))) + def test_nested_scalar_assigns(self): grid = Grid(shape=(4, 4)) diff --git a/tests/test_visitors.py b/tests/test_visitors.py index b5d12f81d5f..b0acb0cdcd2 100644 --- a/tests/test_visitors.py +++ b/tests/test_visitors.py @@ -1,3 +1,5 @@ +from ctypes import c_void_p + import cgen as c import pytest from sympy import Mod @@ -5,11 +7,12 @@ from devito import Eq, Function, Grid, Min, Operator, TimeFunction, sin from devito.ir.equations import DummyEq from devito.ir.iet import ( - Block, Call, Callable, Conditional, Expression, FindApplications, FindNodes, - FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, MapNodes, - Transformer, Uxreplace, printAST + Block, Call, Callable, Conditional, Definition, Expression, FindApplications, + FindNodes, FindSections, FindSymbols, FindWithin, IsPerfectIteration, Iteration, + MapNodes, Transformer, Uxreplace, printAST ) -from devito.types import Array, SpaceDimension, Symbol +from devito.symbolics import ListInitializer +from devito.types import Array, LocalObject, SpaceDimension, Symbol @pytest.fixture(scope="module") @@ -422,3 +425,31 @@ def test_find_apps_nested_calls(exprs, iters): block = iters[0](iters[1](exprs + [call])) assert len(FindApplications().visit(block)) == 1 + + +def test_find_apps_in_definition(): + """ + A Definition carries expressions in the constructor arguments and in the + initializer of the defined object, and both end up in the generated code, + so both must be searched -- otherwise e.g. `generate_macros` would leave + the macros they apply undefined. + """ + s = Symbol(name='s') + d = SpaceDimension(name='d') + + class DummyObject(LocalObject): + dtype = c_void_p + + obj = DummyObject(name='obj', initvalue=ListInitializer([Min(s, 1)])) + assert FindApplications().visit(Definition(obj)) == {Min(s, 1)} + + obj = DummyObject(name='obj', cargs=(Min(s, 2),)) + assert FindApplications().visit(Definition(obj)) == {Min(s, 2)} + + # An Array carries an initializer too + a = Array(name='a', dimensions=(d,), scope='stack', + initvalue=[Min(s, 3), 0]) + assert FindApplications().visit(Definition(a)) == {Min(s, 3)} + + # An object with neither must not trip the search up + assert FindApplications().visit(Definition(DummyObject(name='obj'))) == set()