gh-133395: add option for extension modules to specialize BINARY_OP/SUBSCR, apply to arrays (#133396)

Author: iritkatriel

Include/cpython/object.h

@@ -143,6 +143,11 @@ typedef struct {
  * backwards-compatibility */
 typedef Py_ssize_t printfunc;
 
+/* Specialize a binary op by setting the descriptor pointer */
+struct _PyBinopSpecializationDescr;
+typedef int (*binop_specialize_func)(PyObject *v, PyObject *w, int oparg,
+                                    struct _PyBinopSpecializationDescr **descr);
+
 // If this structure is modified, Doc/includes/typestruct.h should be updated
 // as well.
 struct _typeobject {
@@ -233,6 +238,13 @@ struct _typeobject {
     /* bitset of which type-watchers care about this type */
     unsigned char tp_watched;
 
+    /* callback that may specialize BINARY_OP
+     * this is an experimental API based on the ideas in the paper
+     * Cross Module Quickening - The Curious Case of C Extensions
+     * by Felix Berlakovich and Stefan Brunthaler.
+     */
+    binop_specialize_func tp_binop_specialize;
+
     /* Number of tp_version_tag values used.
      * Set to _Py_ATTR_CACHE_UNUSED if the attribute cache is
      * disabled for this type (e.g. due to custom MRO entries).

Include/internal/pycore_code.h

@@ -482,13 +482,18 @@ adaptive_counter_backoff(_Py_BackoffCounter counter) {
 /* Specialization Extensions */
 
 /* callbacks for an external specialization */
+
+struct _PyBinopSpecializationDescr;
+
 typedef int (*binaryopguardfunc)(PyObject *lhs, PyObject *rhs);
-typedef PyObject *(*binaryopactionfunc)(PyObject *lhs, PyObject *rhs);
+typedef PyObject* (*binaryopactionfunc)(PyObject *lhs, PyObject *rhs);
+typedef void (*binaryopfreefunc)(struct _PyBinopSpecializationDescr *descr);
 
-typedef struct {
+typedef struct _PyBinopSpecializationDescr {
     int oparg;
     binaryopguardfunc guard;
     binaryopactionfunc action;
+    binaryopfreefunc free;
 } _PyBinaryOpSpecializationDescr;
 
 /* Comparison bit masks. */

Include/internal/pycore_opcode_metadata.h

@@ -93,4 +93,5 @@
 #if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x030E0000
 /* New in 3.14 */
 #define Py_tp_token 83
+#define Py_tp_binop_specialize 84
 #endif

Include/internal/pycore_uop_metadata.h

@@ -1776,7 +1776,7 @@ def delx(self): del self.__x
         check((1,2,3), vsize('') + self.P + 3*self.P)
         # type
         # static type: PyTypeObject
-        fmt = 'P2nPI13Pl4Pn9Pn12PIPc'
+        fmt = 'P2nPI13Pl4Pn9Pn12PI3Pc'
         s = vsize(fmt)
         check(int, s)
         typeid = 'n' if support.Py_GIL_DISABLED else ''

Include/typeslots.h

@@ -0,0 +1,2 @@
+Add option for extension modules to specialize ``BINARY_OP`` instructions.
+Applied to ``array`` objects.

Lib/test/test_sys.py

@@ -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>
 
@@ -848,6 +850,10 @@ array_richcompare(PyObject *v, PyObject *w, int op)
     return res;
 }
 
+static int
+array_binop_specialize(PyObject *v, PyObject *w, int oparg,
+                       _PyBinaryOpSpecializationDescr **descr);
+
 static Py_ssize_t
 array_length(PyObject *op)
 {
@@ -2963,6 +2969,8 @@ static PyType_Slot array_slots[] = {
     {Py_tp_alloc, PyType_GenericAlloc},
     {Py_tp_new, array_new},
     {Py_tp_traverse, array_tp_traverse},
+    {Py_tp_token, Py_TP_USE_SPEC},
+    {Py_tp_binop_specialize, array_binop_specialize},
 
     /* as sequence */
     {Py_sq_length, array_length},
@@ -2995,6 +3003,70 @@ static PyType_Spec array_spec = {
     .slots = array_slots,
 };
 
+static inline int
+array_subscr_guard(PyObject *lhs, PyObject *rhs)
+{
+    PyObject *exc = PyErr_GetRaisedException();
+    int ret = PyType_GetBaseByToken(Py_TYPE(lhs), &array_spec, NULL);
+    if (ret < 0) {
+        if (PyErr_ExceptionMatches(PyExc_TypeError)) {
+            PyErr_Clear();
+            ret = 0;
+        }
+    }
+    _PyErr_ChainExceptions1(exc);
+    return ret;
+}
+
+static PyObject *
+array_subscr_action(PyObject *lhs, PyObject *rhs)
+{
+    return array_subscr(lhs, rhs);
+}
+
+static void
+array_subscr_free(_PyBinaryOpSpecializationDescr* descr)
+{
+    if (descr != NULL) {
+        PyMem_Free(descr);
+    }
+}
+
+static int
+array_binop_specialize(PyObject *v, PyObject *w, int oparg,
+                       _PyBinaryOpSpecializationDescr **descr)
+{
+    array_state *state = find_array_state_by_type(Py_TYPE(v));
+
+    if (!array_Check(v, state)) {
+        return 0;
+    }
+
+    *descr = NULL;
+    switch(oparg) {
+        case NB_SUBSCR:
+            if (array_subscr_guard(v, w)) {
+                *descr = (_PyBinaryOpSpecializationDescr*)PyMem_Malloc(
+                             sizeof(_PyBinaryOpSpecializationDescr));
+                if (*descr == NULL) {
+                    PyErr_NoMemory();
+                    return -1;
+                }
+                **descr = (_PyBinaryOpSpecializationDescr) {
+                    .oparg = oparg,
+                    .guard = array_subscr_guard,
+                    .action = array_subscr_action,
+                    .free = array_subscr_free,
+                };
+                return 1;
+            }
+            break;
+    }
+
+    return 0;
+}
+
+
 /*********************** Array Iterator **************************/
 
 /*[clinic input]

Misc/NEWS.d/next/Core_and_Builtins/2025-05-04-19-32-47.gh-issue-133395.VhWWEP.rst

@@ -801,9 +801,19 @@ dummy_func(
             PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
             _PyBinaryOpSpecializationDescr *d = (_PyBinaryOpSpecializationDescr*)descr;
             assert(INLINE_CACHE_ENTRIES_BINARY_OP == 5);
-            assert(d && d->guard);
+            assert(d);
+            assert(d->guard);
             int res = d->guard(left_o, right_o);
-            DEOPT_IF(!res);
+            ERROR_IF(res < 0);
+            if (res == 0) {
+                if (d->free) {
+                    d->free(d);
+                }
+                _PyBinaryOpCache *cache = (_PyBinaryOpCache *)(this_instr+1);
+                write_ptr(cache->external_cache, NULL);
+                this_instr->op.code = BINARY_OP;
+                DEOPT_IF(true);
+            }
         }
 
         pure op(_BINARY_OP_EXTEND, (descr/4, left, right -- res)) {
@@ -816,6 +826,7 @@ dummy_func(
 
             PyObject *res_o = d->action(left_o, right_o);
             DECREF_INPUTS();
+            ERROR_IF(res_o == NULL);
             res = PyStackRef_FromPyObjectSteal(res_o);
         }
 

Modules/arraymodule.c

@@ -2534,7 +2534,7 @@ LONG_FLOAT_ACTION(compactlong_float_multiply, *)
 LONG_FLOAT_ACTION(compactlong_float_true_div, /)
 #undef LONG_FLOAT_ACTION
 
-static _PyBinaryOpSpecializationDescr binaryop_extend_descrs[] = {
+static const _PyBinaryOpSpecializationDescr binaryop_extend_builtins[] = {
     /* long-long arithmetic */
     {NB_OR, compactlongs_guard, compactlongs_or},
     {NB_AND, compactlongs_guard, compactlongs_and},
@@ -2560,14 +2560,41 @@ static int
 binary_op_extended_specialization(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];
+    /* We are currently using this only for NB_SUBSCR, which is not
+     * commutative. Will need to revisit this function when we use
+     * this for operators which are.
+     */
+
+    typedef _PyBinaryOpSpecializationDescr descr_type;
+    size_t size = Py_ARRAY_LENGTH(binaryop_extend_builtins);
+    for (size_t i = 0; i < size; i++) {
+        descr_type *d = (descr_type *)&binaryop_extend_builtins[i];
+        assert(d != NULL);
+        assert(d->guard != NULL);
         if (d->oparg == oparg && d->guard(lhs, rhs)) {
             *descr = d;
             return 1;
         }
     }
+
+    PyTypeObject *lhs_type = Py_TYPE(lhs);
+    if (lhs_type->tp_binop_specialize != NULL) {
+        int ret = lhs_type->tp_binop_specialize(lhs, rhs, oparg, descr);
+        if (ret < 0) {
+            return -1;
+        }
+        if (ret == 1) {
+            if (*descr == NULL) {
+                PyErr_Format(
+                    PyExc_ValueError,
+                    "tp_binop_specialize of '%T' returned 1 with *descr == NULL",
+                    lhs);
+                return -1;
+            }
+            (*descr)->oparg = oparg;
+        }
+        return ret;
+    }
     return 0;
 }