From 88243859f91849d45471b3d8e53c66ffe15cc7dd Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Sat, 31 Oct 2026 23:59:03 +0900 Subject: [PATCH] Feat: add per-function @compiled(fastmath=True) opt-in (#332) --- docs/SPEC.md | 9 ++++ .../src/main/python/typedpython/cgen.py | 33 ++++++++++-- .../src/main/python/typedpython/frontend.py | 53 ++++++++++++------- .../src/main/python/typedpython/ir.py | 3 ++ .../src/test/python/test_frontend.py | 37 +++++++++++++ .../src/test/python/test_pipeline.py | 53 +++++++++++++++++++ 6 files changed, 165 insertions(+), 23 deletions(-) diff --git a/docs/SPEC.md b/docs/SPEC.md index a41dd5b2..1e0697a6 100644 --- a/docs/SPEC.md +++ b/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 diff --git a/python-multiplatform-ksp/src/main/python/typedpython/cgen.py b/python-multiplatform-ksp/src/main/python/typedpython/cgen.py index fb776ce0..940837ed 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/cgen.py +++ b/python-multiplatform-ksp/src/main/python/typedpython/cgen.py @@ -494,10 +494,19 @@ 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: @@ -505,7 +514,7 @@ def visit(name: str) -> int | None: 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) diff --git a/python-multiplatform-ksp/src/main/python/typedpython/frontend.py b/python-multiplatform-ksp/src/main/python/typedpython/frontend.py index a65ea1b9..83fa2083 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/frontend.py +++ b/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 diff --git a/python-multiplatform-ksp/src/main/python/typedpython/ir.py b/python-multiplatform-ksp/src/main/python/typedpython/ir.py index 1176d218..4b28c068 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/ir.py +++ b/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 diff --git a/python-multiplatform-ksp/src/test/python/test_frontend.py b/python-multiplatform-ksp/src/test/python/test_frontend.py index ddcf3fcc..cd631efa 100644 --- a/python-multiplatform-ksp/src/test/python/test_frontend.py +++ b/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 diff --git a/python-multiplatform-ksp/src/test/python/test_pipeline.py b/python-multiplatform-ksp/src/test/python/test_pipeline.py index ed326eb4..18355aec 100644 --- a/python-multiplatform-ksp/src/test/python/test_pipeline.py +++ b/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) +