GH-113464: Generate a more efficient JIT (GH-118512)

2 files changed, 142 insertions, 81 deletions

diff --git a/Tools/jit/_stencils.py b/Tools/jit/_stencils.py
index 9feceb4..6e046df 100644
--- a/Tools/jit/_stencils.py
+++ b/Tools/jit/_stencils.py
@@ -3,6 +3,7 @@
 import dataclasses
 import enum
 import sys
+import typing
 
 import _schema
 
@@ -47,6 +48,73 @@ class HoleValue(enum.Enum):
     ZERO = enum.auto()
 
 
+# Map relocation types to our JIT's patch functions. "r" suffixes indicate that
+# the patch function is relative. "x" suffixes indicate that they are "relaxing"
+# (see comments in jit.c for more info):
+_PATCH_FUNCS = {
+    # aarch64-apple-darwin:
+    "ARM64_RELOC_BRANCH26": "patch_aarch64_26r",
+    "ARM64_RELOC_GOT_LOAD_PAGE21": "patch_aarch64_21rx",
+    "ARM64_RELOC_GOT_LOAD_PAGEOFF12": "patch_aarch64_12x",
+    "ARM64_RELOC_PAGE21": "patch_aarch64_21r",
+    "ARM64_RELOC_PAGEOFF12": "patch_aarch64_12",
+    "ARM64_RELOC_UNSIGNED": "patch_64",
+    # x86_64-pc-windows-msvc:
+    "IMAGE_REL_AMD64_REL32": "patch_x86_64_32rx",
+    # aarch64-pc-windows-msvc:
+    "IMAGE_REL_ARM64_BRANCH26": "patch_aarch64_26r",
+    "IMAGE_REL_ARM64_PAGEBASE_REL21": "patch_aarch64_21rx",
+    "IMAGE_REL_ARM64_PAGEOFFSET_12A": "patch_aarch64_12",
+    "IMAGE_REL_ARM64_PAGEOFFSET_12L": "patch_aarch64_12x",
+    # i686-pc-windows-msvc:
+    "IMAGE_REL_I386_DIR32": "patch_32",
+    "IMAGE_REL_I386_REL32": "patch_x86_64_32rx",
+    # aarch64-unknown-linux-gnu:
+    "R_AARCH64_ABS64": "patch_64",
+    "R_AARCH64_ADD_ABS_LO12_NC": "patch_aarch64_12",
+    "R_AARCH64_ADR_GOT_PAGE": "patch_aarch64_21rx",
+    "R_AARCH64_ADR_PREL_PG_HI21": "patch_aarch64_21r",
+    "R_AARCH64_CALL26": "patch_aarch64_26r",
+    "R_AARCH64_JUMP26": "patch_aarch64_26r",
+    "R_AARCH64_LD64_GOT_LO12_NC": "patch_aarch64_12x",
+    "R_AARCH64_MOVW_UABS_G0_NC": "patch_aarch64_16a",
+    "R_AARCH64_MOVW_UABS_G1_NC": "patch_aarch64_16b",
+    "R_AARCH64_MOVW_UABS_G2_NC": "patch_aarch64_16c",
+    "R_AARCH64_MOVW_UABS_G3": "patch_aarch64_16d",
+    # x86_64-unknown-linux-gnu:
+    "R_X86_64_64": "patch_64",
+    "R_X86_64_GOTPCREL": "patch_32r",
+    "R_X86_64_GOTPCRELX": "patch_x86_64_32rx",
+    "R_X86_64_PC32": "patch_32r",
+    "R_X86_64_REX_GOTPCRELX": "patch_x86_64_32rx",
+    # x86_64-apple-darwin:
+    "X86_64_RELOC_BRANCH": "patch_32r",
+    "X86_64_RELOC_GOT": "patch_x86_64_32rx",
+    "X86_64_RELOC_GOT_LOAD": "patch_x86_64_32rx",
+    "X86_64_RELOC_SIGNED": "patch_32r",
+    "X86_64_RELOC_UNSIGNED": "patch_64",
+}
+# Translate HoleValues to C expressions:
+_HOLE_EXPRS = {
+    HoleValue.CODE: "(uintptr_t)code",
+    HoleValue.CONTINUE: "(uintptr_t)code + sizeof(code_body)",
+    HoleValue.DATA: "(uintptr_t)data",
+    HoleValue.EXECUTOR: "(uintptr_t)executor",
+    # These should all have been turned into DATA values by process_relocations:
+    # HoleValue.GOT: "",
+    HoleValue.OPARG: "instruction->oparg",
+    HoleValue.OPERAND: "instruction->operand",
+    HoleValue.OPERAND_HI: "(instruction->operand >> 32)",
+    HoleValue.OPERAND_LO: "(instruction->operand & UINT32_MAX)",
+    HoleValue.TARGET: "instruction->target",
+    HoleValue.JUMP_TARGET: "instruction_starts[instruction->jump_target]",
+    HoleValue.ERROR_TARGET: "instruction_starts[instruction->error_target]",
+    HoleValue.EXIT_INDEX: "instruction->exit_index",
+    HoleValue.TOP: "instruction_starts[1]",
+    HoleValue.ZERO: "",
+}
+
+
 @dataclasses.dataclass
 class Hole:
     """
@@ -63,19 +131,43 @@ class Hole:
     symbol: str | None
     # ...plus this addend:
     addend: int
+    func: str = dataclasses.field(init=False)
     # Convenience method:
     replace = dataclasses.replace
 
-    def as_c(self) -> str:
-        """Dump this hole as an initialization of a C Hole struct."""
-        parts = [
-            f"{self.offset:#x}",
-            f"HoleKind_{self.kind}",
-            f"HoleValue_{self.value.name}",
-            f"&{self.symbol}" if self.symbol else "NULL",
-            f"{_signed(self.addend):#x}",
-        ]
-        return f"{{{', '.join(parts)}}}"
+    def __post_init__(self) -> None:
+        self.func = _PATCH_FUNCS[self.kind]
+
+    def fold(self, other: typing.Self) -> typing.Self | None:
+        """Combine two holes into a single hole, if possible."""
+        if (
+            self.offset + 4 == other.offset
+            and self.value == other.value
+            and self.symbol == other.symbol
+            and self.addend == other.addend
+            and self.func == "patch_aarch64_21rx"
+            and other.func == "patch_aarch64_12x"
+        ):
+            # These can *only* be properly relaxed when they appear together and
+            # patch the same value:
+            folded = self.replace()
+            folded.func = "patch_aarch64_33rx"
+            return folded
+        return None
+
+    def as_c(self, where: str) -> str:
+        """Dump this hole as a call to a patch_* function."""
+        location = f"{where} + {self.offset:#x}"
+        value = _HOLE_EXPRS[self.value]
+        if self.symbol:
+            if value:
+                value += " + "
+            value += f"(uintptr_t)&{self.symbol}"
+        if _signed(self.addend):
+            if value:
+                value += " + "
+            value += f"{_signed(self.addend):#x}"
+        return f"{self.func}({location}, {value});"
 
 
 @dataclasses.dataclass
@@ -265,6 +357,10 @@ class StencilGroup:
             )
             self.data.body.extend([0] * 8)
 
+    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)}}}"
+
 
 def symbol_to_value(symbol: str) -> tuple[HoleValue, str | None]:
     """
