From 90a319650c04eca9bd956a30cecc81a30161565b Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 09:44:06 +0100 Subject: [PATCH 01/10] move constant subscript folding to CFG --- Lib/test/test_ast/test_ast.py | 10 ----- Lib/test/test_peepholer.py | 2 + Python/ast_opt.c | 21 +-------- Python/flowgraph.c | 84 +++++++++++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 30 deletions(-) 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..92e8284a55cc29 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -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]' ] for e in exprs: with self.subTest(e=e): diff --git a/Python/ast_opt.c b/Python/ast_opt.c index 01e208b88eca8b..7caffa5b26c3bd 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,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_constant_subscr) */ break; case Starred_kind: CALL(astfold_expr, expr_ty, node_->v.Starred.value); diff --git a/Python/flowgraph.c b/Python/flowgraph.c index a0b76050fd4af6..ec6ba700226f8b 100644 --- a/Python/flowgraph.c +++ b/Python/flowgraph.c @@ -21,6 +21,15 @@ return ERROR; \ } +#define RETURN_IF_FOLD_FAIL(X) \ + if ((X) == NULL) { \ + if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) { \ + return ERROR; \ + } \ + PyErr_Clear(); \ + return SUCCESS; \ + } + #define DEFAULT_BLOCK_SIZE 16 typedef _Py_SourceLocation location; @@ -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) +{ + if (PyLong_CheckExact(newconst)) { + int overflow; + long val = PyLong_AsLongAndOverflow(newconst, &overflow); + if (!overflow && val >= 0 && val < 256) { + *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); + 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 @@ -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; } } From 825c75dbac87f0693ca37f7cdd0ac33a9f3bf0ab Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 09:46:22 +0100 Subject: [PATCH 02/10] update comment --- Python/ast_opt.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Python/ast_opt.c b/Python/ast_opt.c index 7caffa5b26c3bd..8884227e66a285 100644 --- a/Python/ast_opt.c +++ b/Python/ast_opt.c @@ -803,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); - /* Subscript folding is now done in flowgraph.c (optimize_constant_subscr) */ + /* Subscript folding is now done in flowgraph.c (optimize_if_const_subscr) */ break; case Starred_kind: CALL(astfold_expr, expr_ty, node_->v.Starred.value); From 5e3eeec87dd8cac3ec86af428c63685c77703cf7 Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 16:22:07 +0100 Subject: [PATCH 03/10] address review --- Python/ast_opt.c | 1 - Python/flowgraph.c | 17 +++++++---------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/Python/ast_opt.c b/Python/ast_opt.c index 8884227e66a285..78d84002d593fb 100644 --- a/Python/ast_opt.c +++ b/Python/ast_opt.c @@ -803,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); - /* Subscript folding is now done in flowgraph.c (optimize_if_const_subscr) */ break; case Starred_kind: CALL(astfold_expr, expr_ty, node_->v.Starred.value); diff --git a/Python/flowgraph.c b/Python/flowgraph.c index ec6ba700226f8b..2f48d69aeada06 100644 --- a/Python/flowgraph.c +++ b/Python/flowgraph.c @@ -21,15 +21,6 @@ return ERROR; \ } -#define RETURN_IF_FOLD_FAIL(X) \ - if ((X) == NULL) { \ - if (PyErr_ExceptionMatches(PyExc_KeyboardInterrupt)) { \ - return ERROR; \ - } \ - PyErr_Clear(); \ - return SUCCESS; \ - } - #define DEFAULT_BLOCK_SIZE 16 typedef _Py_SourceLocation location; @@ -1515,7 +1506,13 @@ optimize_if_const_subscr(basicblock *bb, int n, PyObject *consts, PyObject *cons return ERROR; } PyObject *newconst = PyObject_GetItem(o, key); - RETURN_IF_FOLD_FAIL(newconst); + 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); From bfb913cab3d6e67e78d09fd9f80ad888602dbaef Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 19:47:40 +0100 Subject: [PATCH 04/10] add more tests --- Lib/test/test_peepholer.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 92e8284a55cc29..e1b2d8aa5d465e 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -475,6 +475,40 @@ def test_constant_folding(self): self.assertFalse(instr.opname.startswith('BUILD_')) self.check_lnotab(code) + def test_folding_subscript(self): + tests = [ + # small ints + ('(1, )[0]', True, False), + ('(255, )[0]', True, False), + ('(1, (1, 2))[1][1]', True, False), + ('(1, 2)[2-1]', True, False), + ('(1, (1, 2))[1][2-1]', True, False), + ('(1, (1, 2))[1:6][0][2-1]', True, False), + # regular ints + ('(256, )[0]', False, False), + ('(1, (1, 1000))[1][1]', False, False), + ('(1, 1000)[2-1]', False, False), + ('(1, (1, 1000))[1][2-1]', False, False), + # errors + ('(1, )[1]', True, True), + ('(1, )[-2]', False, True), + ('"a"[1]', True, True), + ('"a"[-2]', False, True), + ('(1, (1, 2))[2:6][0][2-1]', True, True), + ] + for expr, has_small_int, has_error in tests: + with self.subTest(expr=expr, has_small_int=has_small_int, has_error=has_error): + code = compile(expr, '', 'single') + if not has_error: + self.assertNotInBytecode(code, 'BINARY_SUBSCR') + else: + self.assertInBytecode(code, 'BINARY_SUBSCR') + if has_small_int: + self.assertInBytecode(code, 'LOAD_SMALL_INT') + else: + self.assertNotInBytecode(code, 'LOAD_SMALL_INT') + self.check_lnotab(code) + def test_in_literal_list(self): def containtest(): return x in [a, b] From 4c861c21aee93592be5d46d8b0a99b9f0ab95940 Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 19:50:15 +0100 Subject: [PATCH 05/10] remove previous tests --- Lib/test/test_peepholer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index e1b2d8aa5d465e..4b94ac71f79a41 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -463,8 +463,6 @@ 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]' ] for e in exprs: with self.subTest(e=e): From 729e813cae841616bcc70195c4a0b02a101c21b8 Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Sun, 2 Feb 2025 22:21:09 +0100 Subject: [PATCH 06/10] Update Python/flowgraph.c Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com> --- Python/flowgraph.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Python/flowgraph.c b/Python/flowgraph.c index 2f48d69aeada06..c2647ec133f758 100644 --- a/Python/flowgraph.c +++ b/Python/flowgraph.c @@ -1473,7 +1473,7 @@ find_load_const_pair(basicblock *bb, int start, cfg_instr **first, cfg_instr **s /* Determine opcode & oparg for freshly folded constant. */ static int newop_from_folded(PyObject *newconst, PyObject *consts, - PyObject *const_cache, int *newopcode, int *newoparg) + PyObject *const_cache, int *newopcode, int *newoparg) { if (PyLong_CheckExact(newconst)) { int overflow; From a4c8db612f8c10ee9dd8cc9cbf2686d6caf412b1 Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Mon, 3 Feb 2025 13:32:35 +0100 Subject: [PATCH 07/10] add _PY_IS_SMALL_INT macro --- Include/internal/pycore_long.h | 2 ++ Python/codegen.c | 2 +- Python/flowgraph.c | 3 ++- 3 files changed, 5 insertions(+), 2 deletions(-) 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/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 c2647ec133f758..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 @@ -1478,7 +1479,7 @@ newop_from_folded(PyObject *newconst, PyObject *consts, if (PyLong_CheckExact(newconst)) { int overflow; long val = PyLong_AsLongAndOverflow(newconst, &overflow); - if (!overflow && val >= 0 && val < 256) { + if (!overflow && _PY_IS_SMALL_INT(val)) { *newopcode = LOAD_SMALL_INT; *newoparg = val; return SUCCESS; From 5ad6781252ed3bf58471ae1a45b0299ed77d1adf Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Mon, 3 Feb 2025 13:44:32 +0100 Subject: [PATCH 08/10] update `test_folding_subscript` --- Lib/test/test_peepholer.py | 46 ++++++++++++++++++-------------------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 4b94ac71f79a41..f10bbf9f74ce27 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -475,36 +475,34 @@ def test_constant_folding(self): def test_folding_subscript(self): tests = [ - # small ints - ('(1, )[0]', True, False), - ('(255, )[0]', True, False), - ('(1, (1, 2))[1][1]', True, False), - ('(1, 2)[2-1]', True, False), - ('(1, (1, 2))[1][2-1]', True, False), - ('(1, (1, 2))[1:6][0][2-1]', True, False), - # regular ints - ('(256, )[0]', False, False), - ('(1, (1, 1000))[1][1]', False, False), - ('(1, 1000)[2-1]', False, False), - ('(1, (1, 1000))[1][2-1]', False, False), - # errors - ('(1, )[1]', True, True), - ('(1, )[-2]', False, True), - ('"a"[1]', True, True), - ('"a"[-2]', False, True), - ('(1, (1, 2))[2:6][0][2-1]', True, True), + ('(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_small_int, has_error in tests: - with self.subTest(expr=expr, has_small_int=has_small_int, has_error=has_error): + 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') - if has_small_int: - self.assertInBytecode(code, 'LOAD_SMALL_INT') - else: - self.assertNotInBytecode(code, 'LOAD_SMALL_INT') self.check_lnotab(code) def test_in_literal_list(self): From 59423b5450be9de4362363954b94b2214204976c Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Mon, 3 Feb 2025 14:03:05 +0100 Subject: [PATCH 09/10] add `test_constant_folding_small_int` --- Lib/test/test_peepholer.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index f10bbf9f74ce27..0fb2303dc72bf9 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -473,6 +473,28 @@ 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), From a9a1ee65e831f592a86263fea86a5491a4aa4179 Mon Sep 17 00:00:00 2001 From: Yan Yanchii Date: Mon, 3 Feb 2025 14:04:39 +0100 Subject: [PATCH 10/10] fix styling --- Lib/test/test_peepholer.py | 1 - 1 file changed, 1 deletion(-) diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py index 0fb2303dc72bf9..9f2f9350d74661 100644 --- a/Lib/test/test_peepholer.py +++ b/Lib/test/test_peepholer.py @@ -485,7 +485,6 @@ def test_constant_folding_small_int(self): ('(1000, )[0]', None), ('(1 - 2, )[0]', None), ] - for expr, oparg in tests: with self.subTest(expr=expr, oparg=oparg): code = compile(expr, '', 'single')