Add AArch64 stub

Author: brandtbucher

Tools/jit/_optimizers.py

@@ -4,69 +4,34 @@
 import re
 import typing
 
-branches = {}
-for op, nop in [
-    ("ja", "jna"),
-    ("jae", "jnae"),
-    ("jb", "jnb"),
-    ("jbe", "jnbe"),
-    ("jc", "jnc"),
-    ("jcxz", None),
-    ("je", "jne"),
-    ("jecxz", None),
-    ("jg", "jng"),
-    ("jge", "jnge"),
-    ("jl", "jnl"),
-    ("jle", "jnle"),
-    ("jo", "jno"),
-    ("jp", "jnp"),
-    ("jpe", "jpo"),
-    ("jrxz", None),
-    ("js", "jns"),
-    ("jz", "jnz"),
-    ("loop", None),
-    ("loope", None),
-    ("loopne", None),
-    ("loopnz", None),
-    ("loopz", None),
-]:
-    branches[op] = nop
-    if nop:
-        branches[nop] = op
-
-
-def _get_branch(line: str) -> str | None:
-    branch = re.match(rf"\s*({'|'.join(branches)})\s+([\w\.]+)", line)
-    return branch and branch[2]
-
-
-def _invert_branch(line: str, label: str) -> str | None:
-    branch = re.match(rf"\s*({'|'.join(branches)})\s+([\w\.]+)", line)
-    assert branch
-    inverted = branches.get(branch[1])
-    if not inverted:
-        return None
-    line = line.replace(branch[1], inverted, 1)  # XXX
-    line = line.replace(branch[2], label, 1)  # XXX
-    return line
-
-
-def _get_jump(line: str) -> str | None:
-    jump = re.match(r"\s*(?:rex64\s+)?jmpq?\s+\*?([\w\.]+)", line)
-    return jump and jump[1]
-
-
-def _get_label(line: str) -> str | None:
-    label = re.match(r"\s*([\w\.]+):", line)
-    return label and label[1]
-
-
-def _is_return(line: str) -> bool:
-    return re.match(r"\s*ret\s+", line) is not None
-
-
-def _is_noise(line: str) -> bool:
-    return re.match(r"\s*([#\.]|$)", line) is not None
+_RE_NEVER_MATCH = re.compile(r"(?!)")
+
+_X86_BRANCHES = {
+    "ja": "jna",
+    "jae": "jnae",
+    "jb": "jnb",
+    "jbe": "jnbe",
+    "jc": "jnc",
+    "jcxz": None,
+    "je": "jne",
+    "jecxz": None,
+    "jg": "jng",
+    "jge": "jnge",
+    "jl": "jnl",
+    "jle": "jnle",
+    "jo": "jno",
+    "jp": "jnp",
+    "jpe": "jpo",
+    "jrxz": None,
+    "js": "jns",
+    "jz": "jnz",
+    "loop": None,
+    "loope": None,
+    "loopne": None,
+    "loopnz": None,
+    "loopz": None,
+}
+_X86_BRANCHES |= {v: k for k, v in _X86_BRANCHES.items() if v}
 
 
 @dataclasses.dataclass
@@ -91,57 +56,111 @@ def resolve(self) -> typing.Self:
         return self
 
 
-class _Labels(dict):
-    def __missing__(self, key: str) -> _Block:
-        self[key] = _Block(key)
-        return self[key]
-
-
 @dataclasses.dataclass
 class Optimizer:
 
     path: pathlib.Path
     _: dataclasses.KW_ONLY
     prefix: str = ""
     _graph: _Block = dataclasses.field(init=False)
-    _labels: _Labels = dataclasses.field(init=False, default_factory=_Labels)
-
-    _re_branch: typing.ClassVar[re.Pattern[str]]  # Two groups: instruction and target.
-    _re_jump: typing.ClassVar[re.Pattern[str]]  # One group: target.
-    _re_return: typing.ClassVar[re.Pattern[str]]  # No groups.
+    _labels: dict = dataclasses.field(init=False, default_factory=dict)
+    _alignment: typing.ClassVar[int] = 1
+    _branches: typing.ClassVar[dict[str, str | None]] = {}
+    _re_branch: typing.ClassVar[re.Pattern[str]] = (
+        _RE_NEVER_MATCH  # Two groups: instruction and target.
+    )
+    _re_jump: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH  # One group: target.
+    _re_label: typing.ClassVar[re.Pattern[str]] = re.compile(
+        r"\s*(?P<label>[\w\.]+):"
+    )  # One group: label.
+    _re_noise: typing.ClassVar[re.Pattern[str]] = re.compile(
+        r"\s*(?:[#\.]|$)"
+    )  # No groups.
+    _re_return: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH  # No groups.
 
     def __post_init__(self) -> None:
         text = self._preprocess(self.path.read_text())
         self._graph = block = self._new_block()
         for line in text.splitlines():
