plumbing for registering specialization

python · iritkatriel · May 5, 2025 · Apr 24, 2025 · Apr 28, 2025 · Apr 29, 2025
commit 79e45a1ea6c7aa7d6181da2355efdcea0b718c72
diff --git a/Include/internal/pycore_c_array.h b/Include/internal/pycore_c_array.h
@@ -14,7 +14,7 @@ extern "C" {
 
 typedef struct {
     void *array;              /* pointer to the array */
-    int allocated_entries;    /* pointer to the capacity of the array */
+    int allocated_entries;    /* the capacity of the array */
     size_t item_size;         /* size of each element */
     int initial_num_entries;  /* initial allocation size */
 } _Py_c_array_t;

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -491,6 +491,8 @@ typedef struct {
     binaryopactionfunc action;
 } _PyBinaryOpSpecializationDescr;
 
+PyAPI_DATA(int) _Py_Specialize_AddBinaryOpExtention(_PyBinaryOpSpecializationDescr* descr);
+
 /* Comparison bit masks. */
 
 /* Note this evaluates its arguments twice each */

diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h
@@ -8,6 +8,7 @@ extern "C" {
 #endif
 
 #include "pycore_ast_state.h"     // struct ast_state
+#include "pycore_c_array.h"       // _Py_c_array_t
 #include "pycore_llist.h"         // struct llist_node
 #include "pycore_opcode_utils.h"  // NUM_COMMON_CONSTANTS
 #include "pycore_pymath.h"        // _PY_SHORT_FLOAT_REPR
@@ -951,6 +952,10 @@ struct _is {
 #  endif
 #endif
 
+    /* Specialization extensions */
+    Py_ssize_t num_binop_specializer_extentions;
+    _Py_c_array_t binop_specializer_extentions;
+
     /* the initial PyInterpreterState.threads.head */
     _PyThreadStateImpl _initial_thread;
     // _initial_thread should be the last field of PyInterpreterState.

diff --git a/Include/object.h b/Include/object.h
@@ -348,6 +348,8 @@ typedef int (*objobjproc)(PyObject *, PyObject *);
 typedef int (*visitproc)(PyObject *, void *);
 typedef int (*traverseproc)(PyObject *, visitproc, void *);
 
+struct _PyBinaryOpSpecializationDescr;
+typedef int (*binopspecfunc)(PyObject *lhs, PyObject *rhs, int oparg, struct _PyBinaryOpSpecializationDescr* descr);
 
 typedef void (*freefunc)(void *);
 typedef void (*destructor)(PyObject *);

diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
@@ -14,6 +14,8 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_moduleobject.h"  // _PyModule_GetState()
 
+#include "opcode.h"               // binary op opargs (NB_*)
+
 #include <stddef.h>               // offsetof()
 #include <stdbool.h>
 
@@ -3201,6 +3203,33 @@ do {                                                     \
     state->str_ ## string = tmp;                         \
 } while (0)
 
+static inline int
+array_guard(PyObject *lhs, PyObject *rhs)
+{
+     fprintf(stderr, "array_guard\n");
+    return 0;
+}
+
+static PyObject *
+array_action(PyObject *lhs, PyObject *rhs)
+{
+    return NULL;
+}
+
+static int
+array_register_specializations(void)
+{
+    _PyBinaryOpSpecializationDescr descr = {
+        .oparg = NB_MULTIPLY,
+        .guard = array_guard,
+        .action = array_action,
+    };
+    if (_Py_Specialize_AddBinaryOpExtention(&descr) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
 static int
 array_modexec(PyObject *m)
 {
@@ -3240,6 +3269,10 @@ array_modexec(PyObject *m)
     }
     Py_DECREF(res);
 
+    if (array_register_specializations() < 0) {
+        return -1;
+    }
+
     if (PyModule_AddType(m, state->ArrayType) < 0) {
         return -1;
     }
@@ -3252,7 +3285,6 @@ array_modexec(PyObject *m)
     if (PyModule_Add(m, "typecodes", typecodes) < 0) {
         return -1;
     }
-
     return 0;
 }
 

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -873,6 +873,8 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
     interp->ceval.instrumentation_version = 0;
     tstate->eval_breaker = 0;
 
+    _Py_CArray_Fini(&interp->binop_specializer_extentions);
+
     for (int i = 0; i < _PY_MONITORING_UNGROUPED_EVENTS; i++) {
         interp->monitors.tools[i] = 0;
     }

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2556,13 +2556,32 @@ static _PyBinaryOpSpecializationDescr binaryop_extend_descrs[] = {
     {NB_MULTIPLY, compactlong_float_guard, compactlong_float_multiply},
 };
 
+int
+_Py_Specialize_AddBinaryOpExtention(_PyBinaryOpSpecializationDescr* descr)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+    _Py_c_array_t *extensions = &tstate->interp->binop_specializer_extentions;
+    Py_ssize_t idx = tstate->interp->num_binop_specializer_extentions;
+    if (idx == 0) {
+        _Py_CArray_Init(extensions, sizeof(_PyBinaryOpSpecializationDescr), 10);
+    }
+    if (_Py_CArray_EnsureCapacity(extensions, idx) < 0) {
+        return -1;
+    }
+    _PyBinaryOpSpecializationDescr* descrs = (_PyBinaryOpSpecializationDescr*)extensions->array;
+    descrs[idx] = *descr;
+    tstate->interp->num_binop_specializer_extentions++;
+    return 0;
+}
+
 static int
-binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
-                                  _PyBinaryOpSpecializationDescr **descr)
+binary_op_extended_specialization_from_list(
+        _PyBinaryOpSpecializationDescr *descrs, size_t size,
+        PyObject *lhs, PyObject *rhs, int oparg,
+        _PyBinaryOpSpecializationDescr **descr)
 {
-    size_t n = sizeof(binaryop_extend_descrs)/sizeof(_PyBinaryOpSpecializationDescr);
-    for (size_t i = 0; i < n; i++) {
-        _PyBinaryOpSpecializationDescr *d = &binaryop_extend_descrs[i];
+    for (size_t i = 0; i < size; i++) {
+        _PyBinaryOpSpecializationDescr *d = &descrs[i];
         if (d->oparg == oparg && d->guard(lhs, rhs)) {
             *descr = d;
             return 1;
@@ -2571,6 +2590,30 @@ binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
     return 0;
 }
 
 static int
+binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
+                                  _PyBinaryOpSpecializationDescr **descr)
+{
+    if (binary_op_extended_specialization_from_list(
+            binaryop_extend_descrs,
+            sizeof(binaryop_extend_descrs)/sizeof(_PyBinaryOpSpecializationDescr),
+            lhs, rhs, oparg, descr))
+    {
+        return 1;
+    }
+
+    PyThreadState *tstate = PyThreadState_Get();
+    _Py_c_array_t *extensions = &tstate->interp->binop_specializer_extentions;
+    if (binary_op_extended_specialization_from_list(
+            (_PyBinaryOpSpecializationDescr *)extensions->array,
+            tstate->interp->num_binop_specializer_extentions,
+            lhs, rhs, oparg, descr))
+    {
+        return 1;
+    }
+    return 0;
+}
+
 Py_NO_INLINE void
 _Py_Specialize_BinaryOp(_PyStackRef lhs_st, _PyStackRef rhs_st, _Py_CODEUNIT *instr,
                         int oparg, _PyStackRef *locals)