Address code review

Author: corona10

Include/internal/pycore_interp_structs.h

@@ -14,6 +14,7 @@ extern "C" {
 #include "pycore_structs.h"       // PyHamtObject
 #include "pycore_tstate.h"        // _PyThreadStateImpl
 #include "pycore_typedefs.h"      // _PyRuntimeState
+#include "pycore_uop.h"           // struct _PyUOpInstruction
 
 
 #define CODE_MAX_WATCHERS 8
@@ -898,6 +899,10 @@ struct _is {
     struct _stoptheworld_state stoptheworld;
     struct _qsbr_shared qsbr;
 
+#ifdef _Py_TIER2
+    struct _PyUOpInstruction *jit_uop_buffer;
+#endif
+
 #if defined(Py_GIL_DISABLED)
     struct _mimalloc_interp_state mimalloc;
     struct _brc_state brc;  // biased reference counting state

Include/internal/pycore_tstate.h

@@ -12,7 +12,6 @@ extern "C" {
 #include "pycore_freelist_state.h"  // struct _Py_freelists
 #include "pycore_mimalloc.h"        // struct _mimalloc_thread_state
 #include "pycore_qsbr.h"            // struct qsbr
-#include "pycore_uop.h"             // struct _PyUOpInstruction
 
 
 #ifdef Py_GIL_DISABLED
@@ -77,10 +76,6 @@ typedef struct _PyThreadStateImpl {
     Py_ssize_t reftotal;  // this thread's total refcount operations
 #endif
 
-#ifdef _Py_TIER2
-    struct _PyUOpInstruction *jit_uop_buffer;
-#endif
-
 } _PyThreadStateImpl;
 
 #ifdef __cplusplus

Python/optimizer.c

@@ -1280,14 +1280,14 @@ uop_optimize(
 {
     _PyBloomFilter dependencies;
     _Py_BloomFilter_Init(&dependencies);
-    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
-    if (tstate->jit_uop_buffer == NULL) {
-        tstate->jit_uop_buffer = (_PyUOpInstruction *)PyMem_RawMalloc(UOP_MAX_TRACE_LENGTH*sizeof(_PyUOpInstruction));
-        if (tstate->jit_uop_buffer == NULL) {
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp->jit_uop_buffer == NULL) {
+        interp->jit_uop_buffer = (_PyUOpInstruction *)PyMem_RawMalloc(UOP_MAX_TRACE_LENGTH*sizeof(_PyUOpInstruction));
+        if (interp->jit_uop_buffer == NULL) {
             return 0;
         }
     }
-    _PyUOpInstruction *buffer = tstate->jit_uop_buffer;
+    _PyUOpInstruction *buffer = interp->jit_uop_buffer;
     OPT_STAT_INC(attempts);
     char *env_var = Py_GETENV("PYTHON_UOPS_OPTIMIZE");
     bool is_noopt = true;

Python/pystate.c

@@ -556,6 +556,11 @@ init_interpreter(PyInterpreterState *interp,
 #ifdef Py_GIL_DISABLED
     _Py_brc_init_state(interp);
 #endif
+
+#ifdef _Py_TIER2
+     // Ensure the buffer is to be set as NULL.
+    interp->jit_uop_buffer = NULL;
+#endif
     llist_init(&interp->mem_free_queue.head);
     llist_init(&interp->asyncio_tasks_head);
     interp->asyncio_tasks_lock = (PyMutex){0};
@@ -803,6 +808,10 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
 
 #ifdef _Py_TIER2
     _Py_ClearExecutorDeletionList(interp);
+    if (interp->jit_uop_buffer != NULL) {
+        PyMem_RawFree(interp->jit_uop_buffer);
+        interp->jit_uop_buffer = NULL;
+    }
 #endif
     _PyAST_Fini(interp);
     _PyAtExit_Fini(interp);
@@ -1536,11 +1545,6 @@ new_threadstate(PyInterpreterState *interp, int whence)
     }
 #endif
 
-#ifdef _Py_TIER2
-     // Ensure the buffer is to be set as NULL.
-    tstate->jit_uop_buffer = NULL;
-#endif
-
     /* We serialize concurrent creation to protect global state. */
     HEAD_LOCK(interp->runtime);
 
@@ -1731,14 +1735,6 @@ PyThreadState_Clear(PyThreadState *tstate)
     _Py_ClearTLBCIndex((_PyThreadStateImpl *)tstate);
 #endif
 
-#ifdef _Py_TIER2
-    _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
-    if (_tstate->jit_uop_buffer != NULL) {
-        PyMem_RawFree(_tstate->jit_uop_buffer);
-        _tstate->jit_uop_buffer = NULL;
-    }
-#endif
-
     // Merge our queue of pointers to be freed into the interpreter queue.
     _PyMem_AbandonDelayed(tstate);