https://github.com/python/cpython/commit/0664c1af9b29a5af2404e04a522f8e9e175ba05a commit: 0664c1af9b29a5af2404e04a522f8e9e175ba05a branch: main author: Yan Yanchii <[email protected]> committer: Eclips4 <[email protected]> date: 2025-02-04T10:10:55+02:00 summary:
gh-126835: Move constant subscript folding to CFG (#129568) Move folding of constant subscription from AST optimizer to CFG. Co-authored-by: Irit Katriel <[email protected]> files: M Include/internal/pycore_long.h M Lib/test/test_ast/test_ast.py M Lib/test/test_peepholer.py M Python/ast_opt.c M Python/codegen.c M Python/flowgraph.c diff --git a/Include/internal/pycore_long.h b/Include/internal/pycore_long.h index c52eb77692dd6a..df0656a7cb8f0c 100644 --- a/Include/internal/pycore_long.h +++ b/Include/internal/pycore_long.h @@ -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) diff --git a/Lib/test/test_ast/test_ast.py b/Lib/test/test_ast/test_ast.py index c268a1f00f938e..a438c8e81e4fd1 100644 --- a/Lib/test/test_ast/test_ast.py +++ b/Lib/test/test_ast/test_ast.py @@ -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" diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index b5b2b350e77a3b..9f2f9350d74661 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -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] diff --git a/Python/ast_opt.c b/Python/ast_opt.c index 01e208b88eca8b..78d84002d593fb 100644 --- a/Python/ast_opt.c +++ b/Python/ast_opt.c @@ -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. @@ -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); diff --git a/Python/codegen.c b/Python/codegen.c index 0bf9526cdc8435..e9853d7302f67f 100644 --- a/Python/codegen.c +++ b/Python/codegen.c @@ -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; } diff --git a/Python/flowgraph.c b/Python/flowgraph.c index a0b76050fd4af6..9ca7fadb8d7665 100644 --- a/Python/flowgraph.c +++ b/Python/flowgraph.c @@ -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 @@ -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 @@ -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; } } _______________________________________________ Python-checkins mailing list -- [email protected] To unsubscribe send an email to [email protected] https://mail.python.org/mailman3/lists/python-checkins.python.org/ Member address: [email protected]
