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 1500
 
 #define SPECIALIZATION_FAILURE_KINDS 60
 

Include/internal/pycore_jit.h

@@ -15,7 +15,7 @@ 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

@@ -119,7 +119,7 @@ PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyInterpreterState *interp);
 #define JIT_CLEANUP_THRESHOLD 100000
 
 // This is the length of the trace we project initially.
-#define UOP_MAX_TRACE_LENGTH 800
+#define UOP_MAX_TRACE_LENGTH 1600
 
 #define TRACE_STACK_SIZE 5
 

Include/internal/pycore_stackref.h

@@ -70,6 +70,9 @@ static const _PyStackRef PyStackRef_ERROR = { .index = 2 };
 #define PyStackRef_False ((_PyStackRef){ .index = 6 })
 #define PyStackRef_True ((_PyStackRef){ .index = 8 })
 
+#define PyStackRef_ZERO_BITS PyStackRef_NULL
+
+
 #define INITIAL_STACKREF_INDEX 10
 
 static inline int
@@ -264,6 +267,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.

Include/internal/pycore_stats.h

@@ -28,7 +28,7 @@ extern "C" {
 #define GC_STAT_ADD(gen, name, n) do { if (_Py_stats) _Py_stats->gc_stats[(gen)].name += (n); } while (0)
 #define OPT_STAT_INC(name) do { if (_Py_stats) _Py_stats->optimization_stats.name++; } while (0)
 #define OPT_STAT_ADD(name, n) do { if (_Py_stats) _Py_stats->optimization_stats.name += (n); } while (0)
-#define UOP_STAT_INC(opname, name) do { if (_Py_stats) { assert(opname < 512); _Py_stats->optimization_stats.opcode[opname].name++; } } while (0)
+#define UOP_STAT_INC(opname, name) do { if (_Py_stats) { _Py_stats->optimization_stats.opcode[opname].name++; } } while (0)
 #define UOP_PAIR_INC(uopcode, lastuop)                                              \
     do {                                                                            \
         if (lastuop && _Py_stats) {                                                 \