GH-131798: Narrow the return type of _GET_LEN to int (GH-133345)

This commit is contained in:
Nadeshiko Manju 2025-05-21 06:02:50 +08:00 committed by GitHub
parent e4fbfb1288
commit e1c0c451a2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 87 additions and 1 deletions

View file

@ -1925,6 +1925,50 @@ class TestUopsOptimization(unittest.TestCase):
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertNotIn("_GUARD_TOS_INT", uops)
def test_get_len_with_const_tuple(self):
def testfunc(n):
x = 0.0
for _ in range(n):
match (1, 2, 3, 4):
case [_, _, _, _]:
x += 1.0
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(int(res), TIER2_THRESHOLD)
uops = get_opnames(ex)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertNotIn("_GET_LEN", uops)
self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
def test_get_len_with_non_const_tuple(self):
def testfunc(n):
x = 0.0
for _ in range(n):
match object(), object():
case [_, _]:
x += 1.0
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(int(res), TIER2_THRESHOLD)
uops = get_opnames(ex)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertNotIn("_GET_LEN", uops)
self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
def test_get_len_with_non_tuple(self):
def testfunc(n):
x = 0.0
for _ in range(n):
match [1, 2, 3, 4]:
case [_, _, _, _]:
x += 1.0
return x
res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
self.assertEqual(int(res), TIER2_THRESHOLD)
uops = get_opnames(ex)
self.assertNotIn("_GUARD_NOS_INT", uops)
self.assertIn("_GET_LEN", uops)
def test_binary_op_subscr_tuple_int(self):
def testfunc(n):
x = 0

View file

@ -0,0 +1,2 @@
Allow the JIT to remove int guards after ``_GET_LEN`` by setting the return
type to int.

View file

@ -1116,6 +1116,25 @@ dummy_func(void) {
res = sym_new_type(ctx, &PyLong_Type);
}
op(_GET_LEN, (obj -- obj, len)) {
int tuple_length = sym_tuple_length(obj);
if (tuple_length == -1) {
len = sym_new_type(ctx, &PyLong_Type);
}
else {
assert(tuple_length >= 0);
PyObject *temp = PyLong_FromLong(tuple_length);
if (temp == NULL) {
goto error;
}
if (_Py_IsImmortal(temp)) {
REPLACE_OP(this_instr, _LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)temp);
}
len = sym_new_const(ctx, temp);
Py_DECREF(temp);
}
}
op(_GUARD_CALLABLE_LEN, (callable, unused, unused -- callable, unused, unused)) {
PyObject *len = _PyInterpreterState_GET()->callable_cache.len;
if (sym_get_const(ctx, callable) == len) {

View file

@ -1463,8 +1463,29 @@
}
case _GET_LEN: {
JitOptSymbol *obj;
JitOptSymbol *len;
len = sym_new_not_null(ctx);
obj = stack_pointer[-1];
int tuple_length = sym_tuple_length(obj);
if (tuple_length == -1) {
len = sym_new_type(ctx, &PyLong_Type);
}
else {
assert(tuple_length >= 0);
PyObject *temp = PyLong_FromLong(tuple_length);
if (temp == NULL) {
goto error;
}
if (_Py_IsImmortal(temp)) {
REPLACE_OP(this_instr, _LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)temp);
}
len = sym_new_const(ctx, temp);
stack_pointer[0] = len;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());
Py_DECREF(temp);
stack_pointer += -1;
}
stack_pointer[0] = len;
stack_pointer += 1;
assert(WITHIN_STACK_BOUNDS());