-            if label := _get_label(line):
-                block.link = block = self._labels[label]
+            if label := self._parse_label(line):
+                block.link = block = self._lookup_label(label)
+                block.noise.append(line)
+                continue
             elif block.target or not block.fallthrough:
                 block.link = block = self._new_block()
-            if _is_noise(line) or _get_label(line):
+            if self._parse_noise(line):
                 if block.instructions:
                     block.link = block = self._new_block()
                 block.noise.append(line)
                 continue
             block.instructions.append(line)
-            if target := _get_branch(line):
-                block.target = self._labels[target]
+            if target := self._parse_branch(line):
+                block.target = self._lookup_label(target)
                 assert block.fallthrough
-            elif target := _get_jump(line):
-                block.target = self._labels[target]
+            elif target := self._parse_jump(line):
+                block.target = self._lookup_label(target)
                 block.fallthrough = False
-            elif _is_return(line):
+            elif self._parse_return(line):
                 assert not block.target
                 block.fallthrough = False
 
-    def _new_block(self, label: str | None = None) -> _Block:
-        if not label:
-            label = f"{self.prefix}_JIT_LABEL_{len(self._labels)}"
-        assert label not in self._labels, label
-        block = self._labels[label] = _Block(label, [f"{label}:"])
+    @classmethod
+    def _parse_branch(cls, line: str) -> str | None:
+        branch = cls._re_branch.match(line)
+        return branch and branch["target"]
+
+    @classmethod
+    def _parse_jump(cls, line: str) -> str | None:
+        jump = cls._re_jump.match(line)
+        return jump and jump["target"]
+
+    @classmethod
+    def _parse_label(cls, line: str) -> str | None:
+        label = cls._re_label.match(line)
+        return label and label["label"]
+
+    @classmethod
+    def _parse_noise(cls, line: str) -> bool:
+        return cls._re_noise.match(line) is not None
+
+    @classmethod
+    def _parse_return(cls, line: str) -> bool:
+        return cls._re_return.match(line) is not None
+
+    @classmethod
+    def _invert_branch(cls, line: str, target: str) -> str | None:
+        branch = cls._re_branch.match(line)
+        assert branch
+        inverted = cls._branches.get(branch["instruction"])
+        if not inverted:
+            return None
+        (a, b), (c, d) = branch.span("instruction"), branch.span("target")
+        return "".join([line[:a], inverted, line[b:c], target, line[d:]])
+
+    @classmethod
+    def _thread_jump(cls, line: str, target: str) -> str:
+        jump = cls._re_jump.match(line) or cls._re_branch.match(line)
+        assert jump
+        a, b = jump.span("target")
+        return "".join([line[:a], target, line[b:]])
+
+    @staticmethod
+    def _create_jump(target: str) -> str | None:
+        return None
+
+    def _lookup_label(self, label: str) -> _Block:
+        if label not in self._labels:
+            self._labels[label] = _Block(label)
+        return self._labels[label]
+
+    def _new_block(self) -> _Block:
+        label = f"{self.prefix}_JIT_LABEL_{len(self._labels)}"
+        block = self._lookup_label(label)
+        block.noise.append(f"{label}:")
         return block
 
-    def _preprocess(self, text: str) -> str:
+    @staticmethod
+    def _preprocess(text: str) -> str:
         return text
 
     def _blocks(self) -> typing.Generator[_Block, None, None]:
@@ -159,9 +178,11 @@ def _insert_continue_label(self) -> None:
         for end in reversed(list(self._blocks())):
             if end.instructions:
                 break
-        continuation = self._labels[f"{self.prefix}_JIT_CONTINUE"]
+        align = self._new_block()
+        align.noise.append(f"\t.balign\t{self._alignment}")
+        continuation = self._lookup_label(f"{self.prefix}_JIT_CONTINUE")
         continuation.noise.append(f"{continuation.label}:")
-        end.link, continuation.link = continuation, end.link
+        end.link, align.link, continuation.link = align, continuation, end.link
 
     def _mark_hot_blocks(self) -> None:
         predecessors = collections.defaultdict(set)
@@ -170,7 +191,7 @@ def _mark_hot_blocks(self) -> None:
                 predecessors[block.target].add(block)
             if block.fallthrough and block.link:
                 predecessors[block.link].add(block)
-        todo = [self._labels[f"{self.prefix}_JIT_CONTINUE"]]
+        todo = [self._lookup_label(f"{self.prefix}_JIT_CONTINUE")]
         while todo:
             block = todo.pop()
             block.hot = True