diff --git a/Tools/jit/_writer.py b/Tools/jit/_writer.py
index ccd6785..9d11094 100644
--- a/Tools/jit/_writer.py
+++ b/Tools/jit/_writer.py
@@ -1,100 +1,65 @@
 """Utilities for writing StencilGroups out to a C header file."""
 
+import itertools
 import typing
 
-import _schema
 import _stencils
 
 
-def _dump_header() -> typing.Iterator[str]:
-    yield "typedef enum {"
-    for kind in typing.get_args(_schema.HoleKind):
-        yield f"    HoleKind_{kind},"
-    yield "} HoleKind;"
-    yield ""
-    yield "typedef enum {"
-    for value in _stencils.HoleValue:
-        yield f"    HoleValue_{value.name},"
-    yield "} HoleValue;"
-    yield ""
-    yield "typedef struct {"
-    yield "    const size_t offset;"
-    yield "    const HoleKind kind;"
-    yield "    const HoleValue value;"
-    yield "    const void *symbol;"
-    yield "    const uint64_t addend;"
-    yield "} Hole;"
-    yield ""
+def _dump_footer(groups: dict[str, _stencils.StencilGroup]) -> typing.Iterator[str]:
     yield "typedef struct {"
-    yield "    const size_t body_size;"
-    yield "    const unsigned char * const body;"
-    yield "    const size_t holes_size;"
-    yield "    const Hole * const holes;"
-    yield "} Stencil;"
-    yield ""
-    yield "typedef struct {"
-    yield "    const Stencil code;"
-    yield "    const Stencil data;"
+    yield "    void (*emit)("
+    yield "        unsigned char *code, unsigned char *data, _PyExecutorObject *executor,"
+    yield "        const _PyUOpInstruction *instruction, uintptr_t instruction_starts[]);"
+    yield "    size_t code_size;"
+    yield "    size_t data_size;"
     yield "} StencilGroup;"
     yield ""
