gh-134584: Decref elimination for float ops in the JIT (GH-134588)

This PR adds a PyJitRef API to the JIT's optimizer that mimics the _PyStackRef API. This allows it to track references and their stack lifetimes properly. Thus opening up the doorway to refcount elimination in the JIT.
This commit is contained in:
Ken Jin 2025-06-17 23:25:53 +08:00 committed by GitHub
parent 8dd8b5c2f0
commit fba5dded6d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1023 additions and 736 deletions

View file

@ -687,6 +687,52 @@ dummy_func(
ERROR_IF(PyStackRef_IsNull(res));
}
pure op(_BINARY_OP_MULTIPLY_FLOAT__NO_DECREF_INPUTS, (left, right -- res)) {
PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
assert(PyFloat_CheckExact(left_o));
assert(PyFloat_CheckExact(right_o));
STAT_INC(BINARY_OP, hit);
double dres =
((PyFloatObject *)left_o)->ob_fval *
((PyFloatObject *)right_o)->ob_fval;
res = PyStackRef_FromPyObjectSteal(PyFloat_FromDouble(dres));
INPUTS_DEAD();
ERROR_IF(PyStackRef_IsNull(res));
}
pure op(_BINARY_OP_ADD_FLOAT__NO_DECREF_INPUTS, (left, right -- res)) {
PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
assert(PyFloat_CheckExact(left_o));
assert(PyFloat_CheckExact(right_o));
STAT_INC(BINARY_OP, hit);
double dres =
((PyFloatObject *)left_o)->ob_fval +
((PyFloatObject *)right_o)->ob_fval;
res = PyStackRef_FromPyObjectSteal(PyFloat_FromDouble(dres));
INPUTS_DEAD();
ERROR_IF(PyStackRef_IsNull(res));
}
pure op(_BINARY_OP_SUBTRACT_FLOAT__NO_DECREF_INPUTS, (left, right -- res)) {
PyObject *left_o = PyStackRef_AsPyObjectBorrow(left);
PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
assert(PyFloat_CheckExact(left_o));
assert(PyFloat_CheckExact(right_o));
STAT_INC(BINARY_OP, hit);
double dres =
((PyFloatObject *)left_o)->ob_fval -
((PyFloatObject *)right_o)->ob_fval;
res = PyStackRef_FromPyObjectSteal(PyFloat_FromDouble(dres));
INPUTS_DEAD();
ERROR_IF(PyStackRef_IsNull(res));
}
macro(BINARY_OP_MULTIPLY_FLOAT) =
_GUARD_TOS_FLOAT + _GUARD_NOS_FLOAT + unused/5 + _BINARY_OP_MULTIPLY_FLOAT;
macro(BINARY_OP_ADD_FLOAT) =