python · brandtbucher · Sep 1, 2023 · Sep 1, 2023 · Sep 1, 2023 · Sep 5, 2023
diff --git a/Python/jit.c b/Python/jit.c
@@ -118,7 +118,7 @@ size_to_pages(uint64_t size)
     return (size + page_size - 1) / page_size;
 }
 
-static char *emit_trampoline(uint64_t where, char **trampolines);
+// static char *emit_trampoline(uint64_t where, char **trampolines);
 
 static bool
 fits_in_signed_bits(uint64_t value, int bits)
@@ -127,14 +127,14 @@ fits_in_signed_bits(uint64_t value, int bits)
     return (v >= -(1LL << (bits - 1))) && (v < (1LL << (bits - 1)));
 }
 
-static uint64_t
-potential_trampoline_size(const StencilGroup *stencil_group)
-{
-    return stencil_group->text.holes_size * trampoline_stencil_group.text.body_size;
-}
+// static uint64_t
+// potential_trampoline_size(const StencilGroup *stencil_group)
+// {
+//     return stencil_group->text.holes_size * trampoline_stencil_group.text.body_size;
+// }
 
 static void
-patch(char *base, const Hole *hole, uint64_t *patches, char **trampolines)
+patch(char *base, const Hole *hole, uint64_t *patches)
 {
     char *location = base + hole->offset;
     uint64_t value = patches[hole->value] + (uint64_t)hole->symbol + hole->addend;
@@ -150,9 +150,9 @@ patch(char *base, const Hole *hole, uint64_t *patches, char **trampolines)
         case HoleKind_X86_64_RELOC_BRANCH:
         case HoleKind_X86_64_RELOC_GOT:
         case HoleKind_X86_64_RELOC_GOT_LOAD:
-            if (!fits_in_signed_bits(value - (uint64_t)location, 32)) {
-                value = (uint64_t)emit_trampoline(value + 4, trampolines) - 4;  // XXX
-            }
+            // if (!fits_in_signed_bits(value - (uint64_t)location, 32)) {
+            //     value = (uint64_t)emit_trampoline(value + 4, trampolines) - 4;  // XXX
+            // }
             value -= (uint64_t)location;
             assert(fits_in_signed_bits(value, 32));
             *addr = (uint32_t)value;
@@ -174,9 +174,9 @@ patch(char *base, const Hole *hole, uint64_t *patches, char **trampolines)
             return;
         case HoleKind_R_AARCH64_CALL26:
         case HoleKind_R_AARCH64_JUMP26:
-            if (!fits_in_signed_bits(value - (uint64_t)location, 28)) {
-                value = (uint64_t)emit_trampoline(value, trampolines);
-            }
+            // if (!fits_in_signed_bits(value - (uint64_t)location, 28)) {
+            //     value = (uint64_t)emit_trampoline(value, trampolines);
+            // }
             value -= (uint64_t)location;
             assert(((*addr & 0xFC000000) == 0x14000000) ||
                    ((*addr & 0xFC000000) == 0x94000000));
@@ -219,38 +219,38 @@ patch(char *base, const Hole *hole, uint64_t *patches, char **trampolines)
 }
 
 static void
-copy_and_patch(char *base, const Stencil *stencil, uint64_t *patches, char **trampolines)
+copy_and_patch(char *base, const Stencil *stencil, uint64_t *patches)
 {
     memcpy(base, stencil->body, stencil->body_size);
     for (uint64_t i = 0; i < stencil->holes_size; i++) {
-        patch(base, &stencil->holes[i], patches, trampolines);
+        patch(base, &stencil->holes[i], patches);
     }
 }
 
 static void
-emit(const StencilGroup *stencil_group, uint64_t patches[], char **trampolines)
+emit(const StencilGroup *stencil_group, uint64_t patches[])
 {
     char *data = (char *)patches[HoleValue_DATA];
-    copy_and_patch(data, &stencil_group->data, patches, trampolines);
+    copy_and_patch(data, &stencil_group->data, patches);
     char *text = (char *)patches[HoleValue_TEXT];
-    copy_and_patch(text, &stencil_group->text, patches, trampolines);
+    copy_and_patch(text, &stencil_group->text, patches);
 }
 
-static char *
-emit_trampoline(uint64_t where, char **trampolines)
-{
-    assert(trampolines && *trampolines);
-    const StencilGroup *stencil_group = &trampoline_stencil_group;
-    assert(stencil_group->data.body_size == 0);
-    uint64_t patches[] = GET_PATCHES();
-    patches[HoleValue_CONTINUE] = where;
-    patches[HoleValue_TEXT] = (uint64_t)*trampolines;
-    patches[HoleValue_ZERO] = 0;
-    emit(stencil_group, patches, NULL);
-    char *trampoline = *trampolines;
-    *trampolines += stencil_group->text.body_size;
-    return trampoline;
-}
+// static char *
+// emit_trampoline(uint64_t where, char **trampolines)
+// {
+//     assert(trampolines && *trampolines);
+//     const StencilGroup *stencil_group = &trampoline_stencil_group;
+//     assert(stencil_group->data.body_size == 0);
+//     uint64_t patches[] = GET_PATCHES();
+//     patches[HoleValue_CONTINUE] = where;
+//     patches[HoleValue_TEXT] = (uint64_t)*trampolines;
+//     patches[HoleValue_ZERO] = 0;
+//     emit(stencil_group, patches, NULL);
+//     char *trampoline = *trampolines;
+//     *trampolines += stencil_group->text.body_size;
+//     return trampoline;
+// }
 
 static void
 initialize_jit(void)
