Feat: add per-function @compiled(fastmath=True) opt-in (#332)

docs/SPEC.md

@@ -739,6 +739,15 @@ All in `PM/commonTest`, so they run wherever the interpreter loads.
   `continue` and `return` meaning exactly what they mean without chunking (trip counts just below, at and
   above a multiple of the interval, negative steps, empty ranges and ranges spanning the i64 limits all
   behave as CPython's `range`). A function holding an array copy-in or entry-globals snapshot polls too (#152, option (i)): at the top of a loop iteration, once per `TP_POLL_INTERVAL`, it writes its copied lists back (CPython's own stores would already be visible there), ends their critical sections, polls, takes the critical sections again, copies the lists in again and re-checks the entry guards (list lengths, exact element types, entry globals). Unchanged: it continues compiled, which is an execution CPython can produce, since every list operation of another thread now falls before or after this point. Changed: it leaves for the interpreter at that point (an OSR-style exit) and does not deopt or raise: a continuation function generated from the function's own source, holding the rest of the function from that loop iteration (the innermost loop from its current index; each enclosing loop's remaining body, then its remaining iterations; then the statements after), receives every live local (a native value as its Python object, a local array as a new list) and each enclosing `range`'s already evaluated bounds, and the compiled call returns what it returns or raises what it raises. The continuation keeps the original file, line numbers and function name, so a traceback reads as the interpreted function's.
