gh-120619: Tier 2 partial evaluator foundations

.github/CODEOWNERS

@@ -45,6 +45,8 @@ Python/bytecodes.c            @markshannon
 Python/optimizer*.c           @markshannon
 Python/optimizer_analysis.c   @Fidget-Spinner
 Python/optimizer_bytecodes.c  @Fidget-Spinner
+Python/partial_evaluator.c   @Fidget-Spinner
+Python/partial_evaluator_bytecodes.c  @Fidget-Spinner
 Python/symtable.c             @JelleZijlstra @carljm
 Lib/_pyrepl/*                 @pablogsal @lysnikolaou @ambv
 Lib/test/test_patma.py        @brandtbucher

Include/internal/pycore_optimizer.h

@@ -48,8 +48,9 @@ typedef struct {
  *    uint16_t error_target;
  */
 typedef struct {
-    uint16_t opcode:15;
+    uint16_t opcode:14;
     uint16_t format:1;
+    uint16_t is_virtual:1;   // Used for tier2 optimization.
     uint16_t oparg;
     union {
         uint32_t target;
@@ -147,6 +148,15 @@ int _Py_uop_analyze_and_optimize(struct _PyInterpreterFrame *frame,
     _PyUOpInstruction *trace, int trace_len, int curr_stackentries,
     _PyBloomFilter *dependencies);
 
+int
+_Py_uop_partial_evaluate(
+    _PyInterpreterFrame *frame,
+    _PyUOpInstruction *buffer,
+    int length,
+    int curr_stacklen,
+    _PyBloomFilter *dependencies
+);
+
 extern PyTypeObject _PyCounterExecutor_Type;
 extern PyTypeObject _PyCounterOptimizer_Type;
 extern PyTypeObject _PyDefaultOptimizer_Type;
@@ -156,6 +166,8 @@ extern PyTypeObject _PyUOpOptimizer_Type;
 /* Symbols */
 /* See explanation in optimizer_symbols.c */
 
+// Specializer.
+
 struct _Py_UopsSymbol {
     int flags;  // 0 bits: Top; 2 or more bits: Bottom
     PyTypeObject *typ;  // Borrowed reference
@@ -285,6 +297,88 @@ static inline int is_terminator(const _PyUOpInstruction *uop)
     );
 }
 
+// Partial evaluator.
+struct _Py_UopsPESymbol {
+    int flags;  // 0 bits: Top; 2 or more bits: Bottom
+    PyObject *const_val;  // Owned reference (!)
+};
+
+typedef struct _Py_UopsPESymbol _Py_UopsPESymbol;
+
+typedef struct _Py_UopsPESlot {
+    _Py_UopsPESymbol *sym;
+    _PyUOpInstruction *origin_inst; // The instruction this symbol originates from.
+} _Py_UopsPESlot;
+
+struct _Py_UOpsPEAbstractFrame {
+    // Max stacklen
+    int stack_len;
+    int locals_len;
+
+    _Py_UopsPESlot *stack_pointer;
+    _Py_UopsPESlot *stack;
+    _Py_UopsPESlot *locals;
+};
+
+typedef struct _Py_UOpsPEAbstractFrame _Py_UOpsPEAbstractFrame;
+
+typedef struct pe_arena {
+    int sym_curr_number;
+    int sym_max_number;
+    _Py_UopsPESymbol arena[TY_ARENA_SIZE];
+} pe_arena;
+
+struct _Py_UOpsPEContext {
+    char done;
+    char out_of_space;
+    bool contradiction;
+    // The current "executing" frame.
+    _Py_UOpsPEAbstractFrame *frame;
+    _Py_UOpsPEAbstractFrame frames[MAX_ABSTRACT_FRAME_DEPTH];
+    int curr_frame_depth;
+
+    // Arena for the symbolic information.
+    pe_arena sym_arena;
+
+    _Py_UopsPESlot *n_consumed;
+    _Py_UopsPESlot *limit;
+    _Py_UopsPESlot locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
+};
+
+typedef struct _Py_UOpsPEContext _Py_UOpsPEContext;
+
+extern bool _Py_uop_pe_sym_is_null(_Py_UopsPESlot *sym);
+extern bool _Py_uop_pe_sym_is_not_null(_Py_UopsPESlot *sym);
+extern bool _Py_uop_pe_sym_is_const(_Py_UopsPESlot *sym);
+extern PyObject *_Py_uop_pe_sym_get_const(_Py_UopsPESlot *sym);
+extern _Py_UopsPESlot _Py_uop_pe_sym_new_unknown(_Py_UOpsPEContext *ctx);
+extern _Py_UopsPESlot _Py_uop_pe_sym_new_not_null(_Py_UOpsPEContext *ctx);
+extern _Py_UopsPESlot _Py_uop_pe_sym_new_const(_Py_UOpsPEContext *ctx, PyObject *const_val);
+extern _Py_UopsPESlot _Py_uop_pe_sym_new_null(_Py_UOpsPEContext *ctx);
+extern void _Py_uop_pe_sym_set_null(_Py_UOpsPEContext *ctx, _Py_UopsPESlot *sym);
+extern void _Py_uop_pe_sym_set_non_null(_Py_UOpsPEContext *ctx, _Py_UopsPESlot *sym);
+extern void _Py_uop_pe_sym_set_const(_Py_UOpsPEContext *ctx, _Py_UopsPESlot *sym, PyObject *const_val);
+extern bool _Py_uop_pe_sym_is_bottom(_Py_UopsPESlot *sym);
+extern int _Py_uop_pe_sym_truthiness(_Py_UopsPESlot *sym);
+extern void _Py_uop_sym_set_origin_inst_override(_Py_UopsPESlot *sym, _PyUOpInstruction *origin);
+extern _PyUOpInstruction *_Py_uop_sym_get_origin(_Py_UopsPESlot *sym);
+extern bool _Py_uop_sym_is_virtual(_Py_UopsPESlot *sym);
+
+
+extern _Py_UOpsPEAbstractFrame *
+_Py_uop_pe_frame_new(
+    _Py_UOpsPEContext *ctx,
+    PyCodeObject *co,
+    int curr_stackentries,
+    _Py_UopsPESlot *args,
+    int arg_len);
+
+int _Py_uop_pe_frame_pop(_Py_UOpsPEContext *ctx);
+
+extern void _Py_uop_pe_abstractcontext_init(_Py_UOpsPEContext *ctx);
+extern void _Py_uop_pe_abstractcontext_fini(_Py_UOpsPEContext *ctx);
+
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_capi/test_opt.py

@@ -1482,6 +1482,33 @@ def fn(a):
 
         fn(A())
 
+    def test_pe_load_fast_pop_top(self):
+        def thing(a):
+            x = 0
+            for i in range(TIER2_THRESHOLD):
+                i
+            return i
+
+
+        res, ex = self._run_with_optimizer(thing, 1)
+        self.assertEqual(res, 4095)
+        self.assertIsNotNone(ex)
+        self.assertEqual(list(iter_opnames(ex)).count("_POP_TOP"), 0)
+        self.assertTrue(ex.is_valid())
+
+    def test_pe_dead_store_elimination(self):
+        def thing(a):
+            x = 0
+            for i in range(TIER2_THRESHOLD):
+                x = x
+            return i
+
+
+        res, ex = self._run_with_optimizer(thing, 1)
+        self.assertEqual(res, 4095)
+        self.assertIsNotNone(ex)
+        self.assertEqual(list(iter_opnames(ex)).count("_LOAD_FAST_1"), 0)
+        self.assertTrue(ex.is_valid())
     def test_func_guards_removed_or_reduced(self):
         def testfunc(n):
             for i in range(n):

Lib/test/test_generated_cases.py

@@ -38,6 +38,7 @@ def skip_if_different_mount_drives():
     import tier1_generator
     import opcode_metadata_generator
     import optimizer_generator
+    import partial_evaluator_generator
 
 
 def handle_stderr():
@@ -1641,6 +1642,9 @@ def test_escaping_call_next_to_cmacro(self):
 
 
 class TestGeneratedAbstractCases(unittest.TestCase):
+
+    generator = None
+
     def setUp(self) -> None:
         super().setUp()
         self.maxDiff = None
@@ -1676,7 +1680,8 @@ def run_cases_test(self, input: str, input2: str, expected: str):
             temp_input.flush()
 
         with handle_stderr():
-            optimizer_generator.generate_tier2_abstract_from_files(
+            assert self.generator is not None
+            self.generator.generate_tier2_abstract_from_files(
                 [self.temp_input_filename, self.temp_input2_filename],
                 self.temp_output_filename
             )
@@ -1690,6 +1695,9 @@ def run_cases_test(self, input: str, input2: str, expected: str):
         actual = "".join(lines)
         self.assertEqual(actual.strip(), expected.strip())
 
+
+class TestGeneratedOptimizerCases(TestGeneratedAbstractCases):
+    generator = optimizer_generator
     def test_overridden_abstract(self):
         input = """
         pure op(OP, (--)) {
@@ -1790,5 +1798,166 @@ def test_missing_override_failure(self):
             self.run_cases_test(input, input2, output)
 
 
+class TestGeneratedPECases(TestGeneratedAbstractCases):
+    generator = partial_evaluator_generator
+
+    def test_overridden_abstract(self):
+        input = """
+        pure op(OP, (--)) {
+            SPAM();
+        }
+        """
+        input2 = """
+        pure op(OP, (--)) {
+            eggs();
+        }
+        """
+        output = """
+        case OP: {
+            eggs();
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
+    def test_overridden_abstract_args(self):
+        input = """
+        pure op(OP, (arg1 -- out)) {
+            out = SPAM(arg1);
+        }
+        op(OP2, (arg1 -- out)) {
+            out = EGGS(arg1);
+        }
+        """
+        input2 = """
+        op(OP, (arg1 -- out)) {
+            out = EGGS(arg1);
+        }
+        """
+        output = """
+        case OP: {
+            _Py_UopsPESlot arg1;
+            _Py_UopsPESlot out;
+            arg1 = stack_pointer[-1];
+            arg1 = stack_pointer[-1];
+            out = EGGS(arg1);
+            stack_pointer[-1] = out;
+            break;
+        }
+
+        case OP2: {
+            _Py_UopsPESlot arg1;
+            _Py_UopsPESlot out;
+            MATERIALIZE_INST();
+            arg1 = stack_pointer[-1];
+            materialize(&arg1);
+            out = sym_new_not_null(ctx);
+            stack_pointer[-1] = out;
+            break;
+        }
+       """
+        self.run_cases_test(input, input2, output)
+
+    def test_no_overridden_case(self):
+        input = """
+        pure op(OP, (arg1 -- out)) {
+            out = SPAM(arg1);
+        }
+
+        pure op(OP2, (arg1 -- out)) {
+        }
+
+        """
+        input2 = """
+        pure op(OP2, (arg1 -- out)) {
+            out = NULL;
+        }
+        """
+        output = """
+        case OP: {
+            _Py_UopsPESlot arg1;
+            _Py_UopsPESlot out;
+            MATERIALIZE_INST();
+            arg1 = stack_pointer[-1];
+            materialize(&arg1);
+            out = sym_new_not_null(ctx);
+            stack_pointer[-1] = out;
+            break;
+        }
+
+        case OP2: {
+            _Py_UopsPESlot arg1;
+            _Py_UopsPESlot out;
+            arg1 = stack_pointer[-1];
+            out = NULL;
+            stack_pointer[-1] = out;
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
+    def test_missing_override_failure(self):
+        input = """
+        pure op(OP, (arg1 -- out)) {
+            SPAM();
+        }
+        """
+        input2 = """
+        pure op(OTHER, (arg1 -- out)) {
+        }
+        """
+        output = """
+        """
+        with self.assertRaisesRegex(AssertionError, "All abstract uops"):
+            self.run_cases_test(input, input2, output)
+
+
+    def test_validate_inputs(self):
+        input = """
+        pure op(OP, (arg1 --)) {
+            SPAM();
+        }
+        """
+        input2 = """
+        // Non-matching input!
+        pure op(OP, (arg1, arg2 --)) {
+        }
+        """
+        output = """
+        """
+        with self.assertRaisesRegex(AssertionError, "input length don't match"):
+            self.run_cases_test(input, input2, output)
+
+    def test_materialize_inputs(self):
+        input = """
+        pure op(OP2, (arg1, arg2, arg3[oparg] --)) {
+        }
+        """
+        input2 = """
+        pure op(OP2, (arg1, arg2, arg3[oparg] --)) {
+            MATERIALIZE_INPUTS();
+        }
+        """
+        output = """
+        case OP2: {
+            _Py_UopsPESlot *arg3;
+            _Py_UopsPESlot arg2;
+            _Py_UopsPESlot arg1;
+            arg3 = &stack_pointer[-2 - oparg];
+            arg2 = stack_pointer[-2];
+            arg1 = stack_pointer[-1];
+            materialize(&arg1);
+            materialize(&arg2);
+            for (int _i = oparg; --_i >= 0;) {
+                materialize(&arg3[_i]);
+            }
+            stack_pointer += -2 - oparg;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
+
 if __name__ == "__main__":
     unittest.main()