gh-134584: Specialize POP_TOP by reference and type in JIT (GH-135761)

This commit is contained in:
Ken Jin 2025-06-24 00:57:14 +08:00 committed by GitHub
parent 99712c45cc
commit 569fc6870f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 239 additions and 62 deletions

View file

@ -344,6 +344,27 @@ dummy_func(
PyStackRef_XCLOSE(value);
}
op(_POP_TOP_NOP, (value --)) {
assert(PyStackRef_IsNull(value) || (!PyStackRef_RefcountOnObject(value)) ||
_Py_IsImmortal((PyStackRef_AsPyObjectBorrow(value))));
DEAD(value);
}
op(_POP_TOP_INT, (value --)) {
assert(PyLong_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyLong_ExactDealloc);
}
op(_POP_TOP_FLOAT, (value --)) {
assert(PyFloat_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyFloat_ExactDealloc);
}
op(_POP_TOP_UNICODE, (value --)) {
assert(PyUnicode_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyUnicode_ExactDealloc);
}
tier2 op(_POP_TWO, (nos, tos --)) {
PyStackRef_CLOSE(tos);
PyStackRef_CLOSE(nos);

View file

@ -539,6 +539,46 @@
break;
}
case _POP_TOP_NOP: {
_PyStackRef value;
value = stack_pointer[-1];
assert(PyStackRef_IsNull(value) || (!PyStackRef_RefcountOnObject(value)) ||
_Py_IsImmortal((PyStackRef_AsPyObjectBorrow(value))));
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_INT: {
_PyStackRef value;
value = stack_pointer[-1];
assert(PyLong_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyLong_ExactDealloc);
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_FLOAT: {
_PyStackRef value;
value = stack_pointer[-1];
assert(PyFloat_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyFloat_ExactDealloc);
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_UNICODE: {
_PyStackRef value;
value = stack_pointer[-1];
assert(PyUnicode_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
PyStackRef_CLOSE_SPECIALIZED(value, _PyUnicode_ExactDealloc);
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TWO: {
_PyStackRef tos;
_PyStackRef nos;

View file

@ -345,7 +345,7 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
#define sym_new_tuple _Py_uop_sym_new_tuple
#define sym_tuple_getitem _Py_uop_sym_tuple_getitem
#define sym_tuple_length _Py_uop_sym_tuple_length
#define sym_is_immortal _Py_uop_sym_is_immortal
#define sym_is_immortal _Py_uop_symbol_is_immortal
#define sym_is_compact_int _Py_uop_sym_is_compact_int
#define sym_new_compact_int _Py_uop_sym_new_compact_int
#define sym_new_truthiness _Py_uop_sym_new_truthiness

View file

@ -34,7 +34,7 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
#define sym_new_tuple _Py_uop_sym_new_tuple
#define sym_tuple_getitem _Py_uop_sym_tuple_getitem
#define sym_tuple_length _Py_uop_sym_tuple_length
#define sym_is_immortal _Py_uop_sym_is_immortal
#define sym_is_immortal _Py_uop_symbol_is_immortal
#define sym_new_compact_int _Py_uop_sym_new_compact_int
#define sym_is_compact_int _Py_uop_sym_is_compact_int
#define sym_new_truthiness _Py_uop_sym_new_truthiness
@ -534,7 +534,7 @@ dummy_func(void) {
}
op(_LOAD_CONST_INLINE, (ptr/4 -- value)) {
value = PyJitRef_Borrow(sym_new_const(ctx, ptr));
value = sym_new_const(ctx, ptr);
}
op(_LOAD_CONST_INLINE_BORROW, (ptr/4 -- value)) {
@ -542,7 +542,7 @@ dummy_func(void) {
}
op(_POP_TOP_LOAD_CONST_INLINE, (ptr/4, pop -- value)) {
value = PyJitRef_Borrow(sym_new_const(ctx, ptr));
value = sym_new_const(ctx, ptr);
}
op(_POP_TOP_LOAD_CONST_INLINE_BORROW, (ptr/4, pop -- value)) {
@ -561,6 +561,24 @@ dummy_func(void) {
value = PyJitRef_Borrow(sym_new_const(ctx, ptr));
}
op(_POP_TOP, (value -- )) {
PyTypeObject *typ = sym_get_type(value);
if (PyJitRef_IsBorrowed(value) ||
sym_is_immortal(PyJitRef_Unwrap(value)) ||
sym_is_null(value)) {
REPLACE_OP(this_instr, _POP_TOP_NOP, 0, 0);
}
else if (typ == &PyLong_Type) {
REPLACE_OP(this_instr, _POP_TOP_INT, 0, 0);
}
else if (typ == &PyFloat_Type) {
REPLACE_OP(this_instr, _POP_TOP_FLOAT, 0, 0);
}
else if (typ == &PyUnicode_Type) {
REPLACE_OP(this_instr, _POP_TOP_UNICODE, 0, 0);
}
}
op(_COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
assert(oparg > 0);
top = bottom;
@ -803,7 +821,9 @@ dummy_func(void) {
}
op(_RETURN_VALUE, (retval -- res)) {
JitOptRef temp = retval;
// We wrap and unwrap the value to mimic PyStackRef_MakeHeapSafe
// in bytecodes.c
JitOptRef temp = PyJitRef_Wrap(PyJitRef_Unwrap(retval));
DEAD(retval);
SAVE_STACK();
ctx->frame->stack_pointer = stack_pointer;

View file

@ -100,6 +100,47 @@
}
case _POP_TOP: {
JitOptRef value;
value = stack_pointer[-1];
PyTypeObject *typ = sym_get_type(value);
if (PyJitRef_IsBorrowed(value) ||
sym_is_immortal(PyJitRef_Unwrap(value)) ||
sym_is_null(value)) {
REPLACE_OP(this_instr, _POP_TOP_NOP, 0, 0);
}
else if (typ == &PyLong_Type) {
REPLACE_OP(this_instr, _POP_TOP_INT, 0, 0);
}
else if (typ == &PyFloat_Type) {
REPLACE_OP(this_instr, _POP_TOP_FLOAT, 0, 0);
}
else if (typ == &PyUnicode_Type) {
REPLACE_OP(this_instr, _POP_TOP_UNICODE, 0, 0);
}
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_NOP: {
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_INT: {
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_FLOAT: {
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
}
case _POP_TOP_UNICODE: {
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
break;
@ -784,7 +825,7 @@
JitOptRef retval;
JitOptRef res;
retval = stack_pointer[-1];
JitOptRef temp = retval;
JitOptRef temp = PyJitRef_Wrap(PyJitRef_Unwrap(retval));
stack_pointer += -1;
assert(WITHIN_STACK_BOUNDS());
ctx->frame->stack_pointer = stack_pointer;
@ -2660,7 +2701,7 @@
case _LOAD_CONST_INLINE: {
JitOptRef value;
PyObject *ptr = (PyObject *)this_instr->operand0;
value = PyJitRef_Borrow(sym_new_const(ctx, ptr));
value = sym_new_const(ctx, ptr);
stack_pointer[0] = value;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());
@ -2670,7 +2711,7 @@
case _POP_TOP_LOAD_CONST_INLINE: {
JitOptRef value;
PyObject *ptr = (PyObject *)this_instr->operand0;
value = PyJitRef_Borrow(sym_new_const(ctx, ptr));
value = sym_new_const(ctx, ptr);
stack_pointer[-1] = value;
break;
}

View file

@ -668,9 +668,6 @@ _Py_uop_symbol_is_immortal(JitOptSymbol *sym)
if (sym->tag == JIT_SYM_KNOWN_CLASS_TAG) {
return sym->cls.type == &PyBool_Type;
}
if (sym->tag == JIT_SYM_TRUTHINESS_TAG) {
return true;
}
return false;
}