add _PyJitOptState

Author: cocolato

Include/internal/pycore_tstate.h

@@ -51,9 +51,13 @@ typedef struct _PyJitTracerTranslatorState {
     int jump_backward_seen;
 } _PyJitTracerTranslatorState;
 
+typedef struct _PyJitOptState {
+    struct _JitOptContext *opt_context;
+} _PyJitOptState;
+
 typedef struct _PyJitTracerState {
     _PyUOpInstruction *code_buffer;
-    struct _JitOptContext *opt_context;
+    _PyJitOptState opt_state;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
     _PyJitTracerTranslatorState translator_state;

Python/optimizer_analysis.c

@@ -345,14 +345,13 @@ optimize_uops(
     assert(!PyErr_Occurred());
     PyFunctionObject *func = tstate->jit_tracer_state.initial_state.func;
 
-    // Use thread-local JitOptContext to avoid stack overflow
-    JitOptContext *ctx = tstate->jit_tracer_state.opt_context;
+    JitOptContext *ctx = tstate->jit_tracer_state.opt_state.opt_context;
     if (ctx == NULL) {
         ctx = (JitOptContext *)PyMem_RawMalloc(sizeof(JitOptContext));
         if (ctx == NULL) {
             return 0;
         }
-        tstate->jit_tracer_state.opt_context = ctx;
+        tstate->jit_tracer_state.opt_state.opt_context = ctx;
     }
     uint32_t opcode = UINT16_MAX;
 

Python/pystate.c

@@ -1554,7 +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;
+    _tstate->jit_tracer_state.opt_state.opt_context = NULL;
 #endif
     tstate->delete_later = NULL;
 
@@ -1875,11 +1875,9 @@ 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;
+    if (_tstate->jit_tracer_state.opt_state.opt_context != NULL) {
+        PyMem_RawFree(_tstate->jit_tracer_state.opt_state.opt_context);
+        _tstate->jit_tracer_state.opt_state.opt_context = NULL;
     }
 #endif