@@ -275,25 +275,25 @@ _PyJIT_CompileTrace(_PyUOpExecutorObject *executor, _PyUOpInstruction *trace, in
 {
     initialize_jit();
     // First, loop over everything once to find the total compiled size:
-    uint64_t trampoline_size = potential_trampoline_size(&wrapper_stencil_group);
+    // uint64_t trampoline_size = potential_trampoline_size(&wrapper_stencil_group);
     uint64_t text_size = wrapper_stencil_group.text.body_size;
     uint64_t data_size = wrapper_stencil_group.data.body_size;
     for (int i = 0; i < size; i++) {
         _PyUOpInstruction *instruction = &trace[i];
         const StencilGroup *stencil_group = &stencil_groups[instruction->opcode];
-        trampoline_size += potential_trampoline_size(stencil_group);
+        // trampoline_size += potential_trampoline_size(stencil_group);
         text_size += stencil_group->text.body_size;
         data_size += stencil_group->data.body_size;
         assert(stencil_group->text.body_size);
     };
-    uint64_t text_pages = size_to_pages(text_size + trampoline_size);
+    uint64_t text_pages = size_to_pages(text_size);
     uint64_t data_pages = size_to_pages(data_size);
     char *text = alloc(text_pages + data_pages);
     if (text == NULL) {
         return NULL;
     }
     char *data = text + text_pages * page_size;
-    char *head_trampolines = text + text_size;
+    // char *head_trampolines = text + text_size;
     char *head_text = text;
     char *head_data = data;
     // First, the wrapper:
@@ -303,7 +303,7 @@ _PyJIT_CompileTrace(_PyUOpExecutorObject *executor, _PyUOpInstruction *trace, in
     patches[HoleValue_DATA] = (uint64_t)head_data;
     patches[HoleValue_TEXT] = (uint64_t)head_text;
     patches[HoleValue_ZERO] = 0;
-    emit(stencil_group, patches, &head_trampolines);
+    emit(stencil_group, patches);
     head_text += stencil_group->text.body_size;
     head_data += stencil_group->data.body_size;
     // Then, all of the stencils:
@@ -320,7 +320,7 @@ _PyJIT_CompileTrace(_PyUOpExecutorObject *executor, _PyUOpInstruction *trace, in
         patches[HoleValue_TEXT] = (uint64_t)head_text;
        patches[HoleValue_TOP] = (uint64_t)text + wrapper_stencil_group.text.body_size;
         patches[HoleValue_ZERO] = 0;
-        emit(stencil_group, patches, &head_trampolines);
+        emit(stencil_group, patches);
         head_text += stencil_group->text.body_size;
         head_data += stencil_group->data.body_size;
     };

diff --git a/Tools/jit/build.py b/Tools/jit/build.py
@@ -33,7 +33,7 @@
 PYTHON_JIT_STENCILS_H = PYTHON / "jit_stencils.h"
 TOOLS_JIT_TEMPLATE_C = TOOLS_JIT / "template.c"
 
-STUBS = ["trampoline", "wrapper"]
+STUBS = ["wrapper"]
 
 LLVM_VERSION = 16
 
@@ -193,9 +193,6 @@ async def parse(self) -> StencilGroup:
         holes = []
         holes_data = []
         padding = 0
-        while len(self.text) % self.target.alignment:
-            self.text.append(0)
-            padding += 1
         offset_data = 0
         disassembly_data = []
         padding_data = 0
@@ -226,6 +223,68 @@ async def parse(self) -> StencilGroup:
                     newhole.offset, newhole.kind, HoleValue.TEXT, None, addend
                 )
             holes.append(newhole)
+        remaining = []
+        for hole in holes:
+            if hole.kind in {"R_AARCH64_CALL26", "R_AARCH64_JUMP26"} and hole.value is HoleValue.ZERO:
+                base = len(self.text)
+                self.text.extend([0xd2, 0x80, 0x00, 0x08][::-1])
+                self.text.extend([0xf2, 0xa0, 0x00, 0x08][::-1])
+                self.text.extend([0xf2, 0xc0, 0x00, 0x08][::-1])
+                self.text.extend([0xf2, 0xe0, 0x00, 0x08][::-1])
+                self.text.extend([0xd6, 0x1f, 0x01, 0x00][::-1])
+                disassembly.extend(
+                    [  # XXX: Include addend:
+                        f"{base:x}: d2800008      mov     x8, #0x0",
+                        f"{base:016x}:  R_AARCH64_MOVW_UABS_G0_NC    {hole.symbol}",
+                        f"{base + 4:x}: f2a00008      movk    x8, #0x0, lsl #16",
+                        f"{base + 4:016x}:  R_AARCH64_MOVW_UABS_G1_NC    {hole.symbol}",
+                        f"{base + 8:x}: f2c00008      movk    x8, #0x0, lsl #32",
+                        f"{base + 8:016x}:  R_AARCH64_MOVW_UABS_G2_NC    {hole.symbol}",
+                        f"{base + 12:x}: f2e00008      movk    x8, #0x0, lsl #48",
+                        f"{base + 12:016x}:  R_AARCH64_MOVW_UABS_G3       {hole.symbol}",
+                        f"{base + 16:x}: d61f0100      br      x8",
+                    ]
+                )
+                remaining.extend(
+                    [
+                        Hole(
+                            base,
+                            "R_AARCH64_MOVW_UABS_G0_NC",
+                            hole.value,
+                            hole.symbol,
+                            hole.addend
+                        ),
+                        Hole(
+                            base + 4,
+                            "R_AARCH64_MOVW_UABS_G1_NC",
+                            hole.value,
+                            hole.symbol,
+                            hole.addend
+                        ),
+                        Hole(
+                            base + 8,
+                            "R_AARCH64_MOVW_UABS_G2_NC",
+                            hole.value,
+                            hole.symbol,
+                            hole.addend
+                        ),
+                        Hole(
+                            base + 12,
+                            "R_AARCH64_MOVW_UABS_G3",
+                            hole.value,
+                            hole.symbol,
+                            hole.addend
+                        ),
+                    ]
+                )
+                instruction = int.from_bytes(self.text[hole.offset : hole.offset + 4], sys.byteorder)
+                instruction = (instruction & 0xFC000000) | (((base - hole.offset) >> 2) & 0x03FFFFFF)
+                self.text[hole.offset : hole.offset + 4] = instruction.to_bytes(4, sys.byteorder)
+            else:
+                remaining.append(hole)
+        holes = remaining
+        while len(self.text) % self.target.alignment:
+            self.text.append(0)
         for base, relocation in self.data_relocations:
             newhole = self._handle_relocation(base, relocation, self.data)
             if newhole.symbol in self.data_symbols:
@@ -261,8 +320,8 @@ async def parse(self) -> StencilGroup:
             if value_part and not symbol and not addend:
                 addend_part = ""
             else:
-                addend_part = f"&{symbol} + " if symbol else "",
-                addend_part = format_addend(addend)
+                addend_part = f"&{symbol} + " if symbol else ""
+                addend_part += format_addend(addend)
                 if value_part:
                     value_part += " + "
             disassembly_data.append(f"{offset_data:x}: {value_part}{addend_part}")
@@ -559,7 +618,6 @@ def _handle_relocation(
 @dataclasses.dataclass
 class Target:
     pattern: str
-    model: typing.Literal["small", "medium", "large"]
     pyconfig: pathlib.Path
     alignment: int
     prefix: str
@@ -569,47 +627,41 @@ class Target:
 TARGETS = [
     Target(
         pattern=r"aarch64-apple-darwin.*",
-        model="large",
         pyconfig=PYCONFIG_H,
         alignment=8,
         prefix="_",
         parser=MachO,
     ),
     Target(
         pattern=r"aarch64-.*-linux-gnu",
-        model="large",
         pyconfig=PYCONFIG_H,
         alignment=8,
         prefix="",
         parser=ELF,
     ),
     Target(
         pattern=r"i686-pc-windows-msvc",
-        model="small",
         pyconfig=PC_PYCONFIG_H,
         alignment=1,
         prefix="_",
         parser=COFF,
     ),
     Target(
         pattern=r"x86_64-apple-darwin.*",
-        model="medium",
         pyconfig=PYCONFIG_H,
         alignment=1,
         prefix="_",
         parser=MachO,
     ),
     Target(
         pattern=r"x86_64-pc-windows-msvc",
-        model="medium",
         pyconfig=PC_PYCONFIG_H,
         alignment=1,
         prefix="",
         parser=COFF,
     ),
     Target(
         pattern=r"x86_64-.*-linux-gnu",
-        model="medium",
         pyconfig=PYCONFIG_H,
         alignment=1,
         prefix="",
@@ -677,12 +729,7 @@ async def _compile(
             # - "medium": assumes that code resides in the lowest 2GB of memory,
             #   and makes no assumptions about data (not available on aarch64)
             # - "large": makes no assumptions about either code or data
-            # We need 64-bit addresses for data everywhere, but we'd *really*
-            # prefer direct short jumps instead of indirect long ones where
-            # possible. So, we use the "large" code model on aarch64 and the
-            # "medium" code model elsewhere, which gives us correctly-sized
-            # direct jumps and immediate data loads on basically all platforms:
-            f"-mcmodel={self._target.model}",
+            f"-mcmodel=large",
         ]
         await run(self._clang, *flags, "-o", o, c)
         self._stencils_built[opname] = await self._target.parser(

diff --git a/Tools/jit/trampoline.c b/Tools/jit/trampoline.c