python · Fidget-Spinner · Apr 19, 2025 · Apr 19, 2025 · Apr 19, 2025 · Apr 19, 2025
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
@@ -10,6 +10,7 @@ extern "C" {
 
 #include "pycore_typedefs.h"      // _PyInterpreterFrame
 #include "pycore_uop_ids.h"
+#include "pycore_stackref.h"      // _PyStackRef
 #include <stdbool.h>
 
 
@@ -259,7 +260,10 @@ typedef struct _JitOptContext {
 extern bool _Py_uop_sym_is_null(JitOptSymbol *sym);
 extern bool _Py_uop_sym_is_not_null(JitOptSymbol *sym);
 extern bool _Py_uop_sym_is_const(JitOptContext *ctx, JitOptSymbol *sym);
+extern bool _Py_uop_sym_is_safe_const(JitOptContext *ctx, JitOptSymbol *sym);
 extern PyObject *_Py_uop_sym_get_const(JitOptContext *ctx, JitOptSymbol *sym);
+extern JitOptSymbol *_Py_uop_sym_new_const_steal(JitOptContext *ctx, PyObject *const_val);
+extern _PyStackRef _Py_uop_sym_get_const_as_stackref(JitOptContext *ctx, JitOptSymbol *sym);
 extern JitOptSymbol *_Py_uop_sym_new_unknown(JitOptContext *ctx);
 extern JitOptSymbol *_Py_uop_sym_new_not_null(JitOptContext *ctx);
 extern JitOptSymbol *_Py_uop_sym_new_type(

diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h
diff --git a/Lib/test/test_generated_cases.py b/Lib/test/test_generated_cases.py
@@ -1969,12 +1969,12 @@ def run_cases_test(self, input: str, input2: str, expected: str):
 
     def test_overridden_abstract(self):
         input = """
-        pure op(OP, (--)) {
+        op(OP, (--)) {
             SPAM();
         }
         """
         input2 = """
-        pure op(OP, (--)) {
+        op(OP, (--)) {
             eggs();
         }
         """
@@ -1988,7 +1988,7 @@ def test_overridden_abstract(self):
 
     def test_overridden_abstract_args(self):
         input = """
-        pure op(OP, (arg1 -- out)) {
+        op(OP, (arg1 -- out)) {
             out = SPAM(arg1);
         }
         op(OP2, (arg1 -- out)) {
@@ -2021,16 +2021,16 @@ def test_overridden_abstract_args(self):
 
     def test_no_overridden_case(self):
         input = """
-        pure op(OP, (arg1 -- out)) {
+        op(OP, (arg1 -- out)) {
             out = SPAM(arg1);
         }
 
-        pure op(OP2, (arg1 -- out)) {
+        op(OP2, (arg1 -- out)) {
         }
 
         """
         input2 = """
-        pure op(OP2, (arg1 -- out)) {
+        op(OP2, (arg1 -- out)) {
             out = NULL;
         }
         """
@@ -2250,5 +2250,114 @@ def test_validate_uop_unused_size_mismatch(self):
                                     "Inputs must have equal sizes"):
             self.run_cases_test(input, input2, output)
 
+    def test_pure_uop_body_copied_in(self):
+        input = """
+        pure op(OP, (foo -- res)) {
+            res = body(foo);
+        }
+        """
+        input2 = """
+        op(OP, (foo -- res)) {
+            REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
+            res = sym_new_known(ctx, foo);
+        }
+        """
+        output = """
+        case OP: {
+            JitOptSymbol *foo;
+            JitOptSymbol *res;
+            foo = stack_pointer[-1];
+            if (
+                sym_is_safe_const(ctx, foo)
+            ) {
+                JitOptSymbol *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 = body(foo);
+                /* End of uop copied from bytecodes for constant evaluation */
+                res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref));
+                stack_pointer[-1] = res;
+            }
+            else {
+                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_complex(self):
+        input = """
+        pure op(OP, (foo -- res)) {
+            if (foo) {
+                res = body(foo);
+            }
+            else {
+                res = 1;
+            }
+        }
+        """
+        input2 = """
+        op(OP, (foo -- res)) {
+            REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
+            res = sym_new_known(ctx, foo);
+        }
+        """
+        output = """
+        case OP: {
+            JitOptSymbol *foo;
+            JitOptSymbol *res;
+            foo = stack_pointer[-1];
+            if (
+                sym_is_safe_const(ctx, foo)
+            ) {
+                JitOptSymbol *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 = body(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;
+            }
+            else {
+                res = sym_new_known(ctx, foo);
+                stack_pointer[-1] = res;
+            }
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
+    def test_pure_uop_reject_array_effects(self):
+        input = """
+        pure op(OP, (foo[2] -- res)) {
+            if (foo) {
+                res = body(foo);
+            }
+            else {
+                res = 1;
+            }
+        }
+        """
+        input2 = """
+        op(OP, (foo[2] -- res)) {
+            REPLACE_OPCODE_IF_EVALUATES_PURE(foo[2]);
+            res = sym_new_unknown(ctx);
+        }
+        """
+        output = """
+        """
+        with self.assertRaisesRegex(AssertionError,
+                                    "Unsafe to convert a symbol to an array-like StackRef."):
+            self.run_cases_test(input, input2, output)
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2025-04-19-16-22-47.gh-issue-132732.jgqhlF.rst b/Misc/NEWS.d/next/Core_and_Builtins/2025-04-19-16-22-47.gh-issue-132732.jgqhlF.rst
@@ -0,0 +1 @@
+Automatically constant evaluate bytecode operations marked as pure in the JIT optimizer.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Automatically constant evaluate bytecode operations marked as pure in the JIT optimizer.
Fidget-Spinner marked this conversation as resolved. Show resolved Hide resolved