Feat: Implemented type propagation

.gitattributes

@@ -83,6 +83,7 @@ Parser/token.c                                      generated
 Programs/test_frozenmain.h                          generated
 Python/Python-ast.c                                 generated
 Python/generated_cases.c.h                          generated
+Python/tier2_typepropagator.c.h                     generated
 Python/opcode_targets.h                             generated
 Python/stdlib_module_names.h                        generated
 Tools/peg_generator/pegen/grammar_parser.py         generated

Include/cpython/code.h

@@ -56,18 +56,24 @@ typedef struct {
     PyObject *_co_freevars;
 } _PyCoCached;
 
+// Tier 2 types meta interpreter
+typedef struct _PyTier2TypeContext {
+    // points into type_stack, points to one element after the stack
+    PyTypeObject** type_stack_ptr;
+    int type_locals_len;
+    int type_stack_len;
+    PyTypeObject** type_stack;
+    PyTypeObject** type_locals;
+} _PyTier2TypeContext;
+
 // Tier 2 interpreter information
 typedef struct _PyTier2BBMetadata {
     // Index into _PyTier2Info->bb_data
     int id;
-    // Array of types. This corresponds to the fast locals array.
-    int type_context_len;
-    PyTypeObject **type_context;
+    _PyTier2TypeContext *type_context;
     _Py_CODEUNIT *tier2_start;
     // Note, this is the first tier 1 instruction to execute AFTER the BB ends.
     _Py_CODEUNIT *tier1_end;
-    // Type stack state
-    PyTypeObject **types_stack;
 } _PyTier2BBMetadata;
 
 // Bump allocator for basic blocks (overallocated)
@@ -97,13 +103,19 @@ typedef struct _PyTier2Info {
     // will have [[BB_ID1, BB_ID2], [BB_ID3,], [], []]
     // etc.
     int **backward_jump_target_bb_ids;
-    PyTypeObject **types_stack;
     // Max len of bb_data
     int bb_data_len;
     // Current index to write into in bb_data. Incremented after each write.
     // This also assigns the BB ID.
     int bb_data_curr;
     _PyTier2BBMetadata **bb_data;
+
+    // @TODO:
+    //   Potentially optimise _PyTier2TypeContext by allocating the stacksize
+    //   to the size needed for the snapshot, and the type propagation is performed
+    //   on type_metainterpreter_stack_scratch which is allocated only once per
+    //   code object.
+    // PyTypeObject** type_metainterpreter_stack_scratch;
 } _PyTier2Info;
 
 // To avoid repeating ourselves in deepfreeze.py, all PyCodeObject members are

Include/internal/pycore_code.h

@@ -259,7 +259,7 @@ extern _Py_CODEUNIT *_PyCode_Tier2Warmup(struct _PyInterpreterFrame *,
     _Py_CODEUNIT *);
 extern _Py_CODEUNIT *_PyTier2_GenerateNextBB(
     struct _PyInterpreterFrame *frame, uint16_t bb_id, int jumpby,
-    _Py_CODEUNIT **tier1_fallback);
+    _Py_CODEUNIT **tier1_fallback, char gen_bb_requires_pop);
 extern _Py_CODEUNIT *_PyTier2_LocateJumpBackwardsBB(
     struct _PyInterpreterFrame *frame, uint16_t bb_id, int jumpby,
     _Py_CODEUNIT **tier1_fallback);

Objects/codeobject.c

@@ -1672,6 +1672,11 @@ code_tier2_fini(PyCodeObject *co)
     if (co->_tier2_info == NULL) {
         return;
     }
+    // @TODO:
+    //  Write a proper destructor for _PyTier2Info
+    //  and it's children structures.
+    //  Current implementation e.g., doesn't clear
+    //  bb_data
     _PyTier2Info *t2_info = co->_tier2_info;
     t2_info->_entry_bb = NULL;
     if (t2_info->_bb_space != NULL) {
@@ -1690,10 +1695,7 @@ code_tier2_fini(PyCodeObject *co)
         PyMem_Free(t2_info->backward_jump_offsets);
         t2_info->backward_jump_offsets = NULL;
     }
-    if (t2_info->types_stack != NULL) {
-        PyMem_Free(t2_info->types_stack);
-        t2_info->types_stack = NULL;
-    }
+
     t2_info->backward_jump_count = 0;
     if (t2_info->bb_data != NULL && t2_info->bb_data_len > 0) {
         PyMem_Free(t2_info->bb_data);

Python/bytecodes.c

@@ -102,31 +102,31 @@ dummy_func(
             }
         }
 
