Rewrite side exit jumps in JIT

Author: Fidget-Spinner

Include/internal/pycore_jit.h

@@ -20,6 +20,8 @@ typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, _PyStackRef *stack
 int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);
 
+PyAPI_DATA(int) _PyJit_PatchSideExit(_PyExecutorObject *trunk_executor, _PyExitData *exit_p, _PyExecutorObject *side_exit);
+
 #endif  // _Py_JIT
 
 #ifdef __cplusplus

Include/internal/pycore_optimizer.h

@@ -47,6 +47,12 @@ typedef struct _PyExitData {
     uint16_t index;
     _Py_BackoffCounter temperature;
     struct _PyExecutorObject *executor;
+    // Locations into the JIT code that contains
+    // the side exits pointing to this exit.
+    // This allows us to rewrite them to more efficient direct jumps
+    // to side traces.
+    uintptr_t exiting_uop_side_exit_locations[UOP_MAX_SIDE_EXITS_PER_UOP];
+    int num_side_locations_used;
 } _PyExitData;
 
 typedef struct _PyExecutorObject {

Include/internal/pycore_uop_ids.h

@@ -5350,11 +5350,8 @@ dummy_func(
 #ifndef _Py_JIT
             assert(current_executor == (_PyExecutorObject*)executor);
 #endif
-            assert(tstate->jit_exit == NULL || tstate->jit_exit->executor == current_executor);
             tstate->current_executor = (PyObject *)executor;
             if (!current_executor->vm_data.valid) {
-                assert(tstate->jit_exit->executor == current_executor);
-                assert(tstate->current_executor == executor);
                 _PyExecutor_ClearExit(tstate->jit_exit);
                 DEOPT_IF(true);
             }
@@ -5425,6 +5422,13 @@ dummy_func(
                     exit->temperature = restart_backoff_counter(temperature);
                     GOTO_TIER_ONE(optimized < 0 ? NULL : target);
                 }
+            #ifdef _Py_JIT
+                int err = _PyJit_PatchSideExit(previous_executor, exit, executor);
+                if (err < 0) {
+                    exit->temperature = restart_backoff_counter(temperature);
+                    GOTO_TIER_ONE(NULL);
+                }
+            #endif
                 exit->temperature = initial_temperature_backoff_counter();
             }
             assert(tstate->jit_exit == exit);

Python/bytecodes.c

@@ -25,6 +25,7 @@
 #include "pycore_template.h"
 #include "pycore_tuple.h"
 #include "pycore_unicodeobject.h"
+#include "pycore_uop_metadata.h"
 
 #include "pycore_jit.h"
 
@@ -124,6 +125,26 @@ mark_executable(unsigned char *memory, size_t size)
     return 0;
 }
 
+static int
+mark_read_writeable(unsigned char *memory, size_t size)
+{
+    if (size == 0) {
+        return 0;
+    }
+    assert(size % get_page_size() == 0);
+#ifdef MS_WINDOWS
+    int old;
+    int failed = !VirtualProtect(memory, size, PAGE_READWRITE, &old);
+#else
+    int failed = mprotect(memory, size, PROT_WRITE | PROT_READ);
+#endif
+    if (failed) {
+        jit_error("unable to protect executable memory");
+        return -1;
+    }
+    return 0;
+}
+
 // JIT compiler stuff: /////////////////////////////////////////////////////////
 
 #define SYMBOL_MASK_WORDS 4
@@ -138,6 +159,7 @@ typedef struct {
 
 typedef struct {
     trampoline_state trampolines;
+    _PyExecutorObject *executor;
     uintptr_t instruction_starts[UOP_MAX_TRACE_LENGTH];
 } jit_state;
 
@@ -443,6 +465,22 @@ patch_x86_64_32rx(unsigned char *location, uint64_t value)
     patch_32r(location, value);
 }
 
+static void
+register_side_exit(const _PyUOpInstruction *instruction, unsigned char *loc, const _PyUOpInstruction *exit_trace)
+{
+    if (!(_PyUop_Flags[instruction->opcode] & HAS_EXIT_FLAG)) {
+        return;
+    }
+    assert(exit_trace->opcode == _EXIT_TRACE);
+    _PyExitData *exit = (_PyExitData *)exit_trace->operand0;
+    if (exit->num_side_locations_used >= UOP_MAX_SIDE_EXITS_PER_UOP) {
+        return;
+    }
+    exit->exiting_uop_side_exit_locations[exit->num_side_locations_used] = (uintptr_t)loc;
+    exit->num_side_locations_used++;
+}
+
+
 void patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *state);
 
 #include "jit_stencils.h"
@@ -519,6 +557,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     size_t code_size = 0;
     size_t data_size = 0;
     jit_state state = {0};
