Address review comments

Author: markshannon

Misc/NEWS.d/next/Core_and_Builtins/2025-06-13-13-32-16.gh-issue-135379.pAxZgy.rst

@@ -1,6 +1,6 @@
-Implement a limited form of register allocation know as "top of stack
+Implement a limited form of register allocation known as "top of stack
 caching" in the JIT. It works by keeping 0-3 of the top items in the stack
-in registers. The code generator generates multiple versions of thos uops
+in registers. The code generator generates multiple versions of those uops
 that do not escape and are relatively small. During JIT compilation, the
 copy that produces the least memory traffic is selected, spilling or
 reloading values when needed.

Misc/NEWS.d/next/Core_and_Builtins/2025-06-20-16-03-59.gh-issue-135379.eDg89T.rst

@@ -1005,18 +1005,24 @@ count_exits(_PyUOpInstruction *buffer, int length)
     return exit_count;
 }
 
+#define MAX_CACHED_REGISTER 3
+
+/* The number of cached registers at any exit (`EXIT_IF` or `DEOPT_IF`)
+ * This is the number of cached at entries at start, unless the uop is
+ * marked as `exit_depth_is_output` in which case it is the number of
+ * cached entries at the end */
 static int
-get_exit_depth(_PyUOpInstruction *inst)
+get_cached_entries_for_side_exit(_PyUOpInstruction *inst)
 {
-
     // TO DO -- Add another generated table for this?
     int base_opcode = _PyUop_Uncached[inst->opcode];
     assert(base_opcode != 0);
     if (_PyUop_Flags[base_opcode] & HAS_DEOPT_FLAG) {
         return 0;
     }
     int input = -1;
-    for (int i = 0; i < 4; i++) {
+    /* Find number of cached entries at input. */
+    for (int i = 0; i <= MAX_CACHED_REGISTER; i++) {
         if (_PyUop_Caching[base_opcode].opcodes[i] == inst->opcode) {
             input = i;
             break;
@@ -1076,7 +1082,7 @@ prepare_for_execution(_PyUOpInstruction *buffer, int length)
         if (_PyUop_Flags[base_opcode] & (HAS_EXIT_FLAG | HAS_DEOPT_FLAG)) {
             uint16_t base_exit_op = (_PyUop_Flags[base_opcode] & HAS_EXIT_FLAG) ?
                 _EXIT_TRACE : _DEOPT;
-            int exit_depth = get_exit_depth(inst);
+            int exit_depth = get_cached_entries_for_side_exit(inst);
             uint16_t exit_op = _PyUop_Caching[base_exit_op].opcodes[exit_depth];
             int32_t jump_target = target;
             if (is_for_iter_test[base_opcode]) {

Python/optimizer.c

@@ -7,6 +7,8 @@
 
 from parser import Stmt, SimpleStmt, BlockStmt, IfStmt, WhileStmt
 
+MAX_CACHED_REGISTER = 3
+
 @dataclass
 class EscapingCall:
     stmt: SimpleStmt
@@ -1221,13 +1223,13 @@ def is_large(uop: Uop) -> bool:
 
 def get_uop_cache_depths(uop: Uop) -> Iterator[tuple[int, int, int]]:
     if uop.name == "_SPILL_OR_RELOAD":
-        for inputs in range(4):
-            for outputs in range(4):
+        for inputs in range(MAX_CACHED_REGISTER+1):
+            for outputs in range(MAX_CACHED_REGISTER+1):
                 if inputs != outputs:
                     yield inputs, outputs, inputs
         return
     if uop.name == "_EXIT_TRACE":
-        for i in range(4):
+        for i in range(MAX_CACHED_REGISTER+1):
             yield i, 0, i
         return
     if uop.name in ("_START_EXECUTOR", "_JUMP_TO_TOP", "_DEOPT", "_ERROR_POP_N"):
@@ -1256,24 +1258,18 @@ def get_uop_cache_depths(uop: Uop) -> Iterator[tuple[int, int, int]]:
         if item.peek and uop.properties.escapes:
             break
         ideal_outputs += 1
-    if ideal_inputs > 3:
-        ideal_inputs = 3
-    if ideal_outputs > 3:
-        ideal_outputs = 3
+    if ideal_inputs > MAX_CACHED_REGISTER:
+        ideal_inputs = MAX_CACHED_REGISTER
+    if ideal_outputs > MAX_CACHED_REGISTER:
+        ideal_outputs = MAX_CACHED_REGISTER
     if non_decref_escape:
         yield 0, ideal_outputs, 0
         return
-    # If a uop has an exit, we can get in a mess if the stack caching
-    # changes during execution.
-    #if has_exit and ideal_inputs != ideal_outputs:
-    #    n = min(ideal_inputs, ideal_outputs)
-    #    yield n, n
-    #    return
     exit_depth = ideal_outputs if uop.properties.sync_sp else ideal_inputs
     yield ideal_inputs, ideal_outputs, exit_depth
     if uop.properties.escapes or uop.properties.sync_sp or has_array or is_large(uop):
         return
-    if ideal_inputs >= 3 or ideal_outputs >= 3:
+    if ideal_inputs == MAX_CACHED_REGISTER or ideal_outputs == MAX_CACHED_REGISTER:
         return
     inputs, outputs = ideal_inputs, ideal_outputs
     if inputs < outputs:

Tools/cases_generator/analyzer.py

@@ -10,6 +10,7 @@
     analyze_files,
     get_uop_cache_depths,
     Uop,
+    MAX_CACHED_REGISTER,
 )
 from generators_common import (
     DEFAULT_INPUT,
@@ -24,14 +25,15 @@
 DEFAULT_OUTPUT = ROOT / "Include/internal/pycore_uop_metadata.h"
 
 def uop_cache_info(uop: Uop) -> str | None:
+    table_size = MAX_CACHED_REGISTER + 1
     if uop.name == "_SPILL_OR_RELOAD":
         return None
     min_inputs = 4
-    uops = [ "0" ] * 4
+    uops = [ "0" ] * table_size
     for inputs, outputs, exit_depth in get_uop_cache_depths(uop):
         delta = outputs - inputs
         uops[inputs] = f"{uop.name}_r{inputs}{outputs}"
-    for i in range(4):
+    for i in range(table_size):
         if uops[i] != "0":
             max_inputs = i
             if i < min_inputs:
@@ -48,10 +50,10 @@ def generate_names_and_flags(analysis: Analysis, out: CWriter) -> None:
     out.emit("extern int _PyUop_num_popped(int opcode, int oparg);\n\n")
     out.emit("typedef struct _pyuop_info {\n")
     out.emit("int8_t min_input; int8_t max_input; int8_t delta;\n")
-    out.emit("int8_t exit_depth_is_output; uint16_t opcodes[4];\n")
+    out.emit(f"int8_t exit_depth_is_output; uint16_t opcodes[{MAX_CACHED_REGISTER+1}];\n")
     out.emit("} _PyUopCachingInfo;\n")
     out.emit("extern const _PyUopCachingInfo _PyUop_Caching[MAX_UOP_ID+1];\n\n")
-    out.emit("extern const uint16_t _PyUop_SpillsAndReloads[4][4];\n")
+    out.emit(f"extern const uint16_t _PyUop_SpillsAndReloads[{MAX_CACHED_REGISTER+1}][{MAX_CACHED_REGISTER+1}];\n")
     out.emit("extern const uint16_t _PyUop_Uncached[MAX_UOP_REGS_ID+1];\n\n")
     out.emit("#ifdef NEED_OPCODE_METADATA\n")
     out.emit("const uint16_t _PyUop_Flags[MAX_UOP_ID+1] = {\n")
@@ -80,9 +82,9 @@ def generate_names_and_flags(analysis: Analysis, out: CWriter) -> None:
             for inputs, outputs, _ in get_uop_cache_depths(uop):
                 out.emit(f"[{uop.name}_r{inputs}{outputs}] = {uop.name},\n")
     out.emit("};\n\n")
-    out.emit("const uint16_t _PyUop_SpillsAndReloads[4][4] = {\n")
-    for i in range(4):
-        for j in range(4):
+    out.emit(f"const uint16_t _PyUop_SpillsAndReloads[{MAX_CACHED_REGISTER+1}][{MAX_CACHED_REGISTER+1}] = {{\n")
+    for i in range(MAX_CACHED_REGISTER+1):
+        for j in range(MAX_CACHED_REGISTER+1):
             if i != j:
                 out.emit(f"[{i}][{j}] = _SPILL_OR_RELOAD_r{i}{j},\n")
     out.emit("};\n\n")