gh-132732: Automatically constant evaluate pure operations (GH-132733)

This adds a "macro" to the optimizer DSL called "REPLACE_OPCODE_IF_EVALUATES_PURE", which allows automatically constant evaluating a bytecode body if certain inputs have no side effects upon evaluations (such as ints, strings, and floats).


Co-authored-by: Tomas R. <tomas.roun8@gmail.com>
This commit is contained in:
Ken Jin 2025-06-27 19:37:44 +08:00 committed by GitHub
parent c45f4f3ebe
commit 695ab61351
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 706 additions and 122 deletions

View file

@ -2224,5 +2224,202 @@ class TestGeneratedAbstractCases(unittest.TestCase):
"Inputs must have equal sizes"):
self.run_cases_test(input, input2, output)
def test_pure_uop_body_copied_in(self):
# Note: any non-escaping call works.
# In this case, we use PyStackRef_IsNone.
input = """
pure op(OP, (foo -- res)) {
res = PyStackRef_IsNone(foo);
}
"""
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 */
res_stackref = PyStackRef_IsNone(foo);
/* 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_pure_uop_body_copied_in_deopt(self):
# Note: any non-escaping call works.
# In this case, we use PyStackRef_IsNone.
input = """
pure op(OP, (foo -- res)) {
DEOPT_IF(PyStackRef_IsNull(foo));
res = foo;
}
"""
input2 = """
op(OP, (foo -- res)) {
REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
res = 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 (PyStackRef_IsNull(foo)) {
ctx->done = true;
break;
}
res_stackref = foo;
/* 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 = foo;
stack_pointer[-1] = res;
break;
}
"""
self.run_cases_test(input, input2, output)
def test_pure_uop_body_copied_in_error_if(self):
# Note: any non-escaping call works.
# In this case, we use PyStackRef_IsNone.
input = """
pure op(OP, (foo -- res)) {
ERROR_IF(PyStackRef_IsNull(foo));
res = foo;
}
"""
input2 = """
op(OP, (foo -- res)) {
REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
res = 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 (PyStackRef_IsNull(foo)) {
goto error;
}
res_stackref = foo;
/* 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 = foo;
stack_pointer[-1] = res;
break;
}
"""
self.run_cases_test(input, input2, output)
def test_replace_opcode_uop_body_copied_in_complex(self):
input = """
pure op(OP, (foo -- res)) {
if (foo) {
res = PyStackRef_IsNone(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 = PyStackRef_IsNone(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)) {
if (foo) {
res = PyStackRef_IsNone(foo);
}
else {
res = 1;
}
}
"""
input2 = """
op(OP, (foo[2] -- res)) {
REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
res = sym_new_unknown(ctx);
}
"""
output = """
"""
with self.assertRaisesRegex(SyntaxError,
"Pure evaluation cannot take array-like inputs"):
self.run_cases_test(input, input2, output)
if __name__ == "__main__":
unittest.main()