Feat: N-15 step 1 parallel element-wise loops (#277)

docs/SPEC.md

@@ -869,7 +869,7 @@ All in `PM/commonTest`, so they run wherever the interpreter loads.
 - **N-15** TypedPython parallel execution of order-independent loops on free-threaded CPython (#151;
   `docs/design/typedpython-parallel-loops.md`). A compiled `for ... in range(...)` loop may run its
   iterations on several threads only if its result cannot depend on iteration order, in a function the
-  user marked (the spelling of the marking is settled at implementation), and then the
+  user marked with a `# typedpython: parallel` comment line directly above the `def` (above its decorators if any); module-level functions only, and then the
   result, the exception, the partial writes and the deopt behaviour are bit-for-bit those of the
   sequential loop (INTENT §1.4; N-12: every execution is indistinguishable from some valid CPython
   execution).
@@ -895,20 +895,24 @@ All in `PM/commonTest`, so they run wherever the interpreter loads.
     (#153) awaits the owner's decision, and N-15 follows it. Partials combine in chunk order.
   - **Observability.** N-14 holds (a traced call runs the interpreted function, never a region);
     workers are not Python threads; the calling thread stays attached and polls the eval breaker
-    between micro-chunks when the function polls today; list writes become visible only at the
-    existing write-back. Outside the identity (like N-12's cyclic GC timing): the OS thread count, CPU
+    between micro-chunks when the function polls today (step 1 does not poll inside a region: a
+    signal waits until the region ends); the calling thread stays attached for the whole region, so a
+    stop-the-world waits for it too; list writes become visible only at the existing write-back. Outside the identity (like N-12's cyclic GC timing): the OS thread count, CPU
     time versus wall time, the multi-thread `os.fork()` warning while a region is active, and
     profilers seeing non-Python threads.
   - **Threads and cost.** v1 creates workers per region and joins them before the call returns (no
     pool, no thread outlives the call: nothing at fork or interpreter shutdown). Cap
-    `TYPEDPYTHON_THREADS` (1 disables; default 8 or the core count if lower). A region runs only above
-    a measured work threshold (`TP_PAR_MIN_WORK`), measured together with thread creation cost in a
-    quiet window and recorded on #151; below it the sequential code runs. Platforms without threads compile it out.
+    `TYPEDPYTHON_THREADS` (positive integer, 1 disables; default min(cores, 8); hard cap 64). A thread
+    that cannot be created is not an error: its chunk runs on the calling thread. Workers start with all
+    signals blocked. A region runs only above a measured work threshold (`TP_PAR_MIN_WORK`), which counts
+    weighted iterations (the trip count times the loop body's node count); a region also needs at least
+    1024 iterations per thread. The threshold is measured together with thread creation cost in a quiet
+    window and recorded on #151; below it the sequential code runs. Platforms without threads compile it out.
   - **Tests** (written first): differential vs the interpreter (bit-exact, boundaries of trip count and
     chunk, NaN/inf/-0.0), identical bytes for every thread count, failure at first/middle/last
     iteration, negative qualification cases, disabled-path test, reductions, race and GC tests,
     TSan/ASan, N-14, signals, fork and shutdown, and measurement tests for the threshold.
-  `Status: planned (design)`.
+  `Status: step 1 implemented (#277): element-wise F64 stores a[i] = <F64 + - * and negation of constants, locals and b[i]> into list copy-in or local arrays; no reductions, no failure path (eligible loops cannot fail or deopt); the rest planned (design)`.
 
 ---
 

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

@@ -512,6 +512,9 @@ def __init__(self, module: ir.Module):
         self.module = module
         self.display_path = ""
         self.functions = {f.name: f for f in module.functions}
+        # SPEC N-15: any ForRange.par in the module (the module then exports __typedpython_parallel__)
+        self.has_par = any(isinstance(st, ir.ForRange) and st.par is not None
+                           for f in module.functions for st in _walk_stmts(f.body))
         if len(self.functions) != len(module.functions):
             raise CGenError("duplicate compiled function names")
         self.names: list[str] = []
@@ -758,6 +761,8 @@ def generate(module: ir.Module, source_path: Path, display_path: str | None = No
     w("")
     if module.classes:
         w(_class_info_function(mg))
+    if mg.has_par:
+        w(_par_stats_function())
     w(_exec_function(mg))
     w(_state_functions(mg))
     w("static PyModuleDef_Slot tp_slots[] = {")
@@ -929,6 +934,7 @@ def __init__(self, mg: _ModGen, f: ir.Function, fast: bool = False):
         self.depth = 1
         self.temps: list[tuple[str, str]] = []      # (C type, name)
         self.arrays = {p.name: p for p in f.params if isinstance(p, ir.ArrayParam)}
+        self.par_defs: list[str] = []               # outlined region bodies (SPEC N-15), emitted before the impl
         self.local_arrays: dict[str, Type] = {}       # name -> F64_ARRAY / I64_ARRAY
         self.brk: list = []                    # per open loop: goto label for `break`, or None
         self.label_n = 0
@@ -2062,6 +2068,9 @@ def st_ForRange(self, s: ir.ForRange) -> None:
             self.emit(f"    {tn} = ((uint64_t){ts} - (uint64_t){te} - 1u) / {mag} + 1u;")
             sign = "-"
         self.depth += 1
+        par_open = s.par is not 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
@@ -2095,9 +2104,111 @@ def st_ForRange(self, s: ir.ForRange) -> None:
             self.emit("}")
             if label in self.used_labels:
                 self.emit(f"{label}: ;")
+        if par_open:
+            self.emit("}")                  # the sequential `else` of the parallel region
         self.depth -= 1
         self.emit("}")
 
+    def _par_collect(self, e: ir.Expr, scalars: set, arrays: set) -> None:
+        if isinstance(e, ir.Local):
+            scalars.add(e.name)
+        elif isinstance(e, ir.Index):
+            arrays.add(e.array)
+            self._par_collect(e.index, scalars, arrays)
+        elif isinstance(e, ir.ToFloat):
+            self._par_collect(e.operand, scalars, arrays)
+        elif isinstance(e, ir.BinOp):
+            self._par_collect(e.left, scalars, arrays)
+            self._par_collect(e.right, scalars, arrays)
+        elif isinstance(e, ir.UnaryOp):
+            self._par_collect(e.operand, scalars, arrays)
+        elif not isinstance(e, ir.Const):
+            raise CGenError(f"{self.f.name}: {type(e).__name__} in a parallel loop body")
+
+    def _emit_region(self, s: ir.ForRange, ts: str, tn: str, mag: int) -> None:
+        """SPEC N-15: opens `if (<big enough and threads> ) { region } else {`; the caller closes the
+        else with the sequential loop it emits next. The body is the verified `par` loop body
+        (StoreIndex statements of F64 arithmetic), outlined into `tp_pbody_*`, which declares the
+        names it uses as locals copied from a context struct, so the existing statement and
+        expression emitters produce its code unchanged. Nothing in it can fail, deopt or build an
+        exception (the verifier's rule); cgen refuses text that tries."""
+        f, info = self.f, s.par
+        if not f.parallel or not isinstance(info, ir.ParInfo) or s.step.value <= 0:
+            raise CGenError(f"{f.name}: ForRange.par on a loop that cannot be parallel (unmarked "
+                            "function, bad step or no ParInfo)")
+        scalars: set[str] = set()
+        arrays: set[str] = set()
+        for st in s.body:
+            if not isinstance(st, ir.StoreIndex):
+                raise CGenError(f"{f.name}: parallel loop body has a {type(st).__name__}")
+            arrays.add(st.array)
+            self._par_collect(st.value, scalars, arrays)
+        scalars.discard(s.var)
+        for a in arrays:
+            if a not in self.arrays and a not in self.local_arrays:
+                raise CGenError(f"{f.name}: parallel loop array {a} is not an array")
+            if a in self.arrays and self.arrays[a].buffer:
+                raise CGenError(f"{f.name}: parallel loop over buffer array {a}")
+        for x in scalars:
+            if self.types.get(x) not in (Type.F64, Type.I64):
+                raise CGenError(f"{f.name}: parallel loop scalar {x} is not f64/i64")
+        k = len(self.par_defs)
+        tag = _ident("tp_par", f"{f.name}{'_fast' if self.fast else ''}_{k}")
+        ctx_t, body_f = f"{tag}_ctx", f"{tag}_body"
+        fields = ["    int64_t tp_ts;"]
+        decls = ["    int64_t " + _ident("l", s.var) + " = 0;"]
+        sets = []
+        for a in sorted(arrays):
+            at = self.arrays[a].type if a in self.arrays else self.local_arrays[a]
+            ident = _ident("l", a)
+            fields.append(f"    {_ARRAY_CT[at]} *{ident};")
+            if a in self.arrays:                           # a parameter is already a pointer
+                decls.append(f"    {_ARRAY_CT[at]} *{ident} = c->{ident};")
+                sets.append(f"tp_pr.{ident} = {ident};")
+            else:                                          # a local array is a struct: the worker's copy
+                decls.append(f"    {_ARRAY_CT[at]} {ident} = *c->{ident};")
+                sets.append(f"tp_pr.{ident} = &{ident};")
+        for x in sorted(scalars):
+            ident, ct = _ident("l", x), _CT[self.types[x]]
+            fields.append(f"    {ct} {ident};")
+            decls.append(f"    {ct} {ident} = c->{ident};")
+            sets.append(f"tp_pr.{ident} = {ident};")
+        saved = (self.lines, self.temps, self.depth, self.brk)
+        self.lines, self.temps, self.depth, self.brk = [], [], 2, []
+        try:
+            self.stmts(s.body)
+            body_lines, body_temps = self.lines, self.temps
+        finally:
+            self.lines, self.temps, self.depth, self.brk = saved
+        text = "\n".join(body_lines)
+        for banned in ("goto", "tp_rc", "tp_s ", "return", "tp_st", "tp_dict", "PyObject", "tp_release"):
+            if banned in text:
+                raise CGenError(f"{f.name}: parallel loop body emitted `{banned}`")
+        sign_var = _ident("l", s.var)
+        out = [f"typedef struct {{", *fields, f"}} {ctx_t};", "",
+               f"static void {body_f}(const void *tp_c, uint64_t tp_k0, uint64_t tp_k1)", "{",
+               f"    const {ctx_t} *c = (const {ctx_t} *)tp_c;",
+               "    uint64_t tp_k;", *decls]
+        for ct, n in body_temps:
+            out.append(f"    {ct}{'' if ct.endswith('*') else ' '}{n} = 0;")
+        out.append("    (void)c;" + "".join(f" (void){_ident('l', x)};" for x in sorted(scalars)))
+        out.append(f"    for (tp_k = tp_k0; tp_k < tp_k1; tp_k++) {{")
+        out.append(f"        {sign_var} = (int64_t)((uint64_t)c->tp_ts + tp_k * UINT64_C({mag}));")
+        out += body_lines
+        out += ["    }", "}", ""]
+        self.par_defs.append("\n".join(out))
+        w = max(1, info.weight)
+        self.emit(f"if ({tn} >= (uint64_t)((TP_PAR_MIN_WORK + {w - 1}) / {w}) && tp_par_threads() > 1) {{")
+        self.depth += 1
+        self.emit(f"{ctx_t} tp_pr;")
+        self.emit(f"tp_pr.tp_ts = {ts};")
+        for line in sets:
+            self.emit(line)
+        self.emit(f"(void)tp_par_run({body_f}, &tp_pr, {tn});")
+        self.emit(f"{_ident('l', s.var)} = (int64_t)((uint64_t){ts} + ({tn} - 1u) * UINT64_C({mag}));")
+        self.depth -= 1
+        self.emit("} else {")
+
     def st_Return(self, s: ir.Return) -> None:
         rt = self.f.returns
         if rt is Type.NONE:
@@ -2217,8 +2328,13 @@ def generate(self) -> str:
         """The impl (preceded by its fast twin when bounded) and the wrapper."""
         parts = []
         if self.f.name in self.mg.bounded:
-            parts.append(_FnGen(self.mg, self.f, fast=True).impl())
-        parts.append(self.impl())
+            twin = _FnGen(self.mg, self.f, fast=True)
+            twin_text = twin.impl()
+            parts += twin.par_defs
+            parts.append(twin_text)
+        impl_text = self.impl()
+        parts += self.par_defs
+        parts.append(impl_text)
         parts.append(self.wrapper())
         return "\n".join(parts)
 
@@ -2631,6 +2747,23 @@ def _state_functions(mg: _ModGen) -> str:
     return "\n".join(L)
 
 
+def _par_stats_function() -> str:
+    """`__typedpython_parallel__()`: (regions that used workers, workers started, failed thread
+    creations) since import, so a test can prove the parallel path ran (SPEC N-15)."""
+    return "\n".join([
+        "static PyObject *tp_par_stats_fn(PyObject *tp_module, PyObject *tp_unused)",
+        "{",
+        "    int64_t tp_v[3];",
+        "    (void)tp_module; (void)tp_unused;",
+        "    tp_par_stats(tp_v);",
+        "    return Py_BuildValue(\"(LLL)\", (long long)tp_v[0], (long long)tp_v[1], (long long)tp_v[2]);",
+        "}",
+        "",
+        "static PyMethodDef tp_par_stats_def = {\"__typedpython_parallel__\", tp_par_stats_fn, METH_NOARGS, NULL};",
+        "",
+    ])
+
+
 def _class_info_function(mg: _ModGen) -> str:
     """`__typedpython_class_info__()`: {class name: {compiled, captured_tag, tag, current, slow,
     refreshed, watched}}, introspection of the guards (module docstring)."""
@@ -2792,6 +2925,12 @@ def _exec_function(mg: _ModGen) -> str:
             "    if (tp_cf == NULL || PyDict_SetItemString(tp_dict, \"__typedpython_class_info__\", tp_cf) < 0) goto tp_done;",
             "    tp_release(&tp_cf);",
         ]
+    if mg.has_par:
+        L += [
+            "    tp_cf = PyCFunction_NewEx(&tp_par_stats_def, tp_module, NULL);",
+            "    if (tp_cf == NULL || PyDict_SetItemString(tp_dict, \"__typedpython_parallel__\", tp_cf) < 0) goto tp_done;",
+            "    tp_release(&tp_cf);",
+        ]
     L += [
         "    tp_zero = PyLong_FromLong(0);",
         "    if (tp_zero == NULL || PyDict_SetItemString(tp_dict, \"__typedpython_deopts__\", tp_zero) < 0) goto tp_done;",

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

@@ -237,9 +237,28 @@ def lower(path: Path) -> ir.Module:
         else:
             skipped[name] = result
     functions.sort(key=lambda f: f.source_line)
+    marked = _parallel_marked(tree, source)
+    functions = [dataclasses.replace(f, parallel=True) if f.name in marked else f for f in functions]
     return ir.Module(path.stem, tuple(functions), skipped, tuple(class_decls.values()))
 
 
+PARALLEL_MARKER = "# typedpython: parallel"
+
+
+def _parallel_marked(tree: ast.Module, source: str) -> set[str]:
+    """Module-level functions the user marked for SPEC N-15: the comment line `# typedpython: parallel`
+    directly above the `def` (above its decorators if it has any). A comment, like the module marker,
+    so the interpreted module is unchanged and needs no runtime name."""
+    lines = source.splitlines()
+    out: set[str] = set()
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef):
+            first = min([node.lineno] + [d.lineno for d in node.decorator_list])
+            if first >= 2 and lines[first - 2].strip() == PARALLEL_MARKER:
+                out.add(node.name)
+    return out
+
+
 def _module_info(tree: ast.Module) -> _ModuleInfo:
     bindings: dict[str, list[tuple[int, str]]] = {}
 

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

@@ -650,6 +650,22 @@ class While(Stmt):
     body: tuple[Stmt, ...]
 
 
+@dataclass(frozen=True)
+class ParInfo:
+    """Why a `ForRange` may run its iterations on several threads (SPEC N-15, step 1): the body is
+    only `a[var] = <F64 expression>` stores, `stored` are the arrays stored into (each at the loop
+    variable, proven in bounds, so distinct iterations write distinct elements) and `read` the arrays
+    read, all at the loop variable and proven in bounds. The expression nodes are Const, scalar
+    Local (not assigned in the body), ToFloat of an I64, F64 + - * and negation, so no node can fail,
+    deopt or touch a Python object, and `weight` is their count (the cost rule's per-iteration work).
+    Arrays are distinct objects (entry guard), so a store at `var` cannot alias a read of another
+    array."""
+
+    stored: tuple[str, ...]
+    read: tuple[str, ...]
+    weight: int = 1
+
+
 @dataclass(frozen=True)
 class ForRange(Stmt):
     """`for var in range(start, stop, step)`: `step` is a non-zero I64 Const; `start`/`stop` are
@@ -663,6 +679,10 @@ class ForRange(Stmt):
     stop: Expr
     step: Const
     body: tuple[Stmt, ...]
+    # SPEC N-15 (#151, #277): set only by `verify` (like `proven`), which recomputes it from the body
+    # and rejects a loop whose declared value differs (`verify/par-loop`). The front end leaves it
+    # None. Not None: the iterations may run on several C threads (see `ParInfo`). None: sequential.
+    par: ParInfo | None = None
 
 
 @dataclass(frozen=True)
@@ -783,6 +803,9 @@ class Function:
     # rejects a function whose declared entries differ (`verify/virtual-escape`). The front end
     # leaves it empty.
     virtual: tuple[str, ...] = ()
+    # SPEC N-15: the user marked this function with the comment line `# typedpython: parallel`
+    # directly above its `def`. Only a marked function can hold a `ForRange.par`.
+    parallel: bool = False
 
 
 @dataclass(frozen=True)