diff --git a/docs/SPEC.md b/docs/SPEC.md index d9a68a5a..f3b360e0 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -726,7 +726,7 @@ All in `PM/commonTest`, so they run wherever the interpreter loads. runs in chunks of `TP_POLL_INTERVAL` iterations and polls once per chunk (#206), with `break`, `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`). Functions holding an array copy-in or entry-globals snapshot do not poll. + 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. `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 2d1838fc..cd058ece 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/cgen.py +++ b/python-multiplatform-ksp/src/main/python/typedpython/cgen.py @@ -697,6 +697,10 @@ def generate(module: ir.Module, source_path: Path, display_path: str | None = No source_path = Path(source_path) display_path = source_path.name if display_path is None else display_path source = source_path.read_text(encoding="utf-8") + # SPEC N-8 / #152: the continuations of every snapshot holder's loops run at import beside the + # module's own code, so an OSR-style exit finds them in the module dict + from typedpython import resume as _resume + source = source + _resume.resume_sources(module) mg = _ModGen(module) mg.display_path = display_path # the file name of a compiled numpy RuntimeWarning (#220) for f in module.functions: @@ -1060,6 +1064,12 @@ def __init__(self, mg: _ModGen, f: ir.Function, fast: bool = False, rmode: str | self.virtual[x] = next(iter(classes)) for k in range(len(mg.cls(f.name, self.virtual[x]).fields)): self.types[_vfield(x, k)] = Type.OBJ + # SPEC N-8 / #152: a snapshot holder polls with an OSR-style exit when every loop has a + # continuation and every continuation parameter maps to a live C value. Otherwise it keeps + # the #141 rule (no poll), which is correct but starves other threads during its loops. + self.osr_sites: dict[int, ir.Resume] = {} + self.osr_ctx: list[tuple[int, str, str, int]] = [] # enclosing ForRange: (site, var, stop, step) + self.osr = self._osr_plan() # emit primitives @@ -2208,10 +2218,159 @@ def poll(self, every_time: bool = False) -> None: self.emit("if (tp_pc_tick(&tp_pc) && tp_poll_slow(tp_t, tp_st->breaker) != 0) " "{ tp_rc = -1; goto tp_exit; }") + def _osr_plan(self) -> bool: + f = self.f + if not (self.snapshot and f.resumes) or f.virtual or getattr(f, "regions", ()) or self.rmode: + return False + from typedpython import resume as _resume + loops = list(_resume.ir_loops(f.body)) + by_site = {r.site: r for r in f.resumes} + if len(loops) != len(by_site) or any(k not in by_site for k in range(len(loops))): + return False + for k, st in enumerate(loops): + if isinstance(st, ir.ForRange) and st.par is not None: + return False # a parallel region does not poll + if not isinstance(st, (ir.ForRange, ir.While)): + return False + for name in by_site[k].params: + if not self._osr_mappable(name): + return False + self.osr_sites[id(st)] = by_site[k] + self.osr_sites_by_site = by_site + return True + + def _osr_mappable(self, name: str) -> bool: + if name.startswith(("__tp_cur_", "__tp_stop_", "__tp_step_")): + return True + if name in self.arrays or self.local_arrays.get(name) in _ARRAY_CT: + return True + return self.types.get(name) in (Type.I64, Type.F64, Type.BOOL, Type.OBJ, Type.INT) \ + and name not in {g.name for g in self.f.entry_globals} + + def _list_params(self) -> list[ir.ArrayParam]: + return [p for p in self.f.params if isinstance(p, ir.ArrayParam) and p.buffer is None] + + def _osr_value(self, name: str, site: int, cur: str, stop: str = "0", step: int = 0) -> str: + """A new reference to [name]'s value for the continuation of [site] (NULL on failure).""" + for prefix in ("__tp_cur_", "__tp_stop_", "__tp_step_"): + if name.startswith(prefix): + k = int(name[len(prefix):]) + if k == site: + value = {"__tp_cur_": cur, "__tp_stop_": stop, "__tp_step_": str(step)}[prefix] + return f"PyLong_FromLongLong((long long)({value}))" + ctx = next(c for c in self.osr_ctx if c[0] == k) + value = {"__tp_cur_": ctx[1], "__tp_stop_": ctx[2], "__tp_step_": str(ctx[3])}[prefix] + return f"PyLong_FromLongLong((long long)({value}))" + if name in self.arrays: + p = self.arrays[name] + if p.buffer is not None: + return f"Py_NewRef({_ident('l', name)}->view.obj)" + j = self._list_params().index(p) + return f"Py_NewRef(tp_my->lists[{j}])" + if name in self.local_arrays: + sfx = "f64" if self.local_arrays[name] is Type.F64_ARRAY else "i64" + v = _ident("l", name) + return f"tp_list_from_{sfx}({v}.data, {v}.len)" + t = self.types[name] + is_param = any(q.name == name for q in self.f.params) + if t is Type.OBJ: + v = _ident("p" if is_param else "l", name) + return f"Py_NewRef({v} != NULL ? {v} : tp_osent)" + v = _ident("l", name) + if t is Type.INT: + return f"tp_int_box(&{v})" + if t is Type.F64: + return f"PyFloat_FromDouble({v})" + if t is Type.BOOL: + return f"PyBool_FromLong({v})" + return f"PyLong_FromLongLong((long long){v})" + + def snap_poll(self, site: int, cur: str, stop: str = "0", step: int = 0) -> None: + """The poll of a snapshot holder (#152): write the lists back, end their critical section, + poll, take it again, copy the lists in again and re-check every entry guard. Unchanged: + continue. Changed: free the fresh copies, end the section and leave for the continuation of + [site] (the impl returns 2 with tp_my->result set). Only the call its wrapper entered polls.""" + r = self.osr_sites_by_site[site] + lists = self._list_params() + self.emit("if (tp_my != NULL) { /* #152: poll inside a snapshot holder */") + self.depth += 1 + self.emit("int tp_och = 0;") + for j, p in enumerate(lists): + ptr = _ident("l", p.name) + self.emit(f"Py_ssize_t tp_ol{j} = {ptr}->len;") + self.emit(f"if ({_ARRAY_PFX[p.type]}_exit({ptr}) < 0) {{ tp_rc = -1; goto tp_exit; }}") + if lists: + self.emit("tp_cs_end(tp_my->lock);") + self.emit("tp_t->snapshot_depth--;") + self.emit("tp_s = tp_poll_slow(tp_t, tp_st->breaker);") + # Tests only (-DTP_TEST_SNAP_HOOK): a module-level `__tp_test_snap_hook__()` runs here, in the + # window where the lists are released, so a test can change them at an exact poll. + self.emit("#ifdef TP_TEST_SNAP_HOOK") + self.emit("if (tp_s == 0) { PyObject *tp_hk = NULL;") + self.emit(" if (PyDict_GetItemStringRef(tp_dict, \"__tp_test_snap_hook__\", &tp_hk) > 0) {") + self.emit(" PyObject *tp_hr = PyObject_CallNoArgs(tp_hk); Py_DECREF(tp_hk);") + self.emit(" if (tp_hr == NULL) tp_s = -1; else Py_DECREF(tp_hr); } }") + self.emit("#endif") + self.emit("tp_t->snapshot_depth++;") + self.emit("if (tp_s != 0) { tp_rc = -1; goto tp_exit; }") + if lists: + self.emit("tp_cs_begin(tp_my->lock, tp_my->lists[0], tp_my->lists[1]);") + for j, p in enumerate(lists): + ptr = _ident("l", p.name) + self.emit(f"tp_s = {_ARRAY_PFX[p.type]}_enter(tp_my->lists[{j}], {ptr});") + self.emit("if (tp_s < 0) { tp_cs_end(tp_my->lock); tp_rc = -1; goto tp_exit; }") + self.emit(f"if (tp_s != 0 || {ptr}->len != tp_ol{j}) tp_och = 1;") + for g in self.f.entry_globals: + k = self.mg.name_index(g.name) + lv = _ident("l", g.name) + self.emit(f"{{ PyObject *tp_go = tp_global(tp_dict, tp_st->names[{k}]);") + self.emit(" if (tp_go == NULL) {") + self.emit(" if (!PyErr_ExceptionMatches(PyExc_NameError)) { tp_cs_end(tp_my->lock); tp_rc = -1; goto tp_exit; }" + if lists else " if (!PyErr_ExceptionMatches(PyExc_NameError)) { tp_rc = -1; goto tp_exit; }") + self.emit(" PyErr_Clear(); tp_och = 1; }") + self.emit(f" else {{ {_CT[g.type]} tp_gv = 0; tp_s = {_UNBOX[g.type]}(tp_go, &tp_gv); Py_DECREF(tp_go);") + self.emit(f" if (tp_s < 0) {{ {'tp_cs_end(tp_my->lock); ' if lists else ''}tp_rc = -1; goto tp_exit; }}") + self.emit(f" if (tp_s != 0 || memcmp(&tp_gv, &{lv}, sizeof tp_gv) != 0) tp_och = 1; }} }}") + self.emit("if (tp_och) { /* leave for the interpreter at this iteration (OSR-style exit) */") + self.depth += 1 + for j, p in enumerate(lists): + # fresh copies hold no dirty element: this exit only frees them + self.emit(f"(void){_ARRAY_PFX[p.type]}_exit({_ident('l', p.name)});") + if lists: + self.emit("tp_cs_end(tp_my->lock);") + n = len(r.params) + self.emit(f"PyObject *tp_ofn = NULL, *tp_osent = NULL, *tp_oa[{max(n, 1)}] = {{0}};") + self.emit(f"int tp_oe = PyDict_GetItemStringRef(tp_dict, \"{r.name}\", &tp_ofn) <= 0;") + self.emit("if (!tp_oe) tp_oe = PyDict_GetItemStringRef(tp_dict, \"__tp_unbound\", &tp_osent) <= 0;") + for i, name in enumerate(r.params): + self.emit(f"if (!tp_oe) {{ tp_oa[{i}] = {self._osr_value(name, site, cur, stop, step)}; tp_oe = tp_oa[{i}] == NULL; }}") + self.emit("if (!tp_oe) {") + self.emit(" tp_t->snapshot_depth--;") + self.emit(f" tp_my->result = PyObject_Vectorcall(tp_ofn, tp_oa, {n}, NULL);") + self.emit(" tp_t->snapshot_depth++;") + self.emit("}") + self.emit(f"for (int tp_oi = 0; tp_oi < {n}; tp_oi++) Py_XDECREF(tp_oa[tp_oi]);") + self.emit("Py_XDECREF(tp_ofn); Py_XDECREF(tp_osent);") + self.emit("if (tp_oe && !PyErr_Occurred()) PyErr_SetString(PyExc_SystemError, \"typedpython: continuation missing\");") + self.emit("tp_rc = tp_my->result != NULL ? 2 : -1;") + self.emit("goto tp_exit;") + self.depth -= 1 + self.emit("}") + self.depth -= 1 + self.emit("}") + def st_While(self, s: ir.While) -> None: self.emit("for (;;) {") self.depth += 1 - self.poll() + r = self.osr_sites.get(id(s)) if self.osr else None + if r is not None: + self.emit("if (tp_pc_tick(&tp_pc)) {") + self.depth += 1 + self.snap_poll(r.site, "0") + self.depth -= 1 + self.emit("}") + else: + self.poll() c = self._cond(s.cond) self.emit(f"if (!({c})) break;") self.brk.append(None) @@ -2330,8 +2489,9 @@ def st_ForRange(self, s: ir.ForRange) -> None: if par_open: self._emit_region(s, ts, tn, abs(step)) vec = self._vec_plan(s) if step == 1 else None - if self.snapshot: - # no polling in a snapshot holder: a plain loop + r = self.osr_sites.get(id(s)) if self.osr and not par_open else None + if self.snapshot and r is None: + # no polling in a snapshot holder without an OSR-style exit (#141): a plain loop self.emit(f"{tk} = 0;") if vec is not None: self._emit_vec_loop(vec, ts, tk, tn) @@ -2345,7 +2505,16 @@ def st_ForRange(self, s: ir.ForRange) -> None: self.label_n += 1 self.emit(f"for ({tk} = 0; {tk} < {tn}; ) {{") self.depth += 1 - self.poll(every_time=True) + if r is not None: + # Not before the first chunk: a loop shorter than TP_POLL_INTERVAL keeps its whole + # critical section (still an execution CPython can produce); only a long loop detaches. + self.emit(f"if ({tk} != 0) {{") + self.depth += 1 + self.snap_poll(r.site, f"(int64_t)((uint64_t){ts} {sign} {tk} * {mag})", te, step) + self.depth -= 1 + self.emit("}") + else: + self.poll(every_time=True) self.emit(f"{te2} = {tn} - {tk} > (uint64_t)TP_POLL_INTERVAL ? {tk} + (uint64_t)TP_POLL_INTERVAL : {tn};") if vec is not None: self._emit_vec_loop(vec, ts, tk, te2) @@ -2353,7 +2522,12 @@ def st_ForRange(self, s: ir.ForRange) -> None: self.depth += 1 self.brk.append(label) self.emit(f"{var} = (int64_t)((uint64_t){ts} {sign} {tk} * {mag});") + site = self.osr_sites.get(id(s)).site if self.osr and id(s) in self.osr_sites else None + if site is not None: + self.osr_ctx.append((site, var, te, step)) self.stmts(s.body) + if site is not None: + self.osr_ctx.pop() label = self.brk.pop() self.depth -= 1 self.emit("}") @@ -2679,6 +2853,9 @@ def impl(self) -> str: # Eval breaker (#141): a snapshot-holding function never polls and defers others' polls. if self.snapshot: out.append(" tp_t->snapshot_depth++;") + if self.osr: + # #152: only the call the wrapper just entered gets the context; internal callees see NULL + out.append(" tp_osr *tp_my = tp_t->osr; tp_t->osr = NULL;") elif not self.fast: out.append(" if (tp_poll(tp_t, tp_st->breaker) != 0) { tp_rc = -1; goto tp_exit; }") out += prologue @@ -2828,7 +3005,22 @@ def wrapper(self) -> str: if gnames: at = len(args) - (1 if f.returns is not Type.NONE else 0) args[at:at] = gnames + osr = self.osr + if osr: + lst = [i for i, p in arrays] + l0 = f"tp_args[{lst[0]}]" if lst else "NULL" + l1 = f"tp_args[{lst[1]}]" if len(lst) > 1 else "NULL" + lock = "&tp_lock" if lst else "NULL" + L.append(f" tp_osr tp_ox = {{{lock}, {{{l0}, {l1}}}, NULL}};") + L.append(" tp_t->osr = &tp_ox; /* #152: taken by the impl at entry */") L.append(f" tp_s = {_ident('tp_impl', f.name)}(tp_module, tp_ts, tp_t{''.join(', ' + a for a in args)});") + if osr: + L.append(" tp_t->osr = NULL;") + L.append(" if (tp_s == 2) { /* left for the continuation: its result, lists already final */") + L += [" " + e for e in exits] + L.append(" if (tp_x < 0) { Py_CLEAR(tp_ox.result); return NULL; }") + L.append(" return tp_ox.result;") + L.append(" }") if f.pure: L.append(" if (tp_s == 1) goto tp_deopt;") else: diff --git a/python-multiplatform-ksp/src/main/python/typedpython/frontend.py b/python-multiplatform-ksp/src/main/python/typedpython/frontend.py index 7653876f..0da894f1 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/frontend.py +++ b/python-multiplatform-ksp/src/main/python/typedpython/frontend.py @@ -77,7 +77,7 @@ from dataclasses import dataclass, field from pathlib import Path -from typedpython import gate, ir, pyrefly, rebinding, taint +from typedpython import gate, ir, pyrefly, rebinding, resume, taint from typedpython.ir import Type MARKER = "# typedpython: compiled" @@ -1050,11 +1050,13 @@ def run(self) -> ir.Function: pure = not self.impure assert pure or not self.deopts or not self.impure_mode, \ "impure mode emitted a deopting node" - return ir.Function( + fn = ir.Function( name=self.sig.qname, params=params, returns=self.sig.returns, locals=locals_, body=body, pure=pure, may_deopt=self.deopts, source_line=self.fn.lineno, entry_globals=tuple(ir.Param(n, t) for n, t in self.entry_used.items()), owner=self.sig.owner) + # SPEC N-8 (#152): a snapshot holder gets a Python continuation per loop (OSR-style exit). + return dataclasses.replace(fn, resumes=resume.resumes_for(self.fn, fn)) # --- names and types --- diff --git a/python-multiplatform-ksp/src/main/python/typedpython/ir.py b/python-multiplatform-ksp/src/main/python/typedpython/ir.py index 4882949d..368c1c3b 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/ir.py +++ b/python-multiplatform-ksp/src/main/python/typedpython/ir.py @@ -825,6 +825,30 @@ class Function: # which recomputes it and rejects a function whose declared entries differ # (`verify/region-escape`). The front end leaves it empty. regions: tuple[str, ...] = () + # SPEC N-8 (#152, option (i)): one entry per loop of a snapshot-holding function (see `resume.py`), + # in IR pre-order. Empty for a function that holds no snapshot. A loop whose continuation cannot + # be generated (an enclosing `with`/`try`, say) has no entry: its `site` is missing from the tuple. + resumes: tuple["Resume", ...] = () + + +@dataclass(frozen=True) +class Resume: + """A Python continuation of one loop of a snapshot-holding function (SPEC N-8, OSR-style exit). + + `site`: the loop's number, counting every ForRange, While and ForObj of `Function.body` in IR + pre-order: statements in order, `If.then` before `If.orelse`, a loop numbered before the loops + in its body. `name`: `__tp_resume__` (a `.` in a method's name becomes `_`). + `params`: the exact parameter order of the continuation, see `resume.py` for the contract. + `source`: the Python source of the `def`, followed by a line replacing `co_name`/`co_qualname` + with the interpreted function's. `linemap`: for each line of `source` (index = line - 1) the + original (lineno, end_lineno, col_offset, end_col_offset); `resume.resume_sources` uses it so + the continuation's code object carries the original line numbers.""" + + site: int + name: str + params: tuple[str, ...] + source: str + linemap: tuple[tuple[int, int, int, int], ...] = () @dataclass(frozen=True) diff --git a/python-multiplatform-ksp/src/main/python/typedpython/resume.py b/python-multiplatform-ksp/src/main/python/typedpython/resume.py new file mode 100644 index 00000000..0a87c5db --- /dev/null +++ b/python-multiplatform-ksp/src/main/python/typedpython/resume.py @@ -0,0 +1,397 @@ +"""OSR-style exit continuations (SPEC N-8, #152 option (i)). + +A function that holds a snapshot (an array copy-in or entry globals) polls at the top of a loop +iteration. When the re-validation fails it must leave for the interpreter at that point. For every +loop of such a function this module generates a Python function, the continuation, built from the +function's own AST: it runs "the rest of the function from the top of that iteration". + +Contract (the C side builds the call from `ir.Resume.params`): + +* Site numbering: loops of `Function.body` in IR pre-order (see `ir.Resume`). The AST loops are + numbered the same way; if the counts or kinds disagree, no resume is generated at all. +* Parameters, in this order: + 1. every source-level local and parameter that may be bound at the loop, sorted by name. "Source + level" means a name the function's AST binds (parameter, assignment, augmented assignment, + loop target, `with ... as`, `except ... as`, import, def/class name, walrus, `del` target). + Compiler temporaries never appear in the AST, so they are excluded by construction. Left out: + names declared `global`, the site's own loop targets and every enclosing ForRange variable (all + of those are set by the loop header from the hidden state), and a local with no binding + statement before the loop, inside it or inside an enclosing loop (it cannot be bound yet; the + continuation still treats it as a local, so reading it raises UnboundLocalError). + 2. the hidden loop state, outermost loop first and the site last: for a ForRange `k`: + `__tp_cur_`, `__tp_stop_`, `__tp_step_` (the loop variable's value for the iteration + about to run and the already evaluated bounds); for an enclosing ForObj `k`: `__tp_it_` + (the live iterator; its targets are ordinary locals); for the site if it is a ForObj: + `__tp_it_`, `__tp_cur_` (the item just taken from the iterator, before unpacking, which + the continuation processes first); a While has none. `` is that loop's site number. +* A local that is unbound at the call is passed as the module-level singleton `__tp_unbound` + (defined by `resume_sources`); the continuation first `del`s it, so a later read raises + UnboundLocalError exactly as CPython would. +* Returns what the rest of the function returns (falling off the end gives None) or raises. +""" +from __future__ import annotations + +import ast +import copy +from typing import Iterable + +from typedpython import ir + +PRELUDE = '''\ +from itertools import chain as __tp_chain +__tp_unbound = object() +def __tp_resume_def(src, lmap): + import ast, sys + frame = sys._getframe(1) + tree = ast.parse(src) + last = len(lmap) - 1 + for n in ast.walk(tree): + if hasattr(n, "lineno"): + n.lineno, n.end_lineno, n.col_offset, n.end_col_offset = lmap[min(n.lineno - 1, last)] + exec(compile(tree, frame.f_code.co_filename, "exec"), frame.f_globals) +''' + + +def has_snapshot(f: ir.Function) -> bool: + """Same rule as `cgen._FnGen.snapshot` (cgen.py, property `snapshot`, ~line 2186): any + entry global, or any ArrayParam (buffer or not; the property does not look at `buffer`).""" + return bool(f.entry_globals) or any(isinstance(p, ir.ArrayParam) for p in f.params) + + +def resumes_for(fn: ast.FunctionDef, f: ir.Function) -> tuple[ir.Resume, ...]: + """The resumes of `f`, lowered from `fn`; () when `f` holds no snapshot.""" + if not has_snapshot(f): + return () + kinds = [_kind(s) for s in ir_loops(f.body)] + return build(fn, kinds, f.name) + + +def ir_loops(stmts: Iterable[ir.Stmt]): + """Loops in IR pre-order (the numbering of `ir.Resume.site`).""" + for s in stmts: + if isinstance(s, (ir.ForRange, ir.While, ir.ForObj)): + yield s + for fl in s.__dataclass_fields__.values(): + v = getattr(s, fl.name) + if isinstance(v, tuple) and v and all(isinstance(x, ir.Stmt) for x in v): + yield from ir_loops(v) + + +def _kind(s: ir.Stmt) -> str: + return "range" if isinstance(s, ir.ForRange) else "obj" if isinstance(s, ir.ForObj) else "while" + + +# --- the AST side --------------------------------------------------------------------------------- + +class _Site: + def __init__(self, node: ast.stmt, number: int, chain: list, ok: bool) -> None: + self.node, self.number, self.chain, self.ok = node, number, chain, ok + + +def _sub_lists(st: ast.stmt): + for name in ("body", "orelse", "finalbody"): + v = getattr(st, name, None) + if isinstance(v, list) and v and isinstance(v[0], ast.stmt): + yield name, v + for h in getattr(st, "handlers", []): + yield "handler", h.body + for c in getattr(st, "cases", []): + yield "case", c.body + + +def _collect(stmts: list, chain: list, ok: bool, out: list) -> None: + for i, st in enumerate(stmts): + if isinstance(st, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + here = [*chain, (stmts, i)] + if isinstance(st, (ast.For, ast.While)) and not st.orelse: + out.append(_Site(st, len(out), here, ok)) + _collect(st.body, here, ok, out) + continue + if isinstance(st, (ast.For, ast.While)): + out.append(_Site(st, len(out), here, False)) + sub_ok = ok and isinstance(st, (ast.If, ast.For, ast.While)) and not ( + isinstance(st, (ast.For, ast.While)) and st.orelse) + for _, lst in _sub_lists(st): + _collect(lst, here, sub_ok, out) + + +def _names(e: ast.AST, out: set[str]) -> None: + """Names an expression (not a statement) binds in the enclosing function scope.""" + stack = [e] + while stack: + n = stack.pop() + if isinstance(n, ast.Lambda): + continue + if isinstance(n, (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)): + for sub in ast.walk(n): + if isinstance(sub, ast.NamedExpr) and isinstance(sub.target, ast.Name): + out.add(sub.target.id) + continue + if isinstance(n, ast.Name) and isinstance(n.ctx, (ast.Store, ast.Del)): + out.add(n.id) + stack.extend(ast.iter_child_nodes(n)) + + +def _stmt_binds(st: ast.stmt) -> set[str]: + out: set[str] = set() + if isinstance(st, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + out.add(st.name) + return out + if isinstance(st, (ast.Import, ast.ImportFrom)): + for a in st.names: + out.add((a.asname or a.name).split(".")[0]) + return out + for child in ast.iter_child_nodes(st): + if isinstance(child, (ast.stmt, ast.ExceptHandler, ast.match_case)): + continue + _names(child, out) + return out + + +def _own_stmts(stmts: list): + """Every statement of this scope in pre-order, with exception handler names as pseudo bindings.""" + for st in stmts: + yield st + if isinstance(st, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue + for _, lst in _sub_lists(st): + yield from _own_stmts(lst) + + +def _all_binders(fn: ast.FunctionDef) -> tuple[dict[str, list[ast.AST]], set[str], set[str]]: + """name -> binding statements; (global names); (names bound by handlers/patterns: unsupported).""" + binders: dict[str, list[ast.AST]] = {} + globals_: set[str] = set() + odd: set[str] = set() + for st in _own_stmts(fn.body): + if isinstance(st, ast.Global): + globals_ |= set(st.names) + if isinstance(st, ast.Nonlocal): + odd |= set(st.names) + for n in _stmt_binds(st): + binders.setdefault(n, []).append(st) + for h in getattr(st, "handlers", []): + if h.name: + binders.setdefault(h.name, []).append(st) + for c in getattr(st, "cases", []): + for sub in ast.walk(c.pattern): + for attr in ("name", "rest"): + nm = getattr(sub, attr, None) + if isinstance(nm, str): + binders.setdefault(nm, []).append(st) + return binders, globals_, odd + + +def _arg_names(fn: ast.FunctionDef) -> list[str]: + a = fn.args + names = [x.arg for x in [*a.posonlyargs, *a.args, *a.kwonlyargs]] + names += [x.arg for x in (a.vararg, a.kwarg) if x is not None] + return names + + +def _target_names(t: ast.AST) -> list[str]: + return [n.id for n in ast.walk(t) if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store)] + + +def _skeleton(src: str, ref: ast.AST) -> ast.Module: + tree = ast.parse(src) + for n in ast.walk(tree): + if "lineno" in n._attributes: + n.lineno = n.end_lineno = ref.lineno + n.col_offset = ref.col_offset + n.end_col_offset = ref.col_offset + 1 + return tree + + +def _copy_all(stmts) -> list[ast.stmt]: + return [copy.deepcopy(s) for s in stmts] + + +def build(fn: ast.FunctionDef, kinds: list[str], qname: str) -> tuple[ir.Resume, ...]: + """Continuations for every loop of `fn`. `kinds[i]` is "range", "obj" or "while" for loop `i` in + pre-order (from the IR). Returns () when the AST's loops do not match `kinds`.""" + sites: list[_Site] = [] + _collect(fn.body, [], True, sites) + if len(sites) != len(kinds): + return () + for s, k in zip(sites, kinds): + if isinstance(s.node, ast.While) != (k == "while"): + return () + if k == "range" and not (isinstance(s.node, ast.For) and isinstance(s.node.target, ast.Name)): + return () + binders, globals_, odd = _all_binders(fn) + if odd: + return () + out = [] + for s in sites: + if not s.ok: + continue + r = _one(fn, qname, s, sites, kinds, binders, globals_) + if r is not None: + out.append(r) + return tuple(out) + + +def _one(fn, qname, site, sites, kinds, binders, globals_): + n = len(site.chain) + by_node = {id(s.node): s for s in sites} + # loops on the path, outermost first, the site last + loops = [by_node[id(lst[i])] for lst, i in site.chain if id(lst[i]) in by_node] + ref = site.node + + # --- parameters + excluded: set[str] = set(globals_) + hidden: list[str] = [] + for lp in loops: + k, kind = lp.number, kinds[lp.number] + if kind == "range": + excluded.add(lp.node.target.id) + hidden += [f"__tp_cur_{k}", f"__tp_stop_{k}", f"__tp_step_{k}"] + elif kind == "obj": + if lp is site: + excluded |= set(_target_names(lp.node.target)) + hidden += [f"__tp_it_{k}", f"__tp_cur_{k}"] + else: + hidden += [f"__tp_it_{k}"] + in_loop_ids: set[int] = set() + for lp in loops: + in_loop_ids |= {id(x) for x in _own_stmts([lp.node])} + order = {id(st): i for i, st in enumerate(_own_stmts(fn.body))} + site_order = order[id(site.node)] + args = set(_arg_names(fn)) + locals_ = [] + for name in sorted(set(binders) | args): + if name in excluded: + continue + if name in args or any(order[id(b)] < site_order or id(b) in in_loop_ids + for b in binders.get(name, [])): + locals_.append(name) + params = tuple(locals_) + tuple(hidden) + + # --- body + def rest(k: int) -> list[ast.stmt]: + lst, i = site.chain[k] + return [*enter(k), *_copy_all(lst[i + 1:])] + + def enter(k: int) -> list[ast.stmt]: + lst, i = site.chain[k] + node = lst[i] + if k == n - 1: + return site_construct(node) + if isinstance(node, ast.If): + return rest(k + 1) + return enclosing(node, rest(k + 1)) + + def site_construct(node) -> list[ast.stmt]: + kind, num = kinds[site.number], site.number + if kind == "while": + return [copy.deepcopy(node)] + if kind == "range": + skel = _skeleton(f"for _ in range(__tp_cur_{num}, __tp_stop_{num}, __tp_step_{num}): pass", + node) + else: + skel = _skeleton(f"for _ in __tp_chain((__tp_cur_{num},), __tp_it_{num}): pass", node) + loop = skel.body[0] + loop.target = copy.deepcopy(node.target) + loop.body = _copy_all(node.body) + return [loop] + + def enclosing(node, r: list[ast.stmt]) -> list[ast.stmt]: + num = by_node[id(node)].number + kind = kinds[num] + first = f"__tp_first_{num}" + init = _skeleton(f"{first} = True", node).body[0] + flip = _skeleton(f"{first} = False", node).body[0] + body = _copy_all(node.body) + if kind == "range": + skel = _skeleton(f"for _ in range(__tp_cur_{num}, __tp_stop_{num}, __tp_step_{num}):\n" + f" if {first}: pass\n else: pass", node) + loop = skel.body[0] + loop.target = copy.deepcopy(node.target) + loop.body[0].body = [flip, *r] + loop.body[0].orelse = body + return [init, loop] + if kind == "while": + skel = _skeleton(f"while True:\n if {first}: pass\n else: pass", node) + loop = skel.body[0] + loop.test = ast.BoolOp(ast.Or(), [_skeleton(first, node).body[0].value, + copy.deepcopy(node.test)]) + ast.copy_location(loop.test, node) + loop.body[0].body = [flip, *r] + loop.body[0].orelse = body + return [init, loop] + item = f"__tp_item_{num}" + skel = _skeleton( + f"while True:\n if {first}: pass\n else:\n" + f" try:\n {item} = next(__tp_it_{num})\n" + f" except StopIteration:\n break\n" + f" _ = {item}\n del {item}", node) + loop = skel.body[0] + els = loop.body[0].orelse + assign = els[1] + assign.targets = [copy.deepcopy(node.target)] + loop.body[0].body = [flip, *r] + loop.body[0].orelse = [*els, *body] + return [init, loop] + + top, i0 = site.chain[0] + stmts = [*enter(0), *_copy_all(top[i0 + 1:])] + pre: list[ast.stmt] = [] + if globals_: + pre.append(_skeleton("global " + ", ".join(sorted(globals_)), fn).body[0]) + for p in locals_: + pre.extend(_skeleton(f"if {p} is __tp_unbound: del {p}", fn).body) + name = f"__tp_resume_{qname.replace('.', '_')}_{site.number}" + func = _skeleton(f"def {name}({', '.join(params)}): pass", fn).body[0] + func.body = [*pre, *stmts] + ast.fix_missing_locations(func) + text = ast.unparse(func) + short = qname.split(".")[-1] + source = (f"{text}\n{name}.__code__ = {name}.__code__.replace(" + f"co_name={short!r}, co_qualname={qname!r})\n") + return ir.Resume(site.number, name, params, source, _linemap(func, text, fn)) + + +def _entry(st: ast.stmt) -> tuple[int, int, int, int]: + if isinstance(st, (ast.For, ast.While, ast.If, ast.With, ast.Try, ast.FunctionDef)): + end_col = st.col_offset + 1 + for ch in ast.iter_child_nodes(st): + if isinstance(ch, ast.expr) and ch.end_lineno == st.lineno: + end_col = max(end_col, ch.end_col_offset) + return (st.lineno, st.lineno, st.col_offset, end_col) + return (st.lineno, st.end_lineno or st.lineno, st.col_offset, st.end_col_offset or st.col_offset + 1) + + +def _linemap(func: ast.FunctionDef, text: str, fn: ast.FunctionDef): + """Generated line -> original position, by pairing the statements of the unparsed-and-reparsed + function with those of the one we built (same shapes, so same order).""" + reparsed = ast.parse(text).body[0] + a = [s for s in ast.walk(func) if isinstance(s, ast.stmt)] + b = [s for s in ast.walk(reparsed) if isinstance(s, ast.stmt)] + nlines = len(text.splitlines()) + head = (fn.lineno, fn.lineno, fn.col_offset, fn.col_offset + 1) + table: dict[int, tuple[int, int, int, int]] = {1: head} + if [type(x) for x in a] == [type(x) for x in b]: + for x, y in zip(a, b): + table.setdefault(y.lineno, _entry(x)) + table[1] = head + out = [] + cur = head + for line in range(1, nlines + 1): + cur = table.get(line, cur) + out.append(cur) + out.append(cur) # the `__code__` replacement line + return tuple(out) + + +def resume_sources(module: ir.Module) -> str: + """Python source to append to the module's source (what the exec slot runs at import): the + prelude, then every continuation of every function. "" when no function has a resume.""" + parts = [] + for f in module.functions: + for r in f.resumes: + parts.append(f"__tp_resume_def({r.source!r}, {r.linemap!r})\n") + if not parts: + return "" + return "\n" + PRELUDE + "".join(parts) diff --git a/python-multiplatform-ksp/src/main/python/typedpython/runtime/API.md b/python-multiplatform-ksp/src/main/python/typedpython/runtime/API.md index 30b4cfd6..9c30cdad 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/runtime/API.md +++ b/python-multiplatform-ksp/src/main/python/typedpython/runtime/API.md @@ -234,6 +234,13 @@ Same structs as above with `list == NULL` and `dirty == NULL`. 0 ok, -1 MemoryEr Thread state (#125): the wrapper calls `PyThreadState_GetUnchecked()` once and passes it as `tp_ts` to the impl, which passes it to `tp_enter_call`, `tp_depth_room`, `tp_traced` and to its direct compiled callees, so a compiled call and its C-to-C callees do one TLS lookup of the thread state per call. Likewise the per-thread state (call depth, poll countdown, snapshot depth) is one `__thread` struct +`tp_osr` (#152, OSR-style exit): `{ tp_cs *lock; PyObject *lists[2]; PyObject *result; }`. A snapshot-holding +function's wrapper sets `tp_t->osr` to its own context right before calling the impl; the impl takes it at entry +and sets `tp_t->osr` back to NULL. The impl returns 2 when it left for the continuation, with `result` set (strong). + + PyObject *tp_list_from_f64(const double *d, Py_ssize_t n); /* new list; NULL + exception */ + PyObject *tp_list_from_i64(const int64_t *d, Py_ssize_t n); /* new list; NULL + exception */ + `tp_tls`: the wrapper takes `tp_tls_get()` once and passes `tp_t` down. `Py_EnterRecursiveCall`, `Py_LeaveRecursiveCall` and `Py_GetRecursionLimit` have no public tstate-taking variant and are kept. diff --git a/python-multiplatform-ksp/src/main/python/typedpython/runtime/tp_runtime.h b/python-multiplatform-ksp/src/main/python/typedpython/runtime/tp_runtime.h index 3ba26954..ee12da41 100644 --- a/python-multiplatform-ksp/src/main/python/typedpython/runtime/tp_runtime.h +++ b/python-multiplatform-ksp/src/main/python/typedpython/runtime/tp_runtime.h @@ -2749,13 +2749,64 @@ typedef struct { /* All per-thread state of the compiled code in ONE thread-local, so a compiled call looks its address * up once (tp_tls_get, in the wrapper) and passes the pointer down (#125). */ +/* A local array as a new list, for an OSR-style exit (#152): a local array never escapes, so the + * continuation gets a fresh list with the same elements. NULL with an exception on failure. */ +static inline PyObject *tp_list_from_f64(const double *d, Py_ssize_t n) +{ + PyObject *l = PyList_New(n); + Py_ssize_t i; + if (l == NULL) { + return NULL; + } + for (i = 0; i < n; i++) { + PyObject *v = PyFloat_FromDouble(d[i]); + if (v == NULL) { + Py_DECREF(l); + return NULL; + } + PyList_SET_ITEM(l, i, v); + } + return l; +} + +static inline PyObject *tp_list_from_i64(const int64_t *d, Py_ssize_t n) +{ + PyObject *l = PyList_New(n); + Py_ssize_t i; + if (l == NULL) { + return NULL; + } + for (i = 0; i < n; i++) { + PyObject *v = PyLong_FromLongLong((long long)d[i]); + if (v == NULL) { + Py_DECREF(l); + return NULL; + } + PyList_SET_ITEM(l, i, v); + } + return l; +} + +/* OSR-style exit (#152, SPEC N-8). A snapshot-holding function's wrapper fills one of these on its own + * stack and points tp_tls.osr at it right before calling its impl; the impl takes it at entry and sets + * tp_tls.osr back to NULL, so only that one call (not an internal callee, not a recursive call) polls + * with these lists. `lock` and `lists` are the wrapper's: the impl ends and re-takes the critical + * section around each poll. `result` is set (strong) when the impl left for the continuation and + * returned 2; the wrapper then returns it, with its write-back already done. */ +typedef struct tp_osr { + tp_cs *lock; /* the wrapper's critical section, or NULL (no list parameter) */ + PyObject *lists[2]; /* the list arguments in critical-section order (borrowed), or NULL */ + PyObject *result; /* strong, after an exit: what the continuation returned */ +} tp_osr; + typedef struct { tp_depth_state depth; /* compiled frames in flight (see above) */ int poll_countdown; /* eval-breaker poll countdown (#141) */ Py_ssize_t snapshot_depth; /* snapshot-holding functions in flight (#141) */ + tp_osr *osr; /* the OSR context for the impl about to be entered (#152), or NULL */ } tp_tls; -static TP_THREAD_LOCAL tp_tls tp_tls_v = {{0, NULL, 0}, TP_POLL_INTERVAL, 0}; +static TP_THREAD_LOCAL tp_tls tp_tls_v = {{0, NULL, 0}, TP_POLL_INTERVAL, 0, NULL}; static inline tp_tls *tp_tls_get(void) { diff --git a/python-multiplatform-ksp/src/test/python/test_free_threaded.py b/python-multiplatform-ksp/src/test/python/test_free_threaded.py index 7e9ead81..4bc95587 100644 --- a/python-multiplatform-ksp/src/test/python/test_free_threaded.py +++ b/python-multiplatform-ksp/src/test/python/test_free_threaded.py @@ -320,11 +320,13 @@ def go(): def test_shared_list_writers_and_readers_see_whole_calls(extension): - """`fill` writes every element in one call, `edge` reads the first and last: with the list's - critical section held from copy-in to write-back, a reader sees a state between two whole calls, - so first == last, always.""" + """`fill` writes every element in one call, `edge` reads the first and last. A call whose loops stay + below TP_POLL_INTERVAL holds the list's critical section from copy-in to write-back, so a reader + sees a state between two whole calls: first == last, always. (A longer loop polls and releases it + between chunks since #152, as an interpreted loop's list stores are visible one by one; that is + CPython's behaviour, so this test keeps the list below the interval.)""" r = run_child(extension, """ - shared = [0.0] * 20000 + shared = [0.0] * 4000 # below TP_POLL_INTERVAL: one call is one critical section (#152) torn = [] stop = False def writer(v): @@ -354,7 +356,7 @@ def driver(): def test_shared_list_int_writers_and_readers(extension): r = run_child(extension, """ - shared = [0] * 20000 + shared = [0] * 4000 # below TP_POLL_INTERVAL: one call is one critical section (#152) torn = [] stop = False def writer(v): diff --git a/python-multiplatform-ksp/src/test/python/test_osr_exit.py b/python-multiplatform-ksp/src/test/python/test_osr_exit.py new file mode 100644 index 00000000..fbfbdd3c --- /dev/null +++ b/python-multiplatform-ksp/src/test/python/test_osr_exit.py @@ -0,0 +1,244 @@ +"""SPEC N-8 / #152: a snapshot-holding function polls inside long loops and, when a list or an entry +global changed at the poll, leaves for the interpreter at that iteration (an OSR-style exit). + +Deterministic: the module is built with TP_POLL_INTERVAL=32 and TP_TEST_SNAP_HOOK, so a module-level +`__tp_test_snap_hook__()` runs at every poll of a snapshot holder, in the window where its lists are +released, exactly where another thread or a signal handler could run. The reference for each case is +the interpreted function with the same hook called at the same iterations (every 32nd, from the second +chunk on), so compiled and interpreted must agree on every result, list and exception. +""" +from __future__ import annotations + +import textwrap + +import pytest + +from test_cgen import need_runtime, runtime_dir +from typedpython import cbuild, cgen, frontend, verify + +SOURCE = textwrap.dedent(''' + # typedpython: compiled + + SCALE = 2.0 + + + def scale_all(x: list[float]) -> float: + s = 0.0 + for i in range(len(x)): + x[i] = x[i] * SCALE + s += x[i] + return s + + + def nested(x: list[float], n: int) -> int: + hits = 0 + for r in range(n): + for i in range(len(x)): + if x[i] < 0.0: + continue + hits += 1 + if hits > 100000: + break + hits += 1000 + return hits + + + def countdown(x: list[int]) -> int: + k = 0 + total = 0 + while k < len(x): + total += x[k] + k += 1 + return total +''') + +# The interpreted references: the same bodies, with the hook at the iterations the compiled poll runs +# it (the top of every 32nd iteration of a for loop, from the second chunk on; every 32nd iteration of +# a while loop, counting down from the function's entry). +REFERENCE = textwrap.dedent(''' + SCALE = 2.0 + HOOK = None + + + def scale_all(x): + s = 0.0 + for i in range(len(x)): + if i and i % 32 == 0: + HOOK() + x[i] = x[i] * SCALE + s += x[i] + return s + + + def nested(x, n): + hits = 0 + for r in range(n): + for i in range(len(x)): + if i and i % 32 == 0: + HOOK() + if x[i] < 0.0: + continue + hits += 1 + if hits > 100000: + break + hits += 1000 + return hits + + + def countdown(x): + k = 0 + total = 0 + tick = 32 + while True: + tick -= 1 + if tick == 0: + tick = 32 + HOOK() + if not k < len(x): + break + total += x[k] + k += 1 + return total +''') + + +@pytest.fixture(scope="module") +def mods(tmp_path_factory): + need_runtime() + d = tmp_path_factory.mktemp("osr") + path = d / "osrmod.py" + path.write_text(SOURCE) + proved, _ = verify.verify(frontend.lower(path)) + c = cgen.generate(proved, path, path.name) + so = cbuild.build(c, "osrmod", d / "out", runtime_dir=runtime_dir(), + flags=(*cbuild.DEFAULT_FLAGS, "-DTP_POLL_INTERVAL=32", "-DTP_TEST_SNAP_HOOK")) + m = cbuild.load("osrmod", so) + ref: dict = {} + exec(compile(REFERENCE, "reference", "exec"), ref) + return m, ref, c + + +def _hook_at(calls: int, action): + """A hook that runs `action` on its `calls`-th call (1-based) and counts every call.""" + state = {"n": 0} + + def hook(): + state["n"] += 1 + if state["n"] == calls: + action() + return hook, state + + +def _both(mods, fname, make_args, calls, action_for): + """Run compiled and interpreted with a hook firing `action_for(args, module_globals)` at the + `calls`-th poll; return ((result or exception, args), ...) for both.""" + m, ref, _ = mods + out = [] + for kind in ("compiled", "interpreted"): + args = make_args() + g = vars(m) if kind == "compiled" else ref + hook, state = _hook_at(calls, lambda: action_for(args, g)) + if kind == "compiled": + m.__tp_test_snap_hook__ = hook + else: + ref["HOOK"] = hook + try: + res = ("ok", getattr(m if kind == "compiled" else None, fname, None)(*args) + if kind == "compiled" else ref[fname](*args)) + except Exception as e: # compared by type and text + res = ("raise", type(e).__name__, str(e)) + finally: + if kind == "compiled": + del m.__tp_test_snap_hook__ + g["SCALE"] = 2.0 + out.append((res, args, state["n"])) + return out + + +def test_snapshot_holders_get_continuations_and_poll(mods): + _, _, c = mods + assert "__tp_resume_scale_all_0" in c and "tp_osr *tp_my" in c + assert c.count("#152: poll inside a snapshot holder") >= 4 # scale_all, nested (2), countdown + + +@pytest.mark.parametrize("calls", [1, 2, 5]) +def test_an_unchanged_poll_keeps_running_compiled(mods, calls): + m, ref, _ = mods + seen = [] + orig = vars(m)["__tp_resume_scale_all_0"] + m.__tp_resume_scale_all_0 = lambda *a: seen.append(1) or orig(*a) + try: + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), calls, + lambda args, g: None) + finally: + m.__tp_resume_scale_all_0 = orig + assert comp[0] == interp[0] and comp[1] == interp[1] + assert comp[2] == interp[2] == 6 # polls at 32, 64, ..., 192 + assert seen == [] # no exit when nothing changed + + +@pytest.mark.parametrize("calls", [1, 3, 6]) +def test_a_length_change_at_a_poll_continues_interpreted(mods, calls): + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), calls, + lambda args, g: args[0].append(1.5)) + assert comp[0] == interp[0] + assert comp[1] == interp[1] + + +@pytest.mark.parametrize("calls", [1, 4]) +def test_an_element_type_change_at_a_poll_continues_interpreted(mods, calls): + def put_int(args, g): + args[0][100] = 7 # an int where the copy-in requires floats + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), calls, put_int) + assert comp[0] == interp[0] + assert comp[1] == interp[1] + + +def test_a_rebound_entry_global_at_a_poll_continues_interpreted(mods): + def rebind(args, g): + g["SCALE"] = 3.0 + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), 2, rebind) + assert comp[0] == interp[0] + assert comp[1] == interp[1] + + +@pytest.mark.parametrize("calls", [1, 3, 8, 13]) +def test_an_exit_inside_a_nested_loop_keeps_break_and_continue(mods, calls): + def flip(args, g): + x = args[0] + x[5] = -1.0 + x.append(2.0) + comp, interp = _both(mods, "nested", lambda: ([float(i % 7) for i in range(100)], 5), calls, flip) + assert comp[0] == interp[0] + assert comp[1] == interp[1] + + +@pytest.mark.parametrize("calls", [1, 2]) +def test_a_while_loop_exit(mods, calls): + def grow(args, g): + args[0].extend([10, 20, 30]) + comp, interp = _both(mods, "countdown", lambda: (list(range(100)),), calls, grow) + assert comp[0] == interp[0] + assert comp[1] == interp[1] + + +def test_an_exception_raised_in_the_continuation_reaches_the_caller(mods): + def bad(args, g): + args[0][150] = "x" # the interpreted rest raises TypeError on this element + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), 1, bad) + assert comp[0][0] == "raise" and comp[0] == interp[0] + assert comp[1] == interp[1] + + +@pytest.mark.parametrize("calls", [1, 3]) +def test_an_exception_at_a_poll_leaves_cpythons_partial_writes(mods, calls): + """What a signal handler raising KeyboardInterrupt does: the list holds the writes of the + iterations before the poll, written back before the poll ran, as CPython's own stores would be.""" + class Stop(Exception): + pass + + def stop(args, g): + raise Stop("at the poll") + comp, interp = _both(mods, "scale_all", lambda: ([float(i) for i in range(200)],), calls, stop) + assert comp[0] == interp[0] == ("raise", "Stop", "at the poll") + assert comp[1] == interp[1] diff --git a/python-multiplatform-ksp/src/test/python/test_regions.py b/python-multiplatform-ksp/src/test/python/test_regions.py index 3456586d..05afdce9 100644 --- a/python-multiplatform-ksp/src/test/python/test_regions.py +++ b/python-multiplatform-ksp/src/test/python/test_regions.py @@ -337,7 +337,8 @@ def test_gc_and_allocation_counts_are_unchanged_by_a_call_GREEN(built): built.m.tree_sum(8) gc.collect() assert nodes() == n0 # no Node is left behind (counting all objects saw pytest's own) - assert abs(sys.getallocatedblocks() - b0) < 1000 # a leaked tree is ~2300 blocks per call + # one-sided: a leaked tree adds ~2300 blocks per call; a call may also free other tests' garbage + assert sys.getallocatedblocks() - b0 < 1000 def test_tracing_runs_the_interpreted_function_GREEN(built): diff --git a/python-multiplatform-ksp/src/test/python/test_resume.py b/python-multiplatform-ksp/src/test/python/test_resume.py new file mode 100644 index 00000000..14115cc7 --- /dev/null +++ b/python-multiplatform-ksp/src/test/python/test_resume.py @@ -0,0 +1,583 @@ +"""OSR-style exit continuations (SPEC N-8, #152): `typedpython.resume`. + +Pure Python, no C. For each function and EVERY iteration of EVERY loop ("resume point") the test +runs the original up to that point, captures the locals and the loop state, calls the generated +continuation, and requires the same result, the same final list contents and the same exception +(type and text) as the interpreted original run end to end. + +The capture is test-only: the function's AST is copied and a hook is inserted at the top of each +loop body (for `while`, before the condition), see `_Instrument`. Lists are deep-copied at capture +(one memo, so aliasing is kept) and the continuation works on its own copy. +""" +import ast +import copy +import textwrap +import traceback +from types import SimpleNamespace + +import pytest + +from typedpython import frontend, ir, resume + +UNBOUND_NAME = "__tp_unbound" + + +# --- harness -------------------------------------------------------------------------------------- + +def loop_kinds(fn: ast.FunctionDef, force_obj: frozenset = frozenset()) -> list[str]: + """Pre-order kinds, as the front end would give them (`range` for a range call with a name target).""" + kinds: list[str] = [] + + def walk(stmts): + for st in stmts: + if isinstance(st, ast.While): + kinds.append("while") + elif isinstance(st, ast.For): + simple = (isinstance(st.iter, ast.Call) and isinstance(st.iter.func, ast.Name) + and st.iter.func.id == "range" and isinstance(st.target, ast.Name)) + kinds.append("range" if simple and len(kinds) not in force_obj else "obj") + for name in ("body", "orelse"): + walk(getattr(st, name, [])) + + walk(fn.body) + return kinds + + +def instrument(fn: ast.FunctionDef, kinds: list[str]) -> ast.FunctionDef: + """A copy of `fn` with hooks. Written as a plain recursive rewrite (clearer than a visitor).""" + fn = copy.deepcopy(fn) + fn.decorator_list = [] + counter = [0] + + def walk(stmts): + for st in stmts: + if isinstance(st, ast.While): + s = counter[0] + counter[0] += 1 + hook = ast.parse(f"(__hook({s}, locals()) or True)", mode="eval").body + st.test = ast.BoolOp(ast.And(), [hook, st.test]) + walk(st.body) + elif isinstance(st, ast.For): + s = counter[0] + counter[0] += 1 + call = ast.parse(f"{'__rng' if kinds[s] == 'range' else '__it'}({s}, 0)", + mode="eval").body + call.args[1] = st.iter + st.iter = call + st.body.insert(0, ast.parse(f"__hook({s}, locals())").body[0]) + walk(st.body) + else: + for name in ("body", "orelse"): + walk(getattr(st, name, [])) + + walk(fn.body) + ast.fix_missing_locations(fn) + return fn + + +class _RecIter: + def __init__(self, iterable): + self.inner, self.cur = iter(iterable), None + + def __iter__(self): + return self + + def __next__(self): + self.cur = next(self.inner) + return self.cur + + +def build_module(source: str, force_obj=frozenset()): + """(namespace with f and its continuations, fn ast, kinds, resumes).""" + source = textwrap.dedent(source).lstrip("\n") + tree = ast.parse(source) + fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef)) + kinds = loop_kinds(fn, force_obj) + resumes = resume.build(fn, kinds, fn.name) + assert len(resumes) == len(kinds), "every loop must have a continuation" + mod = ir.Module("m", (ir.Function(fn.name, (), ir.Type.NONE, {}, (), True, False, fn.lineno, + resumes=resumes),)) + full = source + resume.resume_sources(mod) + ns: dict = {} + exec(compile(full, "mod.py", "exec"), ns) + # a second namespace, instrumented + inst = ast.Module(body=[instrument(fn, kinds)], type_ignores=[]) + ins_ns: dict = {} + return ns, fn, kinds, resumes, inst, ins_ns, source + + +def outcome(call): + try: + return ("ok", call()) + except BaseException as e: # noqa: BLE001 - the test compares every exception + return ("raise", type(e).__name__, str(e)) + + +def capture_points(inst, kinds, make_args): + """Run the instrumented function; return the list of (site, loc copy, aux copy).""" + points = [] + rng: dict[int, tuple[int, int]] = {} + iters: dict[int, _RecIter] = {} + + def rng_(s, r): + rng[s] = (r.stop, r.step) + return r + + def it_(s, x): + iters[s] = _RecIter(x) + return iters[s] + + def hook(s, loc): + aux = {"rng": dict(rng), + "it": {k: v.inner for k, v in iters.items()}, + "cur": {k: v.cur for k, v in iters.items()}} + loc_c, aux_c = copy.deepcopy((dict(loc), aux)) + points.append((s, loc_c, aux_c)) + + ns = {"__hook": hook, "__rng": rng_, "__it": it_} + exec(compile(ast.fix_missing_locations(inst), "mod.py", "exec"), ns) + args = make_args() + outcome(lambda: ns[inst.body[0].name](*args)) + return points + + +def loop_vars(fn, kinds): + """site -> loop target AST (pre-order, same numbering).""" + out = [] + + def walk(stmts): + for st in stmts: + if isinstance(st, (ast.For, ast.While)): + out.append(st) + for name in ("body", "orelse"): + walk(getattr(st, name, [])) + + walk(fn.body) + return out + + +def check(source, make_args, force_obj=frozenset(), list_params=None, min_points=1): + """Compare every resume point. Returns the number of points checked.""" + ns, fn, kinds, resumes, inst, _, _ = build_module(source, force_obj) + f = ns[fn.name] + sentinel = ns[UNBOUND_NAME] + loops = loop_vars(fn, kinds) + by_site = {r.site: r for r in resumes} + names = [a.arg for a in fn.args.args] + # the reference run + args0 = make_args() + expect = outcome(lambda: f(*args0)) + expect_lists = copy.deepcopy([a for a in args0 if isinstance(a, list)]) + + points = capture_points(inst, kinds, make_args) + assert len(points) >= min_points + for s, loc, aux in points: + r = by_site[s] + vals = [] + for p in r.params: + if p.startswith("__tp_cur_"): + k = int(p.rsplit("_", 1)[1]) + if kinds[k] == "range": + vals.append(loc[loops[k].target.id]) + else: + vals.append(aux["cur"][k]) + elif p.startswith("__tp_stop_"): + vals.append(aux["rng"][int(p.rsplit("_", 1)[1])][0]) + elif p.startswith("__tp_step_"): + vals.append(aux["rng"][int(p.rsplit("_", 1)[1])][1]) + elif p.startswith("__tp_it_"): + vals.append(aux["it"][int(p.rsplit("_", 1)[1])]) + else: + vals.append(loc.get(p, sentinel)) + got = outcome(lambda: ns[r.name](*vals)) + got_lists = [loc[n] for n in names if isinstance(loc.get(n), list)] + assert got == expect, (s, r.params, vals, got, expect) + assert got_lists == expect_lists, (s, got_lists, expect_lists) + return len(points) + + +# --- the functions ---------------------------------------------------------------------------------- + +SINGLE = """ +def f(xs, n): + t = 0 + for i in range(n): + t += xs[i] * 2 + xs[i] = t + return t +""" + +NESTED = """ +def f(xs, n): + total = 0 + for i in range(n): + if i == 2: + continue + for j in range(i, n): + if j == 3: + continue + if xs[j] < 0: + break + total += xs[j] + xs[j] += 1 + total += 100 + if total > 600: + break + xs[i] = total + total += 7 + return total +""" + +WHILE = """ +def f(xs): + k = 0 + while k < len(xs): + xs[k] += k + k += 1 + if k == 3: + continue + xs[0] += 1 + return k +""" + +WHILE_FOR = """ +def f(xs, n): + k = 0 + while k < len(xs): + for i in range(3): + xs[i] += k + if xs[i] > 12: + break + k += 1 + if k == 2: + continue + xs[0] -= 1 + return k +""" + +OBJ = """ +def f(xs, n): + t = 0 + for i in range(n): + xs.append(i) + t += i + if t > 20: + break + return t +""" + +OBJ_NESTED = """ +def f(xs, n): + t = 0 + for i in range(n): + for j in range(i): + xs.append(j) + if j == 3: + continue + t += j + xs.append(-i) + if i == 5: + continue + t += 1000 + return t +""" + +UNPACK = """ +def f(xs, ys): + t = 0 + for a, b in zip(xs, ys): + t += a * b + xs.append(t) + if t > 30: + break + return t +""" + +LATE = """ +def f(xs, n): + for i in range(n): + if i == 2: + late = i * 10 + xs[i] += 1 + return late +""" + +RETURN = """ +def f(xs, n): + for i in range(n): + for j in range(n): + if xs[i] + j == 7: + return (i, j) + xs[i] += 1 + return None +""" + +NEG = """ +def f(xs, n): + t = 0 + for i in range(n, 0, -2): + for j in range(10, i, -3): + t += j * i + xs.append(t) + if t % 7 == 0: + break + t -= 1 + return t +""" + +IF_ELSE = """ +def f(xs, n): + acc = 0 + for i in range(n): + if i % 2: + for j in range(3): + acc += j + xs.append(acc) + else: + acc -= 1 + for j in range(2): + if j == 1: + continue + acc += 10 + acc += 1 + return acc +""" + + +def test_single_for_range(): + assert check(SINGLE, lambda: ([1, 2, 3, 4], 4)) == 4 + + +def test_single_for_range_raising_inside(): + assert check(SINGLE, lambda: ([1, 2, 3, 4], 6)) == 5 # the hook of i=4 runs before the IndexError + + +def test_single_for_range_zero_iterations(): + assert check(SINGLE, lambda: ([1], 0), min_points=0) == 0 + + +@pytest.mark.parametrize("n", [3, 6, 9]) +def test_nested_for_range_with_break_and_continue(n): + assert check(NESTED, lambda: ([3, 1, -1, 4, 1, 5, 9, 2, 6], n)) > n + + +def test_nested_for_range_with_break_in_the_outer_loop(): + assert check(NESTED, lambda: ([300, 200, 100, 4, 1, 5, 9, 2, 6], 9)) > 5 + + +def test_while_with_a_counter(): + assert check(WHILE, lambda: ([5, 5, 5, 5, 5],)) == 6 + + +def test_while_with_an_inner_for(): + assert check(WHILE_FOR, lambda: ([0, 1, 2], 4)) > 8 + + +def test_for_obj_over_a_range_object(): + assert check(OBJ, lambda: ([], 9), force_obj=frozenset({0})) == 7 + + +def test_for_obj_range_object_exhausts_without_break(): + assert check(OBJ, lambda: ([], 4), force_obj=frozenset({0})) == 4 + + +def test_for_obj_nested_outer_object_inner_range(): + assert check(OBJ_NESTED, lambda: ([], 8), force_obj=frozenset({0})) > 8 + + +def test_for_obj_nested_both_objects(): + assert check(OBJ_NESTED, lambda: ([], 8), force_obj=frozenset({0, 1})) > 8 + + +def test_for_obj_nested_inner_object_outer_range(): + assert check(OBJ_NESTED, lambda: ([], 8), force_obj=frozenset({1})) > 8 + + +def test_for_obj_with_unpacking_targets(): + assert check(UNPACK, lambda: ([1, 2, 3, 4, 5], [5, 4, 3, 2, 1])) >= 3 + + +@pytest.mark.parametrize("n", [2, 4]) +def test_local_first_assigned_inside_the_loop(n): + assert check(LATE, lambda: ([0, 0, 0, 0, 0], n)) == n + + +def test_return_inside_the_loop(): + assert check(RETURN, lambda: ([0, 3, 9], 3)) > 3 + + +def test_return_inside_the_loop_not_taken(): + assert check(RETURN, lambda: ([0, 0, 0], 3)) > 3 + + +@pytest.mark.parametrize("force", [frozenset(), frozenset({0}), frozenset({1}), frozenset({0, 1})]) +def test_negative_step(force): + assert check(NEG, lambda: ([], 9), force_obj=force) > 5 + + +def test_loops_in_both_branches_of_an_if(): + assert check(IF_ELSE, lambda: ([], 5)) > 10 + + +def test_hidden_state_and_param_order(): + ns, fn, kinds, resumes, *_ = build_module(NESTED) + assert [r.site for r in resumes] == [0, 1] + assert resumes[0].name == "__tp_resume_f_0" + # sorted locals, then the loop state, outermost first + # `j` is bound inside loop 0 (a previous iteration), so it may be bound at site 0 + assert resumes[0].params == ("j", "n", "total", "xs", "__tp_cur_0", "__tp_stop_0", "__tp_step_0") + assert resumes[1].params == ("n", "total", "xs", "__tp_cur_0", "__tp_stop_0", "__tp_step_0", + "__tp_cur_1", "__tp_stop_1", "__tp_step_1") + # the loop variables are set by the loop headers, not passed + assert "i" not in resumes[1].params and "j" not in resumes[1].params + + +def test_unbound_local_is_passed_as_the_sentinel_and_deleted(): + ns, fn, kinds, resumes, *_ = build_module(LATE) + assert "late" in resumes[0].params and "if late is __tp_unbound:\n del late" in resumes[0].source + with pytest.raises(UnboundLocalError): + ns["__tp_resume_f_0"](ns[UNBOUND_NAME], 2, [0, 0], 1, 2, 1) + + +def test_local_bound_only_after_the_loop_is_not_a_parameter(): + src = """ + def f(xs, n): + for i in range(n): + xs[i] += 1 + after = 1 + return after + """ + resumes = build_module(src)[3] + assert resumes[0].params == ("n", "xs", "__tp_cur_0", "__tp_stop_0", "__tp_step_0") + + +def test_traceback_reads_as_the_interpreted_function(): + src = """ + def f(xs, n): + t = 0 + for i in range(n): + t += 1 + xs[i] = 10 // (3 - i) + return t + """ + ns, fn, kinds, resumes, *_ = build_module(src) + + def frames(call): + try: + call() + except ZeroDivisionError as e: + return [(fs.filename, fs.name, fs.lineno) for fs in traceback.extract_tb(e.__traceback__)][-1] + + direct = frames(lambda: ns["f"]([0] * 6, 6)) + resumed = frames(lambda: ns["__tp_resume_f_0"](6, 1, [0] * 6, 1, 6, 1)) + assert resumed == direct == ("mod.py", "f", 5) + assert ns["__tp_resume_f_0"].__code__.co_qualname == "f" + + +def test_a_global_declaration_is_kept(): + src = """ + G = 0 + def f(xs, n): + global G + for i in range(n): + G += xs[i] + return G + """ + ns, fn, kinds, resumes, *_ = build_module(src) + assert "G" not in resumes[0].params + ns["G"] = 5 + assert ns["__tp_resume_f_0"](3, [1, 2, 3], 1, 3, 1) == 5 + 2 + 3 + + +def test_loop_inside_a_with_has_no_continuation(): + src = """ + def f(xs, n): + with open("/dev/null") as fh: + for i in range(n): + xs[i] += 1 + for k in range(n): + xs[k] += 1 + """ + tree = ast.parse(textwrap.dedent(src)) + fn = tree.body[0] + rs = resume.build(fn, ["range", "range"], "f") + assert [r.site for r in rs] == [1] + + +def test_mismatching_loop_count_gives_nothing(): + tree = ast.parse(textwrap.dedent(SINGLE)) + assert resume.build(tree.body[0], [], "f") == () + + +# --- the front end ------------------------------------------------------------------------------------ + +def lower(tmp_path, source): + path = tmp_path / "mod.py" + path.write_text(textwrap.dedent(source).lstrip("\n")) + return frontend.lower(path) + + +def test_frontend_gives_a_snapshot_holder_one_resume_per_loop(tmp_path): + module = lower(tmp_path, """ + # typedpython: compiled + def f(xs: list[float], n: int) -> float: + t = 0.0 + for i in range(n): + for j in range(i, n): + t += xs[j] + t += 1 + k = 1.0 + while t < 40.0: + a = k * 2.0 + t += a + k += 1.0 + xs[0] = t + return t + """) + assert module.skipped == {} + fn = module.functions[0] + assert any(isinstance(p, ir.ArrayParam) for p in fn.params) + assert [r.site for r in fn.resumes] == [0, 1, 2] + assert [r.name for r in fn.resumes] == ["__tp_resume_f_0", "__tp_resume_f_1", "__tp_resume_f_2"] + # `k` is bound by an assignment after loop 0 and 1 only: not a parameter there + assert "k" not in fn.resumes[0].params and "k" in fn.resumes[2].params + assert "a" in fn.resumes[2].params + # compiler temporaries never appear + assert not any(p.startswith("tp_") or p.startswith("__t") for r in fn.resumes for p in r.params + if not p.startswith("__tp_")) + src = resume.resume_sources(module) + ns = {} + exec(compile(textwrap.dedent(open(tmp_path / "mod.py").read()) + src, "mod.py", "exec"), ns) + # resume loop 1 (inner) at i=1, j=2 of n=4: t so far 0 + xs = [1.0, 2.0, 3.0, 4.0] + got = ns["__tp_resume_f_1"](*[{"xs": xs, "n": 4, "t": 0.0, "k": ns[UNBOUND_NAME], "a": ns[UNBOUND_NAME], + "__tp_cur_0": 1, "__tp_stop_0": 4, "__tp_step_0": 1, + "__tp_cur_1": 2, "__tp_stop_1": 4, "__tp_step_1": 1}[p] + for p in fn.resumes[1].params if p not in ("i", "j")]) + ref = [1.0, 2.0, 3.0, 4.0] + t = 0.0 + for i in range(1, 4): + for j in range(i if i != 1 else 2, 4): + t += ref[j] + t += 1 + k, a = 1.0, 0.0 + while t < 40.0: + a = k * 2.0 + t += a + k += 1.0 + ref[0] = t + assert got == t + assert xs == ref # the while loop ran too + + +def test_frontend_function_without_a_snapshot_has_no_resumes(tmp_path): + module = lower(tmp_path, """ + # typedpython: compiled + def g(n: int) -> int: + t = 0 + for i in range(n): + t += i + return t + """) + assert module.functions[0].resumes == () + assert resume.resume_sources(module) == ""