use s_arena to save slots mapping

Author: cocolato

Include/internal/pycore_optimizer_types.h

@@ -16,6 +16,10 @@ extern "C" {
 
 #define TY_ARENA_SIZE (UOP_MAX_TRACE_LENGTH * 5)
 
+// Maximum slots per object tracked symbolically
+#define MAX_SYMBOLIC_SLOTS_SIZE 16
+#define SLOTS_ARENA_SIZE (MAX_SYMBOLIC_SLOTS_SIZE * 100)
+
 // Need extras for root frame and for overflow frame (see TRACE_STACK_PUSH())
 #define MAX_ABSTRACT_FRAME_DEPTH (16)
 
@@ -77,8 +81,6 @@ typedef struct {
     uint8_t tag;
 } JitOptCompactInt;
 
-#define MAX_SYMBOLIC_SLOTS_SIZE 4
-
 typedef struct {
     uint16_t slot_index;
     uint16_t symbol;
@@ -88,7 +90,7 @@ typedef struct _jit_opt_slots {
     uint8_t tag;
     uint8_t num_slots;
     uint32_t type_version;
-    JitOptSlotMapping slots[MAX_SYMBOLIC_SLOTS_SIZE];
+    JitOptSlotMapping *slots;
 } JitOptSlotsObject;
 
 typedef union _jit_opt_symbol {
@@ -128,6 +130,12 @@ typedef struct ty_arena {
     JitOptSymbol arena[TY_ARENA_SIZE];
 } ty_arena;
 
+typedef struct slots_arena {
+    int slots_curr_number;
+    int slots_max_number;
+    JitOptSlotMapping arena[SLOTS_ARENA_SIZE];
+} slots_arena;
+
 typedef struct _JitOptContext {
     char done;
     char out_of_space;
@@ -142,6 +150,9 @@ typedef struct _JitOptContext {
     // Arena for the symbolic types.
     ty_arena t_arena;
 
+    // Arena for slots mappings.
+    slots_arena s_arena;
+
     JitOptRef *n_consumed;
     JitOptRef *limit;
     JitOptRef locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];

Python/optimizer_symbols.c

@@ -114,7 +114,12 @@ _PyUOpSymPrint(JitOptRef ref)
             printf("<compact_int at %p>", (void *)sym);
             break;
         case JIT_SYM_SLOTS_TAG:
-            printf("<slots[%d] v%u at %p>", sym->slots.num_slots, sym->slots.type_version, (void *)sym);
+            PyTypeObject *slots_type = _PyType_LookupByVersion(sym->slots.type_version);
+            if (slots_type) {
+                printf("<%s slots[%d] v%u at %p>", slots_type->tp_name, sym->slots.num_slots, sym->slots.type_version, (void *)sym);
+            } else {
+                printf("<slots[%d] v%u at %p>", sym->slots.num_slots, sym->slots.type_version, (void *)sym);
+            }
             break;
         default:
             printf("<tag=%d at %p>", sym->tag, (void *)sym);
@@ -890,6 +895,7 @@ _Py_uop_sym_new_slots_object(JitOptContext *ctx, unsigned int type_version)
     }
     res->tag = JIT_SYM_SLOTS_TAG;
     res->slots.num_slots = 0;
+    res->slots.slots = NULL;
     res->slots.type_version = type_version;
     return PyJitRef_Wrap(res);
 }
@@ -899,19 +905,28 @@ _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) {
-        // Linear search through the mapping array
+    if (sym->tag == JIT_SYM_SLOTS_TAG && sym->slots.slots != NULL) {
         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);
             }
         }
     }
 
-    // Not found, return not_null
     return _Py_uop_sym_new_not_null(ctx);
 }
 
+static JitOptSlotMapping *
+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];
+    ctx->s_arena.slots_curr_number += MAX_SYMBOLIC_SLOTS_SIZE;
+    return slots;
+}
+
 void
 _Py_uop_sym_slots_setattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index, JitOptRef value)
 {
@@ -922,31 +937,44 @@ _Py_uop_sym_slots_setattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index
         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;
+        }
     }
     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;
         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;
     }
-
-    if (sym->slots.num_slots >= MAX_SYMBOLIC_SLOTS_SIZE) {
-        return;
+    // Check if have arena space allocated
+    if (sym->slots.slots == NULL) {
+        sym->slots.slots = slots_arena_alloc(ctx);
+        if (sym->slots.slots == NULL) {
+            return;
+        }
     }
-
+    // Check if the slot already exists
     for (int i = 0; i < sym->slots.num_slots; i++) {
         if (sym->slots.slots[i].slot_index == slot_index) {
             sym->slots.slots[i].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
             return;
         }
     }
-
-    int idx = sym->slots.num_slots++;
-    sym->slots.slots[idx].slot_index = slot_index;
-    sym->slots.slots[idx].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
+    // Add new mapping if there's space
+    if (sym->slots.num_slots < MAX_SYMBOLIC_SLOTS_SIZE) {
+        int idx = sym->slots.num_slots++;
+        sym->slots.slots[idx].slot_index = slot_index;
+        sym->slots.slots[idx].symbol = (uint16_t)(PyJitRef_Unwrap(value) - allocation_base(ctx));
+    }
 }
 
 // 0 on success, -1 on error.
@@ -1038,6 +1066,10 @@ _Py_uop_abstractcontext_init(JitOptContext *ctx)
     ctx->t_arena.ty_curr_number = 0;
     ctx->t_arena.ty_max_number = TY_ARENA_SIZE;
 
+    // Setup the arena for slot mappings.
+    ctx->s_arena.slots_curr_number = 0;
+    ctx->s_arena.slots_max_number = SLOTS_ARENA_SIZE;
+
     // Frame setup
     ctx->curr_frame_depth = 0;