Skip to content
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.

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

gh-126835: Move constant subscript folding to CFG #129568

Merged
merged 10 commits into from
Feb 4, 2025
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
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
10 changes: 0 additions & 10 deletions Lib/test/test_ast/test_ast.py
Original file line number Diff line number Diff line change
Expand Up @@ -3279,16 +3279,6 @@ def test_folding_iter(self):

self.assert_ast(code % (left, right), non_optimized_target, optimized_target)

def test_folding_subscript(self):
code = "(1,)[0]"

non_optimized_target = self.wrap_expr(
ast.Subscript(value=ast.Tuple(elts=[ast.Constant(value=1)]), slice=ast.Constant(value=0))
)
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 + 1](): pass"

Expand Down
2 changes: 2 additions & 0 deletions Lib/test/test_peepholer.py
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,8 @@ def test_constant_folding(self):
'(1, 2, -3)',
'(1, 2, -3) * 6',
'lambda x: x in {(3 * -5) + (-1 - 6), (1, -2, 3) * 2, None}',
'(1, )[0]',
'(1, (1, 2))[1][1]'
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
]
for e in exprs:
with self.subTest(e=e):
Expand Down
21 changes: 1 addition & 20 deletions Python/ast_opt.c
Original file line number Diff line number Diff line change
Expand Up @@ -567,25 +567,6 @@ fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
return make_const(node, newval, arena);
}

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

arg = node->v.Subscript.value;
idx = node->v.Subscript.slice;
if (node->v.Subscript.ctx != Load ||
arg->kind != Constant_kind ||
idx->kind != Constant_kind)
{
return 1;
}

newval = PyObject_GetItem(arg->v.Constant.value, idx->v.Constant.value);
return make_const(node, newval, arena);
}

/* Change literal list or set of constants into constant
tuple or frozenset respectively. Change literal list of
non-constants into tuple.
Expand Down Expand Up @@ -822,7 +803,7 @@ astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
case Subscript_kind:
CALL(astfold_expr, expr_ty, node_->v.Subscript.value);
CALL(astfold_expr, expr_ty, node_->v.Subscript.slice);
CALL(fold_subscr, expr_ty, node_);
/* Subscript folding is now done in flowgraph.c (optimize_if_const_subscr) */
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
break;
case Starred_kind:
CALL(astfold_expr, expr_ty, node_->v.Starred.value);
Expand Down
84 changes: 84 additions & 0 deletions Python/flowgraph.c
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,15 @@
return ERROR; \
}

#define RETURN_IF_FOLD_FAIL(X) \
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
if ((X) == NULL) { \
if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) { \
return ERROR; \
} \
PyErr_Clear(); \
return SUCCESS; \
}

#define DEFAULT_BLOCK_SIZE 16

typedef _Py_SourceLocation location;
Expand Down Expand Up @@ -1443,6 +1452,78 @@ optimize_if_const_list_or_set(PyObject *const_cache, cfg_instr* inst, int n, PyO
return SUCCESS;
}

/*
Walk basic block upwards starting from "start" to collect instruction pair
that loads consts skipping NOP's in between.
*/
static bool
find_load_const_pair(basicblock *bb, int start, cfg_instr **first, cfg_instr **second)
{
cfg_instr *second_load_const = NULL;
while (start >= 0) {
cfg_instr *inst = &bb->b_instr[start--];
if (inst->i_opcode == NOP) {
continue;
}
if (!loads_const(inst->i_opcode)) {
return false;
}
if (second_load_const == NULL) {
second_load_const = inst;
continue;
}
*first = inst;
*second = second_load_const;
return true;
}
return false;
}

/* Determine opcode & oparg for freshly folded constant. */
static int
newop_from_folded(PyObject *newconst, PyObject *consts,
PyObject *const_cache, int *newopcode, int *newoparg)
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
{
if (PyLong_CheckExact(newconst)) {
int overflow;
long val = PyLong_AsLongAndOverflow(newconst, &overflow);
if (!overflow && val >= 0 && val < 256) {
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
*newopcode = LOAD_SMALL_INT;
*newoparg = val;
return SUCCESS;
}
}
*newopcode = LOAD_CONST;
*newoparg = add_const(newconst, consts, const_cache);
RETURN_IF_ERROR(*newoparg);
return SUCCESS;
}

static int
optimize_if_const_subscr(basicblock *bb, int n, PyObject *consts, PyObject *const_cache)
WolframAlph marked this conversation as resolved.
Show resolved Hide resolved
{
cfg_instr *subscr = &bb->b_instr[n];
assert(subscr->i_opcode == BINARY_SUBSCR);
cfg_instr *arg, *idx;
if (!find_load_const_pair(bb, n-1, &arg, &idx)) {
return SUCCESS;
}
PyObject *o, *key;
if ((o = get_const_value(arg->i_opcode, arg->i_oparg, consts)) == NULL
|| (key = get_const_value(idx->i_opcode, idx->i_oparg, consts)) == NULL)
{
return ERROR;
}
PyObject *newconst = PyObject_GetItem(o, key);
RETURN_IF_FOLD_FAIL(newconst);
int newopcode, newoparg;
RETURN_IF_ERROR(newop_from_folded(newconst, consts, const_cache, &newopcode, &newoparg));
INSTR_SET_OP1(subscr, newopcode, newoparg);
INSTR_SET_OP0(arg, NOP);
INSTR_SET_OP0(idx, NOP);
return SUCCESS;
}

#define VISITED (-1)

// Replace an arbitrary run of SWAPs and NOPs with an optimal one that has the
Expand Down Expand Up @@ -1948,6 +2029,9 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
INSTR_SET_OP0(inst, NOP);
}
break;
case BINARY_SUBSCR:
RETURN_IF_ERROR(optimize_if_const_subscr(bb, i, consts, const_cache));
break;
}
}

Expand Down
Loading