support track escape

Author: cocolato

Include/internal/pycore_optimizer.h

@@ -44,6 +44,8 @@ typedef struct _JitOptContext {
     JitOptRef *limit;
     JitOptRef locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
     _PyJitUopBuffer out_buffer;
+    // Index of the last escaped uop in out_buffer.
+    int last_escape_index;
 } JitOptContext;
 
 

Include/internal/pycore_optimizer_types.h

@@ -97,13 +97,14 @@ typedef struct {
 typedef struct {
     uint16_t slot_index;
     uint16_t symbol;
-} JitOptSlotMapping;
+} JitOptDescrMapping;
 
 typedef struct _jit_opt_slots {
     uint8_t tag;
     uint8_t num_slots;
+    uint16_t last_modified_index;  // Index in out_buffer when this object was last modified
     uint32_t type_version;
-    JitOptSlotMapping *slots;
+    JitOptDescrMapping *slots;
 } JitOptSlotsObject;
 
 typedef union _jit_opt_symbol {
@@ -147,7 +148,7 @@ typedef struct ty_arena {
 typedef struct slots_arena {
     int slots_curr_number;
     int slots_max_number;
-    JitOptSlotMapping arena[SLOTS_ARENA_SIZE];
+    JitOptDescrMapping arena[SLOTS_ARENA_SIZE];
 } slots_arena;
 
 #ifdef __cplusplus

Python/optimizer_analysis.c

@@ -509,6 +509,10 @@ optimize_uops(
         if (ctx->out_buffer.next == out_ptr) {
             *(ctx->out_buffer.next++) = *this_instr;
         }
+        // Track escapes
+        if (_PyUop_Flags[opcode] & HAS_ESCAPES_FLAG) {
+            ctx->last_escape_index = uop_buffer_length(&ctx->out_buffer);
+        }
         assert(ctx->frame != NULL);
         if (!CURRENT_FRAME_IS_INIT_SHIM()) {
             DPRINTF(3, " stack_level %d\n", STACK_LEVEL());

Python/optimizer_bytecodes.c

@@ -896,10 +896,14 @@ dummy_func(void) {
     }
 
     op(_CHECK_AND_ALLOCATE_OBJECT, (type_version/2, callable, self_or_null, args[oparg] -- callable, self_or_null, args[oparg])) {
-        (void)type_version;
         (void)args;
         callable = sym_new_not_null(ctx);
-        self_or_null = sym_new_not_null(ctx);
+        PyTypeObject *tp = _PyType_LookupByVersion(type_version);
+       if (tp->tp_basicsize > sizeof(PyObject) && !(tp->tp_flags & Py_TPFLAGS_MANAGED_DICT)) {
+            self_or_null = sym_new_slots_object(ctx, type_version);
+        } else {
+            self_or_null = sym_new_not_null(ctx);
+        }
     }
 
     op(_CREATE_INIT_FRAME, (init, self, args[oparg] -- init_frame)) {

Python/optimizer_cases.c.h

@@ -974,6 +974,7 @@ _Py_uop_sym_new_slots_object(JitOptContext *ctx, unsigned int type_version)
     res->slots.num_slots = 0;
     res->slots.slots = NULL;
     res->slots.type_version = type_version;
+    res->slots.last_modified_index = uop_buffer_length(&ctx->out_buffer);
     return PyJitRef_Wrap(res);
 }
 
@@ -982,7 +983,13 @@ _Py_uop_sym_slots_getattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index
 {
     JitOptSymbol *sym = PyJitRef_Unwrap(ref);
 
-    if (sym->tag == JIT_SYM_SLOTS_TAG && sym->slots.slots != NULL) {
+    // Only return tracked slot values if:
+    // 1. Symbol is a slots object with mappings allocated
+    // 2. No escape has occurred since last modification (state is fresh)
+    if (sym->tag == JIT_SYM_SLOTS_TAG &&
+        sym->slots.slots != NULL &&
+        sym->slots.last_modified_index >= ctx->last_escape_index)
+    {
         for (int i = 0; i < sym->slots.num_slots; i++) {
             if (sym->slots.slots[i].slot_index == slot_index) {
                 return PyJitRef_Wrap(allocation_base(ctx) + sym->slots.slots[i].symbol);
@@ -993,13 +1000,13 @@ _Py_uop_sym_slots_getattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index
     return _Py_uop_sym_new_not_null(ctx);
 }
 
-static JitOptSlotMapping *
+static JitOptDescrMapping *
 slots_arena_alloc(JitOptContext *ctx)
 {
     if (ctx->s_arena.slots_curr_number + MAX_SYMBOLIC_SLOTS_SIZE > ctx->s_arena.slots_max_number) {
         return NULL;
     }
-    JitOptSlotMapping *slots = &ctx->s_arena.arena[ctx->s_arena.slots_curr_number];
+    JitOptDescrMapping *slots = &ctx->s_arena.arena[ctx->s_arena.slots_curr_number];
     ctx->s_arena.slots_curr_number += MAX_SYMBOLIC_SLOTS_SIZE;
     return slots;
 }
@@ -1008,30 +1015,18 @@ void
 _Py_uop_sym_slots_setattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index, JitOptRef value)
 {
     JitOptSymbol *sym = PyJitRef_Unwrap(ref);
+    int curr_index = uop_buffer_length(&ctx->out_buffer);
 
-    if (sym->tag == JIT_SYM_TYPE_VERSION_TAG) {
-        uint32_t version = sym->version.version;
-        sym->tag = JIT_SYM_SLOTS_TAG;
-        sym->slots.type_version = version;
-        sym->slots.num_slots = 0;
-        sym->slots.slots = slots_arena_alloc(ctx);
-        if (sym->slots.slots == NULL) {
-            return;
-        }
+    if (sym->tag != JIT_SYM_SLOTS_TAG) {
+        return;
     }
-    else if (sym->tag == JIT_SYM_KNOWN_CLASS_TAG) {
-        uint32_t version = sym->cls.version;
-        sym->tag = JIT_SYM_SLOTS_TAG;
-        sym->slots.type_version = version;
+
+    // Check escape
+    if (sym->slots.last_modified_index < ctx->last_escape_index) {
         sym->slots.num_slots = 0;
-        sym->slots.slots = slots_arena_alloc(ctx);
-        if (sym->slots.slots == NULL) {
-            return;
-        }
-    }
-    else if (sym->tag != JIT_SYM_SLOTS_TAG) {
-        return;
     }
+    // Update the last modified timestamp
+    sym->slots.last_modified_index = curr_index;
     // Check if have arena space allocated
     if (sym->slots.slots == NULL) {
         sym->slots.slots = slots_arena_alloc(ctx);
@@ -1157,6 +1152,7 @@ _Py_uop_abstractcontext_init(JitOptContext *ctx)
     ctx->out_of_space = false;
     ctx->contradiction = false;
     ctx->builtins_watched = false;
+    ctx->last_escape_index = 0;
 }
 
 int