Skip to content

gh-126835: Disable tuple folding in the AST optimizer#128802

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tomasr8 wants to merge 23 commits into from
Closed

gh-126835: Disable tuple folding in the AST optimizer #128802

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
2f475e1
Disable tuple folding in the AST optimizer
tomasr8 Jan 3, 2025
7a96d47
Provisionally fix tests
tomasr8 Jan 3, 2025
6d93343
Tweak tests
Eclips4 Jan 16, 2025
02a38a3
Merge branch 'main' into ast-tuple-folding
Eclips4 Jan 17, 2025
5c05d69
Restore UNPACK_SEQUENCE_TWO_TUPLE
Eclips4 Jan 21, 2025
05f3e62
Update Lib/test/test_peepholer.py
Eclips4 Jan 21, 2025
a3d9790
Merge branch 'main' into ast-tuple-folding
Eclips4 Jan 21, 2025
9f1feb4
Restore test_folding_op by changing const from tuple to a string
Eclips4 Jan 27, 2025
ac50aad
Fold set into frozenset in CFG
Eclips4 Jan 27, 2025
2b39a13
Merge branch 'main' into ast-tuple-folding
Eclips4 Feb 1, 2025
17dffaa
`for` target optimization (set/list -> frozenset/tuple)
Eclips4 Feb 1, 2025
46845f0
Fold list into a tuple as rhs in `for`/`in` operators
Eclips4 Feb 1, 2025
e9631d8
Fix tests
Eclips4 Feb 1, 2025
53dfa93
Merge branch 'main' into ast-tuple-folding
Eclips4 Feb 4, 2025
9754820
Merge branch 'main' into ast-tuple-folding
Eclips4 Feb 14, 2025
477c784
Fix merge artifacts
Eclips4 Feb 14, 2025
6731dfb
Merge branch 'main' into ast-tuple-folding
Eclips4 Feb 22, 2025
aad9fb3
Remove failing tests
Eclips4 Feb 22, 2025
be40093
Restore test_peepholer tests
Eclips4 Feb 22, 2025
5aec965
Add a few tests
Eclips4 Feb 22, 2025
ee69f0f
Restore test
Eclips4 Feb 22, 2025
0801463
Regenerate some files
Eclips4 Feb 22, 2025
a268315
Address review
Eclips4 Feb 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 0 additions & 111 deletions Lib/test/test_ast/test_ast.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -153,22 +153,6 @@ def test_optimization_levels__debug__(self):
self.assertIsInstance(res.body[0].value, ast.Name)
self.assertEqual(res.body[0].value.id, expected)

def test_optimization_levels_const_folding(self):
folded = ('Expr', (1, 0, 1, 6), ('Constant', (1, 0, 1, 6), (1, 2), None))
not_folded = ('Expr', (1, 0, 1, 6),
('Tuple', (1, 0, 1, 6),
[('Constant', (1, 1, 1, 2), 1, None),
('Constant', (1, 4, 1, 5), 2, None)], ('Load',)))

cases = [(-1, not_folded), (0, not_folded), (1, folded), (2, folded)]
for (optval, expected) in cases:
with self.subTest(optval=optval):
tree1 = ast.parse("(1, 2)", optimize=optval)
tree2 = ast.parse(ast.parse("(1, 2)"), optimize=optval)
for tree in [tree1, tree2]:
res = to_tuple(tree.body[0])
self.assertEqual(res, expected)

Comment on lines -156 to -171

@WolframAlphWolframAlphFeb 22, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as with test_compile_ast. Maybe add test for __debug__ instead of removing test entirely?

@WolframAlphWolframAlphFeb 24, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, there is test_optimization_levels__debug__ right above this one, so this one can indeed be gone.

