Implement assembler optimization for AArch64.

Removes a number of unnecessary branches.

Change-Id: I4965fe8d8b79f7d859ff9076d9c53f3ac7f094b2

Author: markshannon

Python/jit.c

@@ -167,11 +167,13 @@ set_bits(uint32_t *loc, uint8_t loc_start, uint64_t value, uint8_t value_start,
 
 // See https://developer.arm.com/documentation/ddi0602/2023-09/Base-Instructions
 // for instruction encodings:
-#define IS_AARCH64_ADD_OR_SUB(I) (((I) & 0x11C00000) == 0x11000000)
-#define IS_AARCH64_ADRP(I)       (((I) & 0x9F000000) == 0x90000000)
-#define IS_AARCH64_BRANCH(I)     (((I) & 0x7C000000) == 0x14000000)
-#define IS_AARCH64_LDR_OR_STR(I) (((I) & 0x3B000000) == 0x39000000)
-#define IS_AARCH64_MOV(I)        (((I) & 0x9F800000) == 0x92800000)
+#define IS_AARCH64_ADD_OR_SUB(I)  (((I) & 0x11C00000) == 0x11000000)
+#define IS_AARCH64_ADRP(I)        (((I) & 0x9F000000) == 0x90000000)
+#define IS_AARCH64_BRANCH(I)      (((I) & 0x7C000000) == 0x14000000)
+#define IS_AARCH64_BRANCH_COND(I) (((I) & 0x7C000000) == 0x54000000)
+#define IS_AARCH64_TEST_AND_BRANCH(I) (((I) & 0x7E000000) == 0x36000000)
+#define IS_AARCH64_LDR_OR_STR(I)  (((I) & 0x3B000000) == 0x39000000)
+#define IS_AARCH64_MOV(I)         (((I) & 0x9F800000) == 0x92800000)
 
 // LLD is a great reference for performing relocations... just keep in
 // mind that Tools/jit/build.py does filtering and preprocessing for us!
@@ -332,6 +334,37 @@ patch_aarch64_21rx(unsigned char *location, uint64_t value)
     patch_aarch64_21r(location, value);
 }
 
+
+// 21-bit relative branch.
+void
+patch_aarch64_19r(unsigned char *location, uint64_t value)
+{
+    uint32_t *loc32 = (uint32_t *)location;
+    assert(IS_AARCH64_BRANCH_COND(*loc32));
+    value -= (uintptr_t)location;
+    // Check that we're not out of range of 21 signed bits:
+    assert((int64_t)value >= -(1 << 20));
+    assert((int64_t)value < (1 << 20));
+    // Since instructions are 4-byte aligned, only use 19 bits:
+    assert(get_bits(value, 0, 2) == 0);
+    set_bits(loc32, 5, value, 2, 19);
+}
+
+// 16-bit relative branch.
+void
+patch_aarch64_14r(unsigned char *location, uint64_t value)
+{
+    uint32_t *loc32 = (uint32_t *)location;
+    assert(IS_AARCH64_TEST_AND_BRANCH(*loc32));
+    value -= (uintptr_t)location;
+    // Check that we're not out of range of 16 signed bits:
+    assert((int64_t)value >= -(1 << 15));
+    assert((int64_t)value < (1 << 15));
+    // Since instructions are 4-byte aligned, only use 14 bits:
+    assert(get_bits(value, 0, 2) == 0);
+    set_bits(loc32, 5, value, 2, 14);
+}
+
 // 28-bit relative branch.
 void
 patch_aarch64_26r(unsigned char *location, uint64_t value)

Tools/jit/_optimizers.py

@@ -39,6 +39,41 @@
 # Update with all of the inverted branches, too:
 _X86_BRANCHES |= {v: k for k, v in _X86_BRANCHES.items() if v}
 
+_AARCH64_COND_CODES = {
+    # https://developer.arm.com/documentation/dui0801/b/CJAJIHAD?lang=en
+    "eq": "ne",
+    "ne": "eq",
+    "lt": "ge",
+    "ge": "lt",
+    "gt": "le",
+    "le": "gt",
+    "vs": "vc",
+    "vc": "vs",
+    "mi": "pl",
+    "pl": "mi",
+    "cs": "cc",
+    "cc": "cs",
+    "hs": "lo",
+    "lo": "hs",
+    "hi": "ls",
+    "ls": "hi",
+}
+# Branches are either b.{cond} or bc.{cond}
+_AARCH64_BRANCHES = {
+    "b." + cond: ("b." + inverse if inverse else None)
+    for (cond, inverse) in _AARCH64_COND_CODES.items()
+} | {
+    "bc." + cond: ("bc." + inverse if inverse else None)
+    for (cond, inverse) in _AARCH64_COND_CODES.items()
+}
+# And four compare and branch instructions
+_AARCH64_BRANCHES |= {
+    "cbz": "cbnz",
+    "cbnz": "cbz",
+    "tbz": "tbnz",
+    "tbnz": "tbz",
+}
+
 
 @dataclasses.dataclass
 class _Block:
@@ -286,8 +321,15 @@ def run(self) -> None:
 class OptimizerAArch64(Optimizer):  # pylint: disable = too-few-public-methods
     """aarch64-apple-darwin/aarch64-pc-windows-msvc/aarch64-unknown-linux-gnu"""
 
+    _branches = _AARCH64_BRANCHES
+    _re_branch = re.compile(
+        rf"\s*(?P<instruction>{'|'.join(_AARCH64_BRANCHES)})\s+(.+,\s+)*(?P<target>[\w.]+)"
+    )
+
     # https://developer.arm.com/documentation/ddi0602/2025-03/Base-Instructions/B--Branch-
     _re_jump = re.compile(r"\s*b\s+(?P<target>[\w.]+)")
+    # https://developer.arm.com/documentation/ddi0602/2025-09/Base-Instructions/RET--Return-from-subroutine-
+    _re_return = re.compile(r"\s*ret\b")
 
 
 class OptimizerX86(Optimizer):  # pylint: disable = too-few-public-methods

Tools/jit/_schema.py

@@ -20,13 +20,15 @@
     "R_AARCH64_ADR_GOT_PAGE",
     "R_AARCH64_ADR_PREL_PG_HI21",
     "R_AARCH64_CALL26",
+    "R_AARCH64_CONDBR19",
     "R_AARCH64_JUMP26",
     "R_AARCH64_ADD_ABS_LO12_NC",
     "R_AARCH64_LD64_GOT_LO12_NC",
     "R_AARCH64_MOVW_UABS_G0_NC",
     "R_AARCH64_MOVW_UABS_G1_NC",
     "R_AARCH64_MOVW_UABS_G2_NC",
     "R_AARCH64_MOVW_UABS_G3",
+    "R_AARCH64_TSTBR14",
     "R_X86_64_64",
     "R_X86_64_GOTPCREL",
     "R_X86_64_GOTPCRELX",

Tools/jit/_stencils.py

@@ -74,12 +74,14 @@ class HoleValue(enum.Enum):
     "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_CONDBR19": "patch_aarch64_19r",
     "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",
+    "R_AARCH64_TSTBR14": "patch_aarch64_14r",
     # x86_64-unknown-linux-gnu:
     "R_X86_64_64": "patch_64",
     "R_X86_64_GOTPCRELX": "patch_x86_64_32rx",