gh-143421: Lazily allocate tracer code and opt buffers (GH-143597)

Author: Fidget-Spinner

Include/internal/pycore_optimizer_types.h

@@ -10,8 +10,8 @@ extern "C" {
 
 #include "pycore_uop.h"  // UOP_MAX_TRACE_LENGTH
 
-// Holds locals, stack, locals, stack ... co_consts (in that order)
-#define MAX_ABSTRACT_INTERP_SIZE 4096
+// Holds locals, stack, locals, stack ... (in that order)
+#define MAX_ABSTRACT_INTERP_SIZE 512
 
 #define TY_ARENA_SIZE (UOP_MAX_TRACE_LENGTH * 5)
 

Include/internal/pycore_tstate.h

@@ -56,8 +56,8 @@ typedef struct _PyJitTracerState {
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
     _PyJitTracerTranslatorState translator_state;
-    JitOptContext opt_context;
-    _PyUOpInstruction code_buffer[UOP_MAX_TRACE_LENGTH];
+    JitOptContext *opt_context;
+    _PyUOpInstruction *code_buffer;
 } _PyJitTracerState;
 
 #endif

Python/optimizer.c

@@ -1025,6 +1025,13 @@ _PyJit_TryInitializeTracing(
     if (oparg > 0xFFFF) {
         return 0;
     }
+    if (_tstate->jit_tracer_state.code_buffer == NULL) {
+        _tstate->jit_tracer_state.code_buffer = (_PyUOpInstruction *)_PyObject_VirtualAlloc(UOP_BUFFER_SIZE);
+        if (_tstate->jit_tracer_state.code_buffer == NULL) {
+            // Don't error, just go to next instruction.
+            return 0;
+        }
+    }
     PyObject *func = PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
     if (func == NULL) {
         return 0;

Python/optimizer_analysis.c

@@ -345,7 +345,15 @@ optimize_uops(
     assert(!PyErr_Occurred());
     PyFunctionObject *func = tstate->jit_tracer_state.initial_state.func;
 
-    JitOptContext *ctx = &tstate->jit_tracer_state.opt_context;
+    JitOptContext *ctx = tstate->jit_tracer_state.opt_context;
+    if (ctx == NULL) {
+        ctx = (JitOptContext *)_PyObject_VirtualAlloc(sizeof(JitOptContext));
+        if (ctx == NULL) {
+            // Don't error, just bail.
+            return 0;
+        }
+        tstate->jit_tracer_state.opt_context = ctx;
+    }
     uint32_t opcode = UINT16_MAX;
 
     // Make sure that watchers are set up

Python/pystate.c

@@ -1553,6 +1553,8 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     init_policy(&_tstate->policy.jit.side_exit_initial_backoff,
                 "PYTHON_JIT_SIDE_EXIT_INITIAL_BACKOFF",
                 SIDE_EXIT_INITIAL_BACKOFF, 0, MAX_BACKOFF);
+    _tstate->jit_tracer_state.code_buffer = NULL;
+    _tstate->jit_tracer_state.opt_context = NULL;
 #endif
     tstate->delete_later = NULL;
 
@@ -1867,6 +1869,18 @@ tstate_delete_common(PyThreadState *tstate, int release_gil)
     assert(tstate_impl->refcounts.values == NULL);
 #endif
 
+#if _Py_TIER2
+    _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
+    if (_tstate->jit_tracer_state.code_buffer != NULL) {
+        _PyObject_VirtualFree(_tstate->jit_tracer_state.code_buffer, UOP_BUFFER_SIZE);
+        _tstate->jit_tracer_state.code_buffer = NULL;
+    }
+    if (_tstate->jit_tracer_state.opt_context != NULL) {
+        _PyObject_VirtualFree(_tstate->jit_tracer_state.opt_context, sizeof(JitOptContext));
+        _tstate->jit_tracer_state.opt_context = NULL;
+    }
+#endif
+
     HEAD_UNLOCK(runtime);
 
     // XXX Unbind in PyThreadState_Clear(), or earlier