-        inst(LOAD_CLOSURE, (-- value)) {
+        inst(LOAD_CLOSURE, (-- value : locals[oparg])) {
             /* We keep LOAD_CLOSURE so that the bytecode stays more readable. */
             value = GETLOCAL(oparg);
             ERROR_IF(value == NULL, unbound_local_error);
             Py_INCREF(value);
         }
 
-        inst(LOAD_FAST_CHECK, (-- value)) {
+        inst(LOAD_FAST_CHECK, (-- value : locals[oparg])) {
             value = GETLOCAL(oparg);
             ERROR_IF(value == NULL, unbound_local_error);
             Py_INCREF(value);
         }
 
-        inst(LOAD_FAST, (-- value)) {
+        inst(LOAD_FAST, (-- value : locals[oparg])) {
             value = GETLOCAL(oparg);
             assert(value != NULL);
             Py_INCREF(value);
         }
 
-        inst(LOAD_CONST, (-- value)) {
+        inst(LOAD_CONST, (-- value : consts[oparg])) {
             value = GETITEM(consts, oparg);
             Py_INCREF(value);
         }
 
-        inst(STORE_FAST, (value --)) {
+        inst(STORE_FAST, (value --), locals[oparg] = *value) {
             SETLOCAL(oparg, value);
         }
 
@@ -303,7 +303,7 @@ dummy_func(
             bb_test = PyLong_CheckExact(left) && (Py_TYPE(left) == Py_TYPE(right));
         }
 
-        u_inst(BINARY_OP_ADD_INT_REST, (left, right -- sum)) {
+        u_inst(BINARY_OP_ADD_INT_REST, (left : PyLong_Type, right : PyLong_Type -- sum : PyLong_Type)) {
             STAT_INC(BINARY_OP, hit);
             sum = _PyLong_Add((PyLongObject *)left, (PyLongObject *)right);
             _Py_DECREF_SPECIALIZED(right, (destructor)PyObject_Free);
@@ -1180,13 +1180,13 @@ dummy_func(
             null = NULL;
         }
 
-        inst(DELETE_FAST, (--)) {
+        inst(DELETE_FAST, (--), locals[oparg] = NULL) {
             PyObject *v = GETLOCAL(oparg);
             ERROR_IF(v == NULL, unbound_local_error);
             SETLOCAL(oparg, NULL);
         }
 
-        inst(MAKE_CELL, (--)) {
+        inst(MAKE_CELL, (--), locals[oparg] = NULL) {
             // "initial" is probably NULL but not if it's an arg (or set
             // via PyFrame_LocalsToFast() before MAKE_CELL has run).
             PyObject *initial = GETLOCAL(oparg);
@@ -2084,6 +2084,10 @@ dummy_func(
                     goto error;
                 }
             }
+            // This gets set so BRANCH_BB knows whether to pop
+            // the type stack (type propagation) when generating the
+            // target BB
+            gen_bb_requires_pop = !jump;
         }
 
         inst(JUMP_IF_TRUE_OR_POP, (cond -- cond if (jump))) {
@@ -2136,6 +2140,10 @@ dummy_func(
                     goto error;
                 }
             }
+            // This gets set so BRANCH_BB knows whether to pop
+            // the type stack (type propagation) when generating the
+            // target BB
+            gen_bb_requires_pop = !jump;
         }
 
         inst(JUMP_BACKWARD_NO_INTERRUPT, (--)) {
@@ -3284,6 +3292,7 @@ dummy_func(
         }
 
         // Tier 2 instructions
+        // Type propagator assumes this doesn't affect type context
         inst(BB_BRANCH, (unused/1 --)) {
             _Py_CODEUNIT *t2_nextinstr = NULL;
             _PyBBBranchCache *cache = (_PyBBBranchCache *)next_instr;
@@ -3293,7 +3302,8 @@ dummy_func(
                 _py_set_opcode(next_instr - 1, BB_BRANCH_IF_FLAG_UNSET);
                 // Generate consequent.
                 t2_nextinstr = _PyTier2_GenerateNextBB(
-                    frame, cache->bb_id, 0, &tier1_fallback);
+                    frame, cache->bb_id, 0, &tier1_fallback, gen_bb_requires_pop);
+                gen_bb_requires_pop = false;
                 if (t2_nextinstr == NULL) {
                     // Fall back to tier 1.
                     next_instr = tier1_fallback;
@@ -3305,7 +3315,8 @@ dummy_func(
                 _py_set_opcode(next_instr - 1, BB_BRANCH_IF_FLAG_SET);
                 // Generate predicate.
                 t2_nextinstr = _PyTier2_GenerateNextBB(
-                    frame, cache->bb_id, oparg, &tier1_fallback);
+                    frame, cache->bb_id, oparg, &tier1_fallback, gen_bb_requires_pop);
+                gen_bb_requires_pop = false;
                 if (t2_nextinstr == NULL) {
                     // Fall back to tier 1.
                     next_instr = tier1_fallback + oparg;
@@ -3327,7 +3338,8 @@ dummy_func(
                 _Py_CODEUNIT *tier1_fallback = NULL;
 
                 t2_nextinstr = _PyTier2_GenerateNextBB(
-                    frame, cache->bb_id, oparg, &tier1_fallback);
+                    frame, cache->bb_id, oparg, &tier1_fallback, gen_bb_requires_pop);
+                gen_bb_requires_pop = false;
                 if (t2_nextinstr == NULL) {
                     // Fall back to tier 1.
                     next_instr = tier1_fallback;
@@ -3358,7 +3370,8 @@ dummy_func(
                 // @TODO: Rewrite TEST intruction above to a JUMP above..
 
                 t2_nextinstr = _PyTier2_GenerateNextBB(
-                    frame, cache->bb_id, oparg, &tier1_fallback);
+                    frame, cache->bb_id, oparg, &tier1_fallback, gen_bb_requires_pop);
+                gen_bb_requires_pop = false;
                 if (t2_nextinstr == NULL) {
                     // Fall back to tier 1.
                     next_instr = tier1_fallback;
@@ -3379,6 +3392,7 @@ dummy_func(
             // Fall through to next instruction.
         }
 
+        // Type propagator assumes this doesn't affect type context
         inst(BB_JUMP_BACKWARD_LAZY, (--)) {
             _Py_CODEUNIT *curr = next_instr - 1;
             _Py_CODEUNIT *t2_nextinstr = NULL;

Python/ceval.c

@@ -729,6 +729,11 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     // true = successor
     // false = alternate
     bool bb_test = true;
+    // For tier2 type propagation, handling of jump instructions with
+    // runtime-dependent stack effect.
+    // This flag is used to determine if the type context of a new bb
+    // requires a stack element to be popped.
+    bool gen_bb_requires_pop = false;
 
     /* WARNING: Because the _PyCFrame lives on the C stack,
      * but can be accessed from a heap allocated object (tstate)