def test_invalid_position_information(self):
invalid_linenos = [
(10, 1), (-10, -11), (10, -11), (-5, -2), (-5, 1)
Expand DownExpand Up@@ -3138,101 +3122,6 @@ def test_folding_format(self):

self.assert_ast(code, non_optimized_target, optimized_target)


def test_folding_tuple(self):
code = "(1,)"

non_optimized_target = self.wrap_expr(ast.Tuple(elts=[ast.Constant(1)]))
optimized_target = self.wrap_expr(ast.Constant(value=(1,)))

self.assert_ast(code, non_optimized_target, optimized_target)

def test_folding_type_param_in_function_def(self):
code = "def foo[%s = (1, 2)](): pass"

unoptimized_tuple = ast.Tuple(elts=[ast.Constant(1), ast.Constant(2)])
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.FunctionDef(
name='foo',
args=ast.arguments(),
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=ast.Constant((1, 2)))]
)
)
non_optimized_target = self.wrap_statement(
ast.FunctionDef(
name='foo',
args=ast.arguments(),
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=unoptimized_tuple)]
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)

def test_folding_type_param_in_class_def(self):
code = "class foo[%s = (1, 2)]: pass"

unoptimized_tuple = ast.Tuple(elts=[ast.Constant(1), ast.Constant(2)])
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.ClassDef(
name='foo',
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=ast.Constant((1, 2)))]
)
)
non_optimized_target = self.wrap_statement(
ast.ClassDef(
name='foo',
body=[ast.Pass()],
type_params=[type_param(name=name, default_value=unoptimized_tuple)]
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)

def test_folding_type_param_in_type_alias(self):
code = "type foo[%s = (1, 2)] = 1"

unoptimized_tuple = ast.Tuple(elts=[ast.Constant(1), ast.Constant(2)])
unoptimized_type_params = [
("T", "T", ast.TypeVar),
("**P", "P", ast.ParamSpec),
("*Ts", "Ts", ast.TypeVarTuple),
]

for type, name, type_param in unoptimized_type_params:
result_code = code % type
optimized_target = self.wrap_statement(
ast.TypeAlias(
name=ast.Name(id='foo', ctx=ast.Store()),
type_params=[type_param(name=name, default_value=ast.Constant((1, 2)))],
value=ast.Constant(value=1),
)
)
non_optimized_target = self.wrap_statement(
ast.TypeAlias(
name=ast.Name(id='foo', ctx=ast.Store()),
type_params=[type_param(name=name, default_value=unoptimized_tuple)],
value=ast.Constant(value=1),
)
)
self.assert_ast(result_code, non_optimized_target, optimized_target)
Comment on lines -3150 to -3234

@WolframAlphWolframAlphFeb 22, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as with test_compile_ast. Maybe not just remove them, but add testing of __debug__ as it is still in ast optimizer?


def test_folding_match_case_allowed_expressions(self):
def get_match_case_values(node):
result = []
Expand Down
11 changes: 5 additions & 6 deletions Lib/test/test_builtin.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -555,7 +555,7 @@ def test_compile_async_generator(self):
self.assertEqual(type(glob['ticker']()), AsyncGeneratorType)

def test_compile_ast(self):
args = ("a*(1,2)", "f.py", "exec")
args = ("a*__debug__", "f.py", "exec")
raw = compile(*args, flags = ast.PyCF_ONLY_AST).body[0]
opt1 = compile(*args, flags = ast.PyCF_OPTIMIZED_AST).body[0]
opt2 = compile(ast.parse(args[0]), *args[1:], flags = ast.PyCF_OPTIMIZED_AST).body[0]
Expand All@@ -566,14 +566,13 @@ def test_compile_ast(self):
self.assertIsInstance(tree.value.left, ast.Name)
self.assertEqual(tree.value.left.id, 'a')

raw_right = raw.value.right # expect Tuple((1, 2))
self.assertIsInstance(raw_right, ast.Tuple)
self.assertListEqual([elt.value for elt in raw_right.elts], [1, 2])
raw_right = raw.value.right
self.assertIsInstance(raw_right, ast.Name)

for opt in [opt1, opt2]:
opt_right = opt.value.right # expect Constant((1,2))
opt_right = opt.value.right
self.assertIsInstance(opt_right, ast.Constant)
self.assertEqual(opt_right.value, (1, 2))
self.assertEqual(opt_right.value, __debug__)

def test_delattr(self):
sys.spam = 1
Expand Down
4 changes: 2 additions & 2 deletions Lib/test/test_compile.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -797,9 +797,9 @@ def check_same_constant(const):
f1, f2 = lambda: "not a name", lambda: ("not a name",)
f3 = lambda x: x in {("not a name",)}
self.assertIs(f1.__code__.co_consts[0],
f2.__code__.co_consts[0][0])
f2.__code__.co_consts[1][0])
self.assertIs(next(iter(f3.__code__.co_consts[1])),
f2.__code__.co_consts[0])
f2.__code__.co_consts[1])

