Perf: one thread-state and one TLS lookup per compiled call (#125)

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

@@ -375,7 +375,7 @@ def _all_exprs(f: ir.Function):
 def _traced_call(k: int) -> str:
     """tp_traced for the function whose interpreted original is `names[k]` (SPEC N-14). The code
     object is a strong reference in the module state, filled at module exec."""
-    return (f"tp_traced(tp_module, (PyCodeObject *)((tp_state *)PyModule_GetState(tp_module))->orig_code[{k}])")
+    return (f"tp_traced(tp_module, tp_ts, (PyCodeObject *)((tp_state *)PyModule_GetState(tp_module))->orig_code[{k}])")
 
 
 def _can_deopt(e: ir.Expr, functions: dict[str, ir.Function]) -> bool:
@@ -887,6 +887,9 @@ def _impl_signature(f: ir.Function, fast: bool = False) -> str:
     params = ["PyObject *tp_module"]
     if fast:
         params.append("int tp_fd")              # compiled frames of this fast tree, itself included
+    else:
+        params.append("PyThreadState *tp_ts")   # the caller's one thread-state lookup (#125)
+    params.append("tp_tls *tp_t")               # its one TLS-address lookup (#125)
     for p in f.params:
         if isinstance(p, ir.ArrayParam):
             params.append(f"{'tp_buf' if p.buffer else _ARRAY_CT[p.type]} *{_ident('l', p.name)}")
@@ -907,7 +910,7 @@ def _impl_signature(f: ir.Function, fast: bool = False) -> str:
 def _impl_call_args(f: ir.Function, fd: str | None = None) -> str:
     """The arguments that forward an impl's own parameters to another impl of the same signature
     (`fd`: the value of a fast twin's `tp_fd`)."""
-    args = ["tp_module"] + ([fd] if fd is not None else [])
+    args = ["tp_module"] + ([fd] if fd is not None else []) + ["tp_t"]
     for p in f.params:
         args.append(_ident("l" if isinstance(p, ir.ArrayParam) or p.type is not Type.OBJ else "p", p.name))
     args += [_ident("l", g.name) for g in f.entry_globals]
@@ -1327,9 +1330,9 @@ def ev_Call(self, e: ir.Call) -> str | None:
         # Depth is counted by the callee itself (tp_enter_call at its entry, issue #57).
         if self.fast:
             # the fast twin's callees are bounded too (_bounded_functions): their fast twins, uncounted
-            self.emit(f"tp_s = {_ident('tp_fimpl', e.function)}(tp_module, tp_fd + 1{''.join(', ' + x for x in args)});")
+            self.emit(f"tp_s = {_ident('tp_fimpl', e.function)}(tp_module, tp_fd + 1, tp_t{''.join(', ' + x for x in args)});")
         else:
-            self.emit(f"tp_s = {_ident('tp_impl', e.function)}(tp_module{''.join(', ' + x for x in args)});")
+            self.emit(f"tp_s = {_ident('tp_impl', e.function)}(tp_module, tp_ts, tp_t{''.join(', ' + x for x in args)});")
         self.depth -= 1
         self.emit("}")
         if ti is not None:
@@ -1393,9 +1396,9 @@ def _redo_call(self, callee: ir.Function, arg_exprs, values: list[str], t: str |
         # In a fast twin the interpreted callee (a Python frame) and whatever it calls must see the
         # frames of this fast tree: count them for exactly the span that can run Python code, from
         # the call to the release of its result (module docstring, "Fast entry").
-        end = "tp_depth_py_end(tp_fd); " if self.fast else ""
+        end = "tp_depth_py_end(tp_t, tp_fd); " if self.fast else ""
         if self.fast:
-            self.emit("tp_depth_py_begin(tp_fd);")
+            self.emit("tp_depth_py_begin(tp_t, tp_fd);")
         if av:
             self.emit(f"{{ PyObject *tp_av[{len(av)}] = {{{', '.join(av)}}};")
             self.emit(f"  {rr} = tp_cg_redo(tp_module, {k}, tp_av, {len(av)}, NULL); }}")
@@ -1419,7 +1422,7 @@ def _redo_call(self, callee: ir.Function, arg_exprs, values: list[str], t: str |
             what = {Type.F64: "a float", Type.BOOL: "True or False", Type.I64: "an int in i64"}[rt]
         self._own_release(rr)
         if self.fast:
-            self.emit("tp_depth_py_end(tp_fd);")
+            self.emit("tp_depth_py_end(tp_t, tp_fd);")
         msg = c_string(f"typedpython: the interpreted redo of {callee.name} did not return {what}")
         self.emit(f"if (tp_s != 0) {{ PyErr_SetString(PyExc_SystemError, {msg}); tp_rc = -1; goto tp_exit; }}")
         self.depth -= 1
@@ -1935,11 +1938,11 @@ def poll(self) -> None:
             # store of a global on every iteration of a short inner loop cost fannkuch ~38%.
             if self.fast:
                 # the poll can run a signal handler or another compiled call: count this fast tree
-                self.emit("if (tp_pc_tick(&tp_pc)) { tp_depth_py_begin(tp_fd); "
-                          "tp_s = tp_poll_slow(tp_st->breaker); tp_depth_py_end(tp_fd); "
+                self.emit("if (tp_pc_tick(&tp_pc)) { tp_depth_py_begin(tp_t, tp_fd); "
+                          "tp_s = tp_poll_slow(tp_t, tp_st->breaker); tp_depth_py_end(tp_t, tp_fd); "
                           "if (tp_s != 0) { tp_rc = -1; goto tp_exit; } }")
                 return
-            self.emit("if (tp_pc_tick(&tp_pc) && tp_poll_slow(tp_st->breaker) != 0) "
+            self.emit("if (tp_pc_tick(&tp_pc) && tp_poll_slow(tp_t, tp_st->breaker) != 0) "
                       "{ tp_rc = -1; goto tp_exit; }")
 
     def st_While(self, s: ir.While) -> None:
@@ -2163,7 +2166,7 @@ def impl(self) -> str:
             out.append(f"    {ct}{sep}{n} = {init};")
         # (a fast twin declares tp_st / tp_dict only when its body reads them: PyModule_GetState and
         # PyModule_GetDict are external calls, and a leaf must stay free of them to inline well)
-        out.append("    (void)tp_s; (void)tp_pc;" + (" (void)tp_fd;" if self.fast else "")
+        out.append("    (void)tp_s; (void)tp_pc;" + (" (void)tp_fd;" if self.fast else " (void)tp_ts;") + " (void)tp_t;"
                    + (" (void)tp_st;" if not self.fast or "tp_st" in body_text else "")
                    + (" (void)tp_dict;" if not self.fast or "tp_dict" in body_text else "")
                    + "".join(f" (void){_ident('l', n)};" for n in self.types if n not in params)
@@ -2184,27 +2187,27 @@ def impl(self) -> str:
                 # Precheck: room for the k frames of this whole call tree, then run its uncounted
                 # twin. Without room, fall through to the counted path below, which raises
                 # RecursionError at exactly the call where the interpreter would.
-                out.append(f"    if (tp_depth_room({k_bound})) {{")
+                out.append(f"    if (tp_depth_room(tp_ts, tp_t, {k_bound})) {{")
                 out.append("        int tp_fr;")
                 out.append("        tp_trace_fast_hit();")
                 out.append(f"        tp_fr = {_ident('tp_fimpl', f.name)}({_impl_call_args(f, '1')});")
-                out.append("        tp_depth_done();")
+                out.append("        tp_depth_done(tp_t);")
                 out.append("        return tp_fr;")
                 out.append("    }")
-            out.append("    if (tp_enter_call() != 0) return -1;")
+            out.append("    if (tp_enter_call(tp_ts, tp_t) != 0) return -1;")
         # Eval breaker (#141): a snapshot-holding function never polls and defers others' polls.
         if self.snapshot:
-            out.append("    tp_snapshot_depth++;")
+            out.append("    tp_t->snapshot_depth++;")
         elif not self.fast:
-            out.append("    if (tp_poll(tp_st->breaker) != 0) { tp_rc = -1; goto tp_exit; }")
+            out.append("    if (tp_poll(tp_t, tp_st->breaker) != 0) { tp_rc = -1; goto tp_exit; }")
         out += prologue
         out += self.lines
         out.append("tp_exit:")
         out += self._own_release_all()
         if self.snapshot:
-            out.append("    tp_snapshot_depth--;")
+            out.append("    tp_t->snapshot_depth--;")
         if not self.fast:
-            out.append("    tp_leave_call();")
+            out.append("    tp_leave_call(tp_t);")
         out.append("    return tp_rc;")
         out.append("}")
         out.append("")
@@ -2225,7 +2228,9 @@ def wrapper(self) -> str:
         n = len(f.params)
         L = [f"static PyObject *{_ident('tp_wrap', f.name)}(PyObject *tp_module, PyObject *const *tp_args, "
              "Py_ssize_t tp_nargs, PyObject *tp_kwnames)", "{",
-             "    int tp_s = 0, tp_x = 0;"]
+             "    int tp_s = 0, tp_x = 0;",
+             "    PyThreadState *tp_ts = PyThreadState_GetUnchecked(); /* the one lookup of this call (#125) */",
+             "    tp_tls *tp_t = tp_tls_get(); /* and its one TLS-address lookup (#125) */"]
         if f.returns not in (Type.NONE,):
             init = "NULL" if f.returns is Type.OBJ else "0"
             L.append(f"    {_CT[f.returns]} tp_r = {init};")
@@ -2333,7 +2338,7 @@ def wrapper(self) -> str:
         if gnames:
             at = len(args) - (1 if f.returns is not Type.NONE else 0)
             args[at:at] = gnames
-        L.append(f"    tp_s = {_ident('tp_impl', f.name)}(tp_module{''.join(', ' + a for a in args)});")
+        L.append(f"    tp_s = {_ident('tp_impl', f.name)}(tp_module, tp_ts, tp_t{''.join(', ' + a for a in args)});")
         if f.pure:
             L.append("    if (tp_s == 1) goto tp_deopt;")
         else:

python-multiplatform-ksp/src/main/python/typedpython/runtime/API.md

@@ -223,8 +223,20 @@ Same structs as above with `list == NULL` and `dirty == NULL`. 0 ok, -1 MemoryEr
     PyObject *tp_tuple(PyObject *const *items, Py_ssize_t n);           /* new tuple; steals nothing */
 
 ## Call depth (every compiled impl function: enter at entry, leave at every exit; issue #57)
-    int  tp_enter_call(void);   /* 0 ok; -1 RecursionError set, nothing counted (the caller must NOT call leave) */
-    void tp_leave_call(void);   /* undoes one successful tp_enter_call; exactly once per success, on ok, error and deopt exits alike */
+    int  tp_enter_call(PyThreadState *ts, tp_tls *t);  /* 0 ok; -1 RecursionError set, nothing counted (the caller must NOT call leave) */
+    void tp_leave_call(tp_tls *t);   /* undoes one successful tp_enter_call; exactly once per success, on ok, error and deopt exits alike */
+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_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.
+
+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 no further thread-state lookup. Likewise the per-thread state (call depth, poll countdown, snapshot depth) is one `__thread` struct
+`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.
+
 `tp_enter_call` makes compiled recursion fail where interpreted recursion fails. It counts, per thread,
 the compiled frames in flight (they have no Python frame), adds the current Python frame depth
 (`PyThreadState_GetFrame` / `PyFrame_GetBack`, public API only; cached while the top frame is the same
@@ -237,42 +249,41 @@ compiled module that sit between two frames of this one are not counted, and ano
 frames never are.
 
 ## Fast entry (bounded functions; issue #147)
-    int  tp_depth_room(Py_ssize_t k);   /* 1 iff tp_python_depth() + count + k <= Py_GetRecursionLimit() */
+    int  tp_depth_room(PyThreadState *ts, tp_tls *t, Py_ssize_t k);   /* 1 iff tp_python_depth(ts) + count + k <= Py_GetRecursionLimit() */
     void tp_trace_fast_hit(void);       /* counts a fast entry; a no-op unless -DTP_TRACE_FAST */
-    void tp_depth_py_begin(int fd);     /* count += fd, around code a fast tree may run (redo, poll) */
-    void tp_depth_py_end(int fd);       /* count -= fd; drops the cached frame at count 0 */
-    void tp_depth_done(void);           /* after a fast tree: drop the cached top frame when count is 0 */
+    void tp_depth_py_begin(tp_tls *t, int fd);     /* count += fd, around code a fast tree may run (redo, poll) */
+    void tp_depth_py_end(tp_tls *t, int fd);       /* count -= fd; drops the cached frame at count 0 */
+    void tp_depth_done(tp_tls *t);           /* after a fast tree: drop the cached top frame when count is 0 */
 A function is *bounded* when it is in no call cycle, every function it Calls is bounded, and its body
 has no node that can run Python code (CallObject, GetAttr, Truth, CompareObj, ObjToFloat, FieldGet,
 FieldSet, New, any OBJ-typed expression except Box of a scalar). A `redo` Call to a bounded callee is allowed. Its k is `1 + max(k(callee))` (1 with no
 callees). cgen emits, per bounded `f`, `static inline int tp_fimpl_f(...)` (the impl body without
 tp_enter_call / tp_leave_call and without the entry poll, Calls to the callees' `tp_fimpl_*`) and, at
-the top of `tp_impl_f` after the Param.cls guards, `if (tp_depth_room(k)) { r = tp_fimpl_f(...);
+the top of `tp_impl_f` after the Param.cls guards, `if (tp_depth_room(tp_ts, tp_t, k)) { r = tp_fimpl_f(...);
 tp_depth_done(); return r; }`. No room: the counted path runs and raises RecursionError exactly where
 it raised before. `tp_enter_call` compares `depth + count + 1`; the fast tree has at most k frames, so
 `depth + count + k` is the same comparison at its deepest frame. Py_EnterRecursiveCall is skipped on
 the fast path: the tree is bounded, so it adds at most k small C frames. A twin takes `int tp_fd`
 (the tree's active frames, itself included; the precheck passes 1, a Call passes `tp_fd + 1`). The
 two points where a fast tree can run Python code, a `redo` Call and the loop poll's `tp_poll_slow`,
-are wrapped in `tp_depth_py_begin(tp_fd)` / `tp_depth_py_end(tp_fd)`, so everything run from there
+are wrapped in `tp_depth_py_begin(tp_t, tp_fd)` / `tp_depth_py_end(tp_t, tp_fd)`, so everything run from there
 (interpreted frames, signal handlers, compiled calls through their precheck) sees the exact count.
 `-DTP_TRACE_FAST`
 (tests only) adds `__tp_trace_fast_hits__()` to the module: the number of entries that took the fast path.
 
 ## Eval breaker (top of every loop iteration and impl entry; issue #141)
-    int  tp_poll(PyObject *breaker);        /* 0 continue; -1 exception set (a signal handler raised) */
-    int  tp_poll_slow(PyObject *breaker);   /* every TP_POLL_INTERVAL polls: call the breaker */
+    int  tp_poll(tp_tls *t, PyObject *breaker);        /* 0 continue; -1 exception set (a signal handler raised) */
+    int  tp_poll_slow(tp_tls *t, PyObject *breaker);   /* every TP_POLL_INTERVAL polls: call the breaker */
     int  tp_pc_tick(int *pc);               /* loop countdown: 1 (and reset) when the slow path is due */
-    Py_ssize_t tp_snapshot_depth;           /* THREAD-LOCAL; ++ at entry, -- at tp_exit of a snapshot-holding impl */
-    int        tp_poll_countdown;           /* THREAD-LOCAL */
+    tp_tls *tp_tls_get(void);               /* address of the one THREAD-LOCAL struct: depth, poll_countdown, snapshot_depth (++ at entry, -- at tp_exit of a snapshot-holding impl) */
 `breaker` is the module state's `lambda: None`. Calling it runs CPython's own eval-breaker handling
 in its RESUME: signals, pending calls, stop-the-world requests of other threads (the free-threaded
 GC), async exceptions and scheduled collections, exactly as between two bytecodes. A function that holds an array copy-in
 (ir.ArrayParam) or an entry-globals snapshot never polls, and while one is active the slow path does
 nothing: the event stays pending and is handled after it returns.
 
 ## Tracing, profiling and monitoring (entry of every function; issue #142, SPEC N-14)
-    int  tp_traced(PyObject *module, PyCodeObject *code);       /* 1: run the interpreted original */
+    int  tp_traced(PyObject *module, PyThreadState *ts, PyCodeObject *code);       /* 1: run the interpreted original */
     PyObject *tp_state.orig_code[];  /* GENERATED, module state: STRONG ref to the original's code, per interned name */
 `tp_traced` is true when the calling thread has a `sys.settrace` / `sys.setprofile` function
 (`PyThreadState.c_tracefunc` / `c_profilefunc`), when any `sys.monitoring` tool has a global event
@@ -440,8 +451,8 @@ static or module-state field must be added here):
 | State | Where | Class | How |
 |---|---|---|---|
 | `tp_depth` (compiled frames in flight, cached frame) | tp_runtime.h static | thread-local | `TP_THREAD_LOCAL` |
-| `tp_poll_countdown` | tp_runtime.h static | thread-local | `TP_THREAD_LOCAL` |
-| `tp_snapshot_depth` | tp_runtime.h static | thread-local | `TP_THREAD_LOCAL`: one thread's array function never silences another thread's polls |
+| `tp_tls_v.poll_countdown` | tp_runtime.h static | thread-local | `TP_THREAD_LOCAL` struct `tp_tls` |
+| `tp_tls_v.snapshot_depth` | tp_runtime.h static | thread-local | `TP_THREAD_LOCAL`: one thread's array function never silences another thread's polls |
 | `tp_pc` (loop countdown) | local of each impl | per call | a C local |
 | `tp_globals_epoch` | tp_runtime.h static | atomic | relaxed loads, acq_rel add in the watcher callback (any thread); the epoch + dict lookup that remember a binding run in the module dict's critical section |
 | `cls_rt[k].tag`, `.bad_tag`, `.gepoch` | module state | atomic | relaxed `tp_atomic_*`; a stale tag is never reused (tags only grow), so a lost or reordered refresh costs one more refresh |