Split JIT stencils into hot and cold code

Author: waterlens

Python/jit.c

@@ -652,20 +652,23 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
 {
     const StencilGroup *group;
     // Loop once to find the total compiled size:
-    size_t code_size = 0;
+    size_t hot_code_size = 0;
+    size_t cold_code_size = 0;
     size_t data_size = 0;
     jit_state state = {0};
     for (size_t i = 0; i < length; i++) {
         const _PyUOpInstruction *instruction = &trace[i];
         group = &stencil_groups[instruction->opcode];
-        state.instruction_starts[i] = code_size;
-        code_size += group->code_size;
+        state.instruction_starts[i] = hot_code_size;
+        hot_code_size += group->hot_code_size;
+        cold_code_size += group->cold_code_size;
         data_size += group->data_size;
         combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
         combine_symbol_mask(group->got_mask, state.got_symbols.mask);
     }
     group = &stencil_groups[_FATAL_ERROR_r00];
-    code_size += group->code_size;
+    hot_code_size += group->hot_code_size;
+    cold_code_size += group->cold_code_size;
     data_size += group->data_size;
     combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
     combine_symbol_mask(group->got_mask, state.got_symbols.mask);
@@ -679,6 +682,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     // Round up to the nearest page:
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
+    size_t code_size = hot_code_size + cold_code_size;
     size_t code_padding = DATA_ALIGN - ((code_size + state.trampolines.size) & (DATA_ALIGN - 1));
     size_t padding = page_size - ((code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size) & (page_size - 1));
     size_t total_size = code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size + padding;
@@ -700,23 +704,27 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     }
     // Loop again to emit the code:
     unsigned char *code = memory;
+    unsigned char *cold = memory + hot_code_size;
     state.trampolines.mem = memory + code_size;
     unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
     assert(trace[0].opcode == _START_EXECUTOR_r00 || trace[0].opcode == _COLD_EXIT_r00 || trace[0].opcode == _COLD_DYNAMIC_EXIT_r00);
     state.got_symbols.mem = data + data_size;
     for (size_t i = 0; i < length; i++) {
         const _PyUOpInstruction *instruction = &trace[i];
         group = &stencil_groups[instruction->opcode];
-        group->emit(code, data, executor, instruction, &state);
-        code += group->code_size;
+        group->emit(code, cold, data, executor, instruction, &state);
+        code += group->hot_code_size;
+        cold += group->cold_code_size;
         data += group->data_size;
     }
     // Protect against accidental buffer overrun into data:
     group = &stencil_groups[_FATAL_ERROR_r00];
-    group->emit(code, data, executor, NULL, &state);
-    code += group->code_size;
+    group->emit(code, cold, data, executor, NULL, &state);
+    code += group->hot_code_size;
+    cold += group->cold_code_size;
     data += group->data_size;
