Fix some issues

`replace` didn't preserve relocation;
root block should be at the start of the block list;
fix the pc-relative addressing of branching between hot and cold sections

Author: waterlens

Tools/jit/_optimizers.py

@@ -383,8 +383,17 @@ def _make_cold_target(self, block: _Block) -> str:
         block.noninstructions.insert(0, f"{label}:")
         return label
 
+    def _make_cross_section_target(self, block: _Block) -> str:
+        real_target = self._make_cold_target(block).removeprefix(
+            self.symbol_prefix
+        )
+        return (
+            f"{self.symbol_prefix}_JIT_CROSS_SECTION_"
+            f"{real_target}_JIT_CROSS_SECTION"
+        )
+
     def _relocation_marker(self, reloc: str, block: _Block) -> Instruction:
-        target = self._make_cold_target(block).removeprefix(
+        target = self._make_cross_section_target(block).removeprefix(
             self.symbol_prefix
         )
         label = f"{self.symbol_prefix}{reloc}_JIT_RELOCATION_{target}:"
@@ -555,15 +564,25 @@ def _partition_hot_cold_blocks(self) -> None:
         cold_blocks = [
             block for layout_hot, unit in units if not layout_hot for block in unit
         ]
-        self._relink_blocks(
-            [
-                *hot_blocks,
-                continuation,
-                cold_start,
-                *cold_blocks,
-                *self._metadata_blocks(),
-            ]
-        )
+        blocks = [
+            *hot_blocks,
+            continuation,
+            cold_start,
+            *cold_blocks,
+            *self._metadata_blocks(),
+        ]
+        if self._root in blocks and blocks[0] is not self._root:
+            assert self._root.label is None
+            assert not self._root.instructions
+            assert self._root.target is None
+            assert self._root.fallthrough
+            blocks.remove(self._root)
+            blocks.insert(0, self._root)
+        self._relink_blocks(blocks)
+        linked_blocks = list(self._blocks())
+        assert linked_blocks[0] is self._root
+        assert continuation in linked_blocks
+        assert cold_start in linked_blocks
 
     def _blocks(self) -> typing.Generator[_Block, None, None]:
         block: _Block | None = self._root
@@ -798,7 +817,7 @@ def _fixup_cross_section_branches(self) -> None:
                 continue
             assert inst.kind is not InstructionKind.SHORT_BRANCH
             block.instructions[-1] = inst.update_target(
-                self._make_cold_target(target)
+                self._make_cross_section_target(target)
             )
 
     def _fixup_external_labels(self) -> None:

Tools/jit/_stencils.py

@@ -1,5 +1,6 @@
 """Core data structures for compiled code templates."""
 
+import copy
 import dataclasses
 import enum
 import sys
@@ -148,6 +149,9 @@ class HoleValue(enum.Enum):
     "IMAGE_REL_AMD64_REL32",
 }
 
+_CROSS_SECTION_PREFIX = "_JIT_CROSS_SECTION_"
+_CROSS_SECTION_SUFFIX = "_JIT_CROSS_SECTION"
+
 @dataclasses.dataclass
 class Hole:
     """
@@ -170,8 +174,16 @@ class Hole:
     func: str = dataclasses.field(init=False)
     offset2: int = -1
     void: bool = False
-    # Convenience method:
-    replace = dataclasses.replace
+
+    def replace(self, **changes: typing.Any) -> typing.Self:
+        field_names = {field.name for field in dataclasses.fields(self)}
+        for name in changes:
+            if name not in field_names:
+                raise TypeError(f"unexpected field name: {name!r}")
+        result = copy.copy(self)
+        for name, value in changes.items():
+            setattr(result, name, value)
+        return result
 
     def __post_init__(self) -> None:
         self.func = _PATCH_FUNCS[self.kind]
@@ -295,6 +307,9 @@ def convert_labels_to_relocations(self) -> None:
                         first = first_in_pair[index]
                         hole.fold(first)
 
+        self.split_code_at_cold_start()
+
+    def split_code_at_cold_start(self) -> None:
         cold_start = self.symbols.get("_JIT_COLD_START")
         if cold_start is None:
             return
@@ -330,6 +345,8 @@ def process_relocations(self, known_symbols: dict[str, int]) -> None:
         """Fix up all GOT and internal relocations for this stencil group."""
         for stencil in self._code_stencils():
             for hole in stencil.holes.copy():
+                if self._resolve_cross_section_symbol(hole):
+                    continue
                 if (
                     hole.kind
                     in {"R_AARCH64_CALL26", "R_AARCH64_JUMP26", "ARM64_RELOC_BRANCH26"}
@@ -379,6 +396,8 @@ def process_relocations(self, known_symbols: dict[str, int]) -> None:
         self.data.pad(8)
         for stencil in self._all_stencils():
             for hole in stencil.holes:
+                if self._resolve_cross_section_symbol(hole):
+                    continue
                 if hole.value is HoleValue.GOT:
                     assert hole.symbol is not None
                     if "_JIT_" in hole.symbol:
@@ -415,6 +434,29 @@ def process_relocations(self, known_symbols: dict[str, int]) -> None:
         self.cold_code.holes.sort(key=lambda hole: hole.offset)
         self.data.holes.sort(key=lambda hole: hole.offset)
 
+    def _resolve_cross_section_symbol(self, hole: Hole) -> bool:
+        if hole.value is not HoleValue.ZERO or hole.symbol is None:
+            return False
+        symbol = hole.symbol
+        if not (
+            symbol.startswith(_CROSS_SECTION_PREFIX)
+            and symbol.endswith(_CROSS_SECTION_SUFFIX)
+        ):
+            return False
+        target = symbol.removeprefix(_CROSS_SECTION_PREFIX).removesuffix(
+            _CROSS_SECTION_SUFFIX
+        )
+        try:
+            value, addend = self.symbols[target]
+        except KeyError:
+            raise ValueError(
+                f"Missing cross-section relocation target {target}"
+            ) from None
+        hole.value = value
+        hole.addend += addend
+        hole.symbol = None
+        return True
+
     def _jit_symbol_table_lookup(self, symbol: str) -> int:
         return len(self.data.body) + self._jit_symbol_table.setdefault(
             symbol, 8 * len(self._jit_symbol_table)