invalidate freed code/function objects used for global promotion

Author: Fidget-Spinner

Include/internal/pycore_interp_structs.h

@@ -947,6 +947,7 @@ struct _is {
     int jit_tracer_code_max_size;
     int jit_tracer_code_curr_size;
     _PyBloomFilter jit_tracer_dependencies;
+    bool jit_tracer_dependencies_still_valid;
     _PyUOpInstruction *jit_tracer_code_buffer;
     _Py_CODEUNIT *jit_tracer_initial_instr;
     int jit_tracer_initial_stack_depth;

Include/internal/pycore_optimizer.h

@@ -375,6 +375,9 @@ void
 _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit);
 
 void _PyJIT_FinalizeTracing(PyThreadState *tstate);
+
+void _Py_JITTracer_InvalidateDependency(PyThreadState *old_tstate, void *obj);
+
 #ifdef __cplusplus
 }
 #endif

Objects/codeobject.c

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

Objects/funcobject.c

@@ -11,7 +11,7 @@
 #include "pycore_setobject.h"     // _PySet_NextEntry()
 #include "pycore_stats.h"
 #include "pycore_weakref.h"       // FT_CLEAR_WEAKREFS()
-
+#include "pycore_optimizer.h"     // _Py_JITTracer_InvalidateDependency
 
 static const char *
 func_event_name(PyFunction_WatchEvent event) {
@@ -1149,6 +1149,9 @@ func_dealloc(PyObject *self)
     if (_PyObject_ResurrectEnd(self)) {
         return;
     }
+#if _Py_TIER2
+    _Py_JITTracer_InvalidateDependency(PyThreadState_GET(), self);
+#endif
     _PyObject_GC_UNTRACK(op);
     FT_CLEAR_WEAKREFS(self, op->func_weakreflist);
     (void)func_clear((PyObject*)op);

Python/optimizer.c

@@ -147,6 +147,11 @@ _PyOptimizer_Optimize(
         interp->compiling = false;
         return 0;
     }
+    // One of our depencies while tracing was invalidated. Not worth compiling.
+    if (!tstate->interp->jit_tracer_dependencies_still_valid) {
+        interp->compiling = false;
+        return 0;
+    }
     _PyExecutorObject *executor;
     int err = uop_optimize(frame, tstate, &executor, progress_needed);
     if (err <= 0) {
@@ -560,9 +565,6 @@ _PyJIT_translate_single_bytecode_to_trace(
     int oparg,
     int jump_taken)
 {
-    if (PyStackRef_IsNull(frame->f_funcobj)) {
-        goto unsupported;
-    }
 
     int is_first_instr = tstate->interp->jit_tracer_initial_instr == this_instr;
     bool progress_needed = (tstate->interp->jit_tracer_initial_chain_depth % MAX_CHAIN_DEPTH) == 0 && is_first_instr;;
@@ -794,12 +796,16 @@ _PyJIT_translate_single_bytecode_to_trace(
                 }
                 if (uop == _PUSH_FRAME || uop == _RETURN_VALUE || uop == _RETURN_GENERATOR || uop == _YIELD_VALUE) {
                     PyCodeObject *new_code = (PyCodeObject *)PyStackRef_AsPyObjectBorrow(frame->f_executable);
-                    PyFunctionObject *new_func = (PyCodeObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
+                    PyFunctionObject *new_func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
                     if (new_func != NULL) {
                         operand = (uintptr_t)new_func;
+                        DPRINTF(2, "Adding %p func to op\n", (void *)operand);
+                        _Py_BloomFilter_Add(dependencies, new_func);
                     }
                     else if (new_code != NULL) {
                         operand = (uintptr_t)new_code | 1;
+                        DPRINTF(2, "Adding %p code to op\n", (void *)operand);
+                        _Py_BloomFilter_Add(dependencies, new_code);
                     }
                     else {
                         operand = 0;
@@ -866,6 +872,7 @@ _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_
     tstate->interp->jit_tracer_initial_stack_depth = curr_stackdepth;
     tstate->interp->jit_tracer_initial_chain_depth = chain_depth;
     tstate->interp->jit_tracer_current_frame = frame;
+    tstate->interp->jit_tracer_dependencies_still_valid = true;
 }
 
 void
@@ -1530,6 +1537,25 @@ _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is
     _Py_Executors_InvalidateAll(interp, is_invalidation);
 }
 
+void
+_Py_JITTracer_InvalidateDependency(PyThreadState *old_tstate, void *obj)
+{
+    _PyBloomFilter obj_filter;
+    _Py_BloomFilter_Init(&obj_filter);
+    _Py_BloomFilter_Add(&obj_filter, obj);
+    HEAD_LOCK(&_PyRuntime);
+
+    PyInterpreterState *interp = old_tstate->interp;
+
+    _Py_FOR_EACH_TSTATE_UNLOCKED(interp, tstate) {
+        if (bloom_filter_may_contain(&tstate->interp->jit_tracer_dependencies, &obj_filter))
+        {
+            tstate->interp->jit_tracer_dependencies_still_valid = false;
+        }
+
+    }
+    HEAD_UNLOCK(&_PyRuntime);
+}
 /* Invalidate all executors */
 void
 _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)