# {0} is converted to a constant frozenset({0}) by the peephole
# optimizer
Expand Down
8 changes: 5 additions & 3 deletions Lib/test/test_opcache.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1665,8 +1665,10 @@ def to_bool_str():
@requires_specialization_ft
def test_unpack_sequence(self):
def unpack_sequence_two_tuple():

t = 1, 2
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
a, b = 1, 2
a, b = t
self.assertEqual(a, 1)
self.assertEqual(b, 2)

Expand All@@ -1677,8 +1679,8 @@ def unpack_sequence_two_tuple():

def unpack_sequence_tuple():
for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
a, = 1,
self.assertEqual(a, 1)
a, b, c, d = 1, 2, 3, 4
self.assertEqual((a, b, c, d), (1, 2, 3, 4))

unpack_sequence_tuple()
self.assert_specialized(unpack_sequence_tuple, "UNPACK_SEQUENCE_TUPLE")
Expand Down
21 changes: 18 additions & 3 deletions Lib/test/test_peepholer.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -154,18 +154,33 @@ def test_folding_of_tuples_of_constants(self):
for line, elem in (
('a = 1,2,3', (1, 2, 3)),
('("a","b","c")', ('a', 'b', 'c')),
('a,b,c = 1,2,3', (1, 2, 3)),
('a,b,c,d = 1,2,3,4', (1, 2, 3, 4)),
('(None, 1, None)', (None, 1, None)),
('((1, 2), 3, 4)', ((1, 2), 3, 4)),
('(1, 2, (3, 4))', (1, 2, (3, 4))),
('()', ()),
('(1, (2, (3, (4, (5,)))))', (1, (2, (3, (4, (5,)))))),
):
with self.subTest(line=line):
code = compile(line,'','single')
self.assertInBytecode(code, 'LOAD_CONST', elem)
self.assertNotInBytecode(code, 'BUILD_TUPLE')
self.check_lnotab(code)

# Long tuples should be folded too.
code = compile(repr(tuple(range(10000))),'','single')
for expr, length in (
('(1, a)', 2),
('(a, b, c)', 3),
('(a, (b, c))', 2),
('(1, [], {})', 3),
):
with self.subTest(expr=expr, length=length):
code = compile(expr, '', 'single')
self.assertInBytecode(code, 'BUILD_TUPLE', length)
self.check_lnotab(code)

# Long tuples should be folded too, but their length should not
# exceed the `STACK_USE_GUIDELINE`

@tomasr8tomasr8Feb 22, 2025

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we perhaps add a test that tuples longer than STACK_USE_GUIDELINE are in fact not folded?

@Eclips4Eclips4Feb 23, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure about it. At the moment, for constant tuples which are longer than STACK_USE_GUIDELINE will be generated following bytecode:

BUILD_LIST 0
Pairs of LOAD_CONST + LIST_APPEND
CALL_INTRINSIC_1 (INTRINSIC_LIST_TO_TUPLE)

Shall we assert that this intrinsic is presented in bytecode?

code = compile(repr(tuple(range(30))),'','single')
Comment on lines +181 to +183

@WolframAlphWolframAlphFeb 22, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How much of a concern is that we can no longer create constant tuples beyond length of STACK_USE_GUIDELINE? @iritkatriel

