gh-131798: constant fold special method lookups in JIT (#148432)

This commit is contained in:
Kumar Aditya 2026-04-14 21:32:23 +05:30 committed by GitHub
parent 11da7d4e21
commit 3cb7eaec85
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 71 additions and 6 deletions

View file

@ -1252,7 +1252,7 @@ const struct opcode_metadata _PyOpcode_opcode_metadata[267] = {
[LOAD_LOCALS] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG },
[LOAD_NAME] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
[LOAD_SMALL_INT] = { true, INSTR_FMT_IB, HAS_ARG_FLAG },
[LOAD_SPECIAL] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
[LOAD_SPECIAL] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG | HAS_RECORDS_VALUE_FLAG },
[LOAD_SUPER_ATTR] = { true, INSTR_FMT_IBC, HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
[LOAD_SUPER_ATTR_ATTR] = { true, INSTR_FMT_IBC, HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
[LOAD_SUPER_ATTR_METHOD] = { true, INSTR_FMT_IBC, HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG | HAS_RECORDS_VALUE_FLAG },
@ -1472,7 +1472,7 @@ _PyOpcode_macro_expansion[256] = {
[LOAD_LOCALS] = { .nuops = 1, .uops = { { _LOAD_LOCALS, OPARG_SIMPLE, 0 } } },
[LOAD_NAME] = { .nuops = 1, .uops = { { _LOAD_NAME, OPARG_SIMPLE, 0 } } },
[LOAD_SMALL_INT] = { .nuops = 1, .uops = { { _LOAD_SMALL_INT, OPARG_SIMPLE, 0 } } },
[LOAD_SPECIAL] = { .nuops = 2, .uops = { { _INSERT_NULL, OPARG_SIMPLE, 0 }, { _LOAD_SPECIAL, OPARG_SIMPLE, 0 } } },
[LOAD_SPECIAL] = { .nuops = 3, .uops = { { _RECORD_TOS_TYPE, OPARG_SIMPLE, 0 }, { _INSERT_NULL, OPARG_SIMPLE, 0 }, { _LOAD_SPECIAL, OPARG_SIMPLE, 0 } } },
[LOAD_SUPER_ATTR_ATTR] = { .nuops = 1, .uops = { { _LOAD_SUPER_ATTR_ATTR, OPARG_SIMPLE, 1 } } },
[LOAD_SUPER_ATTR_METHOD] = { .nuops = 3, .uops = { { _RECORD_NOS, OPARG_SIMPLE, 0 }, { _GUARD_LOAD_SUPER_ATTR_METHOD, OPARG_SIMPLE, 1 }, { _LOAD_SUPER_ATTR_METHOD, OPARG_SIMPLE, 1 } } },
[MAKE_CELL] = { .nuops = 1, .uops = { { _MAKE_CELL, OPARG_SIMPLE, 0 } } },

View file

@ -3435,6 +3435,25 @@ def f(n):
self.assertNotIn("_LOAD_ATTR_METHOD_NO_DICT", uops)
self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
def test_cached_load_special(self):
class CM:
def __enter__(self):
return self
def __exit__(self, *args):
pass
def f(n):
cm = CM()
x = 0
for _ in range(n):
with cm:
x += 1
return x
res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
self.assertIsNotNone(ex)
self.assertEqual(res, TIER2_THRESHOLD)
uops = get_opnames(ex)
self.assertNotIn("_LOAD_SPECIAL", uops)
def test_store_fast_refcount_elimination(self):
def foo(x):
# Since x is known to be

View file

@ -3883,6 +3883,7 @@ dummy_func(
}
macro(LOAD_SPECIAL) =
_RECORD_TOS_TYPE +
_INSERT_NULL +
_LOAD_SPECIAL;

View file

@ -1845,8 +1845,30 @@ dummy_func(void) {
}
op(_LOAD_SPECIAL, (method_and_self[2] -- method_and_self[2])) {
method_and_self[0] = sym_new_not_null(ctx);
method_and_self[1] = sym_new_unknown(ctx);
bool optimized = false;
PyTypeObject *type = sym_get_probable_type(method_and_self[1]);
if (type != NULL) {
PyObject *name = _Py_SpecialMethods[oparg].name;
PyObject *descr = _PyType_Lookup(type, name);
if (descr != NULL && (Py_TYPE(descr)->tp_flags & Py_TPFLAGS_METHOD_DESCRIPTOR)) {
ADD_OP(_GUARD_TYPE_VERSION, 0, type->tp_version_tag);
bool immortal = _Py_IsImmortal(descr) || (type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE);
ADD_OP(immortal ? _LOAD_CONST_INLINE_BORROW : _LOAD_CONST_INLINE,
0, (uintptr_t)descr);
ADD_OP(_SWAP, 3, 0);
ADD_OP(_POP_TOP, 0, 0);
if ((type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE) == 0) {
PyType_Watch(TYPE_WATCHER_ID, (PyObject *)type);
_Py_BloomFilter_Add(dependencies, type);
}
method_and_self[0] = sym_new_const(ctx, descr);
optimized = true;
}
}
if (!optimized) {
method_and_self[0] = sym_new_not_null(ctx);
method_and_self[1] = sym_new_unknown(ctx);
}
}
op(_JUMP_TO_TOP, (--)) {

View file

@ -3555,8 +3555,30 @@
case _LOAD_SPECIAL: {
JitOptRef *method_and_self;
method_and_self = &stack_pointer[-2];
method_and_self[0] = sym_new_not_null(ctx);
method_and_self[1] = sym_new_unknown(ctx);
bool optimized = false;
PyTypeObject *type = sym_get_probable_type(method_and_self[1]);
if (type != NULL) {
PyObject *name = _Py_SpecialMethods[oparg].name;
PyObject *descr = _PyType_Lookup(type, name);
if (descr != NULL && (Py_TYPE(descr)->tp_flags & Py_TPFLAGS_METHOD_DESCRIPTOR)) {
ADD_OP(_GUARD_TYPE_VERSION, 0, type->tp_version_tag);
bool immortal = _Py_IsImmortal(descr) || (type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE);
ADD_OP(immortal ? _LOAD_CONST_INLINE_BORROW : _LOAD_CONST_INLINE,
0, (uintptr_t)descr);
ADD_OP(_SWAP, 3, 0);
ADD_OP(_POP_TOP, 0, 0);
if ((type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE) == 0) {
PyType_Watch(TYPE_WATCHER_ID, (PyObject *)type);
_Py_BloomFilter_Add(dependencies, type);
}
method_and_self[0] = sym_new_const(ctx, descr);
optimized = true;
}
}
if (!optimized) {
method_and_self[0] = sym_new_not_null(ctx);
method_and_self[1] = sym_new_unknown(ctx);
}
break;
}

View file

@ -115,6 +115,7 @@ const _PyOpcodeRecordEntry _PyOpcode_RecordEntries[256] = {
[STORE_ATTR_WITH_HINT] = {1, {_RECORD_TOS_TYPE_INDEX}},
[STORE_ATTR_SLOT] = {1, {_RECORD_TOS_TYPE_INDEX}},
[FOR_ITER_GEN] = {1, {_RECORD_NOS_GEN_FUNC_INDEX}},
[LOAD_SPECIAL] = {1, {_RECORD_TOS_TYPE_INDEX}},
[LOAD_ATTR_METHOD_WITH_VALUES] = {1, {_RECORD_TOS_TYPE_INDEX}},
[LOAD_ATTR_METHOD_NO_DICT] = {1, {_RECORD_TOS_TYPE_INDEX}},
[LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES] = {1, {_RECORD_TOS_TYPE_INDEX}},