Rework AArch64 alignment

Author: brandtbucher

Python/jit.c

@@ -431,8 +431,10 @@ void patch_aarch64_trampoline(unsigned char *location, int ordinal, jit_state *s
 
 #if defined(__aarch64__) || defined(_M_ARM64)
     #define TRAMPOLINE_SIZE 16
+    #define DATA_ALIGN 8
 #else
     #define TRAMPOLINE_SIZE 0
+    #define DATA_ALIGN 0
 #endif
 
 // Generate and patch AArch64 trampolines. The symbols to jump to are stored
@@ -522,7 +524,8 @@ _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 padding = page_size - ((code_size + state.trampolines.size + data_size) & (page_size - 1));
+    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) & (page_size - 1));
     size_t total_size = code_size + state.trampolines.size + data_size  + padding;
     unsigned char *memory = jit_alloc(total_size);
     if (memory == NULL) {
@@ -545,7 +548,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     // Loop again to emit the code:
     unsigned char *code = memory;
     state.trampolines.mem = memory + code_size;
-    unsigned char *data = memory + code_size + state.trampolines.size;
+    unsigned char *data = memory + code_size + state.trampolines.size + code_padding;
     // 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).
@@ -567,7 +570,7 @@ _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction trace[], siz
     code += group->code_size;
     data += group->data_size;
     assert(code == memory + code_size);
-    assert(data == memory + code_size + state.trampolines.size + data_size);
+    assert(data == memory + code_size + state.trampolines.size + code_padding + data_size);
 #ifdef MAP_JIT
     pthread_jit_write_protect_np(1);
 #endif

Tools/jit/_optimizers.py

@@ -84,7 +84,6 @@ class Optimizer:
         r'\s*(?P<label>[\w."$?@]+):'
     )
     # Override everything that follows in subclasses:
-    _alignment: typing.ClassVar[int] = 1
     _branches: typing.ClassVar[dict[str, str | None]] = {}
     # Two groups (instruction and target):
     _re_branch: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH
@@ -197,15 +196,12 @@ def _insert_continue_label(self) -> None:
         #    jmp FOO
         # After:
         #    jmp FOO
-        #    .balign 8
         #    _JIT_CONTINUE:
         # This lets the assembler encode _JIT_CONTINUE jumps at build time!
-        align = _Block()
-        align.noninstructions.append(f"\t.balign\t{self._alignment}")
         continuation = self._lookup_label(f"{self.prefix}_JIT_CONTINUE")
         assert continuation.label
         continuation.noninstructions.append(f"{continuation.label}:")
-        end.link, align.link, continuation.link = align, continuation, end.link
+        end.link, continuation.link = continuation, end.link
 
     def _mark_hot_blocks(self) -> None:
         # Start with the last block, and perform a DFS to find all blocks that
@@ -285,8 +281,6 @@ def run(self) -> None:
 class OptimizerAArch64(Optimizer):  # pylint: disable = too-few-public-methods
     """aarch64-apple-darwin/aarch64-pc-windows-msvc/aarch64-unknown-linux-gnu"""
 
-    # TODO: @diegorusso
-    _alignment = 8
     # https://developer.arm.com/documentation/ddi0602/2025-03/Base-Instructions/B--Branch-
     _re_jump = re.compile(r"\s*b\s+(?P<target>[\w.]+)")