+    state.executor = executor;
     for (size_t i = 0; i < length; i++) {
         const _PyUOpInstruction *instruction = &trace[i];
         group = &stencil_groups[instruction->opcode];
@@ -664,4 +703,47 @@ _PyJIT_Free(_PyExecutorObject *executor)
     }
 }
 
+#ifdef __x86_64__
+static void
+patch_side_exit(unsigned char *location, uint64_t value)
+{
+    uint64_t distance = (value - 0x4) - (uintptr_t)location;
+    // Cannot fit in a 32-bit PC-relative address.
+    if ((int64_t)distance < -(1LL << 31) || (int64_t)distance >= (1LL << 31)) {
+        return;
+    }
+    patch_32r(location, value - 0x4);
+}
+#else
+// TODO AArch64
+static void
+patch_side_exit(unsigned char *location, uint64_t value)
+{
+    (void)location;
+    (void)value;
+    return;
+}
+#endif
+
+int
+_PyJit_PatchSideExit(_PyExecutorObject *trunk_executor, _PyExitData *exit_p, _PyExecutorObject *side_exit)
+{
+    if (exit_p->num_side_locations_used > 0) {
+        if (mark_read_writeable(trunk_executor->jit_code, trunk_executor->jit_size) < 0) {
+            return -1;
+        }
+        uintptr_t new_target = (uintptr_t)side_exit->jit_code;
+        for (int i = 0; i < exit_p->num_side_locations_used; i++) {
+            uintptr_t loc = exit_p->exiting_uop_side_exit_locations[i];
+            assert(loc >= (uintptr_t)trunk_executor->jit_code);
+            assert(loc < (uintptr_t)((unsigned char *)trunk_executor->jit_code + trunk_executor->jit_size));
+            patch_side_exit((unsigned char *)loc, new_target);
+        }
+        if (mark_executable(trunk_executor->jit_code, trunk_executor->jit_size) < 0) {
+            return -1;
+        }
+    }
+    return 0;
+}
+
 #endif  // _Py_JIT

Python/executor_cases.c.h

@@ -1209,6 +1209,7 @@ make_executor_from_uops(_PyUOpInstruction *buffer, int length, const _PyBloomFil
         executor->exits[i].index = i;
         executor->exits[i].temperature = initial_temperature_backoff_counter();
         executor->exits[i].executor = cold;
+        executor->exits[i].num_side_locations_used = 0;
     }
     int next_exit = exit_count-1;
     _PyUOpInstruction *dest = (_PyUOpInstruction *)&executor->trace[length];

Python/jit.c

@@ -21,6 +21,20 @@
 DEFAULT_OUTPUT = ROOT / "Include/internal/pycore_uop_ids.h"
 
 
+def generate_max_side_exit_count(analysis: Analysis, out: CWriter) -> None:
+    max_side_exit_count = 1
+    for uop in analysis.uops.values():
+        if uop.is_viable() and uop.properties.tier != 1:
+            uop_side_exit_count = 0
+            for tok in uop.body.tokens():
+                # Note: this may be fewer than the actual max
+                # in the generated asm. However, it's a good enough estimate.
+                if tok.text == "EXIT_IF":
+                    uop_side_exit_count += 1
+            max_side_exit_count = max(max_side_exit_count, uop_side_exit_count)
+    out.emit(f"#define UOP_MAX_SIDE_EXITS_PER_UOP {max_side_exit_count}\n")
+
+
 def generate_uop_ids(
     filenames: list[str], analysis: Analysis, outfile: TextIO, distinct_namespace: bool
 ) -> None:
@@ -49,6 +63,7 @@ def generate_uop_ids(
                 next_id += 1
 
         out.emit(f"#define MAX_UOP_ID {next_id-1}\n")
+        generate_max_side_exit_count(analysis, out)
 
 
 arg_parser = argparse.ArgumentParser(

Python/optimizer.c

@@ -170,6 +170,10 @@ def fold(self, other: typing.Self, body: bytearray) -> typing.Self | None:
     def as_c(self, where: str) -> str:
         """Dump this hole as a call to a patch_* function."""
         location = f"{where} + {self.offset:#x}"
+        register_jump_target = ""
+        if self.value == HoleValue.JUMP_TARGET and self.kind == "R_X86_64_PLT32":
+            assert self.func == "patch_32r"
+            register_jump_target = f"register_side_exit(instruction, {location}, &state->executor->trace[instruction->jump_target]);\n    "
         value = _HOLE_EXPRS[self.value]
         if self.symbol:
             if value:
@@ -181,7 +185,7 @@ def as_c(self, where: str) -> str:
             value += f"{_signed(self.addend):#x}"
         if self.need_state:
             return f"{self.func}({location}, {value}, state);"
-        return f"{self.func}({location}, {value});"
+        return f"{register_jump_target}{self.func}({location}, {value});"
 
 
 @dataclasses.dataclass