JIT: Implement hot-cold code splitting for improved icache locality

Split JIT stencil code into hot and cold sections to improve instruction
cache performance. Hot code (common execution paths) is laid out
contiguously across all uops in a trace, while cold code (error handling,
rare branches) is placed after all hot code.

Changes:
- _optimizers.py: Physically reorder basic blocks so hot blocks come
  first, cold blocks last. Insert _JIT_COLD_START symbol at boundary.
  Add _make_jump() to each arch-specific optimizer for inserting
  explicit jumps when hot->cold fallthrough edges are broken.
- _stencils.py: Track cold_offset in StencilGroup. Add hot_code_size()
  and cold_code_size() methods. Extract cold offset from _JIT_COLD_START
  symbol in object files.
- _writer.py: Split emit functions to handle hot and cold code bodies
  separately, with adjusted hole patching for cold code offsets. Update
  StencilGroup C struct to include hot_code_size and cold_code_size.
- _targets.py: Call extract_cold_offset() during stencil build.
- jit.c: New memory layout [hot code|cold code|trampolines|data].
  Update _PyJIT_Compile and compile_shim to use split layout.

https://claude.ai/code/session_01QteUqrt8X8Ssxhei1kX2TH

Author: claude

Python/jit.c

@@ -642,23 +642,27 @@ _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);
+    size_t code_size = hot_code_size + cold_code_size;
     // Calculate the size of the trampolines required by the whole trace
     for (size_t i = 0; i < Py_ARRAY_LENGTH(state.trampolines.mask); i++) {
         state.trampolines.size += _Py_popcount32(state.trampolines.mask[i]) * TRAMPOLINE_SIZE;
@@ -684,29 +688,36 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     OPT_STAT_ADD(jit_got_size, state.got_symbols.size);
     OPT_STAT_ADD(jit_padding_size, padding);
     OPT_HIST(total_size, trace_total_memory_hist);
-    // Update the offsets of each instruction:
+    // Update the offsets of each instruction (hot code comes first):
     for (size_t i = 0; i < length; i++) {
         state.instruction_starts[i] += (uintptr_t)memory;
     }
-    // Loop again to emit the code:
+    // Memory layout: [hot code][cold code][trampolines][padding][data][GOT][padding]
+    // Hot code for all uops is laid out contiguously first, improving
+    // instruction cache locality. Cold code (error handling, rare paths)
+    // is placed after all hot code.
     unsigned char *code = memory;
+    unsigned char *cold_code = 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_code, data, executor, instruction, &state);
+        code += group->hot_code_size;
+        cold_code += 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_code, data, executor, NULL, &state);
+    code += group->hot_code_size;
+    cold_code += group->cold_code_size;
     data += group->data_size;
-    assert(code == memory + code_size);
+    assert(code == memory + hot_code_size);
+    assert(cold_code == 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);
@@ -728,14 +739,17 @@ 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);
+    size_t code_size = hot_code_size + cold_code_size;
     // Round up to the nearest page:
     size_t page_size = get_page_size();
     assert((page_size & (page_size - 1)) == 0);
@@ -747,17 +761,20 @@ compile_shim(void)
         return NULL;
     }
     unsigned char *code = memory;
+    unsigned char *cold_code = 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_code, data, &dummy, NULL, &state);
+    code += group->hot_code_size;
+    cold_code += group->cold_code_size;
     data += group->data_size;
-    assert(code == memory + code_size);
+    assert(code == memory + hot_code_size);
+    assert(cold_code == 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

@@ -310,6 +310,72 @@ def _blocks(self) -> typing.Generator[_Block, None, None]:
             yield block
             block = block.link
 
+    def _make_jump(self, target: str) -> Instruction:
+        """Create an unconditional jump instruction to the given label."""
+        raise NotImplementedError
+
+    def _reorder_hot_cold(self) -> None:
+        """Reorder blocks so all hot blocks come first, then all cold blocks.
+
+        This improves instruction cache locality by keeping the common execution
+        paths contiguous in memory. A _JIT_COLD_START label is inserted at the
+        boundary between hot and cold code.
+
+        For any hot block that previously fell through to a cold block, an
+        explicit jump is inserted to maintain correctness.
+        """
+        continuation = self._lookup_label(f"{self.label_prefix}_JIT_CONTINUE")
+        hot_blocks: list[_Block] = []
+        cold_blocks: list[_Block] = []
+        # Blocks at and after _JIT_CONTINUE contain metadata; keep them at end
+        post_continuation: list[_Block] = []
+        past_continuation = False
+        for block in self._blocks():
+            if block is continuation:
+                past_continuation = True
+            if past_continuation:
+                post_continuation.append(block)
+            elif block.hot:
+                hot_blocks.append(block)
+            else:
+                cold_blocks.append(block)
+
+        if not cold_blocks:
+            # Nothing to reorder — all code is hot
+            return
+
+        # For any hot block that falls through to a cold block, insert an
+        # explicit jump to maintain correctness after reordering:
+        for block in hot_blocks:
+            if (
+                block.fallthrough
+                and block.link is not None
+                and not block.link.hot
+                and block.link is not continuation
+            ):
+                # The fallthrough target is cold and will be moved away.
+                # Insert a jump to the cold block's label:
+                target_block = block.link.resolve()
+                if target_block.label:
+                    jump = self._make_jump(target_block.label)
+                    block.instructions.append(jump)
+                    block.fallthrough = False
+
+        # Rebuild the linked list: hot -> cold -> post_continuation
+        all_blocks = hot_blocks + cold_blocks + post_continuation
+        for i, block in enumerate(all_blocks):
+            block.link = all_blocks[i + 1] if i + 1 < len(all_blocks) else None
+
+        # Insert the cold start label between hot and cold sections.
+        # Use the symbol_prefix so it appears in the object file's symbol table:
+        if hot_blocks and cold_blocks:
+            cold_label = f"{self.symbol_prefix}_JIT_COLD_START"
+            cold_blocks[0].noninstructions.insert(0, f"{cold_label}:")
+
+        # Update the root to the first block
+        if all_blocks:
+            self._root = all_blocks[0]
+
     def _body(self) -> str:
         lines = ["#" + line for line in self.text.splitlines()]
         hot = True
@@ -563,6 +629,7 @@ def run(self) -> None:
             self._invert_hot_branches()
             self._remove_redundant_jumps()
             self._remove_unreachable()
+        self._reorder_hot_cold()
         self._fixup_external_labels()
         self._fixup_constants()
         self.path.write_text(self._body())
@@ -580,6 +647,9 @@ class OptimizerAArch64(Optimizer):  # pylint: disable = too-few-public-methods
         rf"\s*(?P<instruction>{'|'.join(_branch_patterns)})\s+(.+,\s+)*(?P<target>[\w.]+)"
     )
 
