diff --git a/devito/arch/archinfo.py b/devito/arch/archinfo.py index c294f5def4..ccaed53cfe 100644 --- a/devito/arch/archinfo.py +++ b/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): diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py index 2b8b6d5080..6f2e4b1b82 100644 --- a/devito/finite_differences/differentiable.py +++ b/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 diff --git a/devito/ir/cgen/printer.py b/devito/ir/cgen/printer.py index 27ee0e3112..a28760f63b 100644 --- a/devito/ir/cgen/printer.py +++ b/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): diff --git a/devito/passes/clusters/derivatives.py b/devito/passes/clusters/derivatives.py index fceeddb8bf..8095e72b3d 100644 --- a/devito/passes/clusters/derivatives.py +++ b/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)) 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..4903ef60a6 100644 --- a/devito/symbolics/inspection.py +++ b/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 diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index f3033aa6e9..02f232e2b8 100644 --- a/devito/tools/dtypes_lowering.py +++ b/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) diff --git a/tests/test_arch.py b/tests/test_arch.py index 22d084cdc1..305887d2e2 100644 --- a/tests/test_arch.py +++ b/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') diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index 1e6ca94223..ee93c00488 100644 --- a/tests/test_derivatives.py +++ b/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)) diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py index 305bc23f8c..ef0b6f41e8 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,10 @@ 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 DefFunction, 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 +82,27 @@ 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([dtype, np.float32]) == np.float32 + assert infer_dtype([pickle.loads(pickle.dumps(dtype))], vector=True) == dtype + assert (infer_dtype([dtype, np.float64], vector=True) == + dtypes_vector_mapper[np.float64, 2]) + assert infer_dtype([dtype, np.complex64], vector=True) == np.complex64 + + f = Function(name='f', grid=Grid(shape=(4,))) + for expr in (f + v, f*v, (f + v)**2): + assert expr.dtype == extract_dtype(expr) == dtype + assert (f + ComponentAccess(v, 0)).dtype == np.float32 + + @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], @@ -245,6 +269,15 @@ def test_printer_registry() -> None: assert get_printer(CPrinter, np.float16) is float16 +@pytest.mark.parametrize('cls', [CPrinter, CXXPrinter]) +@pytest.mark.parametrize('dtype,suffix', [(np.float32, 'F'), (np.float64, '')]) +def test_vector_operand_precision(cls, dtype, suffix): + printer = cls({'dtype': dtypes_vector_mapper[dtype, 2]}) + a = Symbol(name='a', dtype=dtype) + expr = DefFunction('make_float2', [a + sympy.Float('0.1'), sympy.Float('0.2')]) + assert printer.doprint(expr) == f'make_float2(a + 1.0e-1{suffix},2.0e-1{suffix})' + + @pytest.mark.parametrize('dtype', [np.complex64, np.complex128]) def test_complex_override(dtype: np.dtype[np.complexfloating]) -> None: """ 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]