+  **Floating-point exactness and `@compiled(fastmath=True)` (#153, #332, owner decision 2026-10-25):**
+  by default (`fastmath=False`), floating-point compilation is bit-identical to CPython (`#pragma STDC
+  FP_CONTRACT OFF` in the generated C, `-ffp-contract=off -fno-fast-math` on GCC/Clang, `/fp:precise`
+  on MSVC). A module-level function marked `@compiled(fastmath=True)` opts into floating-point
+  reassociation and FMA contraction for that function only (never on by default, never applied to an
+  unmarked function or another function in the same module, and never inlined into a non-fastmath
+  caller without fastmath): **this is the one path in TypedPython that is explicitly NOT
+  CPython-identical** (a sum or fused multiply-add in a `@compiled(fastmath=True)` function may round
+  differently from the interpreter so float reductions and FMA loops can vectorise).
   `Status: implemented` for desktop host builds
   (`test_pipeline.py`, `test_cgen*.py`, `test_recursion.py`, `test_eval_breaker.py`, `test_fast_entry.py`,
   `test_free_threaded.py`; free-threaded 3.15t, #159). Not yet: shipping compiled modules in an app through

python-multiplatform-ksp/src/main/python/typedpython/cgen.py

@@ -494,18 +494,27 @@ def _runs_no_python(f: ir.Function) -> bool:
 def _bounded_functions(functions: dict[str, ir.Function]) -> dict[str, int]:
     """name -> k for every bounded function (module docstring, "Fast entry"): in no call cycle,
     every callee bounded, `_runs_no_python`. k(f) = 1 + max(k(callee)), 1 without callees: the most
-    compiled frames a call of `f` can have in flight at once."""
+    compiled frames a call of `f` can have in flight at once.
+    A `fastmath` function is never given an `inline` fast twin (`tp_fimpl_*`) when some other
+    function in the module calls it without `fastmath`, so a non-fastmath caller can never inline
+    a fastmath body (#332)."""
     callees: dict[str, set[str]] = {}
     for name, f in functions.items():
         callees[name] = {e.function for e in _all_exprs(f) if isinstance(e, ir.Call)}
+    called_from_exact = {
+        c
+        for caller_name, cs in callees.items()
+        if not functions[caller_name].fastmath
+        for c in cs
+    }
     k: dict[str, int | None] = {}                 # None: not bounded (also while on the DFS stack)
 
     def visit(name: str) -> int | None:
         if name in k:
             return k[name]
         k[name] = None                            # on the stack: reaching it again is a cycle
         f = functions[name]
-        if not _runs_no_python(f):
+        if not _runs_no_python(f) or (f.fastmath and name in called_from_exact):
             return None
         deepest = 0
         for c in callees[name]:
@@ -743,6 +752,21 @@ def generate(module: ir.Module, source_path: Path, display_path: str | None = No
     w("#if !defined(_MSC_VER) || defined(__clang__)    /* cl.exe warns on an unknown pragma */")
     w("#pragma STDC FP_CONTRACT OFF")
     w("#endif")
+    if any(f.fastmath for f in module.functions):
+        w("/* SPEC N-8 (#153, #332): per-function `@compiled(fastmath=True)` opt-in (NOT CPython-identical). */")
+        w("#if defined(__clang__)")
+        w("#define TP_FASTMATH_BEGIN _Pragma(\"clang fp reassociate(on) contract(fast)\")")
+        w("#define TP_FASTMATH_END _Pragma(\"clang fp reassociate(off) contract(off)\")")
+        w("#elif defined(__GNUC__)")
+        w("#define TP_FASTMATH_BEGIN _Pragma(\"GCC push_options\") _Pragma(\"GCC optimize(\\\"fast-math\\\")\")")
+        w("#define TP_FASTMATH_END _Pragma(\"GCC pop_options\")")
+        w("#elif defined(_MSC_VER)")
+        w("#define TP_FASTMATH_BEGIN __pragma(float_control(precise, off, push)) __pragma(fp_contract(on))")
+        w("#define TP_FASTMATH_END __pragma(fp_contract(off)) __pragma(float_control(pop))")
+        w("#else")
+        w("#define TP_FASTMATH_BEGIN")
+        w("#define TP_FASTMATH_END")
+        w("#endif")
     w("")
     w(f"static const char tp_source[] =\n    {c_source_init(source)};")
     w(f"static const char tp_source_path[] = {c_string(display_path)};")
@@ -2892,7 +2916,8 @@ def generate(self) -> str:
         """The impl (preceded by its fast twin when bounded) and the wrapper."""
         parts = []
         if self.rmode:
-            return self.impl()
+            text = self.impl()
+            return f"TP_FASTMATH_BEGIN\n{text}TP_FASTMATH_END\n" if self.f.fastmath else text
         if self.f.name in self.mg.bounded:
             twin = _FnGen(self.mg, self.f, fast=True)
             twin_text = twin.impl()
@@ -2901,6 +2926,8 @@ def generate(self) -> str:
         impl_text = self.impl()
         parts += self.par_defs
         parts.append(impl_text)
+        if self.f.fastmath:
+            parts = ["TP_FASTMATH_BEGIN", *parts, "TP_FASTMATH_END"]
         parts.append(self.wrapper())
         return "\n".join(parts)
 

python-multiplatform-ksp/src/main/python/typedpython/frontend.py

@@ -191,7 +191,8 @@ def lower(path: Path) -> ir.Module:
     skipped.update(class_skips)
     signatures: dict[str, _Signature] = {}
     marked: set[str] = set()
-    for node, reason, parallel in _targets(tree, source, info):
+    fastmath_marked: set[str] = set()
+    for node, reason, parallel, fastmath in _targets(tree, source, info):
         if reason is not None:
             skipped[node.name] = reason
             continue
@@ -203,6 +204,8 @@ def lower(path: Path) -> ir.Module:
         else:
             if parallel:
                 marked.add(node.name)
+            if fastmath:
+                fastmath_marked.add(node.name)
     for owner, node, reason in _method_targets(tree, info):
         q = f"{owner}.{node.name}"
         if reason is not None:
@@ -240,7 +243,11 @@ def lower(path: Path) -> ir.Module:
         else:
             skipped[name] = result
     functions.sort(key=lambda f: f.source_line)
-    functions = [dataclasses.replace(f, parallel=True) if f.name in marked else f for f in functions]
+    functions = [
+        dataclasses.replace(f, parallel=f.name in marked, fastmath=f.name in fastmath_marked)
+        if (f.name in marked or f.name in fastmath_marked) else f
+        for f in functions
+    ]
     return ir.Module(path.stem, tuple(functions), skipped, tuple(class_decls.values()))
 
 
@@ -423,30 +430,34 @@ def _is_compiled_decorator(d: ast.expr) -> bool:
     )
 
 
-def _compiled_decorator_options(d: ast.expr) -> tuple[bool, str | None]:
-    """Parse `@compiled` or `@compiled(parallel=...)` -> (parallel, error_reason_or_None)."""
+def _compiled_decorator_options(d: ast.expr) -> tuple[bool, bool, str | None]:
+    """Parse `@compiled` or `@compiled(parallel=..., fastmath=...)` -> (parallel, fastmath, error_reason_or_None)."""
     if isinstance(d, ast.Name):
-        return False, None
+        return False, False, None
     assert isinstance(d, ast.Call)
     if d.args:
-        return False, (f"line {d.lineno}: `@compiled` takes keyword arguments only "
-                       f"(got positional `{_src(d.args[0])}`)")
+        return False, False, (f"line {d.lineno}: `@compiled` takes keyword arguments only "
+                              f"(got positional `{_src(d.args[0])}`)")
     parallel = False
+    fastmath = False
     seen: set[str] = set()
     for kw in d.keywords:
         if kw.arg is None:
-            return False, f"line {d.lineno}: `**{_src(kw.value)}` in `@compiled(...)` is not supported"
-        if kw.arg != "parallel":
-            return False, (f"line {d.lineno}: `@compiled` keyword argument `{kw.arg}` is not "
-                           "supported (only `parallel=True` or `parallel=False`)")
+            return False, False, f"line {d.lineno}: `**{_src(kw.value)}` in `@compiled(...)` is not supported"
+        if kw.arg not in ("parallel", "fastmath"):
+            return False, False, (f"line {d.lineno}: `@compiled` keyword argument `{kw.arg}` is not "
+                                  "supported (only `parallel` and `fastmath` with literal `True` or `False`)")
         if kw.arg in seen:
-            return False, f"line {d.lineno}: `@compiled` keyword argument `{kw.arg}` is repeated"
+            return False, False, f"line {d.lineno}: `@compiled` keyword argument `{kw.arg}` is repeated"
         seen.add(kw.arg)
         if not (isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, bool)):
-            return False, (f"line {d.lineno}: `@compiled(parallel=...)` requires a literal `True` "
-                           f"or `False`, got `{_src(kw.value)}`")
-        parallel = kw.value.value
-    return parallel, None
+            return False, False, (f"line {d.lineno}: `@compiled({kw.arg}=...)` requires a literal `True` "
+                                  f"or `False`, got `{_src(kw.value)}`")
+        if kw.arg == "parallel":
+            parallel = kw.value.value
+        else:
+            fastmath = kw.value.value
+    return parallel, fastmath, None
 
 
 def is_opted_in(path: Path) -> bool:
@@ -614,8 +625,8 @@ def _class_of(kind: str, classes: dict[str, ir.ClassDecl]) -> tuple[str, bool] |
 
 def _targets(
     tree: ast.Module, source: str, info: _ModuleInfo,
-) -> list[tuple[ast.FunctionDef | ast.AsyncFunctionDef, str | None, bool]]:
-    found: list[tuple[ast.FunctionDef | ast.AsyncFunctionDef, str | None, bool]] = []
+) -> list[tuple[ast.FunctionDef | ast.AsyncFunctionDef, str | None, bool, bool]]:
+    found: list[tuple[ast.FunctionDef | ast.AsyncFunctionDef, str | None, bool, bool]] = []
     for node in tree.body:
         if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
             continue
@@ -624,6 +635,7 @@ def _targets(
             continue
         reason = None
         parallel = False
+        fastmath = False
         if marked and info.bound("compiled"):
             line = info.bindings["compiled"][0][0]
             reason = (f"line {node.lineno}: `compiled` is bound in this module (line {line}), "
@@ -633,18 +645,19 @@ def _targets(
                       "(only `@compiled`)")
         else:
             for d in marked:
-                par, r = _compiled_decorator_options(d)
+                par, fm, r = _compiled_decorator_options(d)
                 if r is not None:
                     reason = r
                     break
                 parallel = parallel or par
+                fastmath = fastmath or fm
         if reason is None:
             if isinstance(node, ast.AsyncFunctionDef):
                 reason = f"line {node.lineno}: async function is not supported"
             elif len(info.bindings.get(node.name, [])) > 1:
                 lines = ", ".join(str(line) for line, _ in info.bindings[node.name])
                 reason = f"line {node.lineno}: `{node.name}` is bound more than once (lines {lines})"
-        found.append((node, reason, parallel))
+        found.append((node, reason, parallel, fastmath))
     return found
 
 

python-multiplatform-ksp/src/main/python/typedpython/ir.py

@@ -820,6 +820,9 @@ class Function:
     # SPEC N-15: the user marked this function `@compiled(parallel=True)` (#358).
     # Only a marked function can hold a `ForRange.par`.
     parallel: bool = False
+    # SPEC N-8 (#153, #332): the user marked this function `@compiled(fastmath=True)` (reassociation
+    # and FMA contraction enabled for this function only; NOT CPython-identical).
+    fastmath: bool = False
     # SPEC N-12b (#268): the producers whose result this pure function keeps inside one statement
     # (`check(make(d))`), so each such call builds its objects in a region. Set only by `verify`,
     # which recomputes it and rejects a function whose declared entries differ

python-multiplatform-ksp/src/test/python/test_frontend.py

@@ -677,6 +677,43 @@ def bad(x: float) -> float:
     assert set(m.skipped) == {"bad"} and "parallel" in m.skipped["bad"]
 
 
+def test_compiled_fastmath_option_and_validation(tmp_path):
+    m = lower(tmp_path, """
+        @compiled(fastmath=True)
+        def fm(x: float) -> float:
+            return x
+
+        @compiled(fastmath=False)
+        def exact(x: float) -> float:
+            return -x
+
+        @compiled(parallel=True, fastmath=True)
+        def both(x: float) -> float:
+            return x * 2.0
+
+        @compiled
+        def plain(x: float) -> float:
+            return x + 1.0
+
+        @compiled(fastmath=1)
+        def bad_val(x: float) -> float:
+            return x
+
+        @compiled(other=True)
+        def bad_kw(x: float) -> float:
+            return x
+    """)
+    assert [(f.name, f.parallel, f.fastmath) for f in m.functions] == [
+        ("fm", False, True),
+        ("exact", False, False),
+        ("both", True, True),
+        ("plain", False, False),
+    ]
+    assert set(m.skipped) == {"bad_val", "bad_kw"}
+    assert "fastmath" in m.skipped["bad_val"]
+    assert "fastmath" in m.skipped["bad_kw"] and "parallel" in m.skipped["bad_kw"]
+
+
 def test_the_former_module_pragma_comment_compiles_nothing(tmp_path):
     m = lower(tmp_path, """
         # typedpython: compiled

python-multiplatform-ksp/src/test/python/test_pipeline.py

@@ -185,3 +185,56 @@ def test_a_fresh_local_does_not_exempt_the_callee_class_guard(tmp_path):
     c = mod.Cell(None, None)
     mod.seta(c, 7)
     assert c.a == 7
+
+
+def test_default_fp_flags_stay_exact_and_fastmath_is_per_function_only(write, tmp_path):
+    """#153 / #332: by default, `-ffp-contract=off` and `-fno-fast-math` on GCC/Clang and
+    `/fp:precise` on MSVC are pinned so floating-point arithmetic is bit-identical to CPython.
+    `@compiled(fastmath=True)` opts in for that function only (not CPython-identical), leaving
+    sibling `@compiled` functions in the same module exact."""
+    from typedpython import cbuild, cgen, frontend, verify
+
+    assert "-ffp-contract=off" in cbuild.DEFAULT_FLAGS
+    assert "-fno-fast-math" in cbuild.DEFAULT_FLAGS
+    assert "/fp:precise" in cbuild.MSVC_DEFAULT_FLAGS
+    assert not any("fast" in f for f in cbuild.MSVC_DEFAULT_FLAGS)
+
+    path = write(
+        "fp_mod.py",
+        """
+        @compiled
+        def fma_exact(a: float, b: float, c: float) -> float:
+            return a * b + c
+
+        @compiled(fastmath=True)
+        def fma_fast(a: float, b: float, c: float) -> float:
+            return a * b + c
+
+        @compiled
+        def calls_fast_and_adds(a: float, b: float, c: float) -> float:
+            # caller is exact: its own `a * b + c` must not contract even though it calls `fma_fast`
+            _ = fma_fast(1.0, 2.0, 3.0)
+            return a * b + c
+        """,
+    )
+    proved, _ = verify.verify(frontend.lower(path))
+    c_text = cgen.generate(proved, path)
+    assert "#pragma STDC FP_CONTRACT OFF" in c_text
+    assert "\nTP_FASTMATH_BEGIN\n" in c_text and "\nTP_FASTMATH_END\n" in c_text
+    # Only `fma_fast` is wrapped in standalone TP_FASTMATH_BEGIN / TP_FASTMATH_END lines:
+    assert c_text.count("\nTP_FASTMATH_BEGIN\n") == 1
+    assert c_text.count("\nTP_FASTMATH_END\n") == 1
+
+    res = pipeline.compile_module(path, tmp_path / "out", cost_rule=False)
+    mod = pipeline.load(res)
+    ref = interpreted(path)
+    # 1 + 2**-27 squared: exact double product rounds to 1 + 2**-26, while FMA(a, a, -1.0) keeps 2**-54
+    a = 1.0 + 2.0**-27
+    c = -1.0
+    want_exact = ref.fma_exact(a, a, c)
+    assert mod.fma_exact(a, a, c) == want_exact
+    assert mod.calls_fast_and_adds(a, a, c) == want_exact
+    # fma_fast runs and returns a float (on targets with FMA it equals 2**-26 + 2**-54; either is valid under fastmath)
+    got_fast = mod.fma_fast(a, a, c)
+    assert isinstance(got_fast, float)
+