-    assert(code == memory + code_size);
+    assert(code == memory + hot_code_size);
+    assert(cold == memory + code_size);
     assert(data == memory + code_size + state.trampolines.size + code_padding + data_size);
     if (mark_executable(memory, total_size)) {
         jit_free(memory, total_size);
@@ -738,17 +746,20 @@ compile_shim(void)
 {
     _PyExecutorObject dummy;
     const StencilGroup *group;
-    size_t code_size = 0;
+    size_t hot_code_size = 0;
+    size_t cold_code_size = 0;
     size_t data_size = 0;
     jit_state state = {0};
     group = &shim;
-    code_size += group->code_size;
+    hot_code_size += group->hot_code_size;
+    cold_code_size += group->cold_code_size;
     data_size += group->data_size;
     combine_symbol_mask(group->trampoline_mask, state.trampolines.mask);
     combine_symbol_mask(group->got_mask, state.got_symbols.mask);
     // Round up to the nearest page:
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
+    size_t code_size = hot_code_size + cold_code_size;
     size_t code_padding = DATA_ALIGN - ((code_size + state.trampolines.size) & (DATA_ALIGN - 1));
     size_t padding = page_size - ((code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size) & (page_size - 1));
     size_t total_size = code_size + state.trampolines.size + code_padding + data_size + state.got_symbols.size + padding;
@@ -757,17 +768,20 @@ compile_shim(void)
         return NULL;
     }
     unsigned char *code = memory;
+    unsigned char *cold = memory + hot_code_size;
     state.trampolines.mem = memory + code_size;
     unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
     state.got_symbols.mem = data + data_size;
     // Compile the shim, which handles converting between the native
     // calling convention and the calling convention used by jitted code
     // (which may be different for efficiency reasons).
     group = &shim;
-    group->emit(code, data, &dummy, NULL, &state);
-    code += group->code_size;
+    group->emit(code, cold, data, &dummy, NULL, &state);
+    code += group->hot_code_size;
+    cold += group->cold_code_size;
     data += group->data_size;
-    assert(code == memory + code_size);
+    assert(code == memory + hot_code_size);
+    assert(cold == memory + code_size);
     assert(data == memory + code_size + state.trampolines.size + code_padding + data_size);
     if (mark_executable(memory, total_size)) {
         jit_free(memory, total_size);

Tools/jit/_optimizers.py

@@ -65,6 +65,7 @@
 }
 # MyPy doesn't understand that a invariant variable can be initialized by a covariant value
 CUSTOM_AARCH64_BRANCH19: str | None = "CUSTOM_AARCH64_BRANCH19"
+CUSTOM_AARCH64_BRANCH26: str | None = "ARM64_RELOC_BRANCH26"
 
 _AARCH64_SHORT_BRANCHES = {
     "tbz": "tbnz",
@@ -216,6 +217,7 @@ class Optimizer:
     label_index: int = 0
     _cold_start: _Block | None = dataclasses.field(init=False, default=None)
     _jump_name = "<Not supported>"
+    _jump_reloc: typing.ClassVar[str | None] = None
 
     def __post_init__(self) -> None:
         # Split the code into a linked list of basic blocks. A basic block is an
@@ -363,13 +365,31 @@ def _make_label(self) -> str:
         self.label_index += 1
         return label
 
+    def _make_symbol_label(self, name: str) -> str:
+        label = f"{self.symbol_prefix}{name}_{self.label_index}"
+        self.label_index += 1
+        return label
+
     def _ensure_label(self, block: _Block) -> str:
         if block.label is None:
             block.label = self._make_label()
             self._labels[block.label] = block
             block.noninstructions.insert(0, f"{block.label}:")
         return block.label
 
+    def _make_cold_target(self, block: _Block) -> str:
+        label = self._make_symbol_label("_JIT_COLD")
+        self._labels[label] = block
+        block.noninstructions.insert(0, f"{label}:")
+        return label
+
+    def _relocation_marker(self, reloc: str, block: _Block) -> Instruction:
+        target = self._make_cold_target(block).removeprefix(
+            self.symbol_prefix
+        )
+        label = f"{self.symbol_prefix}{reloc}_JIT_RELOCATION_{target}:"
+        return Instruction(InstructionKind.OTHER, "", label, None, None)
+
     def _make_jump(self, target: _Block, *, hot: bool) -> _Block:
         label = self._ensure_label(target)
         return _Block(
@@ -755,6 +775,32 @@ def _remove_unreachable(self) -> None:
         assert seen_continuation
         assert seen_cold_start
 
+    def _fixup_cross_section_branches(self) -> None:
+        layout_blocks = set(self._layout_blocks())
+        for block in self._layout_blocks():
+            target = block.target
+            if target is None or target not in layout_blocks:
+                continue
+            if self._same_layout_section(block, target):
+                continue
+            inst = block.instructions[-1]
+            if inst.kind == InstructionKind.LONG_BRANCH:
+                reloc = self._branches[inst.name][1]
+                if reloc is not None:
+                    block.instructions[-1] = self._relocation_marker(reloc, target)
+                    block.instructions.append(inst.update_target("0"))
+                    continue
+            elif inst.kind == InstructionKind.JUMP and self._jump_reloc is not None:
+                block.instructions[-1] = self._relocation_marker(
+                    self._jump_reloc, target
+                )
+                block.instructions.append(inst.update_target("0"))
+                continue
+            assert inst.kind is not InstructionKind.SHORT_BRANCH
+            block.instructions[-1] = inst.update_target(
+                self._make_cold_target(target)
+            )
+
     def _fixup_external_labels(self) -> None:
         if self._supports_external_relocations:
             # Nothing to fix up
@@ -803,6 +849,7 @@ def run(self) -> None:
             self._remove_redundant_jumps()
             self._remove_unreachable()
         self._fixup_external_labels()
+        self._fixup_cross_section_branches()
         self._fixup_constants()
         self._validate()
         self.path.write_text(self._body())
@@ -814,6 +861,7 @@ class OptimizerAArch64(Optimizer):  # pylint: disable = too-few-public-methods
     _branches = _AARCH64_BRANCHES
     _short_branches = _AARCH64_SHORT_BRANCHES
     _jump_name = "b"
+    _jump_reloc = CUSTOM_AARCH64_BRANCH26
     # Mach-O does not support the 19 bit branch locations needed for branch reordering
     _supports_external_relocations = False
     _branch_patterns = [name.replace(".", r"\.") for name in _AARCH64_BRANCHES]