move JitOptContext to _PyThreadStateImpl

Author: cocolato

Include/internal/pycore_optimizer.h

@@ -84,6 +84,7 @@ PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyInterpreterState *interp);
 #define JIT_CLEANUP_THRESHOLD 1000
 
 int _Py_uop_analyze_and_optimize(
+    struct _PyThreadStateImpl *tstate,
     PyFunctionObject *func,
     _PyUOpInstruction *trace, int trace_len, int curr_stackentries,
     _PyBloomFilter *dependencies);

Include/internal/pycore_tstate.h

@@ -53,6 +53,7 @@ typedef struct _PyJitTracerTranslatorState {
 
 typedef struct _PyJitTracerState {
     _PyUOpInstruction *code_buffer;
+    struct _JitOptContext *opt_context;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
     _PyJitTracerTranslatorState translator_state;

Python/optimizer.c

@@ -1484,6 +1484,7 @@ uop_optimize(
     OPT_STAT_INC(traces_created);
     if (!is_noopt) {
         length = _Py_uop_analyze_and_optimize(
+            _tstate,
             _tstate->jit_tracer_state.initial_state.func,
             buffer,length,
             curr_stackentries, dependencies);

Python/optimizer_analysis.c

@@ -18,6 +18,7 @@
 #include "pycore_opcode_metadata.h"
 #include "pycore_opcode_utils.h"
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
+#include "pycore_tstate.h"        // _PyThreadStateImpl
 #include "pycore_uop_metadata.h"
 #include "pycore_long.h"
 #include "pycore_interpframe.h"  // _PyFrame_GetCode
@@ -334,6 +335,7 @@ _Py_opt_assert_within_stack_bounds(
 /* >0 (length) for success, 0 for not ready, clears all possible errors. */
 static int
 optimize_uops(
+    _PyThreadStateImpl *tstate,
     PyFunctionObject *func,
     _PyUOpInstruction *trace,
     int trace_len,
@@ -343,8 +345,15 @@ optimize_uops(
 {
     assert(!PyErr_Occurred());
 
-    JitOptContext context;
-    JitOptContext *ctx = &context;
+    // Use thread-local JitOptContext to avoid stack overflow
+    JitOptContext *ctx = tstate->jit_tracer_state.opt_context;
+    if (ctx == NULL) {
+        ctx = (JitOptContext *)PyMem_RawMalloc(sizeof(JitOptContext));
+        if (ctx == NULL) {
+            return 0;
+        }
+        tstate->jit_tracer_state.opt_context = ctx;
+    }
     uint32_t opcode = UINT16_MAX;
 
     // Make sure that watchers are set up
@@ -574,6 +583,7 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
 //  > 0 - length of optimized trace
 int
 _Py_uop_analyze_and_optimize(
+    _PyThreadStateImpl *tstate,
     PyFunctionObject *func,
     _PyUOpInstruction *buffer,
     int length,
@@ -584,7 +594,7 @@ _Py_uop_analyze_and_optimize(
     OPT_STAT_INC(optimizer_attempts);
 
     length = optimize_uops(
-         func, buffer,
+         tstate, func, buffer,
          length, curr_stacklen, dependencies);
 
     if (length == 0) {

Python/pystate.c

@@ -1554,6 +1554,7 @@ init_threadstate(_PyThreadStateImpl *_tstate,
                 "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;
 
@@ -1874,6 +1875,12 @@ tstate_delete_common(PyThreadState *tstate, int release_gil)
         _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) {
+        // Ensure any resources in opt_context are cleaned up
+        _Py_uop_abstractcontext_fini(_tstate->jit_tracer_state.opt_context);
+        PyMem_RawFree(_tstate->jit_tracer_state.opt_context);
+        _tstate->jit_tracer_state.opt_context = NULL;
+    }
 #endif
 
     HEAD_UNLOCK(runtime);