From d18afa038ab3c3103fd34fbbcb1f9e9a50f92520 Mon Sep 17 00:00:00 2001 From: Kumar Aditya Date: Sun, 12 Apr 2026 15:39:09 +0530 Subject: [PATCH 1/2] JIT constant fold special method lookup --- Include/internal/pycore_opcode_metadata.h | 4 ++-- Lib/test/test_capi/test_opt.py | 19 +++++++++++++++++ Python/bytecodes.c | 1 + Python/optimizer_bytecodes.c | 25 +++++++++++++++++++++-- Python/optimizer_cases.c.h | 25 +++++++++++++++++++++-- Python/record_functions.c.h | 1 + 6 files changed, 69 insertions(+), 6 deletions(-) diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h index 9bd78f67c6e2d4..9a5fab5cd8256b 100644 --- a/Include/internal/pycore_opcode_metadata.h +++ b/Include/internal/pycore_opcode_metadata.h @@ -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 }, @@ -1471,7 +1471,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 } } }, diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index f11413cc625422..0122c037d898d3 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -3316,6 +3316,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 diff --git a/Python/bytecodes.c b/Python/bytecodes.c index 34e3885a93c8bd..8ee725ec05fbd3 100644 --- a/Python/bytecodes.c +++ b/Python/bytecodes.c @@ -3887,6 +3887,7 @@ dummy_func( } macro(LOAD_SPECIAL) = + _RECORD_TOS_TYPE + _INSERT_NULL + _LOAD_SPECIAL; diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index c12a4f4131bc7e..95de600a0ccd04 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -1682,8 +1682,29 @@ 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); + 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 && _PyType_HasFeature(Py_TYPE(descr), 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); + } + } + else { + method_and_self[0] = sym_new_not_null(ctx); + method_and_self[1] = sym_new_unknown(ctx); + } } op(_JUMP_TO_TOP, (--)) { diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 2db2c87cb3610b..f9d81bd28db5ac 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -3452,8 +3452,29 @@ 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); + 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 && _PyType_HasFeature(Py_TYPE(descr), 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); + } + } + else { + method_and_self[0] = sym_new_not_null(ctx); + method_and_self[1] = sym_new_unknown(ctx); + } break; } diff --git a/Python/record_functions.c.h b/Python/record_functions.c.h index 02b8538bc902b5..861b02cc5d75fa 100644 --- a/Python/record_functions.c.h +++ b/Python/record_functions.c.h @@ -106,6 +106,7 @@ const uint8_t _PyOpcode_RecordFunctionIndices[256] = { [STORE_ATTR_WITH_HINT] = _RECORD_TOS_TYPE_INDEX, [STORE_ATTR_SLOT] = _RECORD_TOS_TYPE_INDEX, [FOR_ITER_GEN] = _RECORD_NOS_GEN_FUNC_INDEX, + [LOAD_SPECIAL] = _RECORD_TOS_TYPE_INDEX, [LOAD_ATTR_METHOD_WITH_VALUES] = _RECORD_TOS_TYPE_INDEX, [LOAD_ATTR_METHOD_NO_DICT] = _RECORD_TOS_TYPE_INDEX, [LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES] = _RECORD_TOS_TYPE_INDEX, From 9da1b9d21bcffb1f970f2a23cec644fa04b67764 Mon Sep 17 00:00:00 2001 From: Kumar Aditya Date: Mon, 13 Apr 2026 16:06:45 +0530 Subject: [PATCH 2/2] refactor it a bit --- Python/optimizer_bytecodes.c | 4 +++- Python/optimizer_cases.c.h | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index 95de600a0ccd04..c15e43d283df66 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -1682,6 +1682,7 @@ dummy_func(void) { } op(_LOAD_SPECIAL, (method_and_self[2] -- method_and_self[2])) { + bool optimized = false; PyTypeObject *type = sym_get_probable_type(method_and_self[1]); if (type != NULL) { PyObject *name = _Py_SpecialMethods[oparg].name; @@ -1699,9 +1700,10 @@ dummy_func(void) { _Py_BloomFilter_Add(dependencies, type); } method_and_self[0] = sym_new_const(ctx, descr); + optimized = true; } } - else { + if (!optimized) { method_and_self[0] = sym_new_not_null(ctx); method_and_self[1] = sym_new_unknown(ctx); } diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index f9d81bd28db5ac..26b09246c8c911 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -3452,6 +3452,7 @@ case _LOAD_SPECIAL: { JitOptRef *method_and_self; method_and_self = &stack_pointer[-2]; + bool optimized = false; PyTypeObject *type = sym_get_probable_type(method_and_self[1]); if (type != NULL) { PyObject *name = _Py_SpecialMethods[oparg].name; @@ -3469,9 +3470,10 @@ _Py_BloomFilter_Add(dependencies, type); } method_and_self[0] = sym_new_const(ctx, descr); + optimized = true; } } - else { + if (!optimized) { method_and_self[0] = sym_new_not_null(ctx); method_and_self[1] = sym_new_unknown(ctx); }