GH-144179: Use recorded values to make optimizer more robust (GH-144437)

* Add three new symbol kinds
* Do not smuggle code object in _PUSH_FRAME operand
* Fix small bug in predicate analysis
This commit is contained in:
Mark Shannon 2026-02-05 08:58:41 +00:00 committed by GitHub
parent b6d8aa436b
commit b53fc7caa6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 2062 additions and 1507 deletions

View file

@ -30,6 +30,7 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
#define sym_set_compact_int(SYM) _Py_uop_sym_set_compact_int(ctx, SYM)
#define sym_is_bottom _Py_uop_sym_is_bottom
#define frame_new _Py_uop_frame_new
#define frame_new_from_symbol _Py_uop_frame_new_from_symbol
#define frame_pop _Py_uop_frame_pop
#define sym_new_tuple _Py_uop_sym_new_tuple
#define sym_tuple_getitem _Py_uop_sym_tuple_getitem
@ -40,6 +41,11 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
#define sym_new_truthiness _Py_uop_sym_new_truthiness
#define sym_new_predicate _Py_uop_sym_new_predicate
#define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
#define sym_set_recorded_type(SYM, TYPE) _Py_uop_sym_set_recorded_type(ctx, SYM, TYPE)
#define sym_set_recorded_value(SYM, VAL) _Py_uop_sym_set_recorded_value(ctx, SYM, VAL)
#define sym_set_recorded_gen_func(SYM, VAL) _Py_uop_sym_set_recorded_gen_func(ctx, SYM, VAL)
#define sym_get_probable_func_code _Py_uop_sym_get_probable_func_code
#define sym_get_probable_value _Py_uop_sym_get_probable_value
extern int
optimize_to_bool(
@ -337,14 +343,23 @@ dummy_func(void) {
GETLOCAL(this_instr->operand0) = sym_new_null(ctx);
}
op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem -- new_frame)) {
assert((this_instr + 1)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging(this_instr + 1);
if (co == NULL) {
ctx->done = true;
break;
op(_BINARY_OP_SUBSCR_CHECK_FUNC, (container, unused -- container, unused, getitem)) {
getitem = sym_new_not_null(ctx);
PyTypeObject *tp = sym_get_type(container);
if (tp == NULL) {
PyObject *c = sym_get_probable_value(container);
if (c != NULL) {
tp = Py_TYPE(c);
}
}
_Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0);
if (tp != NULL) {
PyObject *getitem_o = ((PyHeapTypeObject *)tp)->_spec_cache.getitem;
sym_set_recorded_value(getitem, getitem_o);
}
}
op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem -- new_frame)) {
_Py_UOpsAbstractFrame *f = frame_new_from_symbol(ctx, getitem, 0, NULL, 0);
if (f == NULL) {
break;
}
@ -556,17 +571,10 @@ dummy_func(void) {
}
op(_COMPARE_OP_STR, (left, right -- res, l, r)) {
int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK);
if (cmp_mask == COMPARE_EQ_MASK) {
res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ);
}
else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) {
res = sym_new_predicate(ctx, left, right, JIT_PRED_NE);
}
else {
res = sym_new_type(ctx, &PyBool_Type);
}
/* Cannot use predicate optimization here, as `a == b`
* does not imply that `a` is equivalent to `b`. `a` may be
* mortal, while `b` is immortal */
res = sym_new_type(ctx, &PyBool_Type);
l = left;
r = right;
REPLACE_OPCODE_IF_EVALUATES_PURE(left, right, res);
@ -814,12 +822,8 @@ dummy_func(void) {
op(_LOAD_ATTR_PROPERTY_FRAME, (fget/4, owner -- new_frame)) {
// + 1 for _SAVE_RETURN_OFFSET
assert((this_instr + 2)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging(this_instr + 2);
if (co == NULL) {
ctx->done = true;
break;
}
// FIX ME -- This needs a version check and function watcher
PyCodeObject *co = (PyCodeObject *)((PyFunctionObject *)fget)->func_code;
_Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0);
if (f == NULL) {
break;
@ -872,14 +876,6 @@ dummy_func(void) {
op(_INIT_CALL_PY_EXACT_ARGS, (callable, self_or_null, args[oparg] -- new_frame)) {
int argcount = oparg;
assert((this_instr + 2)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 2));
if (co == NULL) {
ctx->done = true;
break;
}
assert(!PyJitRef_IsNull(self_or_null));
assert(args != NULL);
if (sym_is_not_null(self_or_null)) {
@ -889,9 +885,9 @@ dummy_func(void) {
}
if (sym_is_null(self_or_null) || sym_is_not_null(self_or_null)) {
new_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, args, argcount));
new_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, callable, 0, args, argcount));
} else {
new_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, NULL, 0));
new_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, callable, 0, NULL, 0));
}
}
@ -902,36 +898,15 @@ dummy_func(void) {
}
op(_PY_FRAME_GENERAL, (callable, self_or_null, args[oparg] -- new_frame)) {
assert((this_instr + 2)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 2));
if (co == NULL) {
ctx->done = true;
break;
}
new_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, NULL, 0));
new_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, callable, 0, NULL, 0));
}
op(_PY_FRAME_KW, (callable, self_or_null, args[oparg], kwnames -- new_frame)) {
assert((this_instr + 2)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 2));
if (co == NULL) {
ctx->done = true;
break;
}
new_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, NULL, 0));
new_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, callable, 0, NULL, 0));
}
op(_PY_FRAME_EX, (func_st, null, callargs_st, kwargs_st -- ex_frame)) {
assert((this_instr + 2)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 2));
if (co == NULL) {
ctx->done = true;
break;
}
ex_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, NULL, 0));
ex_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, func_st, 0, NULL, 0));
}
op(_CHECK_AND_ALLOCATE_OBJECT, (type_version/2, callable, self_or_null, args[oparg] -- callable, self_or_null, args[oparg])) {
@ -954,8 +929,7 @@ dummy_func(void) {
ctx->frame = shim;
ctx->curr_frame_depth++;
assert((this_instr + 1)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 1));
init_frame = PyJitRef_WrapInvalid(frame_new(ctx, co, 0, args-1, oparg+1));
init_frame = PyJitRef_WrapInvalid(frame_new_from_symbol(ctx, init, 0, args-1, oparg+1));
}
op(_RETURN_VALUE, (retval -- res)) {
@ -964,7 +938,9 @@ dummy_func(void) {
DEAD(retval);
SAVE_STACK();
ctx->frame->stack_pointer = stack_pointer;
PyCodeObject *returning_code = get_code_with_logging(this_instr);
assert(this_instr[1].opcode == _RECORD_CODE);
PyCodeObject *returning_code = (PyCodeObject *)this_instr[1].operand0;
assert(PyCode_Check(returning_code));
if (returning_code == NULL) {
ctx->done = true;
break;
@ -973,8 +949,8 @@ dummy_func(void) {
if (ctx->curr_frame_depth >= 2) {
PyCodeObject *expected_code = ctx->frames[ctx->curr_frame_depth - 2].code;
if (expected_code == returning_code) {
assert((this_instr + 1)->opcode == _GUARD_IP_RETURN_VALUE);
REPLACE_OP((this_instr + 1), _NOP, 0, 0);
assert(this_instr[2].opcode == _GUARD_IP_RETURN_VALUE);
REPLACE_OP((this_instr + 2), _NOP, 0, 0);
}
}
if (frame_pop(ctx, returning_code, returning_stacklevel)) {
@ -989,7 +965,9 @@ dummy_func(void) {
op(_RETURN_GENERATOR, ( -- res)) {
SYNC_SP();
ctx->frame->stack_pointer = stack_pointer;
PyCodeObject *returning_code = get_code_with_logging(this_instr);
assert(this_instr[1].opcode == _RECORD_CODE);
PyCodeObject *returning_code = (PyCodeObject *)this_instr[1].operand0;
assert(PyCode_Check(returning_code));
if (returning_code == NULL) {
ctx->done = true;
break;
@ -1009,7 +987,9 @@ dummy_func(void) {
DEAD(retval);
SAVE_STACK();
ctx->frame->stack_pointer = stack_pointer;
PyCodeObject *returning_code = get_code_with_logging(this_instr);
assert(this_instr[1].opcode == _RECORD_CODE);
PyCodeObject *returning_code = (PyCodeObject *)this_instr[1].operand0;
assert(PyCode_Check(returning_code));
if (returning_code == NULL) {
ctx->done = true;
break;
@ -1035,14 +1015,8 @@ dummy_func(void) {
}
}
op(_FOR_ITER_GEN_FRAME, (unused, unused -- unused, unused, gen_frame)) {
assert((this_instr + 1)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 1));
if (co == NULL) {
ctx->done = true;
break;
}
_Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
op(_FOR_ITER_GEN_FRAME, (iter, unused -- iter, unused, gen_frame)) {
_Py_UOpsAbstractFrame *new_frame = frame_new_from_symbol(ctx, iter, 1, NULL, 0);
if (new_frame == NULL) {
ctx->done = true;
break;
@ -1051,14 +1025,8 @@ dummy_func(void) {
gen_frame = PyJitRef_WrapInvalid(new_frame);
}
op(_SEND_GEN_FRAME, (unused, v -- unused, gen_frame)) {
assert((this_instr + 1)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 1));
if (co == NULL) {
ctx->done = true;
break;
}
_Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
op(_SEND_GEN_FRAME, (receiver, v -- receiver, gen_frame)) {
_Py_UOpsAbstractFrame *new_frame = frame_new_from_symbol(ctx, receiver, 1, NULL, 0);
if (new_frame == NULL) {
ctx->done = true;
break;
@ -1067,9 +1035,8 @@ dummy_func(void) {
gen_frame = PyJitRef_WrapInvalid(new_frame);
}
op(_CHECK_STACK_SPACE, (unused, unused, unused[oparg] -- unused, unused, unused[oparg])) {
assert((this_instr + 4)->opcode == _PUSH_FRAME);
PyCodeObject *co = get_code_with_logging((this_instr + 4));
op(_CHECK_STACK_SPACE, (callable, unused, unused[oparg] -- callable, unused, unused[oparg])) {
PyCodeObject *co = sym_get_probable_func_code(callable);
if (co == NULL) {
ctx->done = true;
break;
@ -1089,22 +1056,12 @@ dummy_func(void) {
ctx->frame = (_Py_UOpsAbstractFrame *)PyJitRef_Unwrap(new_frame);
ctx->curr_frame_depth++;
stack_pointer = ctx->frame->stack_pointer;
uint64_t operand = this_instr->operand0;
if (operand == 0) {
ctx->done = true;
break;
}
if (!(operand & 1)) {
PyFunctionObject *func = (PyFunctionObject *)operand;
// No need to re-add to dependencies here. Already
// handled by the tracer.
ctx->frame->func = func;
}
// Fixed calls don't need IP guards.
if ((this_instr-1)->opcode == _CREATE_INIT_FRAME) {
assert((this_instr+1)->opcode == _GUARD_IP__PUSH_FRAME);
REPLACE_OP(this_instr+1, _NOP, 0, 0);
}
assert(ctx->frame->locals != NULL);
}
op(_UNPACK_SEQUENCE, (seq -- values[oparg], top[0])) {
@ -1653,6 +1610,32 @@ dummy_func(void) {
ss = sub_st;
}
op(_RECORD_TOS, (tos -- tos)) {
sym_set_recorded_value(tos, (PyObject *)this_instr->operand0);
}
op(_RECORD_TOS_TYPE, (tos -- tos)) {
PyTypeObject *tp = (PyTypeObject *)this_instr->operand0;
sym_set_recorded_type(tos, tp);
}
op(_RECORD_NOS, (nos, tos -- nos, tos)) {
sym_set_recorded_value(nos, (PyObject *)this_instr->operand0);
}
op(_RECORD_4OS, (value, _3os, nos, tos -- value, _3os, nos, tos)) {
sym_set_recorded_value(value, (PyObject *)this_instr->operand0);
}
op(_RECORD_CALLABLE, (func, self, args[oparg] -- func, self, args[oparg])) {
sym_set_recorded_value(func, (PyObject *)this_instr->operand0);
}
op(_RECORD_NOS_GEN_FUNC, (nos, tos -- nos, tos)) {
PyFunctionObject *func = (PyFunctionObject *)this_instr->operand0;
assert(func == NULL || PyFunction_Check(func));
sym_set_recorded_gen_func(nos, func);
}
// END BYTECODES //