Add a constant pool to all executors

Author: Fidget-Spinner

Include/internal/pycore_optimizer.h

@@ -52,6 +52,7 @@ typedef struct _PyExitData {
 typedef struct _PyExecutorObject {
     PyObject_VAR_HEAD
     const _PyUOpInstruction *trace;
+    PyObject *constant_pool;
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
     uint32_t exit_count;
     uint32_t code_size;
@@ -98,7 +99,7 @@ PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyInterpreterState *interp);
 
 int _Py_uop_analyze_and_optimize(_PyInterpreterFrame *frame,
     _PyUOpInstruction *trace, int trace_len, int curr_stackentries,
-    _PyBloomFilter *dependencies);
+    _PyBloomFilter *dependencies, PyObject **constant_pool_ptr);
 
 extern PyTypeObject _PyUOpExecutor_Type;
 
@@ -278,6 +279,7 @@ typedef struct _JitOptContext {
     bool contradiction;
      // Has the builtins dict been watched?
     bool builtins_watched;
+    PyObject *constant_pool;
     // The current "executing" frame.
     _Py_UOpsAbstractFrame *frame;
     _Py_UOpsAbstractFrame frames[MAX_ABSTRACT_FRAME_DEPTH];
@@ -324,8 +326,9 @@ extern bool _Py_uop_sym_is_compact_int(JitOptRef sym);
 extern JitOptRef _Py_uop_sym_new_compact_int(JitOptContext *ctx);
 extern void _Py_uop_sym_set_compact_int(JitOptContext *ctx,  JitOptRef sym);
 
-extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
+extern int _Py_uop_abstractcontext_init(JitOptContext *ctx);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);
+extern int _Py_uop_promote_to_constant_pool(JitOptContext *ctx, PyObject *obj);
 
 extern _Py_UOpsAbstractFrame *_Py_uop_frame_new(
     JitOptContext *ctx,

Lib/test/test_capi/test_opt.py

@@ -2594,8 +2594,7 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
 
-        # For now... until we constant propagate it away.
-        self.assertIn("_BINARY_OP", uops)
+        self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
 
     def test_jitted_code_sees_changed_globals(self):
         "Issue 136154: Check that jitted code spots the change in the globals"

Python/optimizer.c

@@ -100,7 +100,7 @@ insert_executor(PyCodeObject *code, _Py_CODEUNIT *instr, int index, _PyExecutorO
 }
 
 static _PyExecutorObject *
-make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFilter *dependencies);
+make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFilter *dependencies, PyObject *constant_pool);
 
 static int
 uop_optimize(_PyInterpreterFrame *frame, _Py_CODEUNIT *instr,
@@ -388,6 +388,7 @@ static int
 executor_traverse(PyObject *o, visitproc visit, void *arg)
 {
     _PyExecutorObject *executor = _PyExecutorObject_CAST(o);
+    Py_VISIT(executor->constant_pool);
     for (uint32_t i = 0; i < executor->exit_count; i++) {
         Py_VISIT(executor->exits[i].executor);
     }
@@ -1115,13 +1116,15 @@ prepare_for_execution(_PyUOpInstruction *buffer, int length)
 /* Executor side exits */
 
 static _PyExecutorObject *
-allocate_executor(int exit_count, int length)
+allocate_executor(int exit_count, int length, PyObject *constant_pool)
 {
     int size = exit_count*sizeof(_PyExitData) + length*sizeof(_PyUOpInstruction);
     _PyExecutorObject *res = PyObject_GC_NewVar(_PyExecutorObject, &_PyUOpExecutor_Type, size);
     if (res == NULL) {
         return NULL;
     }
+    // Transfer ownership
+    res->constant_pool = constant_pool;
     res->trace = (_PyUOpInstruction *)(res->exits + exit_count);
     res->code_size = length;
     res->exit_count = exit_count;
@@ -1196,10 +1199,10 @@ sanity_check(_PyExecutorObject *executor)
  * and not a NOP.
  */
 static _PyExecutorObject *
-make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFilter *dependencies)
+make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFilter *dependencies, PyObject *constant_pool)
 {
     int exit_count = count_exits(buffer, length);
-    _PyExecutorObject *executor = allocate_executor(exit_count, length);
+    _PyExecutorObject *executor = allocate_executor(exit_count, length, constant_pool);
     if (executor == NULL) {
         return NULL;
     }
@@ -1293,6 +1296,7 @@ uop_optimize(
     _PyBloomFilter dependencies;
     _Py_BloomFilter_Init(&dependencies);
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    PyObject *constant_pool = NULL;
     if (interp->jit_uop_buffer == NULL) {
         interp->jit_uop_buffer = (_PyUOpInstruction *)_PyObject_VirtualAlloc(UOP_BUFFER_SIZE);
         if (interp->jit_uop_buffer == NULL) {
@@ -1316,7 +1320,7 @@ uop_optimize(
     if (!is_noopt) {
         length = _Py_uop_analyze_and_optimize(frame, buffer,
                                            length,
-                                           curr_stackentries, &dependencies);
+                                           curr_stackentries, &dependencies, &constant_pool);
         if (length <= 0) {
             return length;
         }
@@ -1339,7 +1343,7 @@ uop_optimize(
     OPT_HIST(effective_trace_length(buffer, length), optimized_trace_length_hist);
     length = prepare_for_execution(buffer, length);
     assert(length <= UOP_MAX_TRACE_LENGTH);
-    _PyExecutorObject *executor = make_executor_from_uops(buffer, length,  &dependencies);
+    _PyExecutorObject *executor = make_executor_from_uops(buffer, length,  &dependencies, constant_pool);
     if (executor == NULL) {
         return -1;
     }
@@ -1512,7 +1516,7 @@ _PyExecutor_GetColdExecutor(void)
     if (interp->cold_executor != NULL) {
         return interp->cold_executor;
     }
-    _PyExecutorObject *cold = allocate_executor(0, 1);
+    _PyExecutorObject *cold = allocate_executor(0, 1, NULL);
     if (cold == NULL) {
         Py_FatalError("Cannot allocate core JIT code");
     }
@@ -1575,6 +1579,8 @@ executor_clear(PyObject *op)
     unlink_executor(executor);
     executor->vm_data.valid = 0;
 
+    Py_CLEAR(executor->constant_pool);
+
     /* It is possible for an executor to form a reference
      * cycle with itself, so decref'ing a side exit could
      * free the executor unless we hold a strong reference to it

Python/optimizer_analysis.c

@@ -153,6 +153,8 @@ incorrect_keys(PyObject *obj, uint32_t version)
     (INST)->oparg = ARG;            \
     (INST)->operand0 = OPERAND;
 
+#define PROMOTE_TO_CONSTANT_POOL _Py_uop_promote_to_constant_pool
+
 /* Shortened forms for convenience, used in optimizer_bytecodes.c */
 #define sym_is_not_null _Py_uop_sym_is_not_null
 #define sym_is_const _Py_uop_sym_is_const
@@ -290,7 +292,8 @@ optimize_uops(
     _PyUOpInstruction *trace,
     int trace_len,
     int curr_stacklen,
-    _PyBloomFilter *dependencies
+    _PyBloomFilter *dependencies,
+    PyObject **constant_pool_ptr
 )
 {
     assert(!PyErr_Occurred());
@@ -310,9 +313,13 @@ optimize_uops(
         interp->type_watchers[TYPE_WATCHER_ID] = type_watcher_callback;
     }
 
-    _Py_uop_abstractcontext_init(ctx);
+    if (_Py_uop_abstractcontext_init(ctx)) {
+        return 0;
+    }
+
     _Py_UOpsAbstractFrame *frame = _Py_uop_frame_new(ctx, (PyCodeObject *)func->func_code, curr_stacklen, NULL, 0);
     if (frame == NULL) {
+        _Py_uop_abstractcontext_fini(ctx);
         return 0;
     }
     frame->func = func;
@@ -367,6 +374,7 @@ optimize_uops(
 
     /* Either reached the end or cannot optimize further, but there
      * would be no benefit in retrying later */
+    *constant_pool_ptr = Py_NewRef(ctx->constant_pool);
     _Py_uop_abstractcontext_fini(ctx);
     if (first_valid_check_stack != NULL) {
         assert(first_valid_check_stack->opcode == _CHECK_STACK_SPACE);
@@ -522,14 +530,15 @@ _Py_uop_analyze_and_optimize(
     _PyUOpInstruction *buffer,
     int length,
     int curr_stacklen,
-    _PyBloomFilter *dependencies
+    _PyBloomFilter *dependencies,
+    PyObject **constant_pool_ptr
 )
 {
     OPT_STAT_INC(optimizer_attempts);
 
     length = optimize_uops(
         _PyFrame_GetFunction(frame), buffer,
-        length, curr_stacklen, dependencies);
+        length, curr_stacklen, dependencies, constant_pool_ptr);
 
     if (length == 0) {
         return length;