Skip to content

Commit

Permalink
pythongh-126835: Move constant subscript folding to CFG (python#129568)
Browse files Browse the repository at this point in the history
Move folding of constant subscription from AST optimizer to CFG.

Co-authored-by: Irit Katriel <[email protected]>
  • Loading branch information
WolframAlph and iritkatriel authored Feb 4, 2025
1 parent bb5c687 commit 0664c1a
Show file tree
Hide file tree
Showing 6 changed files with 138 additions and 31 deletions.
2 changes: 2 additions & 0 deletions Include/internal/pycore_long.h
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ PyAPI_FUNC(void) _PyLong_ExactDealloc(PyObject *self);
# error "_PY_NSMALLPOSINTS must be greater than or equal to 257"
#endif

#define _PY_IS_SMALL_INT(val) ((val) >= 0 && (val) < 256 && (val) < _PY_NSMALLPOSINTS)

// Return a reference to the immortal zero singleton.
// The function cannot return NULL.
static inline PyObject* _PyLong_GetZero(void)
Expand Down
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
53 changes: 53 additions & 0 deletions Lib/test/test_peepholer.py
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,59 @@ def test_constant_folding(self):
self.assertFalse(instr.opname.startswith('BUILD_'))
self.check_lnotab(code)

def test_constant_folding_small_int(self):
tests = [
# subscript
('(0, )[0]', 0),
('(1 + 2, )[0]', 3),
('(2 + 2 * 2, )[0]', 6),
('(1, (1 + 2 + 3, ))[1][0]', 6),
('(255, )[0]', 255),
('(256, )[0]', None),
('(1000, )[0]', None),
('(1 - 2, )[0]', None),
]
for expr, oparg in tests:
with self.subTest(expr=expr, oparg=oparg):
code = compile(expr, '', 'single')
if oparg is not None:
self.assertInBytecode(code, 'LOAD_SMALL_INT', oparg)
else:
self.assertNotInBytecode(code, 'LOAD_SMALL_INT')
self.check_lnotab(code)

def test_folding_subscript(self):
tests = [
('(1, )[0]', False),
('(1, )[-1]', False),
('(1 + 2, )[0]', False),
('(1, (1, 2))[1][1]', False),
('(1, 2)[2-1]', False),
('(1, (1, 2))[1][2-1]', False),
('(1, (1, 2))[1:6][0][2-1]', False),
('"a"[0]', False),
('("a" + "b")[1]', False),
('("a" + "b", )[0][1]', False),
('("a" * 10)[9]', False),
('(1, )[1]', True),
('(1, )[-2]', True),
('"a"[1]', True),
('"a"[-2]', True),
('("a" + "b")[2]', True),
('("a" + "b", )[0][2]', True),
('("a" + "b", )[1][0]', True),
('("a" * 10)[10]', True),
('(1, (1, 2))[2:6][0][2-1]', True),
]
for expr, has_error in tests:
with self.subTest(expr=expr, has_error=has_error):
code = compile(expr, '', 'single')
if not has_error:
self.assertNotInBytecode(code, 'BINARY_SUBSCR')
else:
self.assertInBytecode(code, 'BINARY_SUBSCR')
self.check_lnotab(code)

def test_in_literal_list(self):
def containtest():
return x in [a, b]
Expand Down
20 changes: 0 additions & 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,6 @@ 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_);
break;
case Starred_kind:
CALL(astfold_expr, expr_ty, node_->v.Starred.value);
Expand Down
2 changes: 1 addition & 1 deletion Python/codegen.c
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ codegen_addop_load_const(compiler *c, location loc, PyObject *o)
if (PyLong_CheckExact(o)) {
int overflow;
long val = PyLong_AsLongAndOverflow(o, &overflow);
if (!overflow && val >= 0 && val < 256 && val < _PY_NSMALLPOSINTS) {
if (!overflow && _PY_IS_SMALL_INT(val)) {
ADDOP_I(c, loc, LOAD_SMALL_INT, val);
return SUCCESS;
}
Expand Down
82 changes: 82 additions & 0 deletions Python/flowgraph.c
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include "pycore_compile.h"
#include "pycore_intrinsics.h"
#include "pycore_pymem.h" // _PyMem_IsPtrFreed()
#include "pycore_long.h" // _PY_IS_SMALL_INT()

#include "pycore_opcode_utils.h"
#include "pycore_opcode_metadata.h" // OPCODE_HAS_ARG, etc
Expand Down Expand Up @@ -1443,6 +1444,84 @@ 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)
{
if (PyLong_CheckExact(newconst)) {
int overflow;
long val = PyLong_AsLongAndOverflow(newconst, &overflow);
if (!overflow && _PY_IS_SMALL_INT(val)) {
*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)
{
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);
if (newconst == NULL) {
if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) {
return ERROR;
}
PyErr_Clear();
return SUCCESS;
}
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 +2027,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

0 comments on commit 0664c1a

Please sign in to comment.