gh-120619: Tier 2 partial evaluation pass foundations

Include/internal/pycore_optimizer.h

@@ -160,10 +160,13 @@ struct _Py_UopsSymbol {
     PyTypeObject *typ;  // Borrowed reference
     PyObject *const_val;  // Owned reference (!)
     unsigned int type_version; // currently stores type version
+    int locals_idx;
+    char is_static;  // used for binding-time analysis
 };
 
 #define UOP_FORMAT_TARGET 0
 #define UOP_FORMAT_JUMP 1
+void _PyUOpPrint(const _PyUOpInstruction*);
 
 static inline uint32_t uop_get_target(const _PyUOpInstruction *inst)
 {
@@ -199,14 +202,21 @@ static inline uint16_t uop_get_error_target(const _PyUOpInstruction *inst)
 
 typedef struct _Py_UopsSymbol _Py_UopsSymbol;
 
+typedef struct _Py_UopsLocalsPlusSlot {
+    _Py_UopsSymbol *sym;
+    char is_virtual;
+} _Py_UopsLocalsPlusSlot;
+
 struct _Py_UOpsAbstractFrame {
     // Max stacklen
     int stack_len;
     int locals_len;
 
-    _Py_UopsSymbol **stack_pointer;
-    _Py_UopsSymbol **stack;
-    _Py_UopsSymbol **locals;
+    _Py_UopsLocalsPlusSlot *stack_pointer;
+    _Py_UopsLocalsPlusSlot *stack;
+    _Py_UopsLocalsPlusSlot *locals;
+
+    void *instr_ptr;
 };
 
 typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
@@ -229,34 +239,39 @@ struct _Py_UOpsContext {
     // Arena for the symbolic types.
     ty_arena t_arena;
 
-    _Py_UopsSymbol **n_consumed;
-    _Py_UopsSymbol **limit;
-    _Py_UopsSymbol *locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
+    _Py_UopsLocalsPlusSlot *n_consumed;
+    _Py_UopsLocalsPlusSlot *limit;
+    _Py_UopsLocalsPlusSlot locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
+
+    _PyUOpInstruction *trace_dest;
+    int n_trace_dest;
 };
 
 typedef struct _Py_UOpsContext _Py_UOpsContext;
 
-extern bool _Py_uop_sym_is_null(_Py_UopsSymbol *sym);
-extern bool _Py_uop_sym_is_not_null(_Py_UopsSymbol *sym);
-extern bool _Py_uop_sym_is_const(_Py_UopsSymbol *sym);
-extern PyObject *_Py_uop_sym_get_const(_Py_UopsSymbol *sym);
-extern _Py_UopsSymbol *_Py_uop_sym_new_unknown(_Py_UOpsContext *ctx);
-extern _Py_UopsSymbol *_Py_uop_sym_new_not_null(_Py_UOpsContext *ctx);
-extern _Py_UopsSymbol *_Py_uop_sym_new_type(
+extern bool _Py_uop_sym_is_null(_Py_UopsLocalsPlusSlot sym);
+extern bool _Py_uop_sym_is_not_null(_Py_UopsLocalsPlusSlot sym);
+extern bool _Py_uop_sym_is_const(_Py_UopsLocalsPlusSlot sym);
+extern PyObject *_Py_uop_sym_get_const(_Py_UopsLocalsPlusSlot sym);
+extern _Py_UopsLocalsPlusSlot _Py_uop_sym_new_unknown(_Py_UOpsContext *ctx);
+extern _Py_UopsLocalsPlusSlot _Py_uop_sym_new_not_null(_Py_UOpsContext *ctx);
+extern _Py_UopsLocalsPlusSlot _Py_uop_sym_new_type(
     _Py_UOpsContext *ctx, PyTypeObject *typ);
-extern _Py_UopsSymbol *_Py_uop_sym_new_const(_Py_UOpsContext *ctx, PyObject *const_val);
-extern _Py_UopsSymbol *_Py_uop_sym_new_null(_Py_UOpsContext *ctx);
-extern bool _Py_uop_sym_has_type(_Py_UopsSymbol *sym);
-extern bool _Py_uop_sym_matches_type(_Py_UopsSymbol *sym, PyTypeObject *typ);
-extern bool _Py_uop_sym_matches_type_version(_Py_UopsSymbol *sym, unsigned int version);
-extern void _Py_uop_sym_set_null(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym);
-extern void _Py_uop_sym_set_non_null(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym);
-extern void _Py_uop_sym_set_type(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, PyTypeObject *typ);
-extern bool _Py_uop_sym_set_type_version(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, unsigned int version);
-extern void _Py_uop_sym_set_const(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, PyObject *const_val);
-extern bool _Py_uop_sym_is_bottom(_Py_UopsSymbol *sym);
-extern int _Py_uop_sym_truthiness(_Py_UopsSymbol *sym);
-extern PyTypeObject *_Py_uop_sym_get_type(_Py_UopsSymbol *sym);
+extern _Py_UopsLocalsPlusSlot _Py_uop_sym_new_const(_Py_UOpsContext *ctx, PyObject *const_val);
+extern _Py_UopsLocalsPlusSlot _Py_uop_sym_new_null(_Py_UOpsContext *ctx);
+extern bool _Py_uop_sym_has_type(_Py_UopsLocalsPlusSlot sym);
+extern bool _Py_uop_sym_matches_type(_Py_UopsLocalsPlusSlot sym, PyTypeObject *typ);
+extern bool _Py_uop_sym_matches_type_version(_Py_UopsLocalsPlusSlot sym, unsigned int version);
+extern void _Py_uop_sym_set_locals_idx(_Py_UopsLocalsPlusSlot sym, int locals_idx);
+extern int _Py_uop_sym_get_locals_idx(_Py_UopsLocalsPlusSlot sym);
+extern void _Py_uop_sym_set_null(_Py_UOpsContext *ctx, _Py_UopsLocalsPlusSlot sym);
+extern void _Py_uop_sym_set_non_null(_Py_UOpsContext *ctx, _Py_UopsLocalsPlusSlot sym);
+extern void _Py_uop_sym_set_type(_Py_UOpsContext *ctx, _Py_UopsLocalsPlusSlot sym, PyTypeObject *typ);
+extern bool _Py_uop_sym_set_type_version(_Py_UOpsContext *ctx, _Py_UopsLocalsPlusSlot sym, unsigned int version);
+extern void _Py_uop_sym_set_const(_Py_UOpsContext *ctx, _Py_UopsLocalsPlusSlot sym, PyObject *const_val);
+extern bool _Py_uop_sym_is_bottom(_Py_UopsLocalsPlusSlot sym);
+extern int _Py_uop_sym_truthiness(_Py_UopsLocalsPlusSlot sym);
+extern PyTypeObject *_Py_uop_sym_get_type(_Py_UopsLocalsPlusSlot sym);
 
 
 extern void _Py_uop_abstractcontext_init(_Py_UOpsContext *ctx);
@@ -266,7 +281,7 @@ extern _Py_UOpsAbstractFrame *_Py_uop_frame_new(
     _Py_UOpsContext *ctx,
     PyCodeObject *co,
     int curr_stackentries,
-    _Py_UopsSymbol **args,
+    _Py_UopsLocalsPlusSlot *args,
     int arg_len);
 extern int _Py_uop_frame_pop(_Py_UOpsContext *ctx);
 

Lib/test/test_capi/test_opt.py

@@ -1481,6 +1481,33 @@ def fn(a):
 
         fn(A())
 
+    def test_pe_load_fast_pop_top(self):
+        def thing(a):
+            x = 0
+            for i in range(20):
+                i
+            return i
+
+
+        res, ex = self._run_with_optimizer(thing, 1)
+        self.assertEqual(res, 19)
+        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(20):
+                x = x
+            return i
+
+
+        res, ex = self._run_with_optimizer(thing, 1)
+        self.assertEqual(res, 19)
+        self.assertIsNotNone(ex)
+        self.assertEqual(list(iter_opnames(ex)).count("_LOAD_FAST_1"), 0)
+        self.assertTrue(ex.is_valid())
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_generated_cases.py

@@ -1264,15 +1264,15 @@ def test_overridden_abstract_args(self):
         """
         output = """
         case OP: {
-            _Py_UopsSymbol *arg1;
-            _Py_UopsSymbol *out;
+            _Py_UopsLocalsPlusSlot arg1;
+            _Py_UopsLocalsPlusSlot out;
             eggs();
             stack_pointer[-1] = out;
             break;
         }
 
         case OP2: {
-            _Py_UopsSymbol *out;
+            _Py_UopsLocalsPlusSlot out;
             out = sym_new_not_null(ctx);
             stack_pointer[-1] = out;
             break;
@@ -1296,15 +1296,15 @@ def test_no_overridden_case(self):
         """
         output = """
         case OP: {
-            _Py_UopsSymbol *out;
+            _Py_UopsLocalsPlusSlot out;
             out = sym_new_not_null(ctx);
             stack_pointer[-1] = out;
             break;
         }
 
         case OP2: {
-            _Py_UopsSymbol *arg1;
-            _Py_UopsSymbol *out;
+            _Py_UopsLocalsPlusSlot arg1;
+            _Py_UopsLocalsPlusSlot out;
             stack_pointer[-1] = out;
             break;
         }

Makefile.pre.in

@@ -1958,7 +1958,7 @@ Objects/mimalloc/page.o: $(srcdir)/Objects/mimalloc/page-queue.c
 regen-cases: \
         regen-opcode-ids regen-opcode-targets regen-uop-ids regen-opcode-metadata-py \
 		regen-generated-cases regen-executor-cases regen-optimizer-cases \
-		regen-opcode-metadata regen-uop-metadata
+		regen-partial-evaluator-cases regen-opcode-metadata regen-uop-metadata
 
 .PHONY: regen-opcode-ids
 regen-opcode-ids:
@@ -2004,6 +2004,16 @@ regen-optimizer-cases:
 	    $(srcdir)/Python/bytecodes.c
 	$(UPDATE_FILE) $(srcdir)/Python/optimizer_cases.c.h $(srcdir)/Python/optimizer_cases.c.h.new
 
+.PHONY: regen-partial-evaluator-cases
+regen-partial-evaluator-cases:
+	$(PYTHON_FOR_REGEN) $(srcdir)/Tools/cases_generator/optimizer_generator.py \
+	    -o $(srcdir)/Python/partial_evaluator_cases.c.h.new \
+        $(srcdir)/Python/optimizer_bytecodes.c \
+	    $(srcdir)/Python/partial_evaluator_bytecodes.c \
+	    $(srcdir)/Python/bytecodes.c
+	$(UPDATE_FILE) $(srcdir)/Python/partial_evaluator_cases.c.h $(srcdir)/Python/partial_evaluator_cases.c.h.new
+
+
 .PHONY: regen-opcode-metadata
 regen-opcode-metadata:
 	$(PYTHON_FOR_REGEN) $(srcdir)/Tools/cases_generator/opcode_metadata_generator.py \
@@ -2041,7 +2051,8 @@ Python/optimizer.o: \
 Python/optimizer_analysis.o: \
 		$(srcdir)/Include/internal/pycore_opcode_metadata.h \
 		$(srcdir)/Include/internal/pycore_optimizer.h \
-		$(srcdir)/Python/optimizer_cases.c.h
+		$(srcdir)/Python/optimizer_cases.c.h \
+		$(srcdir)/Python/partial_evaluator_cases.c.h
 
 Python/frozen.o: $(FROZEN_FILES_OUT)
 

Misc/NEWS.d/next/Core and Builtins/2024-09-04-03-36-48.gh-issue-120619.yE7lQb.rst

@@ -0,0 +1 @@
+Set up a tier 2 partial evaluation pass. Patch by Ken Jin.