diff --git a/fp_arena/__init__.py b/fp_arena/__init__.py index a1a931f..6375dbc 100644 --- a/fp_arena/__init__.py +++ b/fp_arena/__init__.py @@ -6,65 +6,77 @@ Importing this package registers the FP-Arena types (e.g. ``float32sr``, ``float64sr``) into DaCe without modifying DaCe itself, so they can be used like any built-in DaCe scalar type. Call :func:`enable_fp_arena_extensions` on an -SDFG before compiling it to pull in the matching C++ headers. +SDFG before compiling it to attach the DaCe environments that pull in the +matching C++ headers and link flags. """ from fp_arena.dtypes import ( + FP_ARENA_TYPECLASSES, Float32sr, Float64sr, float32sr, float64sr, mpfr, register, - FP_ARENA_TYPECLASSES, ) +from fp_arena.environments import INCLUDE_DIR, MPFR, FPArenaSR from fp_arena.extensions import ( - enable_fp_arena_extensions, - enable_auto_extensions, + attach_environments, disable_auto_extensions, disable_fast_math, - inject_headers, + enable_auto_extensions, + enable_fp_arena_extensions, + patch_aligned_heap_allocation, + patch_memcpy_copies, precise_math, + required_environments, uses_fp_arena_types, - fp_arena_global_code, - INCLUDE_DIR, ) -from fp_arena.transformations.change_fp_types import change_fptype from fp_arena.transformations.change_and_propagate_fp_types import ( DEFAULT_PROMOTION_RULES, change_and_propagate_fp_types, ) +from fp_arena.transformations.change_fp_types import change_fptype # Register the types and the SDFG convenience method on import (idempotent). register() +patch_aligned_heap_allocation() +patch_memcpy_copies() + import dace as _dace if not hasattr(_dace.SDFG, "enable_fp_arena_extensions"): - _dace.SDFG.enable_fp_arena_extensions = lambda self: enable_fp_arena_extensions(self) + _dace.SDFG.enable_fp_arena_extensions = lambda self: enable_fp_arena_extensions( + self + ) # Automatically enable FP-Arena for any SDFG that uses its types, so the # explicit call above becomes optional. Disable with disable_auto_extensions(). enable_auto_extensions() __all__ = [ + "DEFAULT_PROMOTION_RULES", + "FP_ARENA_TYPECLASSES", + "INCLUDE_DIR", + "MPFR", + "FPArenaSR", "Float32sr", "Float64sr", + "attach_environments", + "change_and_propagate_fp_types", + "change_fptype", + "disable_auto_extensions", + "disable_fast_math", + "enable_auto_extensions", + "enable_fp_arena_extensions", "float32sr", "float64sr", "mpfr", - "register", - "FP_ARENA_TYPECLASSES", - "enable_fp_arena_extensions", - "enable_auto_extensions", - "disable_auto_extensions", - "disable_fast_math", - "inject_headers", + "patch_aligned_heap_allocation", + "patch_memcpy_copies", "precise_math", + "register", + "required_environments", "uses_fp_arena_types", - "fp_arena_global_code", - "INCLUDE_DIR", - "change_fptype", - "change_and_propagate_fp_types", - "DEFAULT_PROMOTION_RULES", ] diff --git a/fp_arena/environments.py b/fp_arena/environments.py new file mode 100644 index 0000000..d991fef --- /dev/null +++ b/fp_arena/environments.py @@ -0,0 +1,62 @@ +# Copyright 2019-2026 ETH Zurich and the FP-Arena authors. All rights reserved. +""" +DaCe library environments for the FP-Arena C++ runtime. +""" + +from __future__ import annotations + +import os +from typing import ClassVar + +import dace.library + +INCLUDE_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "runtime", "include" +) + + +@dace.library.environment +class FPArenaSR: + """ + The header-only stochastic-rounding value types (``fp_arena::float32sr`` and ``fp_arena::float64sr``). + """ + + cmake_minimum_version: ClassVar[str | None] = None + cmake_packages: ClassVar[list] = [] + cmake_variables: ClassVar[dict] = {} + cmake_includes: ClassVar[list] = [INCLUDE_DIR] + cmake_libraries: ClassVar[list] = [] + cmake_compile_flags: ClassVar[list] = [] + cmake_link_flags: ClassVar[list] = [] + cmake_files: ClassVar[list] = [] + + headers: ClassVar[dict] = { + "frame": ["fp_arena/float32sr.h", "fp_arena/float64sr.h"], + "cuda": ["fp_arena/float32sr.h", "fp_arena/float64sr.h"], + } + state_fields: ClassVar[list] = [] + init_code: ClassVar[str] = "" + finalize_code: ClassVar[str] = "" + dependencies: ClassVar[list] = [] + + +@dace.library.environment +class MPFR: + """ + The ``dace::mpfr

