gh-144145: Track nullness of properties in the Tier 2 JIT optimizer

Include/internal/pycore_optimizer.h

@@ -37,10 +37,15 @@ typedef struct _JitOptContext {
     // Arena for the symbolic types.
     ty_arena t_arena;
 
+    // Arena for the slots mappings.
+    slots_arena s_arena;
+
     JitOptRef *n_consumed;
     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;
 
 
@@ -295,6 +300,9 @@ extern JitOptRef _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef value,
 extern bool _Py_uop_sym_is_compact_int(JitOptRef sym);
 extern JitOptRef _Py_uop_sym_new_compact_int(JitOptContext *ctx);
 extern void _Py_uop_sym_set_compact_int(JitOptContext *ctx,  JitOptRef sym);
+extern JitOptRef _Py_uop_sym_new_slots_object(JitOptContext *ctx, unsigned int type_version);
+extern JitOptRef _Py_uop_sym_slots_getattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index);
+extern void _Py_uop_sym_slots_setattr(JitOptContext *ctx, JitOptRef ref, uint16_t slot_index, JitOptRef value);
 extern JitOptRef _Py_uop_sym_new_predicate(JitOptContext *ctx, JitOptRef lhs_ref, JitOptRef rhs_ref, JitOptPredicateKind kind);
 extern void _Py_uop_sym_apply_predicate_narrowing(JitOptContext *ctx, JitOptRef sym, bool branch_is_true);
 

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)
 
@@ -41,6 +45,7 @@ typedef enum _JitSymType {
     JIT_SYM_TRUTHINESS_TAG = 9,
     JIT_SYM_COMPACT_INT = 10,
     JIT_SYM_PREDICATE_TAG = 11,
+    JIT_SYM_SLOTS_TAG = 12,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -89,6 +94,19 @@ typedef struct {
     uint8_t tag;
 } JitOptCompactInt;
 
+typedef struct {
+    uint16_t slot_index;
+    uint16_t symbol;
+} 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;
+    JitOptDescrMapping *slots;
+} JitOptSlotsObject;
+
 typedef union _jit_opt_symbol {
     uint8_t tag;
     JitOptKnownClass cls;
@@ -97,6 +115,7 @@ typedef union _jit_opt_symbol {
     JitOptTuple tuple;
     JitOptTruthiness truthiness;
     JitOptCompactInt compact;
+    JitOptSlotsObject slots;
     JitOptPredicate predicate;
 } JitOptSymbol;
 
@@ -126,6 +145,11 @@ typedef struct ty_arena {
     JitOptSymbol arena[TY_ARENA_SIZE];
 } ty_arena;
 
+typedef struct slots_arena {
+    int slots_curr_number;
+    int slots_max_number;
+    JitOptDescrMapping arena[SLOTS_ARENA_SIZE];
+} slots_arena;
 
 #ifdef __cplusplus
 }

Python/optimizer_analysis.c

@@ -247,6 +247,9 @@ add_op(JitOptContext *ctx, _PyUOpInstruction *this_instr,
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_slots_object _Py_uop_sym_new_slots_object
+#define sym_slots_getattr _Py_uop_sym_slots_getattr
+#define sym_slots_setattr _Py_uop_sym_slots_setattr
 #define sym_new_predicate _Py_uop_sym_new_predicate
 #define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
@@ -506,6 +509,10 @@ optimize_uops(
         if (ctx->out_buffer.next == out_ptr) {
             *(ctx->out_buffer.next++) = *this_instr;
         }
+        // Track escapes
+        if (_PyUop_Flags[out_ptr->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

@@ -38,6 +38,9 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_slots_object _Py_uop_sym_new_slots_object
+#define sym_slots_getattr _Py_uop_sym_slots_getattr
+#define sym_slots_setattr _Py_uop_sym_slots_setattr
 #define sym_new_predicate _Py_uop_sym_new_predicate
 #define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
@@ -125,8 +128,7 @@ dummy_func(void) {
     }
 
     op(_STORE_ATTR_SLOT, (index/1, value, owner -- o)) {
-        (void)index;
-        (void)value;
+        sym_slots_setattr(ctx, owner, (uint16_t)index, value);
         o = owner;
     }
 
@@ -710,8 +712,7 @@ dummy_func(void) {
     }
 
     op(_LOAD_ATTR_SLOT, (index/1, owner -- attr, o)) {
-        attr = sym_new_not_null(ctx);
-        (void)index;
+        attr = sym_slots_getattr(ctx, owner, (uint16_t)index);
         o = owner;
     }
 
@@ -895,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 != NULL && 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)) {