Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Commit 0664c1a

Browse filesBrowse files
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
Copy full SHA for 0664c1a

File tree

6 files changed

+138
-31
lines changed
Filter options

6 files changed

+138
-31
lines changed

‎Include/internal/pycore_long.h

Copy file name to clipboardExpand all lines: Include/internal/pycore_long.h
+2Lines 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

Copy file name to clipboardExpand all lines: Lib/test/test_ast/test_ast.py
-10Lines 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

Copy file name to clipboardExpand all lines: Lib/test/test_peepholer.py
+53Lines 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+
('("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

Copy file name to clipboardExpand all lines: Python/ast_opt.c
-20Lines 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

Copy file name to clipboardExpand all lines: Python/codegen.c
+1-1Lines 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

Copy file name to clipboardExpand all lines: Python/flowgraph.c
+82Lines 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)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.