`` wrapper and the ``libmpfr`` it calls into. + """ + + cmake_minimum_version: ClassVar[str | None] = None + cmake_packages: ClassVar[list] = [] + cmake_variables: ClassVar[dict] = {} + cmake_includes: ClassVar[list] = [INCLUDE_DIR] + cmake_libraries: ClassVar[list] = ["mpfr"] + cmake_compile_flags: ClassVar[list] = [] + cmake_link_flags: ClassVar[list] = [] + cmake_files: ClassVar[list] = [] + + headers: ClassVar[dict] = {"frame": ["fp_arena/mpfr.h"]} + state_fields: ClassVar[list] = [] + init_code: ClassVar[str] = "" + finalize_code: ClassVar[str] = "" + dependencies: ClassVar[list] = [] diff --git a/fp_arena/experiment/registry.py b/fp_arena/experiment/registry.py index 92143f0..2e29296 100644 --- a/fp_arena/experiment/registry.py +++ b/fp_arena/experiment/registry.py @@ -34,7 +34,7 @@ def to_typeclass(key: str) -> dace.dtypes.typeclass: def is_mpfr(key: str) -> bool: - """Whether ``key`` names an MPFR precision (which requires linking libmpfr).""" + """Whether ``key`` names an MPFR precision.""" return _MPFR_KEY.match(key) is not None diff --git a/fp_arena/experiment/retarget.py b/fp_arena/experiment/retarget.py index a0cdd49..5ba711e 100644 --- a/fp_arena/experiment/retarget.py +++ b/fp_arena/experiment/retarget.py @@ -54,13 +54,6 @@ def _validate_pins(sdfg: dace.SDFG, pin_map: PrecisionMap) -> None: raise ValueError(f"Pinned array {name!r} is not a floating-point array") -def _ensure_mpfr_linked() -> None: - """Add the MPFR library to DaCe's CPU link line.""" - libs = dace.Config.get("compiler", "cpu", "libs") or "" - if "mpfr" not in libs.split(): - dace.Config.append("compiler", "cpu", "libs", value=" mpfr") - - def apply_precision( sdfg: dace.SDFG, pin_map: PrecisionMap, @@ -72,8 +65,6 @@ def apply_precision( if not pin_map: return _validate_pins(sdfg, pin_map) - if any(registry.is_mpfr(key) for key in pin_map.values()): - _ensure_mpfr_linked() typed = {name: registry.to_typeclass(key) for name, key in pin_map.items()} change_and_propagate_fp_types(sdfg, typed, promotion_rules) diff --git a/fp_arena/extensions.py b/fp_arena/extensions.py index f49867d..ee6a24a 100644 --- a/fp_arena/extensions.py +++ b/fp_arena/extensions.py @@ -16,19 +16,16 @@ Toggle with :func:`enable_auto_extensions` / :func:`disable_auto_extensions`. * Explicit: call :func:`enable_fp_arena_extensions` on an SDFG (it also strips fast-math from DaCe's config as a persistent default). - -The headers are referenced by absolute path, so no ``-I`` flag or config change -is required -- this keeps the integration fully non-invasive. """ -import os +from collections.abc import Iterator from contextlib import contextmanager import dace from dace.config import Config -#: Absolute path to the bundled C++ include root (``.../runtime/include``). -INCLUDE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runtime", "include") +from fp_arena.dtypes import mpfr +from fp_arena.environments import MPFR, FPArenaSR #: Fast-math flags removed when FP-Arena is enabled (incompatible with the #: exact IEEE rounding that stochastic rounding depends on). ``Config.get`` @@ -39,48 +36,138 @@ #: DaCe compiler-argument config paths that may carry fast-math flags. _COMPILER_ARG_PATHS = (("compiler", "cpu", "args"), ("compiler", "cuda", "args")) -#: Headers injected into generated code. -_HEADERS = ( - os.path.join(INCLUDE_DIR, "fp_arena", "float32sr.h"), - os.path.join(INCLUDE_DIR, "fp_arena", "float64sr.h"), - os.path.join(INCLUDE_DIR, "fp_arena", "mpfr.h"), +#: The FP-Arena environments, each with the C-type substrings that require it +#: (``fp_arena::`` for the SR types, ``dace::mpfr`` for MPFR). +_ENVIRONMENTS = ( + (FPArenaSR, ("fp_arena::",)), + (MPFR, ("dace::mpfr",)), ) -#: Backends whose global-code section receives the includes (CPU frame + CUDA). -_BACKENDS = ("frame", "cuda") -#: Marker so the includes are only injected once per SDFG. -_GUARD = "// fp_arena extensions enabled" +def _type_strings(sdfg: dace.SDFG) -> Iterator[str]: + """ + :param sdfg: the SDFG to inspect. + :returns: an iterator over every string in ``sdfg`` and its nested SDFGs that + may name a C type: data descriptor C types and tasklet bodies. + """ + from dace.sdfg import nodes as _dnodes -#: Substrings identifying an FP-Arena C type (fp_arena:: for SR types, dace::mpfr for mpfr). -_CTYPE_MARKERS = ("fp_arena::", "dace::mpfr") + for nested in sdfg.all_sdfgs_recursive(): + for desc in nested.arrays.values(): + yield getattr(desc.dtype, "ctype", "") or "" + for state in nested.states(): + for node in state.nodes(): + if isinstance(node, _dnodes.Tasklet): + try: + yield node.code.as_string + except AttributeError: + yield str(node.code) -def fp_arena_global_code() -> str: +def required_environments(sdfg: dace.SDFG) -> set[str]: """ - :returns: the C++ ``#include`` block (absolute paths) for the SR headers. + :param sdfg: the SDFG to inspect. + :returns: the full class paths (DaCe's environment identifiers) of the + FP-Arena environments the types used by ``sdfg`` need -- empty if it uses + none of them. """ - includes = "\n".join('#include "%s"' % h for h in _HEADERS) - return "%s\n%s\n" % (_GUARD, includes) + strings = list(_type_strings(sdfg)) + return { + env.full_class_path() + for env, markers in _ENVIRONMENTS + if any(marker in s for s in strings for marker in markers) + } -def inject_headers(sdfg: dace.SDFG) -> dace.SDFG: +def attach_environments(sdfg: dace.SDFG) -> dace.SDFG: """ - Inject the SR header ``#include``s into ``sdfg``'s global code (CPU and CUDA). - Idempotent. + Attach the FP-Arena environments that ``sdfg`` needs to its code nodes. - :param sdfg: the SDFG to inject into. + :param sdfg: the SDFG to attach to. :returns: the same SDFG, for chaining. """ - code = fp_arena_global_code() - for backend in _BACKENDS: - existing = sdfg.global_code.get(backend) - if existing is not None and _GUARD in existing.code: - continue - sdfg.append_global_code(code, backend) + from dace.sdfg import nodes as _dnodes + + envs = required_environments(sdfg) + if not envs: + return sdfg + for nested in sdfg.all_sdfgs_recursive(): + for state in nested.states(): + for node in state.nodes(): + if isinstance(node, _dnodes.CodeNode): + node.environments = frozenset(node.environments) | envs return sdfg +#: Whether the aligned-allocation patch is installed. +_aligned_patch_installed = False + + +def patch_aligned_heap_allocation() -> bool: + """Route mpfr heap arrays through plain ``new[]`` / ``delete[]``, which run destructors.""" + + global _aligned_patch_installed + if _aligned_patch_installed: + return False + + from dace.codegen.targets import cpu, experimental_cpu + + original = cpu.use_aligned_operator_new + + def _use_aligned_operator_new(desc) -> bool: + if isinstance(desc.dtype, mpfr): + return False + return original(desc) + + cpu.use_aligned_operator_new = _use_aligned_operator_new + experimental_cpu.use_aligned_operator_new = _use_aligned_operator_new + + _aligned_patch_installed = True + return True + + +#: Copy implementations that emit a raw ``memcpy``, which shallow-copies mpfr's limb pointer. +_MEMCPY_IMPLEMENTATIONS = frozenset( + {"MemcpyCPU", "MemcpyCUDA1D", "MemcpyCUDA2D", "MemcpyCUDANDStrided"} +) + +#: Whether the element-wise copy patch is installed. +_copy_patch_installed = False + + +def patch_memcpy_copies() -> bool: + """Route mpfr array copies through element-wise assignment, which deep-copies.""" + + global _copy_patch_installed + if _copy_patch_installed: + return False + + from dace.libraries.standard.nodes import copy_node + from dace.transformation.passes.canonicalize import finalize + + original = copy_node.select_copy_implementation + + def _select_copy_implementation(node, parent_state) -> str: + impl = original(node, parent_state) + if impl not in _MEMCPY_IMPLEMENTATIONS: + return impl + _, inp, in_subset, _, _, out_subset = node.validate( + parent_state.sdfg, parent_state, allow_cross_storage=True + ) + if not isinstance(inp.dtype, mpfr): + return impl + single = ( + in_subset.num_elements_exact() == 1 and out_subset.num_elements_exact() == 1 + ) + return "Tasklet" if single else "MappedTasklet" + + copy_node.select_copy_implementation = _select_copy_implementation + finalize.select_copy_implementation = _select_copy_implementation + + _copy_patch_installed = True + return True + + def _strip_fast_math(args: str) -> str: """:returns: ``args`` with any fast-math flag removed.""" return " ".join(tok for tok in args.split() if tok not in _FAST_MATH_FLAGS) @@ -128,18 +215,18 @@ def precise_math(): def enable_fp_arena_extensions(sdfg: dace.SDFG) -> dace.SDFG: """ - Make ``sdfg`` compile with the FP-Arena types by injecting the SR header - includes into its generated global code, and remove ``-ffast-math`` from - DaCe's compiler flags (persistently; see :func:`disable_fast_math`). + Make ``sdfg`` compile with the FP-Arena types by attaching the FP-Arena + environments and remove ``-ffast-math`` from DaCe's compiler flags (persistently; + see :func:`disable_fast_math`). - Idempotent. Applies to the CPU and CUDA backends. With automatic enablement - active (the default), calling this explicitly is optional. + Idempotent. With automatic enablement active (the default), calling this + explicitly is optional. :param sdfg: the SDFG to enable FP-Arena types for. :returns: the same SDFG, for chaining. """ disable_fast_math() - inject_headers(sdfg) + attach_environments(sdfg) return sdfg @@ -149,21 +236,7 @@ def uses_fp_arena_types(sdfg: dace.SDFG) -> bool: :returns: ``True`` if any data descriptor or tasklet body in ``sdfg`` or its nested SDFGs references an FP-Arena C type, ``False`` otherwise. """ - from dace.sdfg import nodes as _dnodes - for nested in sdfg.all_sdfgs_recursive(): - for desc in nested.arrays.values(): - if any(m in (getattr(desc.dtype, "ctype", "") or "") for m in _CTYPE_MARKERS): - return True - for state in nested.states(): - for node in state.nodes(): - if isinstance(node, _dnodes.Tasklet): - try: - code_str = node.code.as_string - except AttributeError: - code_str = str(node.code) - if any(m in code_str for m in _CTYPE_MARKERS): - return True - return False + return bool(required_environments(sdfg)) #: Whether the automatic wrappers are currently installed. @@ -181,11 +254,11 @@ def enable_auto_extensions(): Two wrappers are installed, each at the layer its concern belongs to: - * Header injection (a code-generation concern) wraps the single codegen - entry point ``dace.codegen.codegen.generate_code``, so headers are added - on every path -- ``compile``, ``SDFG.generate_code``, or a direct codegen - call. ``compile`` code-generates a deep copy, so the caller's SDFG object - is left untouched. + * Attaching the environments (a code-generation concern) wraps the single + codegen entry point ``dace.codegen.codegen.generate_code``, so they are + attached on every path -- ``compile``, ``SDFG.generate_code``, or a direct + codegen call. ``compile`` code-generates a deep copy, so the caller's SDFG + object is left untouched. * Removing ``-ffast-math`` (a build concern) wraps ``SDFG.compile``, scoped to the build of FP-Arena SDFGs only. @@ -199,8 +272,7 @@ def enable_auto_extensions(): _original_generate_code = codegen.generate_code def _generate_code(sdfg, *args, **kwargs): - if uses_fp_arena_types(sdfg): - inject_headers(sdfg) + attach_environments(sdfg) return _original_generate_code(sdfg, *args, **kwargs) codegen.generate_code = _generate_code @@ -223,6 +295,7 @@ def disable_auto_extensions(): if not _auto_installed: return from dace.codegen import codegen + codegen.generate_code = _original_generate_code dace.SDFG.compile = _original_compile _original_generate_code = None diff --git a/tests/test_auto_enable.py b/tests/test_auto_enable.py index c7b3efe..4b1f27a 100644 --- a/tests/test_auto_enable.py +++ b/tests/test_auto_enable.py @@ -5,9 +5,9 @@ without an explicit call. """ +import dace import numpy as np import pytest -import dace from dace.codegen.exceptions import CompilationError import fp_arena @@ -21,9 +21,14 @@ def _cast_sdfg_without_enable(n: int) -> dace.SDFG: state = sdfg.add_state() a = state.add_read("A") c = state.add_write("C") - me, mx = state.add_map("cast", dict(i=f"0:{n}")) - tasklet = state.add_tasklet("cast", {"inp"}, {"out"}, "out = static_cast(inp);", - dace.Language.CPP) + me, mx = state.add_map("cast", {"i": f"0:{n}"}) + tasklet = state.add_tasklet( + "cast", + {"inp"}, + {"out"}, + "out = static_cast(inp);", + dace.Language.CPP, + ) state.add_memlet_path(a, me, tasklet, dst_conn="inp", memlet=dace.Memlet("A[i]")) state.add_memlet_path(tasklet, mx, c, src_conn="out", memlet=dace.Memlet("C[i]")) return sdfg @@ -38,13 +43,33 @@ def test_auto_enable_runs_without_explicit_enable(): assert np.all(C == np.float32(0.5)) -def test_headers_injected_at_codegen_not_compile(): - # Header injection is a codegen concern: invoking codegen directly (no - # compile) must already contain the FP-Arena includes. +def test_headers_included_at_codegen_not_compile(): + # Attaching the environment is a codegen concern: invoking codegen directly + # (no compile) must already contain the FP-Arena includes, by a path + # relative to the environment's include root rather than an absolute one. sdfg = _cast_sdfg_without_enable(8) generated = "\n".join(co.code for co in sdfg.generate_code()) - assert "fp_arena extensions enabled" in generated - assert "float32sr.h" in generated + assert '#include "fp_arena/float32sr.h"' in generated + assert fp_arena.INCLUDE_DIR not in generated + + +def test_environment_attached_to_code_nodes(): + sdfg = _cast_sdfg_without_enable(8) + assert fp_arena.required_environments(sdfg) == {"fp_arena.environments.FPArenaSR"} + + fp_arena.attach_environments(sdfg) + tasklets = [ + n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.Tasklet) + ] + assert tasklets + assert all("fp_arena.environments.FPArenaSR" in n.environments for n in tasklets) + + +def test_mpfr_environment_carries_the_link_flag(): + sdfg = dace.SDFG("mpfr_env") + sdfg.add_array("A", [8], dace.mpfr(128)) + assert fp_arena.required_environments(sdfg) == {"fp_arena.environments.MPFR"} + assert "mpfr" in fp_arena.MPFR.cmake_libraries def test_uses_fp_arena_types_detection(): diff --git a/tests/test_mpfr.py b/tests/test_mpfr.py index e9cb21d..1ca7bc3 100644 --- a/tests/test_mpfr.py +++ b/tests/test_mpfr.py @@ -1,12 +1,10 @@ -import numpy as np import dace -import fp_arena +import numpy as np + from fp_arena.transformations.change_and_propagate_fp_types import ( change_and_propagate_fp_types, ) -dace.Config.append("compiler", "cpu", "libs", value="mpfr") - def test_sdfg_scalar_compute(): """Compute 1/3 at 128-bit precision and return as double."""