+    def _make_jump(self, target: str) -> Instruction:
+        return Instruction(InstructionKind.JUMP, "b", f"\tb {target}", target)
+
     # https://developer.arm.com/documentation/ddi0406/b/Application-Level-Architecture/Instruction-Details/Alphabetical-list-of-instructions/BL--BLX--immediate-
     _re_call = re.compile(r"\s*blx?\s+(?P<target>[\w.]+)")
     # https://developer.arm.com/documentation/ddi0602/2025-03/Base-Instructions/B--Branch-
@@ -644,6 +714,9 @@ class OptimizerX86(Optimizer):  # pylint: disable = too-few-public-methods
     _re_branch = re.compile(
         rf"\s*(?P<instruction>{'|'.join(_X86_BRANCHES)})\s+(?P<target>[\w.]+)"
     )
+
+    def _make_jump(self, target: str) -> Instruction:
+        return Instruction(InstructionKind.JUMP, "jmp", f"\tjmp {target}", target)
     # https://www.felixcloutier.com/x86/call
     _re_call = re.compile(r"\s*callq?\s+(?P<target>[\w.]+)")
     # https://www.felixcloutier.com/x86/jmp

Tools/jit/_stencils.py

@@ -256,6 +256,9 @@ class StencilGroup:
 
     code: Stencil = dataclasses.field(default_factory=Stencil, init=False)
     data: Stencil = dataclasses.field(default_factory=Stencil, init=False)
+    # Byte offset within code.body where cold code begins.
+    # If 0, there is no hot/cold split (all code is hot):
+    cold_offset: int = dataclasses.field(default=0, init=False)
     symbols: dict[int | str, tuple[HoleValue, int]] = dataclasses.field(
         default_factory=dict, init=False
     )
@@ -265,6 +268,14 @@ class StencilGroup:
     _trampolines: set[int] = dataclasses.field(default_factory=set, init=False)
     _got_entries: set[int] = dataclasses.field(default_factory=set, init=False)
 
+    def extract_cold_offset(self) -> None:
+        """Extract the cold code offset from the _JIT_COLD_START symbol."""
+        cold_start = self.symbols.get("_JIT_COLD_START")
+        if cold_start is not None:
+            value, offset = cold_start
+            assert value is HoleValue.CODE
+            self.cold_offset = offset
+
     def convert_labels_to_relocations(self) -> None:
         for name, hole_plus in self.symbols.items():
             if isinstance(name, str) and "_JIT_RELOCATION_" in name:
@@ -420,9 +431,26 @@ def _get_trampoline_mask(self) -> str:
     def _get_got_mask(self) -> str:
         return self._get_symbol_mask(self._got_entries)
 
+    def hot_code_size(self) -> int:
+        """Size of the hot code section (or total code if no cold split)."""
+        if self.cold_offset:
+            return self.cold_offset
+        return len(self.code.body)
+
+    def cold_code_size(self) -> int:
+        """Size of the cold code section."""
+        if self.cold_offset:
+            return len(self.code.body) - self.cold_offset
+        return 0
+
     def as_c(self, opname: str) -> str:
         """Dump this hole as a StencilGroup initializer."""
-        return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {self._get_trampoline_mask()}, {self._get_got_mask()}}}"
+        return (
+            f"{{emit_{opname}, {self.hot_code_size()}, "
+            f"{self.cold_code_size()}, "
+            f"{len(self.data.body)}, "
+            f"{self._get_trampoline_mask()}, {self._get_got_mask()}}}"
+        )
 
 
 def symbol_to_value(symbol: str) -> tuple[HoleValue, str | None]:

Tools/jit/_targets.py

@@ -219,6 +219,7 @@ async def _build_stencils(self) -> dict[str, _stencils.StencilGroup]:
                     tasks.append(group.create_task(coro, name=opname))
         stencil_groups = {task.get_name(): task.result() for task in tasks}
         for stencil_group in stencil_groups.values():
+            stencil_group.extract_cold_offset()
             stencil_group.convert_labels_to_relocations()
             stencil_group.process_relocations(self.known_symbols)
         return stencil_groups