@@ -185,43 +206,41 @@ def _invert_hot_branches(self) -> None:
             if (
                 block.fallthrough
                 and block.target
-                and block.link
                 and block.target.hot
+                and block.link
                 and not block.link.hot
             ):
-                # Turn...
-                #         branch hot
-                # ...into..
-                #         opposite-branch ._JIT_LABEL_N
-                #         jmp hot
-                #     ._JIT_LABEL_N:
                 label_block = self._new_block()
-                inverted = _invert_branch(block.instructions[-1], label_block.label)
-                if inverted is None:
+                branch = block.instructions[-1]
+                inverted = self._invert_branch(branch, label_block.label)
+                jump = self._create_jump(block.target.label)
+                if inverted is None or jump is None:
                     continue
                 jump_block = self._new_block()
-                jump_block.instructions.append(f"\tjmp\t{block.target.label}")
+                jump_block.instructions.append(jump)
                 jump_block.target = block.target
                 jump_block.fallthrough = False
                 block.instructions[-1] = inverted
                 block.target = label_block
-                label_block.link = block.link
-                jump_block.link = label_block
-                block.link = jump_block
+                block.link, jump_block.link, label_block.link = (
+                    jump_block,
+                    label_block,
+                    block.link,
+                )
 
     def _thread_jumps(self) -> None:
         for block in self._blocks():
             while block.target:
-                label = block.target.label
                 target = block.target.resolve()
                 if (
                     not target.fallthrough
                     and target.target
                     and len(target.instructions) == 1
                 ):
-                    block.instructions[-1] = block.instructions[-1].replace(
-                        label, target.target.label
-                    )  # XXX
+                    jump = block.instructions[-1]
+                    block.instructions[-1] = self._thread_jump(
+                        jump, target.target.label
+                    )
                     block.target = target.target
                 else:
                     break
@@ -273,14 +292,24 @@ def run(self) -> None:
         self.path.write_text("\n".join(self._lines()))
 
 
+class OptimizerAArch64(Optimizer):
+    # TODO: @diegorusso
+    _alignment = 8
+
+
 class OptimizerX86(Optimizer):
 
+    _branches = _X86_BRANCHES
     _re_branch = re.compile(
-        rf"\s*(?P<instruction>{'|'.join(branches)})\s+(?P<target>[\w\.]+)"
+        rf"\s*(?P<instruction>{'|'.join(_X86_BRANCHES)})\s+(?P<target>[\w\.]+)"
     )
     _re_jump = re.compile(r"\s*jmp\s+(?P<target>[\w\.]+)")
     _re_return = re.compile(r"\s*ret\b")
 
+    @staticmethod
+    def _create_jump(target: str) -> str:
+        return f"\tjmp\t{target}"
+
 
 class OptimizerX86Windows(OptimizerX86):
 

Tools/jit/_stencils.py

@@ -206,23 +206,6 @@ def pad(self, alignment: int) -> None:
             self.disassembly.append(f"{offset:x}: {' '.join(['00'] * padding)}")
         self.body.extend([0] * padding)
 
-    # def add_nops(self, nop: bytes, alignment: int) -> None:
-    #     """Add NOPs until there is alignment. Fail if it is not possible."""
-    #     offset = len(self.body)
-    #     nop_size = len(nop)
-
-    #     # Calculate the gap to the next multiple of alignment.
-    #     gap = -offset % alignment
-    #     if gap:
-    #         if gap % nop_size == 0:
-    #             count = gap // nop_size
-    #             self.body.extend(nop * count)
-    #         else:
-    #             raise ValueError(
-    #                 f"Cannot add nops of size '{nop_size}' to a body with "
-    #                 f"offset '{offset}' to align with '{alignment}'"
-    #             )
-
 
 @dataclasses.dataclass
 class StencilGroup:
@@ -240,9 +223,7 @@ class StencilGroup:
     _got: dict[str, int] = dataclasses.field(default_factory=dict, init=False)
     _trampolines: set[int] = dataclasses.field(default_factory=set, init=False)
 
-    def process_relocations(
-        self, known_symbols: dict[str, int], *, alignment: int = 1, nop: bytes = b""
-    ) -> None:
+    def process_relocations(self, known_symbols: dict[str, int]) -> None:
         """Fix up all GOT and internal relocations for this stencil group."""
         for hole in self.code.holes.copy():
             if (
@@ -262,7 +243,6 @@ def process_relocations(
                 self._trampolines.add(ordinal)
                 hole.addend = ordinal
                 hole.symbol = None
-        # self.code.add_nops(nop=nop, alignment=alignment)
         self.data.pad(8)
         for stencil in [self.code, self.data]:
             for hole in stencil.holes: