GH-144651: Optimize the new uops added when recording values during tracing. (GH-144948)

* Handle dependencies in the optimizer, not the tracer
* Strengthen some checks to avoid relying on optimizer for correctness

Author: markshannon

Include/internal/pycore_optimizer.h

@@ -22,6 +22,10 @@ typedef struct _PyJitUopBuffer {
     _PyUOpInstruction *end;
 } _PyJitUopBuffer;
 
+typedef struct _JitOptRefBuffer {
+    JitOptRef *used;
+    JitOptRef *end;
+} _JitOptRefBuffer;
 
 typedef struct _JitOptContext {
     char done;
@@ -37,10 +41,15 @@ typedef struct _JitOptContext {
     // Arena for the symbolic types.
     ty_arena t_arena;
 
-    JitOptRef *n_consumed;
-    JitOptRef *limit;
-    JitOptRef locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
+    /* To do -- We could make this more space efficient
+     * by using a single array and growing the stack and
+     * locals toward each other. */
+    _JitOptRefBuffer locals;
+    _JitOptRefBuffer stack;
+    JitOptRef locals_array[ABSTRACT_INTERP_LOCALS_SIZE];
+    JitOptRef stack_array[ABSTRACT_INTERP_STACK_SIZE];
     _PyJitUopBuffer out_buffer;
+    _PyBloomFilter *dependencies;
 } JitOptContext;
 
 
@@ -83,13 +92,11 @@ typedef struct _PyJitTracerInitialState {
 } _PyJitTracerInitialState;
 
 typedef struct _PyJitTracerPreviousState {
-    bool dependencies_still_valid;
     int instr_oparg;
     int instr_stacklevel;
     _Py_CODEUNIT *instr;
     PyCodeObject *instr_code; // Strong
     struct _PyInterpreterFrame *instr_frame;
-    _PyBloomFilter dependencies;
     PyObject *recorded_value; // Strong, may be NULL
 } _PyJitTracerPreviousState;
 
@@ -303,25 +310,24 @@ extern void _Py_uop_sym_set_recorded_type(JitOptContext *ctx, JitOptRef sym, PyT
 extern void _Py_uop_sym_set_recorded_gen_func(JitOptContext *ctx, JitOptRef ref, PyFunctionObject *value);
 extern PyCodeObject *_Py_uop_sym_get_probable_func_code(JitOptRef sym);
 extern PyObject *_Py_uop_sym_get_probable_value(JitOptRef sym);
+extern JitOptRef *_Py_uop_sym_set_stack_depth(JitOptContext *ctx, int stack_depth, JitOptRef *current_sp);
 
-extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
+extern void _Py_uop_abstractcontext_init(JitOptContext *ctx, _PyBloomFilter *dependencies);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);
 
 extern _Py_UOpsAbstractFrame *_Py_uop_frame_new(
     JitOptContext *ctx,
     PyCodeObject *co,
-    int curr_stackentries,
     JitOptRef *args,
     int arg_len);
 
 extern _Py_UOpsAbstractFrame *_Py_uop_frame_new_from_symbol(
     JitOptContext *ctx,
     JitOptRef callable,
-    int curr_stackentries,
     JitOptRef *args,
     int arg_len);
 
-extern int _Py_uop_frame_pop(JitOptContext *ctx, PyCodeObject *co, int curr_stackentries);
+extern int _Py_uop_frame_pop(JitOptContext *ctx, PyCodeObject *co);
 
 PyAPI_FUNC(PyObject *) _Py_uop_symbols_test(PyObject *self, PyObject *ignored);
 
@@ -357,8 +363,6 @@ PyAPI_FUNC(void) _PyJit_FinalizeTracing(PyThreadState *tstate, int err);
 void _PyPrintExecutor(_PyExecutorObject *executor, const _PyUOpInstruction *marker);
 void _PyJit_TracerFree(_PyThreadStateImpl *_tstate);
 
-void _PyJit_Tracer_InvalidateDependency(PyThreadState *old_tstate, void *obj);
-
 #ifdef _Py_TIER2
 typedef void (*_Py_RecordFuncPtr)(_PyInterpreterFrame *frame, _PyStackRef *stackpointer, int oparg, PyObject **recorded_value);
 PyAPI_DATA(const _Py_RecordFuncPtr) _PyOpcode_RecordFunctions[];

Include/internal/pycore_optimizer_types.h

@@ -11,8 +11,9 @@ extern "C" {
 #include <stdbool.h>
 #include "pycore_uop.h"  // UOP_MAX_TRACE_LENGTH
 
-// Holds locals, stack, locals, stack ... (in that order)
-#define MAX_ABSTRACT_INTERP_SIZE 512
+#define ABSTRACT_INTERP_STACK_SIZE 256
+#define ABSTRACT_INTERP_LOCALS_SIZE 512
+
 
 #define TY_ARENA_SIZE (UOP_MAX_TRACE_LENGTH * 5)
 
@@ -138,6 +139,7 @@ typedef struct _Py_UOpsAbstractFrame {
     // Max stacklen
     int stack_len;
     int locals_len;
+    bool caller; // We have made a call from this frame during the trace
     PyFunctionObject *func;
     PyCodeObject *code;
 

Include/internal/pycore_uop_ids.h

@@ -110,6 +110,7 @@ def f{n}():
             for exe in executors[:i]:
                 self.assertTrue(exe.is_valid())
 
+    @unittest.skipIf(os.getenv("PYTHON_UOPS_OPTIMIZE") == "0", "Needs uop optimizer to run.")
     def test_uop_optimizer_invalidation(self):
         # Generate a new function at each call
         ns = {}

Include/internal/pycore_uop_metadata.h

@@ -2433,7 +2433,6 @@ code_dealloc(PyObject *self)
         PyMem_Free(co_extra);
     }
 #ifdef _Py_TIER2
-    _PyJit_Tracer_InvalidateDependency(tstate, self);
     if (co->co_executors != NULL) {
         clear_executors(co);
     }

Lib/test/test_capi/test_opt.py

@@ -262,7 +262,6 @@ framelocalsproxy_setitem(PyObject *self, PyObject *key, PyObject *value)
 
 #if _Py_TIER2
         _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), co, 1);
