Elias/dev

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

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<P>`` 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] = []

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
 
 

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)