diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py index 2b8b6d5080..6803e9f0f8 100644 --- a/devito/finite_differences/differentiable.py +++ b/devito/finite_differences/differentiable.py @@ -898,7 +898,7 @@ class RealComplexPart(ComplexPart): @cached_property def dtype(self): - dtype = extract_dtype(self) + dtype = extract_dtype(self.args[0]) return dtype(0).real.__class__ @@ -1321,6 +1321,10 @@ def base(self): def pivot(self): return self.base.subs({d: 0 for d in self.dimensions}) + @cached_property + def dtype(self): + return extract_dtype(self.expr) + @property def weights(self): return self._weights diff --git a/devito/ir/clusters/cluster.py b/devito/ir/clusters/cluster.py index bba54fbafc..4c41fd7f49 100644 --- a/devito/ir/clusters/cluster.py +++ b/devito/ir/clusters/cluster.py @@ -15,7 +15,8 @@ from devito.mpi.reduction_scheme import DistReduce from devito.symbolics import estimate_cost, uxreplace from devito.tools import ( - CacheInstances, as_tuple, cached_hash, filter_ordered, flatten, infer_dtype + CacheInstances, as_tuple, cached_hash, dtypes_vector_mapper, filter_ordered, + flatten, infer_dtype ) from devito.types import ( CriticalRegion, Fence, Indexed, PhaseMarker, TensorMove, ThreadArrive, ThreadCommit, @@ -340,6 +341,8 @@ def dtype(self): # E.g. `i.dtype` is a ctypes pointer, which has no dtype equivalent pass + # A Cluster's precision is independent of its expressions' vector lengths + dtypes = {dtypes_vector_mapper.get_base_dtype(i, i) for i in dtypes} return infer_dtype(dtypes) @cached_property diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 8f3c02f983..3c98874fc1 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -635,7 +635,11 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs): # Process data-carrier overrides for p in overrides: - args.update(p._arg_values(estimate_memory=estimate_memory, **kwargs)) + values = p._arg_values(estimate_memory=estimate_memory, **kwargs) + # A replacement may supply defaults for other arguments, such as + # a ConditionalDimension's factor. Explicit overrides take precedence + args.update({k: v for k, v in values.items() + if k == p.name or k not in kwargs}) try: args.reduce_inplace() except ValueError as e: diff --git a/devito/passes/clusters/derivatives.py b/devito/passes/clusters/derivatives.py index fceeddb8bf..3658e0051b 100644 --- a/devito/passes/clusters/derivatives.py +++ b/devito/passes/clusters/derivatives.py @@ -6,9 +6,9 @@ from devito.finite_differences import IndexDerivative, Weights from devito.ir import Backward, Forward, Interval, IterationSpace, Queue from devito.passes.clusters.fusion import fuse -from devito.symbolics import BasicWrapperMixin, reuse_if_untouched, search, uxreplace -from devito.symbolics.inspection import sympy_dtype -from devito.tools import infer_dtype, timed_pass +from devito.symbolics import (BasicWrapperMixin, ListInitializer, reuse_if_untouched, + search, uxreplace) +from devito.tools import dtype_len, extract_dtype, infer_dtype, timed_pass from devito.types import Eq, Inc, Indexed, Symbol, Temp __all__ = ['lower_index_derivatives'] @@ -173,19 +173,22 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs): # The temporary that will hold the result of the IndexDerivative computation # NOTE: created before recursing so that we ultimately get a sound ordering - dtype = sympy_dtype(ideriv) - try: - s = reusables.pop() - assert np.can_cast(s.dtype, dtype) - except KeyError: + dtype = ideriv.dtype + s = next((s for s in reusables if np.can_cast(s.dtype, dtype)), None) + if s is None: name = sregistry.make_name(prefix='r') s = Temp(name=name, dtype=dtype) + else: + reusables.remove(s) # Go inside `expr` and recursively lower any nested IndexDerivatives expr, processed = _core(expr, c, ispace1, weights, reusables, mapper, **kwargs) # Finally inject the lowered IndexDerivative if init is not None: + # Broadcast scalar initializers; folded centre terms may already be vectors + if (n := dtype_len(dtype)) > 1 and dtype_len(extract_dtype(init)) == 1: + init = ListInitializer(*([init]*n)) expr0 = Eq(s, init) processed.insert(0, c.rebuild(exprs=expr0, ispace=ispace)) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fe9a78b65b..08c5bb368c 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -593,6 +593,10 @@ def _(i, mapper, sregistry): @abstract_object.register(Array) def _(i, mapper, sregistry): + # Constant arrays have global storage shared with other Callables + if i._mem_constant: + return + name = sregistry.make_name(prefix=i._symbol_prefix) if i.initvalue is not None: diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 57d536bf1c..48933e0ccd 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -13,7 +13,7 @@ from devito.symbolics.extended_sympy import CallFromPointer, Cast, DefFunction, Reserved from devito.symbolics.queries import q_routine from devito.tools import as_tuple, is_integer, prod -from devito.tools.dtypes_lowering import infer_dtype +from devito.tools.dtypes_lowering import dtypes_vector_mapper, infer_dtype __all__ = ['compare_ops', 'estimate_cost', 'has_integer_args', 'sympy_dtype'] @@ -311,7 +311,7 @@ def has_integer_args(*args): def sympy_dtype(expr, base=None, default=None, smin=None): """ - Infer the dtype of the expression. + Infer the scalar precision of the expression. """ if expr is None: return default @@ -324,6 +324,7 @@ def sympy_dtype(expr, base=None, default=None, smin=None): if not dtypes or not np.issubdtype(base, np.complexfloating): dtypes.update({base} - {None}) + dtypes = {dtypes_vector_mapper.get_base_dtype(i, i) for i in dtypes} dtype = infer_dtype(dtypes) # Promote if we missed complex number, i.e f + I diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index f3033aa6e9..4f2ef2b9a7 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -93,7 +93,7 @@ def add_dtype(self, field_name, count): def get_base_dtype(self, v, default=None): for (base_dtype, _), dtype in self.items(): - if dtype is v: + if dtype == v: return base_dtype if default is not None: @@ -364,7 +364,12 @@ def infer_dtype(dtypes): * In the case of multiple floating dtypes, return the dtype with highest precision; * If there's at least one floating dtype, ignore any integer dtypes. + + Preserve vector length when combining compatible vector and scalar operands. """ + dtypes = set(dtypes) + n = max(map(dtype_len, dtypes), default=1) + # Resolve the vector types, if any dtypes = {dtypes_vector_mapper.get_base_dtype(i, i) for i in dtypes} # Only keep number types @@ -373,19 +378,32 @@ def infer_dtype(dtypes): fdtypes = {i for i in dtypes if np.issubdtype(i, np.floating) or np.issubdtype(i, np.complexfloating)} if len(fdtypes) > 1: - return max(fdtypes, key=lambda i: np.dtype(i).itemsize) + dtype = max(fdtypes, key=lambda i: np.dtype(i).itemsize) elif len(fdtypes) == 1: - return fdtypes.pop() + dtype = fdtypes.pop() elif len(dtypes) == 1: - return dtypes.pop() + dtype = dtypes.pop() else: # E.g., mixed integer arithmetic - return max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None) + dtype = max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None) + + return dtypes_vector_mapper[dtype, n] if n > 1 else dtype def extract_dtype(expr): - """Extract the "winning" dtype from an expression""" - dtypes = {getattr(e, 'dtype', None) - for e in expr.free_symbols} + """ + Extract the result dtype, retaining typed vector operands such as float2. + """ + # Arithmetic nodes may expose only their Functions' scalar precision + if not (expr.is_Add or expr.is_Mul or expr.is_Pow): + try: + dtype = expr.dtype + except AttributeError: + pass + else: + if dtype is not None: + return dtype + + dtypes = {extract_dtype(i) for i in expr.args} - {None} - return infer_dtype(dtypes - {None}) + return infer_dtype(dtypes) diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 1e6ca94223..6305ebd6aa 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -14,7 +14,12 @@ IndexSum, Weights, interp_for_fd ) from devito.finite_differences.finite_difference import _PRECISION -from devito.symbolics import indexify, retrieve_indexed +from devito.ir import Cluster, SymbolRegistry +from devito.passes.clusters.derivatives import (_lower_index_derivatives, + deriv_schedule_registry) +from devito.symbolics import DefFunction, ListInitializer, indexify, retrieve_indexed +from devito.tools import dtypes_vector_mapper +from devito.types import Array, Temp from devito.types.dimension import StencilDimension from devito.types.grid import DerivedSubDomain from devito.warnings import DevitoWarning @@ -1191,6 +1196,60 @@ def test_index_derivative(self): assert IndexDerivative(vi0*w, {x: i}, properties=idxder1.properties) == vi1 + @pytest.mark.parametrize('width', [1, 2]) + @pytest.mark.parametrize('reuse', [False, True]) + @pytest.mark.parametrize('initial', [0, 1]) + def test_index_derivative_vector(self, width, reuse, initial, monkeypatch): + grid = Grid((10,)) + x, = grid.dimensions + i = StencilDimension('i', 0, 2) + dtype = dtypes_vector_mapper[np.float32, width] + v = Temp(name='v', dtype=dtype) + w = Weights(name='w', dimensions=i, initvalue=[-1., 0., 1.]) + prop = IndexDerivativeProperty('initial') + monkeypatch.setitem(deriv_schedule_registry, prop, + lambda e: (sympify(initial), e)) + ideriv = IndexDerivative(w[i]*v, {x: i}, properties=(prop,)) + + assert ideriv._rebuild() == ideriv + assert ideriv.dtype == dtype + replacement = Temp(name='replacement', dtype=np.float64) + assert ideriv.xreplace({v: replacement}).dtype == np.float64 + + if reuse: + out = Temp(name='out', dtype=dtype) + else: + out = Array(name='out', dimensions=(x,), shape=(10,), dtype=dtype)[x] + + # Exercise the generic reduction lowering without a vector-arithmetic backend + clusters, _, _ = _lower_index_derivatives( + [Cluster(Eq(out, ideriv))], sregistry=SymbolRegistry(), subs={} + ) + init = clusters[0].exprs[0] + increment = clusters[1].exprs[0] + assert init.lhs.dtype == dtype + assert init.rhs == (initial if width == 1 else + ListInitializer(*([initial]*width))) + assert increment.lhs == init.lhs + assert (init.lhs == out) == reuse + + def test_index_derivative_incompatible_reuse(self): + grid = Grid((10,)) + x, = grid.dimensions + i = StencilDimension('i', 0, 2) + v = Temp(name='v', dtype=np.float32) + out = Temp(name='out', dtype=dtypes_vector_mapper[np.float32, 2]) + w = Weights(name='w', dimensions=i, initvalue=[-1., 0., 1.]) + ideriv = IndexDerivative(w[i]*v, {x: i}) + + # A vector result cannot be reused for a scalar reduction in one component + expr = Eq(out, DefFunction('pack', (ideriv, 0))) + clusters, _, _ = _lower_index_derivatives( + [Cluster(expr)], sregistry=SymbolRegistry(), subs={} + ) + init = clusters[0].exprs[0] + assert init.lhs != out and init.lhs.dtype == np.float32 + def test_dx2(self): grid = Grid(shape=(4, 4)) diff --git a/tests/test_dimension.py b/tests/test_dimension.py index 3e28210fe5..d5397b75e4 100644 --- a/tests/test_dimension.py +++ b/tests/test_dimension.py @@ -2074,6 +2074,28 @@ def test_factor_and_condition(self): for t in range(buffer_size): assert np.all(usaved.data[t] == t*factor + bounds[0] - 1) + @pytest.mark.parametrize('legacy', [False, True]) + @pytest.mark.parametrize('array', [False, True]) + @pytest.mark.parametrize('constant', [False, True]) + def test_override_function_and_factor(self, legacy, array, constant): + """An explicit factor takes precedence over a replacement's defaults.""" + grid = Grid(shape=(4, 4)) + factor = Constant(name='factor', dtype=np.int32, value=4) if legacy else 4 + cd = ConditionalDimension(name='ts', parent=grid.time_dim, factor=factor) + u = TimeFunction(name='u', grid=grid, time_order=0, space_order=0) + usave = TimeFunction(name='usave', grid=grid, time_dim=cd, save=4, + space_order=0) + replacement = usave.func(name='replacement') + target = replacement.data if array else replacement + override = Constant(name='override', dtype=np.int32, value=2) if constant else 2 + + op = Operator([Eq(usave, u), Eq(u, u + 1)]) + op.apply(usave=target, time_M=7, **{cd.symbolic_factor.name: override}) + + assert np.all(replacement.data == np.arange(4).reshape(4, 1, 1)*2) + assert np.all(usave.data == 0) + assert cd.factor_data == 4 + @pytest.mark.parametrize('factor', [1, 2, 3]) @pytest.mark.parametrize('relation', [Or, 'strict']) def test_factor_and_condeq_overlap(self, factor, relation): diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py index 305bc23f8c..abab01efd7 100644 --- a/tests/test_dtypes.py +++ b/tests/test_dtypes.py @@ -1,3 +1,4 @@ +import pickle from pathlib import Path import numpy as np @@ -18,8 +19,11 @@ from devito.passes.iet.languages.CXX import CXXPrinter from devito.passes.iet.languages.openacc import AccBB, AccPrinter from devito.passes.iet.languages.openmp import CXXOmpBB, OmpBB +from devito.symbolics import sympy_dtype from devito.symbolics.extended_dtypes import ctypes_vector_mapper -from devito.tools import dtype_to_cstr +from devito.tools import (dtype_to_cstr, dtypes_vector_mapper, extract_dtype, + infer_dtype) +from devito.types import ComponentAccess from devito.types.basic import Basic, Scalar, Symbol from devito.types.dense import TimeFunction from devito.types.sparse import SparseTimeFunction @@ -79,6 +83,19 @@ def kw_id(kwargs): return "-".join(f'{k}' for k in kwargs.values()) +def test_vector_inference(): + dtype = dtypes_vector_mapper[np.float32, 2] + v = Symbol(name='v', dtype=dtype) + a = Symbol(name='a', dtype=np.float32) + expr = 2*v + a + + assert extract_dtype(expr) == dtype + assert extract_dtype(ComponentAccess(v, 0)) == np.float32 + assert sympy_dtype(expr) == np.float32 + assert infer_dtype([pickle.loads(pickle.dumps(dtype)), np.float32]) == dtype + assert infer_dtype([dtype, np.float64]) == dtypes_vector_mapper[np.float64, 2] + + @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) @pytest.mark.parametrize('kwargs', _configs, ids=kw_id) def test_dtype_mapping(dtype: np.dtype[np.inexact], kwargs: dict[str, str], diff --git a/tests/test_iet.py b/tests/test_iet.py index 99a94ed24c..dc7117cba7 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -564,6 +564,21 @@ def make_foo(name, start): assert [i.name for i in calls] == ['foo0', 'foo0'] +@pytest.mark.parametrize('shared', [False, True]) +def test_reuse_efuncs_constant_arrays(shared): + d = Dimension(name='d') + a = Array(name='a', dimensions=d, shape=(2,), scope='constant', + initvalue=(1, 2)) + b = a if shared else a._rebuild(name='b') + foo0 = make_callable('foo0', Call('consume', (a,))) + foo1 = make_callable('foo1', Call('consume', (b,))) + root = make_callable('root', List(body=[Call('foo0'), Call('foo1')])) + efuncs = reuse_efuncs(root, {i.name: i for i in (root, foo0, foo1)}) + + assert len(efuncs) == (2 if shared else 3) + assert FindNodes(Call).visit(efuncs['foo0'])[0].arguments == (a,) + + def test_sorted_efuncs_natural_numbering(): names = ['kernel3', 'foo0', 'kernel11', 'kernel1', 'kernel2'] efuncs = [make_callable(name, List()) for name in names]