compiler: Misc dtype lowering tweaks

devito/arch/archinfo.py

@@ -1174,9 +1174,12 @@ def supports(self, query, language=None):
         elif query == 'async-pipe' and cc >= 80:
             # Asynchronous pipeline loads -- introduced in Ampere
             return True
-        elif query in ('tma', 'thread-block-cluster') and cc >= 90:  # noqa: SIM103
+        elif query in ('tma', 'thread-block-cluster') and cc >= 90:
             # Tensor Memory Accelerator -- introduced in Hopper
             return True
+        elif query == 'fp32x2' and cc >= 100:  # noqa: SIM103
+            # Paired FP32 arithmetic -- introduced in Blackwell
+            return True
         else:
             return False
 
@@ -1208,7 +1211,12 @@ def supports(self, query, language=None):
 
 
 class Blackwell(Hopper):
-    pass
+
+    def supports(self, query, language=None):
+        if query == 'fp32x2':
+            return language == 'cuda'
+        else:
+            return super().supports(query, language)
 
 
 class AmdDevice(Device):

devito/finite_differences/differentiable.py

@@ -22,7 +22,7 @@
 from devito.logger import warning
 from devito.tools import (
     Pickable, Tag, as_tuple, extract_dtype, filter_ordered, flatten, frozendict,
-    infer_dtype, is_integer, is_number, memoized_func, split
+    infer_dtype, infer_dtype_from_args, is_integer, is_number, memoized_func, split
 )
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 from devito.types.basic import AbstractFunction, Indexed
@@ -591,6 +591,11 @@ def __new__(cls, *args, **kwargs):
 
         return obj
 
+    @cached_property
+    def dtype(self):
+        # Typed operands may carry vector types or casts absent from `_functions`
+        return infer_dtype_from_args(self.args)
+
     def subs(self, *args, **kwargs):
         return self.func(
             *[getattr(a, 'subs', lambda x: a)(*args, **kwargs)  # noqa: B023
@@ -898,7 +903,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 +1326,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/cgen/printer.py

@@ -20,7 +20,8 @@
 from devito.symbolics.inspection import has_integer_args, sympy_dtype
 from devito.symbolics.queries import q_leaf
 from devito.tools import (
-    ctypes_to_cstr, ctypes_vector_mapper, dtype_to_ctype, memoized_func
+    ctypes_to_cstr, ctypes_vector_mapper, dtype_to_ctype, dtypes_vector_mapper,
+    memoized_func
 )
 from devito.types.basic import AbstractFunction
 
@@ -62,9 +63,12 @@ class BasePrinter(CodePrinter):
     @property
     def dtype(self):
         try:
-            return self._settings['dtype'].nptype
+            dtype = self._settings['dtype'].nptype
         except AttributeError:
-            return self._settings['dtype']
+            dtype = self._settings['dtype']
+
+        # Use scalar precision in vector operands, e.g. `make_float2(a + 0.1F, b)`
+        return dtypes_vector_mapper.get_base_dtype(dtype, dtype)
 
     @property
     def compiler(self):

devito/passes/clusters/derivatives.py

@@ -6,9 +6,10 @@
 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 +174,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

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

devito/tools/dtypes_lowering.py

@@ -15,8 +15,8 @@
            'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype',
            'dtype_to_cstr', 'dtype_to_ctype', 'infer_datasize', 'dtype_to_mpitype',
            'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper',
-           'is_external_ctype', 'infer_dtype', 'extract_dtype', 'CustomDtype',
-           'mpi4py_mapper', 'as_fp64_decimal']
+           'is_external_ctype', 'infer_dtype', 'infer_dtype_from_args', 'extract_dtype',
+           'CustomDtype', 'mpi4py_mapper', 'as_fp64_decimal']
 
 
 def as_fp64_decimal(v):
@@ -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:
@@ -357,14 +357,19 @@ def is_numpy_dtype(dtype):
         return False
 
 
-def infer_dtype(dtypes):
+def infer_dtype(dtypes, vector=False):
     """
     Given a set of np.dtypes, return the "winning" dtype:
 
         * 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.
+
+    If `vector=True`, preserve vector length when a matching vector dtype exists.
     """
+    dtypes = set(dtypes)
+    n = max(map(dtype_len, dtypes), default=1) if vector else 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,37 @@ 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.get((dtype, n), 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.
+    """
+    try:
+        dtype = expr.dtype
+    except AttributeError:
+        pass
+    else:
+        if dtype is not None:
+            return dtype
+
+    return infer_dtype_from_args(expr.args)
+
+
+def infer_dtype_from_args(args):
+    """
+    Infer the result dtype from expression operands, retaining vector types.
+    """
+    dtypes = {extract_dtype(i) for i in args} - {None}
 
-    return infer_dtype(dtypes - {None})
+    return infer_dtype(dtypes, vector=True)

tests/test_arch.py

@@ -5,7 +5,7 @@
 from packaging.version import Version
 
 from devito import configuration, switchconfig
-from devito.arch import BLACKWELL, NVIDIAX
+from devito.arch import BLACKWELL, NVIDIAX, platform_registry
 from devito.arch.archinfo import AppleArm, get_apple_platform
 from devito.arch.compiler import (
     CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version
@@ -108,3 +108,22 @@ def test_cuda_target(monkeypatch, platform, cc, target):
 
     flags = [i for i in compiler.cflags if i.startswith(('-arch=', '-gencode='))]
     assert flags == [target]
+
+
+@pytest.mark.parametrize('language', [None, 'cuda', 'openacc', 'hip', 'sycl'])
+@pytest.mark.parametrize('platform,cc,expected', [
+    ('nvidiaX', None, False),
+    ('nvidiaX', 90, False),
+    ('nvidiaX', 100, True),
+    ('nvidiaX', 120, True),
+    ('cpu64', 100, False),
+    ('volta', 100, True),
+    ('ampere', 100, True),
+    ('hopper', 100, True),
+    ('blackwell', None, True),
+    ('blackwell', 90, True)
+])
+def test_fp32x2_support(platform, cc, expected, language, monkeypatch):
+    monkeypatch.setattr('devito.arch.archinfo.get_nvidia_cc', lambda: cc)
+    platform = platform_registry[platform]()
+    assert platform.supports('fp32x2', language) is (expected and language == 'cuda')

tests/test_derivatives.py

@@ -14,7 +14,11 @@
     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
+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 +1195,56 @@ 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])
+    def test_index_derivative_vector(self, width, reuse):
+        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.])
+        ideriv = IndexDerivative(w[i]*v, {x: i})
+
+        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
+        expected = 0 if width == 1 else ListInitializer(0, 0)
+        assert init.rhs == expected
+        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))