8000 gh-126835: Move constant subscript folding to CFG (#129568) · python/cpython@0664c1a · GitHub
[go: up one dir, main page]
Skip to content

Commit 0664c1a

Browse files
WolframAlphiritkatriel
WolframAlph
and
iritkatriel
authored
gh-126835: Move constant subscript folding to CFG (#129568)
Move folding of constant subscription from AST optimizer to CFG. Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>
1 parent bb5c687 commit 0664c1a

File tree

6 files changed

+138
-31
lines changed

6 files changed

+138
-31
lines changed

‎Include/internal/pycore_long.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ PyAPI_FUNC(void) _PyLong_ExactDealloc(PyObject *self);
6565
# error "_PY_NSMALLPOSINTS must be greater than or equal to 257"
6666
#endif
6767

68+
#define _PY_IS_SMALL_INT(val) ((val) >= 0 && (val) < 256 && (val) < _PY_NSMALLPOSINTS)
69+
6870
// Return a reference to the immortal zero singleton.
6971
// The function cannot return NULL.
7072
static inline PyObject* _PyLong_GetZero(void)

‎Lib/test/test_ast/test_ast.py

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3279,16 +3279,6 @@ def test_folding_iter(self):
32793279

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

3282-
def test_folding_subscript(self):
3283-
code = "(1,)[0]"
3284-
3285-
non_optimized_target = self.wrap_expr(
3286-
ast.Subscript(value=ast.Tuple(elts=[ast.Constant(value=1)]), slice=ast.Constant(value=0))
3287-
)
3288-
optimized_target = self.wrap_expr(ast.Constant(value=1))
3289-
3290-
self.assert_ast(code, non_optimized_target, optimized_target)
3291-
32923282
def test_folding_type_param_in_function_def(self):
32933283
code = "def foo[%s = 1 + 1](): pass"
32943284

‎Lib/test/test_peepholer.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,59 @@ def test_constant_folding(self):
473473
self.assertFalse(instr.opname.startswith('BUILD_'))
474474
self.check_lnotab(code)
475475

476+
def test_constant_folding_small_int(self):
477+
tests = [
478+
# subscript
479+
('(0, )[0]', 0),
480+
('(1 + 2, )[0]', 3),
481+
('(2 + 2 * 2, )[0]', 6),
482+
('(1, (1 + 2 + 3, ))[1][0]', 6),
483+
('(255, )[0]', 255),
484+
('(256, )[0]', None),
485+
('(1000, )[0]', None),
486+
('(1 - 2, )[0]', None),
487+
]
488+
for expr, oparg in tests:
489+
with self.subTest(expr=expr, oparg=oparg):
490+
code = compile(expr, '', 'single')
491+
if oparg is not None:
492+
self.assertInBytecode(code, 'LOAD_SMALL_INT', oparg)
493+
else:
494+
self.assertNotInBytecode(code, 'LOAD_SMALL_INT')
495+
self.check_lnotab(code)
496+
497+
def test_folding_subscript(self):
498+
tests = [
499+
('(1, )[0]', False),
500+
('(1, )[-1]', False),
501+
('(1 + 2, )[0]', False),
502+
('(1, (1, 2))[1][1]', False),
503+
('(1, 2)[2-1]', False),
504+
('(1, (1, 2))[1][2-1]', False),
505+
('(1, (1, 2))[1:6][0][2-1]', False),
506+
('"a"[0]', False),
507+
('("a" + "b")[1]', False),
508+
(< A3DB span class=pl-s>'("a" + "b", )[0][1]', False),
509+
('("a" * 10)[9]', False),
510+
('(1, )[1]', True),
511+
('(1, )[-2]', True),
512+
('"a"[1]', True),
513+
('"a"[-2]', True),
514+
('("a" + "b")[2]', True),
515+
('("a" + "b", )[0][2]', True),
516+
('("a" + "b", )[1][0]', True),
517+
('("a" * 10)[10]', True),
518+
('(1, (1, 2))[2:6][0][2-1]', True),
519+
]
520+
for expr, has_error in tests:
521+
with self.subTest(expr=expr, has_error=has_error):
522+
code = compile(expr, '', 'single')
523+
if not has_error:
524+
self.assertNotInBytecode(code, 'BINARY_SUBSCR')
525+
else:
526+
self.assertInBytecode(code, 'BINARY_SUBSCR')
527+
self.check_lnotab(code)
528+
476529
def test_in_literal_list(self):
477530
def containtest():
478531
return x in [a, b]

‎Python/ast_opt.c

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -567,25 +567,6 @@ fold_tuple(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
567567
return make_const(node, newval, arena);
568568
}
569569

570-
static int
571-
fold_subscr(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
572-
{
573-
PyObject *newval;
574-
expr_ty arg, idx;
575-
576-
arg = node->v.Subscript.value;
577-
idx = node->v.Subscript.slice;
578-
if (node->v.Subscript.ctx != Load ||
579-
arg->kind != Constant_kind ||
580-
idx->kind != Constant_kind)
581-
{
582-
return 1;
583-
}
584-
585-
newval = PyObject_GetItem(arg->v.Constant.value, idx->v.Constant.value);
586-
return make_const(node, newval, arena);
587-
}
588-
589570
/* Change literal list or set of constants into constant
590571
tuple or frozenset respectively. Change literal list of
591572
non-constants into tuple.
@@ -822,7 +803,6 @@ astfold_expr(expr_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
822803
case Subscript_kind:
823804
CALL(astfold_expr, expr_ty, node_->v.Subscript.value);
824805
CALL(astfold_expr, expr_ty, node_->v.Subscript.slice);
825-
CALL(fold_subscr, expr_ty, node_);
826806
break;
827807
case Starred_kind:
828808
CALL(astfold_expr, expr_ty, node_->v.Starred.value);

‎Python/codegen.c

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -284,7 +284,7 @@ codegen_addop_load_const(compiler *c, location loc, PyObject *o)
284284
if (PyLong_CheckExact(o)) {
285285
int overflow;
286286
long val = PyLong_AsLongAndOverflow(o, &overflow);
287-
if (!overflow && val >= 0 && val < 256 && val < _PY_NSMALLPOSINTS) {
287+
if (!overflow && _PY_IS_SMALL_INT(val)) {
288288
ADDOP_I(c, loc, LOAD_SMALL_INT, val);
289289
return SUCCESS;
290290
}

‎Python/flowgraph.c

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include "pycore_compile.h"
77
#include "pycore_intrinsics.h"
88
#include "pycore_pymem.h" // _PyMem_IsPtrFreed()
9+
#include "pycore_long.h" // _PY_IS_SMALL_INT()
910

1011
#include "pycore_opcode_utils.h"
1112
#include "pycore_opcode_metadata.h" // OPCODE_HAS_ARG, etc
@@ -1443,6 +1444,84 @@ optimize_if_const_list_or_set(PyObject *const_cache, cfg_instr* inst, int n, PyO
14431444
return SUCCESS;
14441445
}
14451446

1447+
/*
1448+
Walk basic block upwards starting from "start" to collect instruction pair
1449+
that loads consts skipping NOP's in between.
1450+
*/
1451+
static bool
1452+
find_load_const_pair(basicblock *bb, int start, cfg_instr **first, cfg_instr **second)
1453+
{
1454+
cfg_instr *second_load_const = NULL;
1455+
while (start >= 0) {
1456+
cfg_instr *inst = &bb->b_instr[start--];
1457+
if (inst->i_opcode == NOP) {
1458+
continue;
1459+
}
1460+
if (!loads_const(inst->i_opcode)) {
1461+
return false;
1462+
}
1463+
if (second_load_const == NULL) {
1464+
second_load_const = inst;
1465+
continue;
1466+
}
1467+
*first = inst;
1468+
*second = second_load_const;
1469+
return true;
1470+
}
1471+
return false;
1472+
}
1473+
1474+
/* Determine opcode & oparg for freshly folded constant. */
1475+
static int
1476+
newop_from_folded(PyObject *newconst, PyObject *consts,
1477+
PyObject *const_cache, int *newopcode, int *newoparg)
1478+
{
1479+
if (PyLong_CheckExact(newconst)) {
1480+
int overflow;
1481+
long val = PyLong_AsLongAndOverflow(newconst, &overflow);
1482+
if (!overflow && _PY_IS_SMALL_INT(val)) {
1483+
*newopcode = LOAD_SMALL_INT;
1484+
*newoparg = val;
1485+
return SUCCESS;
1486+
}
1487+
}
1488+
*newopcode = LOAD_CONST;
1489+
*newoparg = add_const(newconst, consts, const_cache);
1490+
RETURN_IF_ERROR(*newoparg);
1491+
return SUCCESS;
1492+
}
1493+
1494+
static int
1495+
optimize_if_const_subscr(basicblock *bb, int n, PyObject *consts, PyObject *const_cache)
1496+
{
1497+
cfg_instr *subscr = &bb->b_instr[n];
1498+
assert(subscr->i_opcode == BINARY_SUBSCR);
1499+
cfg_instr *arg, *idx;
1500+
if (!find_load_const_pair(bb, n-1, &arg, &idx)) {
1501+
return SUCCESS;
1502+
}
1503+
PyObject *o, *key;
1504+
if ((o = get_const_value(arg->i_opcode, arg->i_oparg, consts)) == NULL
1505+
|| (key = get_const_value(idx->i_opcode, idx->i_oparg, consts)) == NULL)
1506+
{
1507+
return ERROR;
1508+
}
1509+
PyObject *newconst = PyObject_GetItem(o, key);
1510+
if (newconst == NULL) {
1511+
if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) {
1512+
return ERROR;
1513+
}
1514+
PyErr_Clear();
1515+
return SUCCESS;
1516+
}
1517+
int newopcode, newoparg;
1518+
RETURN_IF_ERROR(newop_from_folded(newconst, consts, const_cache, &newopcode, &newoparg));
1519+
INSTR_SET_OP1(subscr, newopcode, newoparg);
1520+
INSTR_SET_OP0(arg, NOP);
1521+
INSTR_SET_OP0(idx, NOP);
1522+
return SUCCESS;
1523+
}
1524+
14461525
#define VISITED (-1)
14471526

14481527
// Replace an arbitrary run of SWAPs and NOPs with an optimal one that has the
@@ -1948,6 +2027,9 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
19482027
INSTR_SET_OP0(inst, NOP);
19492028
}
19502029
break;
2030+
case BINARY_SUBSCR:
2031+
RETURN_IF_ERROR(optimize_if_const_subscr(bb, i, consts, const_cache));
2032+
break;
19512033
}
19522034
}
19532035

0 commit comments

Comments
 (0)
0