Skip to content

gh-136183: Deal with escapes in JIT optimizer's constant evaluator #136184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 2, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .github/workflows/jit.yml
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ on:
- '**jit**'
- 'Python/bytecodes.c'
- 'Python/optimizer*.c'
- 'Python/executor_cases.c.h'
- 'Python/optimizer_cases.c.h'
- '!Python/perf_jit_trampoline.c'
- '!**/*.md'
- '!**/*.ini'
Expand All @@ -13,6 +15,8 @@ on:
- '**jit**'
- 'Python/bytecodes.c'
- 'Python/optimizer*.c'
- 'Python/executor_cases.c.h'
- 'Python/optimizer_cases.c.h'
- '!Python/perf_jit_trampoline.c'
- '!**/*.md'
- '!**/*.ini'
Expand Down
15 changes: 15 additions & 0 deletions Lib/test/test_capi/test_opt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
2 ** 65

testfunc(TIER2_THRESHOLD)

ex = get_first_executor(testfunc)
self.assertIsNotNone(ex)
uops = get_opnames(ex)

# For now... until we constant propagate it away.
self.assertIn("_BINARY_OP", uops)


def global_identity(x):
return x

Expand Down
47 changes: 47 additions & 0 deletions Lib/test/test_generated_cases.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
6 changes: 3 additions & 3 deletions Python/optimizer_cases.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 7 additions & 4 deletions Tools/cases_generator/generators_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
1 change: 1 addition & 0 deletions Tools/cases_generator/optimizer_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading