compiler: Misc dtype lowering tweaks

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

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

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:

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))
 

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:

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

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)

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))
 

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):

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],

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]