GH-104584: Plugin optimizer API (GH-105100)

Author: markshannon

Include/Python.h

@@ -105,5 +105,6 @@
 #include "fileutils.h"
 #include "cpython/pyfpe.h"
 #include "tracemalloc.h"
+#include "cpython/optimizer.h"
 
 #endif /* !Py_PYTHON_H */

Include/cpython/code.h

@@ -76,6 +76,13 @@ typedef struct {
     int8_t line_delta;
 } _PyCoLineInstrumentationData;
 
+
+typedef struct {
+    int size;
+    int capacity;
+    struct _PyExecutorObject *executors[1];
+} _PyExecutorArray;
+
 /* Main data structure used for instrumentation.
  * This is allocated when needed for instrumentation
  */
@@ -153,6 +160,7 @@ typedef struct {
     PyObject *co_qualname;        /* unicode (qualname, for reference) */      \
     PyObject *co_linetable;       /* bytes object that holds location info */  \
     PyObject *co_weakreflist;     /* to support weakrefs to code objects */    \
+    _PyExecutorArray *co_executors;      /* executors from optimizer */        \
     _PyCoCached *_co_cached;      /* cached co_* attributes */                 \
     uint64_t _co_instrumentation_version; /* current instrumentation version */  \
     _PyCoMonitoringData *_co_monitoring; /* Monitoring data */                 \

Include/cpython/optimizer.h

@@ -0,0 +1,54 @@
+
+#ifndef Py_LIMITED_API
+#ifndef Py_OPTIMIZER_H
+#define Py_OPTIMIZER_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct {
+    uint8_t opcode;
+    uint8_t oparg;
+} _PyVMData;
+
+typedef struct _PyExecutorObject {
+    PyObject_HEAD
+    /* WARNING: execute consumes a reference to self. This is necessary to allow executors to tail call into each other. */
+    struct _PyInterpreterFrame *(*execute)(struct _PyExecutorObject *self, struct _PyInterpreterFrame *frame, PyObject **stack_pointer);
+    _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
+    /* Data needed by the executor goes here, but is opaque to the VM */
+} _PyExecutorObject;
+
+typedef struct _PyOptimizerObject _PyOptimizerObject;
+
+typedef _PyExecutorObject *(*optimize_func)(_PyOptimizerObject* self, PyCodeObject *code, _Py_CODEUNIT *instr);
+
+typedef struct _PyOptimizerObject {
+    PyObject_HEAD
+    optimize_func optimize;
+    uint16_t resume_threshold;
+    uint16_t backedge_threshold;
+    /* Data needed by the optimizer goes here, but is opaque to the VM */
+} _PyOptimizerObject;
+
+PyAPI_FUNC(int) PyUnstable_Replace_Executor(PyCodeObject *code, _Py_CODEUNIT *instr, _PyExecutorObject *executor);
+
+PyAPI_FUNC(void) PyUnstable_SetOptimizer(_PyOptimizerObject* optimizer);
+
+PyAPI_FUNC(_PyOptimizerObject *) PyUnstable_GetOptimizer(void);
+
+struct _PyInterpreterFrame *
+_PyOptimizer_BackEdge(struct _PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNIT *dest, PyObject **stack_pointer);
+
+extern _PyOptimizerObject _PyOptimizer_Default;
+
+/* For testing */
+PyAPI_FUNC(PyObject *)PyUnstable_Optimizer_NewCounter(void);
+
+#define OPTIMIZER_BITS_IN_COUNTER 4
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_OPTIMIZER_H */
+#endif /* Py_LIMITED_API */

Include/internal/pycore_code.h

@@ -487,6 +487,7 @@ extern int _Py_Instrument(PyCodeObject *co, PyInterpreterState *interp);
 
 extern int _Py_GetBaseOpcode(PyCodeObject *code, int offset);
 
+extern int _PyInstruction_GetLength(PyCodeObject *code, int offset);
 
 #ifdef __cplusplus
 }

Include/internal/pycore_interp.h

@@ -160,6 +160,9 @@ struct _is {
     struct types_state types;
     struct callable_cache callable_cache;
     PyCodeObject *interpreter_trampoline;
+    _PyOptimizerObject *optimizer;
+    uint16_t optimizer_resume_threshold;
+    uint16_t optimizer_backedge_threshold;
 
     _Py_Monitors monitors;
     bool f_opcode_trace_set;

Include/internal/pycore_opcode.h

@@ -70,6 +70,7 @@ typedef struct _object_stats {
     uint64_t type_cache_dunder_hits;
     uint64_t type_cache_dunder_misses;
     uint64_t type_cache_collisions;
+    uint64_t optimization_attempts;
 } ObjectStats;
 
 typedef struct _stats {

Include/opcode.h

@@ -445,8 +445,9 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12b1 3529 (Inline list/dict/set comprehensions)
 #     Python 3.12b1 3530 (Shrink the LOAD_SUPER_ATTR caches)
 #     Python 3.12b1 3531 (Add PEP 695 changes)
+#     Python 3.13a1 3550 (Plugin optimizer support)
 
-#     Python 3.13 will start with 3550
+#     Python 3.14 will start with 3600
 
 #     Please don't copy-paste the same pre-release tag for new entries above!!!
 #     You should always use the *upcoming* tag. For example, if 3.12a6 came out
@@ -461,7 +462,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3531).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3550).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Include/pystats.h

@@ -232,6 +232,9 @@ def pseudo_op(name, op, real_ops):
 def_op('LOAD_FROM_DICT_OR_DEREF', 176)
 hasfree.append(176)
 
+# Optimizer hook
+def_op('ENTER_EXECUTOR', 230)
+
 # Instrumented instructions
 MIN_INSTRUMENTED_OPCODE = 237
 
@@ -486,6 +489,9 @@ def pseudo_op(name, op, real_ops):
     "SEND": {
         "counter": 1,
     },
+    "JUMP_BACKWARD": {
+        "counter": 1,
+    },
 }
 
 _inline_cache_entries = [

Lib/importlib/_bootstrap_external.py

@@ -1916,6 +1916,19 @@ def func():
         names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
         self.do_test(func, names)
 
+class TestOptimizerAPI(unittest.TestCase):
+
+    def test_counter_optimizer(self):
+        opt = _testinternalcapi.get_counter_optimizer()
+        self.assertEqual(opt.get_count(), 0)
+        try:
+            _testinternalcapi.set_optimizer(opt)
+            self.assertEqual(opt.get_count(), 0)
+            for _ in range(1000):
+                pass
+            self.assertEqual(opt.get_count(), 1000)
+        finally:
+            _testinternalcapi.set_optimizer(None)
 
 if __name__ == "__main__":
     unittest.main()