-        _PyJit_Tracer_InvalidateDependency(_PyThreadState_GET(), co);
 #endif
 
         _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);

Modules/_testinternalcapi/test_cases.c.h

@@ -12,7 +12,7 @@
 #include "pycore_setobject.h"     // _PySet_NextEntry()
 #include "pycore_stats.h"
 #include "pycore_weakref.h"       // FT_CLEAR_WEAKREFS()
-#include "pycore_optimizer.h"     // _PyJit_Tracer_InvalidateDependency
+#include "pycore_optimizer.h"     // _Py_Executors_InvalidateDependency
 
 static const char *
 func_event_name(PyFunction_WatchEvent event) {

Objects/codeobject.c

@@ -3125,10 +3125,10 @@ dummy_func(
             assert(executor->vm_data.code == code);
             assert(executor->vm_data.valid);
             assert(tstate->current_executor == NULL);
-            /* If the eval breaker is set then stay in tier 1.
-             * This avoids any potentially infinite loops
-             * involving _RESUME_CHECK */
-            if (_Py_atomic_load_uintptr_relaxed(&tstate->eval_breaker) & _PY_EVAL_EVENTS_MASK) {
+            /* If the eval breaker is set, or instrumentation is needed, then stay in tier 1.
+             * This avoids any potentially infinite loops involving _RESUME_CHECK */
+            uintptr_t iversion = FT_ATOMIC_LOAD_UINTPTR_ACQUIRE(code->_co_instrumentation_version);
+            if (_Py_atomic_load_uintptr_relaxed(&tstate->eval_breaker) != iversion) {
                 opcode = executor->vm_data.opcode;
                 oparg = (oparg & ~255) | executor->vm_data.oparg;
                 next_instr = this_instr;
@@ -5616,9 +5616,9 @@ dummy_func(
             HANDLE_PENDING_AND_DEOPT_IF(_Py_emscripten_signal_clock == 0);
             _Py_emscripten_signal_clock -= Py_EMSCRIPTEN_SIGNAL_HANDLING;
 #endif
+            uintptr_t iversion = FT_ATOMIC_LOAD_UINTPTR_ACQUIRE(_PyFrame_GetCode(frame)->_co_instrumentation_version);
             uintptr_t eval_breaker = _Py_atomic_load_uintptr_relaxed(&tstate->eval_breaker);
-            HANDLE_PENDING_AND_DEOPT_IF(eval_breaker & _PY_EVAL_EVENTS_MASK);
-            assert(tstate->tracing || eval_breaker == FT_ATOMIC_LOAD_UINTPTR_ACQUIRE(_PyFrame_GetCode(frame)->_co_instrumentation_version));
+            HANDLE_PENDING_AND_DEOPT_IF(eval_breaker != iversion);
         }
 
         tier2 op(_COLD_EXIT, ( -- )) {
@@ -5668,9 +5668,9 @@ dummy_func(
             Py_UNREACHABLE();
         }
 
-        tier2 op(_GUARD_CODE, (version/2 -- )) {
+        tier2 op(_GUARD_CODE_VERSION, (version/2 -- )) {
             PyObject *code = PyStackRef_AsPyObjectBorrow(frame->f_executable);
-            EXIT_IF(code == Py_None);
+            assert(PyCode_Check(code));
             EXIT_IF(((PyCodeObject *)code)->co_version != version);
         }