feat: bigger CACHE entries for everything, CACHE for JUMP_BACKWARD (p…

…ython#51)

* feat: store bb_test flag in frame to reload during resume

* feat: bigger CACHE entries for everything, CACHE for JUMP_BACKWARD

* nit: increase overallocate factor even more

* feat: remove bb_flag/bb_test and use only the frame's

Include/internal/pycore_code.h

@@ -12,11 +12,17 @@ typedef struct {
     // The LSB indicates whether the bb branch is a type guard or not.
     // To get the actual BB ID, do a right bit shift by one.
     uint16_t bb_id_tagged;
+    // Forward jump if required since not all successor BBs are fall-through.
     uint16_t successor_jumpby;
+    // Function pointers to trace.
+    uint16_t consequent_trace[4];
+    uint16_t alternative_trace[4];
 } _PyBBBranchCache;
 
 #define INLINE_CACHE_ENTRIES_BB_BRANCH CACHE_ENTRIES(_PyBBBranchCache)
 
+#define INLINE_CACHE_ENTRIES_JUMP_BACKWARD CACHE_ENTRIES(_PyBBBranchCache)
+
 /* PEP 659
  * Specialization and quickening structs and helper functions
  */
@@ -103,6 +109,7 @@ typedef struct {
 
 #define INLINE_CACHE_ENTRIES_FOR_ITER CACHE_ENTRIES(_PyForIterCache)
 
+
 typedef struct {
     uint16_t counter;
 } _PySendCache;
@@ -269,8 +276,8 @@ extern int _PyStaticCode_Init(PyCodeObject *co);
 //   requires a stack element to be popped.
 #define BB_TEST(gen_bb_is_successor, gen_bb_requires_pop) \
     (((gen_bb_is_successor) << 4) | (gen_bb_requires_pop))
-#define BB_TEST_IS_SUCCESSOR(bb_test) ((bb_test) >> 4)
-#define BB_TEST_GET_N_REQUIRES_POP(bb_test) ((bb_test) & 0b1111)
+#define BB_TEST_IS_SUCCESSOR(frame) ((frame->bb_test) >> 4)
+#define BB_TEST_GET_N_REQUIRES_POP(bb_flag) ((bb_flag) & 0b1111)
 
 extern _Py_CODEUNIT *_PyCode_Tier2Warmup(struct _PyInterpreterFrame *,
     _Py_CODEUNIT *);

Include/internal/pycore_frame.h

@@ -63,6 +63,8 @@ typedef struct _PyInterpreterFrame {
     uint16_t yield_offset;
     char owner;
     bool is_tier2;
+    // For restoring state after entering the JIT code.
+    char bb_test;
     /* Locals and stack and unboxed bit mask */
     PyObject *localsplus[1];
 } _PyInterpreterFrame;
@@ -138,6 +140,7 @@ _PyFrame_Initialize(
         frame->is_tier2 = false;
         frame->prev_instr = _PyCode_CODE(code) - 1;
     }
+    frame->bb_test = BB_TEST(0, 0);
     frame->yield_offset = 0;
     frame->owner = FRAME_OWNED_BY_THREAD;
 

Lib/dis.py

@@ -62,10 +62,10 @@
     if uop.startswith('BB_BRANCH') or uop.startswith('BB_JUMP'):
         if uop.startswith('BB_JUMP'):
             _bb_jumps.append(uop_opcode)
-        _inline_cache_entries[uop_opcode] = 2
+        _inline_cache_entries[uop_opcode] = sum(_cache_format['BB_BRANCH'].values())
         _uop_hasoparg.append(uop_opcode)
     if uop.startswith('BB_TEST_ITER'):
-        _inline_cache_entries[uop_opcode] = 1      
+        _inline_cache_entries[uop_opcode] = 1
 
 deoptmap = {
     specialized: base for base, family in _specializations.items() for specialized in family

Lib/importlib/_bootstrap_external.py

@@ -436,6 +436,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a6 3520 (Remove PREP_RERAISE_STAR, add CALL_INTRINSIC_2)
 #     Python 3.12a7 3521 (Shrink the LOAD_GLOBAL caches)
 #     Python 3.12a7 3523 (Convert COMPARE_AND_BRANCH back to COMPARE_OP)
+#     Python 3.12a7 3524 (pylbbv: Add support for tier 2 JIT)
 
 #     Python 3.13 will start with 3550
 
@@ -452,7 +453,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3523).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3524).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Lib/opcode.py

@@ -446,8 +446,17 @@ def pseudo_op(name, op, real_ops):
     },
     "BB_BRANCH" : {
         "bb_id": 1,
-        "forward_jumpby": 1
+        "forward_jumpby": 1,
+        "consequent_trace": 4,
+        "alternative_trace": 4,
     },
+    # Keep in sync with BB_BRANCH for simplicity's sake
+    "JUMP_BACKWARD": {
+        "bb_id": 1,
+        "forward_jumpby": 1,
+        "consequent_trace": 4,
+        "alternative_trace": 4,
+    }
 }
 
 _inline_cache_entries = [