GH-135379: Top of stack caching for the JIT.

Include/cpython/pystats.h

@@ -29,7 +29,7 @@
 #  error "this header file must not be included directly"
 #endif
 
-#define PYSTATS_MAX_UOP_ID 1024
+#define PYSTATS_MAX_UOP_ID 2000
 
 #define SPECIALIZATION_FAILURE_KINDS 60
 

Include/internal/pycore_jit.h

@@ -18,7 +18,10 @@ extern "C" {
 
 #ifdef _Py_JIT
 
-typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate);
+typedef _Py_CODEUNIT *(*jit_func)(
+    _PyInterpreterFrame *frame, _PyStackRef *stack_pointer, PyThreadState *tstate,
+    _PyStackRef _tos_cache0, _PyStackRef _tos_cache1, _PyStackRef _tos_cache2
+);
 
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);

Include/internal/pycore_optimizer.h

@@ -24,19 +24,20 @@ typedef struct _PyExecutorLinkListNode {
 typedef struct {
     uint8_t opcode;
     uint8_t oparg;
-    uint8_t valid:1;
-    uint8_t linked:1;
-    uint8_t chain_depth:6;  // Must be big enough for MAX_CHAIN_DEPTH - 1.
+    uint8_t valid;
+    uint8_t linked;
+    uint8_t chain_depth;  // Must be big enough for MAX_CHAIN_DEPTH - 1.
     bool warm;
-    int index;           // Index of ENTER_EXECUTOR (if code isn't NULL, below).
+    int32_t index;           // Index of ENTER_EXECUTOR (if code isn't NULL, below).
     _PyBloomFilter bloom;
     _PyExecutorLinkListNode links;
     PyCodeObject *code;  // Weak (NULL if no corresponding ENTER_EXECUTOR).
 } _PyVMData;
 
 typedef struct _PyExitData {
     uint32_t target;
-    uint16_t index:14;
+    uint16_t index:12;
+    uint16_t stack_cache:2;
     uint16_t is_dynamic:1;
     uint16_t is_control_flow:1;
     _Py_BackoffCounter temperature;
@@ -344,17 +345,6 @@ extern _PyExecutorObject *_PyExecutor_GetColdDynamicExecutor(void);
 
 PyAPI_FUNC(void) _PyExecutor_ClearExit(_PyExitData *exit);
 
-static inline int is_terminator(const _PyUOpInstruction *uop)
-{
-    int opcode = uop->opcode;
-    return (
-        opcode == _EXIT_TRACE ||
-        opcode == _DEOPT ||
-        opcode == _JUMP_TO_TOP ||
-        opcode == _DYNAMIC_EXIT
-    );
-}
-
 extern void _PyExecutor_Free(_PyExecutorObject *self);
 
 PyAPI_FUNC(int) _PyDumpExecutors(FILE *out);

Include/internal/pycore_stackref.h

@@ -76,6 +76,8 @@ static const _PyStackRef PyStackRef_ERROR = { .index = (1 << Py_TAGGED_SHIFT) };
 
 #define INITIAL_STACKREF_INDEX (5 << Py_TAGGED_SHIFT)
 
+#define PyStackRef_ZERO_BITS PyStackRef_NULL
+
 static inline _PyStackRef
 PyStackRef_Wrap(void *ptr)
 {
@@ -369,6 +371,10 @@ PyStackRef_IsNullOrInt(_PyStackRef ref);
 
 static const _PyStackRef PyStackRef_ERROR = { .bits = Py_TAG_INVALID };
 
+/* For use in the JIT to clear an unused value.
+ * PyStackRef_ZERO_BITS has no meaning and should not be used other than by the JIT. */
+static const _PyStackRef PyStackRef_ZERO_BITS = { .bits = 0 };
+
 /* Wrap a pointer in a stack ref.
  * The resulting stack reference is not safe and should only be used
  * in the interpreter to pass values from one uop to another.
@@ -922,6 +928,17 @@ _PyThreadState_PopCStackRef(PyThreadState *tstate, _PyCStackRef *ref)
     PyStackRef_XCLOSE(ref->ref);
 }
 
+static inline _PyStackRef
+_PyThreadState_PopCStackRefSteal(PyThreadState *tstate, _PyCStackRef *ref)
+{
+#ifdef Py_GIL_DISABLED
+    _PyThreadStateImpl *tstate_impl = (_PyThreadStateImpl *)tstate;
+    assert(tstate_impl->c_stack_refs == ref);
+    tstate_impl->c_stack_refs = ref->next;
+#endif
+    return ref->ref;
+}
+
 #ifdef Py_GIL_DISABLED
 
 static inline int

Include/internal/pycore_uop.h

@@ -38,9 +38,9 @@ typedef struct _PyUOpInstruction{
 // This is the length of the trace we translate initially.
 #ifdef Py_DEBUG
     // With asserts, the stencils are a lot larger
-#define UOP_MAX_TRACE_LENGTH 1000
+#define UOP_MAX_TRACE_LENGTH 2000
 #else
-#define UOP_MAX_TRACE_LENGTH 3000
+#define UOP_MAX_TRACE_LENGTH 5000
 #endif
 #define UOP_BUFFER_SIZE (UOP_MAX_TRACE_LENGTH * sizeof(_PyUOpInstruction))