@tomasr8tomasr8Feb 24, 2025

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's probably not ideal in case you have some kind of large constant lookup table for instance. As Kirill pointed out, this would get compiled to a bunch of LOAD_CONST + LIST_APPEND which is probably much slower

@WolframAlphWolframAlphFeb 24, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is not a hard pattern to detect, right? Maybe we could fold it anyway? @Eclips4

@tomasr8tomasr8Feb 24, 2025

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems quite predictable:

>>>deffoo(): ... return (1,2,3, ... ,31) ... >>>dis.dis(foo) 1RESUME02BUILD_LIST0LOAD_SMALL_INT1LIST_APPEND1LOAD_SMALL_INT2LIST_APPEND1 ...
LOAD_SMALL_INT31LIST_APPEND1CALL_INTRINSIC_16 (INTRINSIC_LIST_TO_TUPLE)

(Could also be LOAD_CONST instead of LOAD_SMALL_INT here I suppose)

@WolframAlphWolframAlphFeb 24, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly. I think we can easily fold it.

@tomasr8tomasr8Feb 24, 2025

Copy link
Copy Markdown
MemberAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we could check with is_const_tuple if the tuple is constant and if it is, ignore STACK_USE_GUIDELINE because we know it'll be folded by flowgraph anyway.

@WolframAlphWolframAlphFeb 24, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it's a good idea. We would be creating dependency between both.

@WolframAlphWolframAlphFeb 24, 2025•
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So if we are going to fold constant tuple beyond STACK_USE_GUIDELINE, maybe we could also add this optimization to literal lists & sets? It would be consistent with previous Python version as we are not folding these cases anymore after we migrated them to CFG.

self.assertNotInBytecode(code, 'BUILD_TUPLE')
# One LOAD_CONST for the tuple, one for the None return value
load_consts = [instr for instr in dis.get_instructions(code)
Expand Down
18 changes: 9 additions & 9 deletions Programs/test_frozenmain.h
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

35 changes: 0 additions & 35 deletions Python/ast_opt.c
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -278,40 +278,6 @@ fold_binop(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
return 1;
}

static PyObject*
make_const_tuple(asdl_expr_seq *elts)
{
for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
if (e->kind != Constant_kind) {
return NULL;
}
}

PyObject *newval = PyTuple_New(asdl_seq_LEN(elts));
if (newval == NULL) {
return NULL;
}

for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
PyObject *v = e->v.Constant.value;
PyTuple_SET_ITEM(newval, i, Py_NewRef(v));
}
return newval;
}

static int
fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
{
PyObject *newval;

if (node->v.Tuple.ctx != Load)
return 1;

newval = make_const_tuple(node->v.Tuple.elts);
return make_const(node, newval, arena);
}

static int astfold_mod(mod_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
static int astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state);
Expand DownExpand Up@@ -505,7 +471,6 @@ astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
break;
case Tuple_kind:
CALL_SEQ(astfold_expr, expr, node_->v.Tuple.elts);
CALL(fold_tuple, expr_ty, node_);
break;
case Name_kind:
if (node_->v.Name.ctx == Load &&
Expand Down
15 changes: 15 additions & 0 deletions Python/codegen.c
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1693,11 +1693,26 @@ codegen_typealias(compiler *c, stmt_ty s)
return SUCCESS;
}

static bool
is_const_tuple(asdl_expr_seq *elts)
{
for (Py_ssize_t i = 0; i < asdl_seq_LEN(elts); i++) {
expr_ty e = (expr_ty)asdl_seq_GET(elts, i);
if (e->kind != Constant_kind) {
return false;
}
}
return true;
}

/* Return false if the expression is a constant value except named singletons.
Return true otherwise. */
static bool
check_is_arg(expr_ty e)
{
if (e->kind == Tuple_kind) {
return !is_const_tuple(e->v.Tuple.elts);
}
if (e->kind != Constant_kind) {
return true;
}
Expand Down