-
-
-def _dump_footer(opnames: typing.Iterable[str]) -> typing.Iterator[str]:
-    yield "#define INIT_STENCIL(STENCIL) {                         \\"
-    yield "    .body_size = Py_ARRAY_LENGTH(STENCIL##_body) - 1,   \\"
-    yield "    .body = STENCIL##_body,                             \\"
-    yield "    .holes_size = Py_ARRAY_LENGTH(STENCIL##_holes) - 1, \\"
-    yield "    .holes = STENCIL##_holes,                           \\"
-    yield "}"
-    yield ""
-    yield "#define INIT_STENCIL_GROUP(OP) {     \\"
-    yield "    .code = INIT_STENCIL(OP##_code), \\"
-    yield "    .data = INIT_STENCIL(OP##_data), \\"
-    yield "}"
+    yield f"static const StencilGroup trampoline = {groups['trampoline'].as_c('trampoline')};"
     yield ""
-    yield "static const StencilGroup stencil_groups[512] = {"
-    for opname in opnames:
+    yield "static const StencilGroup stencil_groups[MAX_UOP_ID + 1] = {"
+    for opname, group in sorted(groups.items()):
         if opname == "trampoline":
             continue
-        yield f"    [{opname}] = INIT_STENCIL_GROUP({opname}),"
+        yield f"    [{opname}] = {group.as_c(opname)},"
     yield "};"
-    yield ""
-    yield "static const StencilGroup trampoline = INIT_STENCIL_GROUP(trampoline);"
-    yield ""
-    yield "#define GET_PATCHES() { \\"
-    for value in _stencils.HoleValue:
-        yield f"    [HoleValue_{value.name}] = (uintptr_t)0xBADBADBADBADBADB, \\"
-    yield "}"
 
 
 def _dump_stencil(opname: str, group: _stencils.StencilGroup) -> typing.Iterator[str]:
-    yield f"// {opname}"
+    yield "void"
+    yield f"emit_{opname}("
+    yield "    unsigned char *code, unsigned char *data, _PyExecutorObject *executor,"
+    yield "    const _PyUOpInstruction *instruction, uintptr_t instruction_starts[])"
+    yield "{"
     for part, stencil in [("code", group.code), ("data", group.data)]:
         for line in stencil.disassembly:
-            yield f"// {line}"
+            yield f"    // {line}"
         if stencil.body:
-            size = len(stencil.body) + 1
-            yield f"static const unsigned char {opname}_{part}_body[{size}] = {{"
+            yield f"    const unsigned char {part}_body[{len(stencil.body)}] = {{"
             for i in range(0, len(stencil.body), 8):
                 row = " ".join(f"{byte:#04x}," for byte in stencil.body[i : i + 8])
-                yield f"    {row}"
-            yield "};"
-        else:
-            yield f"static const unsigned char {opname}_{part}_body[1];"
-        if stencil.holes:
-            size = len(stencil.holes) + 1
-            yield f"static const Hole {opname}_{part}_holes[{size}] = {{"
-            for hole in stencil.holes:
-                yield f"    {hole.as_c()},"
-            yield "};"
-        else:
-            yield f"static const Hole {opname}_{part}_holes[1];"
+                yield f"        {row}"
+            yield "    };"
+    # Data is written first (so relaxations in the code work properly):
+    for part, stencil in [("data", group.data), ("code", group.code)]:
+        if stencil.body:
+            yield f"    memcpy({part}, {part}_body, sizeof({part}_body));"
+        skip = False
+        stencil.holes.sort(key=lambda hole: hole.offset)
+        for hole, pair in itertools.zip_longest(stencil.holes, stencil.holes[1:]):
+            if skip:
+                skip = False
+                continue
+            if pair and (folded := hole.fold(pair)):
+                skip = True
+                hole = folded
+            yield f"    {hole.as_c(part)}"
+    yield "}"
     yield ""
 
 
 def dump(groups: dict[str, _stencils.StencilGroup]) -> typing.Iterator[str]:
     """Yield a JIT compiler line-by-line as a C header file."""
-    yield from _dump_header()
-    for opname, group in groups.items():
+    for opname, group in sorted(groups.items()):
         yield from _dump_stencil(opname, group)
     yield from _dump_footer(groups)