From 2daa0e41500bc5dcb110faf9dee3a27b1964f7f4 Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Tue, 1 Jul 2025 23:34:37 +0800 Subject: [PATCH 1/5] JIT optimizer: Ignore escaping functions if inputs are pure --- Python/optimizer_cases.c.h | 6 +++--- Tools/cases_generator/generators_common.py | 11 +++++++---- Tools/cases_generator/optimizer_generator.py | 1 + 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 8d30df3aa7d429..fbd007d79c1333 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -2723,9 +2723,6 @@ PyObject *lhs_o = PyStackRef_AsPyObjectBorrow(lhs); PyObject *rhs_o = PyStackRef_AsPyObjectBorrow(rhs); assert(_PyEval_BinaryOps[oparg]); - stack_pointer[-2] = res; - stack_pointer += -1; - assert(WITHIN_STACK_BOUNDS()); PyObject *res_o = _PyEval_BinaryOps[oparg](lhs_o, rhs_o); if (res_o == NULL) { JUMP_TO_LABEL(error); @@ -2733,6 +2730,9 @@ res_stackref = PyStackRef_FromPyObjectSteal(res_o); /* End of uop copied from bytecodes for constant evaluation */ res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref)); + stack_pointer[-2] = res; + stack_pointer += -1; + assert(WITHIN_STACK_BOUNDS()); break; } bool lhs_int = sym_matches_type(lhs, &PyLong_Type); diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py index 47de205c0e9120..4c210fbf8d28e9 100644 --- a/Tools/cases_generator/generators_common.py +++ b/Tools/cases_generator/generators_common.py @@ -106,8 +106,9 @@ class Emitter: out: CWriter labels: dict[str, Label] _replacers: dict[str, ReplacementFunctionType] + cannot_escape: bool - def __init__(self, out: CWriter, labels: dict[str, Label]): + def __init__(self, out: CWriter, labels: dict[str, Label], cannot_escape: bool = False): self._replacers = { "EXIT_IF": self.exit_if, "DEOPT_IF": self.deopt_if, @@ -127,6 +128,7 @@ def __init__(self, out: CWriter, labels: dict[str, Label]): } self.out = out self.labels = labels + self.cannot_escape = cannot_escape def dispatch( self, @@ -238,7 +240,8 @@ def decref_inputs( next(tkn_iter) self._print_storage("DECREF_INPUTS", storage) try: - storage.close_inputs(self.out) + if not self.cannot_escape: + storage.close_inputs(self.out) except StackError as ex: raise analysis_error(ex.args[0], tkn) except Exception as ex: @@ -476,7 +479,7 @@ def emit_SimpleStmt( reachable = True tkn = stmt.contents[-1] try: - if stmt in uop.properties.escaping_calls: + if stmt in uop.properties.escaping_calls and not self.cannot_escape: escape = uop.properties.escaping_calls[stmt] if escape.kills is not None: self.stackref_kill(escape.kills, storage, True) @@ -513,7 +516,7 @@ def emit_SimpleStmt( self.out.emit(tkn) else: self.out.emit(tkn) - if stmt in uop.properties.escaping_calls: + if stmt in uop.properties.escaping_calls and not self.cannot_escape: self.emit_reload(storage) return reachable, None, storage except StackError as ex: diff --git a/Tools/cases_generator/optimizer_generator.py b/Tools/cases_generator/optimizer_generator.py index 4556b6d5a74f37..81ae534bddae5c 100644 --- a/Tools/cases_generator/optimizer_generator.py +++ b/Tools/cases_generator/optimizer_generator.py @@ -245,6 +245,7 @@ def __init__(self, out: CWriter, labels: dict[str, Label], original_uop: Uop, st outp.name: self.emit_stackref_override for outp in self.original_uop.stack.outputs } self._replacers = {**self._replacers, **overrides} + self.cannot_escape = True def emit_to_with_replacement( self, From a78813d53ef40ae2e16100885d31fa72d066f4e9 Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Tue, 1 Jul 2025 23:37:39 +0800 Subject: [PATCH 2/5] add a test --- Lib/test/test_capi/test_opt.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index e4c9a463855a69..eef789fccb685b 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -2451,6 +2451,21 @@ def testfunc(n): self.assertNotIn("_GUARD_TOS_FLOAT", uops) self.assertNotIn("_GUARD_NOS_FLOAT", uops) + def test_binary_op_constant_evaluate(self): + def testfunc(n): + for _ in range(n): + a = 2 ** 65 + + testfunc(TIER2_THRESHOLD) + + ex = get_first_executor(testfunc) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + # For now.. until we constant propagte it away. + self.assertIn("_BINARY_OP", uops) + + def global_identity(x): return x From e648644a3f2f0ff078a09fbe5bbd6566496c8f52 Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Tue, 1 Jul 2025 23:43:24 +0800 Subject: [PATCH 3/5] Add test to cases generator --- Lib/test/test_capi/test_opt.py | 2 +- Lib/test/test_generated_cases.py | 47 ++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index eef789fccb685b..cb5c9562c2bea2 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -2454,7 +2454,7 @@ def testfunc(n): def test_binary_op_constant_evaluate(self): def testfunc(n): for _ in range(n): - a = 2 ** 65 + 2 ** 65 testfunc(TIER2_THRESHOLD) diff --git a/Lib/test/test_generated_cases.py b/Lib/test/test_generated_cases.py index eb01328b6ea946..81d4e39f5be1ee 100644 --- a/Lib/test/test_generated_cases.py +++ b/Lib/test/test_generated_cases.py @@ -2398,6 +2398,53 @@ def test_replace_opcode_uop_body_copied_in_complex(self): """ self.run_cases_test(input, input2, output) + def test_replace_opcode_escaping_uop_body_copied_in_complex(self): + input = """ + pure op(OP, (foo -- res)) { + if (foo) { + res = ESCAPING_CODE(foo); + } + else { + res = 1; + } + } + """ + input2 = """ + op(OP, (foo -- res)) { + REPLACE_OPCODE_IF_EVALUATES_PURE(foo); + res = sym_new_known(ctx, foo); + } + """ + output = """ + case OP: { + JitOptRef foo; + JitOptRef res; + foo = stack_pointer[-1]; + if ( + sym_is_safe_const(ctx, foo) + ) { + JitOptRef foo_sym = foo; + _PyStackRef foo = sym_get_const_as_stackref(ctx, foo_sym); + _PyStackRef res_stackref; + /* Start of uop copied from bytecodes for constant evaluation */ + if (foo) { + res_stackref = ESCAPING_CODE(foo); + } + else { + res_stackref = 1; + } + /* End of uop copied from bytecodes for constant evaluation */ + res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref)); + stack_pointer[-1] = res; + break; + } + res = sym_new_known(ctx, foo); + stack_pointer[-1] = res; + break; + } + """ + self.run_cases_test(input, input2, output) + def test_replace_opocode_uop_reject_array_effects(self): input = """ pure op(OP, (foo[2] -- res)) { From 75ad8b1b40a5969dee9cd93f76d6de38992218f5 Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Wed, 2 Jul 2025 01:05:18 +0800 Subject: [PATCH 4/5] Add executor_cases to CI triggers --- .github/workflows/jit.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/jit.yml b/.github/workflows/jit.yml index 116e0c1e945e38..74bb11126b7922 100644 --- a/.github/workflows/jit.yml +++ b/.github/workflows/jit.yml @@ -5,6 +5,7 @@ on: - '**jit**' - 'Python/bytecodes.c' - 'Python/optimizer*.c' + - 'Python/executor_cases.c.h' - '!Python/perf_jit_trampoline.c' - '!**/*.md' - '!**/*.ini' @@ -13,6 +14,7 @@ on: - '**jit**' - 'Python/bytecodes.c' - 'Python/optimizer*.c' + - 'Python/executor_cases.c.h' - '!Python/perf_jit_trampoline.c' - '!**/*.md' - '!**/*.ini' From 0724a09e78e337a916dc089eadea4e9239e48c5e Mon Sep 17 00:00:00 2001 From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com> Date: Wed, 2 Jul 2025 13:34:56 +0800 Subject: [PATCH 5/5] Address review --- .github/workflows/jit.yml | 2 ++ Lib/test/test_capi/test_opt.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/jit.yml b/.github/workflows/jit.yml index 74bb11126b7922..947badff816c09 100644 --- a/.github/workflows/jit.yml +++ b/.github/workflows/jit.yml @@ -6,6 +6,7 @@ on: - 'Python/bytecodes.c' - 'Python/optimizer*.c' - 'Python/executor_cases.c.h' + - 'Python/optimizer_cases.c.h' - '!Python/perf_jit_trampoline.c' - '!**/*.md' - '!**/*.ini' @@ -15,6 +16,7 @@ on: - 'Python/bytecodes.c' - 'Python/optimizer*.c' - 'Python/executor_cases.c.h' + - 'Python/optimizer_cases.c.h' - '!Python/perf_jit_trampoline.c' - '!**/*.md' - '!**/*.ini' diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index cb5c9562c2bea2..7be1c9eebb3bf9 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -2462,7 +2462,7 @@ def testfunc(n): self.assertIsNotNone(ex) uops = get_opnames(ex) - # For now.. until we constant propagte it away. + # For now... until we constant propagate it away. self.assertIn("_BINARY_OP", uops)