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

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/typeslots.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

Lib/test/test_sys.py

@@ -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 ''

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

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

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>
 
@@ -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]

Python/bytecodes.c

@@ -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);
  }
 

Python/specialize.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;
 }