gh-131798: optimize through keyword and bound method calls in the JIT (GH-148466)

Co-authored-by: Ken Jin <kenjin4096@gmail.com>
This commit is contained in:
Kumar Aditya
2026-04-13 18:44:48 +05:30
committed by GitHub
parent 10d275fdf8
commit 88e378cc1c
8 changed files with 1174 additions and 1001 deletions
+4 -4
View File
@@ -1134,9 +1134,9 @@ const struct opcode_metadata _PyOpcode_opcode_metadata[267] = {
[CALL_INTRINSIC_2] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
[CALL_ISINSTANCE] = { true, INSTR_FMT_IXC00, HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
[CALL_KW] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG | HAS_NEEDS_GUARD_IP_FLAG },
[CALL_KW_BOUND_METHOD] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG | HAS_NEEDS_GUARD_IP_FLAG },
[CALL_KW_BOUND_METHOD] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG | HAS_NEEDS_GUARD_IP_FLAG | HAS_RECORDS_VALUE_FLAG },
[CALL_KW_NON_PY] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_EVAL_BREAK_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
[CALL_KW_PY] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG | HAS_NEEDS_GUARD_IP_FLAG },
[CALL_KW_PY] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG | HAS_SYNC_SP_FLAG | HAS_NEEDS_GUARD_IP_FLAG | HAS_RECORDS_VALUE_FLAG },
[CALL_LEN] = { true, INSTR_FMT_IXC00, HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
[CALL_LIST_APPEND] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
[CALL_METHOD_DESCRIPTOR_FAST] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_EVAL_BREAK_FLAG | HAS_EXIT_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
@@ -1383,9 +1383,9 @@ _PyOpcode_macro_expansion[256] = {
[CALL_INTRINSIC_1] = { .nuops = 2, .uops = { { _CALL_INTRINSIC_1, OPARG_SIMPLE, 0 }, { _POP_TOP, OPARG_SIMPLE, 0 } } },
[CALL_INTRINSIC_2] = { .nuops = 3, .uops = { { _CALL_INTRINSIC_2, OPARG_SIMPLE, 0 }, { _POP_TOP, OPARG_SIMPLE, 0 }, { _POP_TOP, OPARG_SIMPLE, 0 } } },
[CALL_ISINSTANCE] = { .nuops = 3, .uops = { { _GUARD_THIRD_NULL, OPARG_SIMPLE, 3 }, { _GUARD_CALLABLE_ISINSTANCE, OPARG_SIMPLE, 3 }, { _CALL_ISINSTANCE, OPARG_SIMPLE, 3 } } },
[CALL_KW_BOUND_METHOD] = { .nuops = 6, .uops = { { _CHECK_PEP_523, OPARG_SIMPLE, 1 }, { _CHECK_METHOD_VERSION_KW, 2, 1 }, { _EXPAND_METHOD_KW, OPARG_SIMPLE, 3 }, { _PY_FRAME_KW, OPARG_SIMPLE, 3 }, { _SAVE_RETURN_OFFSET, OPARG_SAVE_RETURN_OFFSET, 3 }, { _PUSH_FRAME, OPARG_SIMPLE, 3 } } },
[CALL_KW_BOUND_METHOD] = { .nuops = 7, .uops = { { _RECORD_CALLABLE_KW, OPARG_SIMPLE, 0 }, { _CHECK_PEP_523, OPARG_SIMPLE, 1 }, { _CHECK_METHOD_VERSION_KW, 2, 1 }, { _EXPAND_METHOD_KW, OPARG_SIMPLE, 3 }, { _PY_FRAME_KW, OPARG_SIMPLE, 3 }, { _SAVE_RETURN_OFFSET, OPARG_SAVE_RETURN_OFFSET, 3 }, { _PUSH_FRAME, OPARG_SIMPLE, 3 } } },
[CALL_KW_NON_PY] = { .nuops = 3, .uops = { { _CHECK_IS_NOT_PY_CALLABLE_KW, OPARG_SIMPLE, 3 }, { _CALL_KW_NON_PY, OPARG_SIMPLE, 3 }, { _CHECK_PERIODIC_AT_END, OPARG_REPLACED, 3 } } },
[CALL_KW_PY] = { .nuops = 6, .uops = { { _CHECK_PEP_523, OPARG_SIMPLE, 1 }, { _CHECK_FUNCTION_VERSION_KW, 2, 1 }, { _CHECK_RECURSION_REMAINING, OPARG_SIMPLE, 3 }, { _PY_FRAME_KW, OPARG_SIMPLE, 3 }, { _SAVE_RETURN_OFFSET, OPARG_SAVE_RETURN_OFFSET, 3 }, { _PUSH_FRAME, OPARG_SIMPLE, 3 } } },
[CALL_KW_PY] = { .nuops = 7, .uops = { { _RECORD_CALLABLE_KW, OPARG_SIMPLE, 0 }, { _CHECK_PEP_523, OPARG_SIMPLE, 1 }, { _CHECK_FUNCTION_VERSION_KW, 2, 1 }, { _CHECK_RECURSION_REMAINING, OPARG_SIMPLE, 3 }, { _PY_FRAME_KW, OPARG_SIMPLE, 3 }, { _SAVE_RETURN_OFFSET, OPARG_SAVE_RETURN_OFFSET, 3 }, { _PUSH_FRAME, OPARG_SIMPLE, 3 } } },
[CALL_LEN] = { .nuops = 5, .uops = { { _GUARD_NOS_NULL, OPARG_SIMPLE, 3 }, { _GUARD_CALLABLE_LEN, OPARG_SIMPLE, 3 }, { _CALL_LEN, OPARG_SIMPLE, 3 }, { _POP_TOP, OPARG_SIMPLE, 3 }, { _POP_TOP, OPARG_SIMPLE, 3 } } },
[CALL_LIST_APPEND] = { .nuops = 6, .uops = { { _GUARD_CALLABLE_LIST_APPEND, OPARG_SIMPLE, 3 }, { _GUARD_NOS_NOT_NULL, OPARG_SIMPLE, 3 }, { _GUARD_NOS_LIST, OPARG_SIMPLE, 3 }, { _CALL_LIST_APPEND, OPARG_SIMPLE, 3 }, { _POP_TOP, OPARG_SIMPLE, 3 }, { _POP_TOP, OPARG_SIMPLE, 3 } } },
[CALL_METHOD_DESCRIPTOR_FAST] = { .nuops = 5, .uops = { { _GUARD_CALLABLE_METHOD_DESCRIPTOR_FAST, OPARG_SIMPLE, 3 }, { _CALL_METHOD_DESCRIPTOR_FAST, OPARG_SIMPLE, 3 }, { _POP_TOP_OPARG, OPARG_SIMPLE, 3 }, { _POP_TOP, OPARG_SIMPLE, 3 }, { _CHECK_PERIODIC_AT_END, OPARG_REPLACED, 3 } } },
+996 -995
View File
File diff suppressed because it is too large Load Diff
+4
View File
@@ -410,6 +410,7 @@ const uint32_t _PyUop_Flags[MAX_UOP_ID+1] = {
[_RECORD_3OS_GEN_FUNC] = HAS_RECORDS_VALUE_FLAG,
[_RECORD_4OS] = HAS_RECORDS_VALUE_FLAG,
[_RECORD_CALLABLE] = HAS_ARG_FLAG | HAS_RECORDS_VALUE_FLAG,
[_RECORD_CALLABLE_KW] = HAS_ARG_FLAG | HAS_RECORDS_VALUE_FLAG,
[_RECORD_BOUND_METHOD] = HAS_ARG_FLAG | HAS_RECORDS_VALUE_FLAG,
[_RECORD_CODE] = HAS_RECORDS_VALUE_FLAG,
};
@@ -5665,6 +5666,7 @@ const char *const _PyOpcode_uop_name[MAX_UOP_REGS_ID+1] = {
[_RECORD_4OS] = "_RECORD_4OS",
[_RECORD_BOUND_METHOD] = "_RECORD_BOUND_METHOD",
[_RECORD_CALLABLE] = "_RECORD_CALLABLE",
[_RECORD_CALLABLE_KW] = "_RECORD_CALLABLE_KW",
[_RECORD_CODE] = "_RECORD_CODE",
[_RECORD_NOS] = "_RECORD_NOS",
[_RECORD_NOS_GEN_FUNC] = "_RECORD_NOS_GEN_FUNC",
@@ -6626,6 +6628,8 @@ int _PyUop_num_popped(int opcode, int oparg)
return 0;
case _RECORD_CALLABLE:
return 0;
case _RECORD_CALLABLE_KW:
return 0;
case _RECORD_BOUND_METHOD:
return 0;
case _RECORD_CODE:
+45
View File
@@ -5253,6 +5253,51 @@ class TestUopsOptimization(unittest.TestCase):
PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
self.assertEqual(result[0].rc, 0, result)
def test_call_kw(self):
def func(a):
return int(a) * 42
def testfunc(n):
x = 0
for _ in range(n):
x += func(a=1)
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(res, 42 * TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_PUSH_FRAME", uops)
self.assertIn("_CHECK_FUNCTION_VERSION_KW", uops)
# Check the optimizer has optmized through the function call
# by promoting global `int` to a constant.
self.assertNotIn("_LOAD_GLOBAL_BUILTINS", uops)
self.assertIn("_CALL_BUILTIN_CLASS", uops)
def test_call_kw_bound_method(self):
class C:
def method(self, a, b):
return int(a) + int(b)
def testfunc(n):
obj = C()
x = 0
meth = obj.method
for _ in range(n):
x += meth(a=1, b=2)
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(res, 3 * TIER2_THRESHOLD)
self.assertIsNotNone(ex)
uops = get_opnames(ex)
self.assertIn("_PUSH_FRAME", uops)
self.assertIn("_CHECK_METHOD_VERSION_KW", uops)
# Check the optimizer has optmized through the function call
# by promoting global `int` to a constant.
self.assertNotIn("_LOAD_GLOBAL_BUILTINS", uops)
self.assertIn("_CALL_BUILTIN_CLASS", uops)
def test_func_version_guarded_on_change(self):
def testfunc(n):
for i in range(n):
+6
View File
@@ -5249,6 +5249,7 @@ dummy_func(
}
macro(CALL_KW_PY) =
_RECORD_CALLABLE_KW +
unused/1 + // Skip over the counter
_CHECK_PEP_523 +
_CHECK_FUNCTION_VERSION_KW +
@@ -5279,6 +5280,7 @@ dummy_func(
}
macro(CALL_KW_BOUND_METHOD) =
_RECORD_CALLABLE_KW +
unused/1 + // Skip over the counter
_CHECK_PEP_523 +
_CHECK_METHOD_VERSION_KW +
@@ -6157,6 +6159,10 @@ dummy_func(
RECORD_VALUE(PyStackRef_AsPyObjectBorrow(func));
}
tier2 op(_RECORD_CALLABLE_KW, (func, self, args[oparg], kwnames -- func, self, args[oparg], kwnames)) {
RECORD_VALUE(PyStackRef_AsPyObjectBorrow(func));
}
tier2 op(_RECORD_BOUND_METHOD, (callable, self, args[oparg] -- callable, self, args[oparg])) {
PyObject *callable_o = PyStackRef_AsPyObjectBorrow(callable);
if (Py_TYPE(callable_o) == &PyMethod_Type) {
+51
View File
@@ -1022,6 +1022,18 @@ dummy_func(void) {
_Py_BloomFilter_Add(dependencies, func);
}
op(_CHECK_FUNCTION_VERSION_KW, (func_version/2, callable, unused, unused[oparg], unused -- callable, unused, unused[oparg], unused)) {
PyObject *func = sym_get_probable_value(callable);
if (func == NULL || !PyFunction_Check(func) || ((PyFunctionObject *)func)->func_version != func_version) {
ctx->contradiction = true;
ctx->done = true;
break;
}
// Guarded on this, so it can be promoted.
sym_set_const(callable, func);
_Py_BloomFilter_Add(dependencies, func);
}
op(_CHECK_METHOD_VERSION, (func_version/2, callable, null, unused[oparg] -- callable, null, unused[oparg])) {
if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyMethod_Type)) {
PyMethodObject *method = (PyMethodObject *)sym_get_const(ctx, callable);
@@ -1045,6 +1057,29 @@ dummy_func(void) {
sym_set_type(callable, &PyMethod_Type);
}
op(_CHECK_METHOD_VERSION_KW, (func_version/2, callable, null, unused[oparg], unused -- callable, null, unused[oparg], unused)) {
if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyMethod_Type)) {
PyMethodObject *method = (PyMethodObject *)sym_get_const(ctx, callable);
assert(PyMethod_Check(method));
ADD_OP(_CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
uop_buffer_last(&ctx->out_buffer)->operand1 = (uintptr_t)method->im_func;
}
else {
// Guarding on the bound method, safe to promote.
PyObject *bound_method = sym_get_probable_value(callable);
if (bound_method != NULL && Py_TYPE(bound_method) == &PyMethod_Type) {
PyMethodObject *method = (PyMethodObject *)bound_method;
PyObject *func = method->im_func;
if (PyFunction_Check(func) &&
((PyFunctionObject *)func)->func_version == func_version) {
_Py_BloomFilter_Add(dependencies, func);
sym_set_const(callable, bound_method);
}
}
}
sym_set_type(callable, &PyMethod_Type);
}
op(_CHECK_FUNCTION_EXACT_ARGS, (callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
assert(sym_matches_type(callable, &PyFunction_Type));
if (sym_is_const(ctx, callable)) {
@@ -1092,6 +1127,18 @@ dummy_func(void) {
}
}
op(_EXPAND_METHOD_KW, (callable, self_or_null, unused[oparg], unused -- callable, self_or_null, unused[oparg], unused)) {
if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyMethod_Type)) {
PyMethodObject *method = (PyMethodObject *)sym_get_const(ctx, callable);
callable = sym_new_const(ctx, method->im_func);
self_or_null = sym_new_const(ctx, method->im_self);
}
else {
callable = sym_new_not_null(ctx);
self_or_null = sym_new_not_null(ctx);
}
}
op(_MAYBE_EXPAND_METHOD, (callable, self_or_null, args[oparg] -- callable, self_or_null, args[oparg])) {
(void)args;
callable = sym_new_not_null(ctx);
@@ -2261,6 +2308,10 @@ dummy_func(void) {
sym_set_recorded_value(func, (PyObject *)this_instr->operand0);
}
op(_RECORD_CALLABLE_KW, (func, self, args[oparg], kwnames -- func, self, args[oparg], kwnames)) {
sym_set_recorded_value(func, (PyObject *)this_instr->operand0);
}
op(_RECORD_BOUND_METHOD, (callable, self, args[oparg] -- callable, self, args[oparg])) {
sym_set_recorded_value(callable, (PyObject *)this_instr->operand0);
}
+55
View File
@@ -4612,14 +4612,62 @@
}
case _CHECK_FUNCTION_VERSION_KW: {
JitOptRef callable;
callable = stack_pointer[-3 - oparg];
uint32_t func_version = (uint32_t)this_instr->operand0;
PyObject *func = sym_get_probable_value(callable);
if (func == NULL || !PyFunction_Check(func) || ((PyFunctionObject *)func)->func_version != func_version) {
ctx->contradiction = true;
ctx->done = true;
break;
}
sym_set_const(callable, func);
_Py_BloomFilter_Add(dependencies, func);
break;
}
case _CHECK_METHOD_VERSION_KW: {
JitOptRef callable;
callable = stack_pointer[-3 - oparg];
uint32_t func_version = (uint32_t)this_instr->operand0;
if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyMethod_Type)) {
PyMethodObject *method = (PyMethodObject *)sym_get_const(ctx, callable);
assert(PyMethod_Check(method));
ADD_OP(_CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
uop_buffer_last(&ctx->out_buffer)->operand1 = (uintptr_t)method->im_func;
}
else {
PyObject *bound_method = sym_get_probable_value(callable);
if (bound_method != NULL && Py_TYPE(bound_method) == &PyMethod_Type) {
PyMethodObject *method = (PyMethodObject *)bound_method;
PyObject *func = method->im_func;
if (PyFunction_Check(func) &&
((PyFunctionObject *)func)->func_version == func_version) {
_Py_BloomFilter_Add(dependencies, func);
sym_set_const(callable, bound_method);
}
}
}
sym_set_type(callable, &PyMethod_Type);
break;
}
case _EXPAND_METHOD_KW: {
JitOptRef self_or_null;
JitOptRef callable;
self_or_null = stack_pointer[-2 - oparg];
callable = stack_pointer[-3 - oparg];
if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyMethod_Type)) {
PyMethodObject *method = (PyMethodObject *)sym_get_const(ctx, callable);
callable = sym_new_const(ctx, method->im_func);
self_or_null = sym_new_const(ctx, method->im_self);
}
else {
callable = sym_new_not_null(ctx);
self_or_null = sym_new_not_null(ctx);
}
stack_pointer[-3 - oparg] = callable;
stack_pointer[-2 - oparg] = self_or_null;
break;
}
@@ -5248,6 +5296,13 @@
break;
}
case _RECORD_CALLABLE_KW: {
JitOptRef func;
func = stack_pointer[-3 - oparg];
sym_set_recorded_value(func, (PyObject *)this_instr->operand0);
break;
}
case _RECORD_BOUND_METHOD: {
JitOptRef callable;
callable = stack_pointer[-2 - oparg];
+13 -2
View File
@@ -69,6 +69,13 @@ void _PyOpcode_RecordFunction_CALLABLE(_PyInterpreterFrame *frame, _PyStackRef *
Py_INCREF(*recorded_value);
}
void _PyOpcode_RecordFunction_CALLABLE_KW(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, int oparg, PyObject **recorded_value) {
_PyStackRef func;
func = stack_pointer[-3 - oparg];
*recorded_value = (PyObject *)PyStackRef_AsPyObjectBorrow(func);
Py_INCREF(*recorded_value);
}
void _PyOpcode_RecordFunction_BOUND_METHOD(_PyInterpreterFrame *frame, _PyStackRef *stack_pointer, int oparg, PyObject **recorded_value) {
_PyStackRef callable;
callable = stack_pointer[-2 - oparg];
@@ -90,7 +97,8 @@ void _PyOpcode_RecordFunction_CODE(_PyInterpreterFrame *frame, _PyStackRef *stac
#define _RECORD_NOS_GEN_FUNC_INDEX 4
#define _RECORD_CALLABLE_INDEX 5
#define _RECORD_BOUND_METHOD_INDEX 6
#define _RECORD_4OS_INDEX 7
#define _RECORD_CALLABLE_KW_INDEX 7
#define _RECORD_4OS_INDEX 8
const uint8_t _PyOpcode_RecordFunctionIndices[256] = {
[TO_BOOL_ALWAYS_TRUE] = _RECORD_TOS_TYPE_INDEX,
[BINARY_OP_SUBSCR_GETITEM] = _RECORD_NOS_INDEX,
@@ -123,10 +131,12 @@ const uint8_t _PyOpcode_RecordFunctionIndices[256] = {
[CALL_METHOD_DESCRIPTOR_O] = _RECORD_CALLABLE_INDEX,
[CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS] = _RECORD_CALLABLE_INDEX,
[CALL_METHOD_DESCRIPTOR_NOARGS] = _RECORD_CALLABLE_INDEX,
[CALL_KW_PY] = _RECORD_CALLABLE_KW_INDEX,
[CALL_KW_BOUND_METHOD] = _RECORD_CALLABLE_KW_INDEX,
[CALL_EX_PY] = _RECORD_4OS_INDEX,
};
const _Py_RecordFuncPtr _PyOpcode_RecordFunctions[8] = {
const _Py_RecordFuncPtr _PyOpcode_RecordFunctions[9] = {
[0] = NULL,
[_RECORD_TOS_TYPE_INDEX] = _PyOpcode_RecordFunction_TOS_TYPE,
[_RECORD_NOS_INDEX] = _PyOpcode_RecordFunction_NOS,
@@ -134,5 +144,6 @@ const _Py_RecordFuncPtr _PyOpcode_RecordFunctions[8] = {
[_RECORD_NOS_GEN_FUNC_INDEX] = _PyOpcode_RecordFunction_NOS_GEN_FUNC,
[_RECORD_CALLABLE_INDEX] = _PyOpcode_RecordFunction_CALLABLE,
[_RECORD_BOUND_METHOD_INDEX] = _PyOpcode_RecordFunction_BOUND_METHOD,
[_RECORD_CALLABLE_KW_INDEX] = _PyOpcode_RecordFunction_CALLABLE_KW,
[_RECORD_4OS_INDEX] = _PyOpcode_RecordFunction_4OS,
};