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

This commit is contained in:
Irit Katriel 2025-05-05 17:46:56 +01:00 committed by GitHub
parent 07f416a3f0
commit 082dbf7788
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 172 additions and 45 deletions

View file

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

View file

@ -1147,26 +1147,7 @@
break;
}
case _GUARD_BINARY_OP_EXTEND: {
_PyStackRef right;
_PyStackRef left;
right = stack_pointer[-1];
left = stack_pointer[-2];
PyObject *descr = (PyObject *)CURRENT_OPERAND0();
PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
_PyBinaryOpSpecializationDescr *d = (_PyBinaryOpSpecializationDescr*)descr;
assert(INLINE_CACHE_ENTRIES_BINARY_OP == 5);
assert(d && d->guard);
_PyFrame_SetStackPointer(frame, stack_pointer);
int res = d->guard(left_o, right_o);
stack_pointer = _PyFrame_GetStackPointer(frame);
if (!res) {
UOP_STAT_INC(uopcode, miss);
JUMP_TO_JUMP_TARGET();
}
break;
}
/* _GUARD_BINARY_OP_EXTEND is not a viable micro-op for tier 2 because it uses the 'this_instr' variable */
case _BINARY_OP_EXTEND: {
_PyStackRef right;
@ -1193,6 +1174,9 @@
stack_pointer = _PyFrame_GetStackPointer(frame);
stack_pointer += -2;
assert(WITHIN_STACK_BOUNDS());
if (res_o == NULL) {
JUMP_TO_ERROR();
}
res = PyStackRef_FromPyObjectSteal(res_o);
stack_pointer[0] = res;
stack_pointer += 1;

View file

@ -283,14 +283,30 @@
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);
_PyFrame_SetStackPointer(frame, stack_pointer);
int res = d->guard(left_o, right_o);
stack_pointer = _PyFrame_GetStackPointer(frame);
if (!res) {
UPDATE_MISS_STATS(BINARY_OP);
assert(_PyOpcode_Deopt[opcode] == (BINARY_OP));
JUMP_TO_PREDICTED(BINARY_OP);
if (res < 0) {
JUMP_TO_LABEL(error);
}
if (res == 0) {
if (d->free) {
_PyFrame_SetStackPointer(frame, stack_pointer);
d->free(d);
stack_pointer = _PyFrame_GetStackPointer(frame);
}
_PyBinaryOpCache *cache = (_PyBinaryOpCache *)(this_instr+1);
_PyFrame_SetStackPointer(frame, stack_pointer);
write_ptr(cache->external_cache, NULL);
stack_pointer = _PyFrame_GetStackPointer(frame);
this_instr->op.code = BINARY_OP;
if (true) {
UPDATE_MISS_STATS(BINARY_OP);
assert(_PyOpcode_Deopt[opcode] == (BINARY_OP));
JUMP_TO_PREDICTED(BINARY_OP);
}
}
}
/* Skip -4 cache entry */
@ -315,6 +331,9 @@
stack_pointer = _PyFrame_GetStackPointer(frame);
stack_pointer += -2;
assert(WITHIN_STACK_BOUNDS());
if (res_o == NULL) {
JUMP_TO_LABEL(error);
}
res = PyStackRef_FromPyObjectSteal(res_o);
}
stack_pointer[0] = res;

View file

@ -554,9 +554,7 @@
break;
}
case _GUARD_BINARY_OP_EXTEND: {
break;
}
/* _GUARD_BINARY_OP_EXTEND is not a viable micro-op for tier 2 */
case _BINARY_OP_EXTEND: {
JitOptSymbol *res;

View file

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