Do some textual assembly magic

brandtbucher · Jun 27, 2025 · Jun 12, 2025 · Jun 12, 2025 · Jun 12, 2025 · Jun 19, 2025
commit 560607812df24abd2e2a48a1911bed2c61834686
@@ -17,8 +17,6 @@ class HoleValue(enum.Enum):
 
     # The base address of the machine code for the current uop (exposed as _JIT_ENTRY):
     CODE = enum.auto()
-    # The base address of the machine code for the next uop (exposed as _JIT_CONTINUE):
-    CONTINUE = enum.auto()
     # The base address of the read-only data for this uop:
     DATA = enum.auto()
     # The address of the current executor (exposed as _JIT_EXECUTOR):
@@ -97,7 +95,6 @@ class HoleValue(enum.Enum):
 # 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:
@@ -209,63 +206,22 @@ 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}'"
-                )
-
-    def remove_jump(self) -> None:
-        """Remove a zero-length continuation jump, if it exists."""
-        hole = max(self.holes, key=lambda hole: hole.offset)
-        match hole:
-            case Hole(
-                offset=offset,
-                kind="IMAGE_REL_AMD64_REL32",
-                value=HoleValue.GOT,
-                symbol="_JIT_CONTINUE",
-                addend=-4,
-            ) as hole:
-                # jmp qword ptr [rip]
-                jump = b"\x48\xff\x25\x00\x00\x00\x00"
-                offset -= 3
-            case Hole(
-                offset=offset,
-                kind="IMAGE_REL_I386_REL32" | "R_X86_64_PLT32" | "X86_64_RELOC_BRANCH",
-                value=HoleValue.CONTINUE,
-                symbol=None,
-                addend=addend,
-            ) as hole if (
-                _signed(addend) == -4
-            ):
-                # jmp 5
-                jump = b"\xe9\x00\x00\x00\x00"
-                offset -= 1
-            case Hole(
-                offset=offset,
-                kind="R_AARCH64_JUMP26",
-                value=HoleValue.CONTINUE,
-                symbol=None,
-                addend=0,
-            ) as hole:
-                # b #4
-                jump = b"\x00\x00\x00\x14"
-            case _:
-                return
-        if self.body[offset:] == jump:
-            self.body = self.body[:offset]
-            self.holes.remove(hole)
+    # 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
@@ -306,8 +262,7 @@ def process_relocations(
                 self._trampolines.add(ordinal)
                 hole.addend = ordinal
                 hole.symbol = None
-        self.code.remove_jump()
-        self.code.add_nops(nop=nop, alignment=alignment)
+        # self.code.add_nops(nop=nop, alignment=alignment)
         self.data.pad(8)
         for stencil in [self.code, self.data]:
             for hole in stencil.holes:

@@ -3,6 +3,7 @@
 import asyncio
 import dataclasses
 import hashlib
+import itertools
 import json
 import os
 import pathlib
@@ -34,6 +35,223 @@
     "_R", _schema.COFFRelocation, _schema.ELFRelocation, _schema.MachORelocation
 )
 
+inverted_branches = {}
+for op, nop in [
+    ("ja", "jna"),
+    ("jae", "jnae"),
+    ("jb", "jnb"),
+    ("jbe", "jnbe"),
+    ("jc", "jnc"),
+    ("je", "jne"),
+    ("jg", "jng"),
+    ("jge", "jnge"),
+    ("jl", "jnl"),
+    ("jle", "jnle"),
+    ("jo", "jno"),
+    ("jp", "jnp"),
+    ("js", "jns"),
+    ("jz", "jnz"),
+    ("jpe", "jpo"),
+    ("jcxz", None),
+    ("jecxz", None),
+    ("jrxz", None),
+    ("loop", None),
+    ("loope", None),
+    ("loopne", None),
+    ("loopnz", None),
+    ("loopz", None),
+]:
+    inverted_branches[op] = nop
+    if nop is not None:
+        inverted_branches[nop] = op
+
+
+@dataclasses.dataclass
+class _Line:
+    text: str
+    hot: bool = dataclasses.field(init=False, default=False)
+    predecessors: list["_Line"] = dataclasses.field(
+        init=False, repr=False, default_factory=list
+    )
+
+
+@dataclasses.dataclass
+class _Label(_Line):
+    label: str
+
+
+@dataclasses.dataclass
+class _Jump(_Line):
+    target: _Label | None
+
+    def remove(self) -> None:
+        self.text = ""
+        if self.target is not None:
+            self.target.predecessors.remove(self)
+            if not self.target.predecessors:
+                self.target.text = ""
+            self.target = None
+
+    def update(self, target: _Label) -> None:
+        assert self.target is not None
+        self.target.predecessors.remove(self)
+        assert self.target.label in self.text
+        self.text = self.text.replace(self.target.label, target.label)
+        self.target = target
+        self.target.predecessors.append(self)
+
+
+@dataclasses.dataclass
+class _Branch(_Line):
+    op: str
+    target: _Label
+    fallthrough: _Line | None = None
+
+    def update(self, target: _Label) -> None:
+        assert self.target is not None
+        self.target.predecessors.remove(self)
+        assert self.target.label in self.text
+        self.text = self.text.replace(self.target.label, target.label)
+        self.target = target
+        self.target.predecessors.append(self)
+
+    def invert(self, jump: _Jump) -> bool:
+        inverted = inverted_branches[self.op]
+        if inverted is None or jump.target is None:
+            return False
+        assert self.op in self.text
+        self.text = self.text.replace(self.op, inverted)
+        old_target = self.target
+        self.update(jump.target)
+        jump.update(old_target)
+        return True
+
+
+@dataclasses.dataclass
+class _Return(_Line):
+    pass
+
+
+@dataclasses.dataclass
+class _Noise(_Line):
+    pass
+
+
+def _branch(
+    line: str, use_label: typing.Callable[[str], _Label]
+) -> tuple[str, _Label] | None:
+    branch = re.match(rf"\s*({'|'.join(inverted_branches)})\s+([\w\.]+)", line)
+    return branch and (branch.group(1), use_label(branch.group(2)))
+
+
+def _jump(line: str, use_label: typing.Callable[[str], _Label]) -> _Label | None:
+    if jump := re.match(r"\s*jmp\s+([\w\.]+)", line):
+        return use_label(jump.group(1))
+    return None
+
+
+def _label(line: str, use_label: typing.Callable[[str], _Label]) -> _Label | None:
+    label = re.match(r"\s*([\w\.]+):", line)
+    return label and use_label(label.group(1))
+
+
+def _return(line: str) -> bool:
+    return re.match(r"\s*ret\s+", line) is not None
+
+
+def _noise(line: str) -> bool:
+    return re.match(r"\s*[#\.]|\s*$", line) is not None
+
+
+def _apply_asm_transformations(path: pathlib.Path) -> None:
+    labels = {}
+
+    def use_label(label: str) -> _Label:
+        if label not in labels:
+            labels[label] = _Label("", label)
+        return labels[label]
+
+    def new_line(text: str) -> _Line:
+        if branch := _branch(text, use_label):
+            op, label = branch
+            line = _Branch(text, op, label)
+            label.predecessors.append(line)
+            return line
+        if label := _jump(text, use_label):
+            line = _Jump(text, label)
+            label.predecessors.append(line)
+            return line
+        if line := _label(text, use_label):
+            assert line.text == ""
+            line.text = text
+            return line
+        if _return(text):
+            return _Return(text)
+        if _noise(text):
+            return _Noise(text)
+        return _Line(text)
+
+    # Build graph:
+    lines = []
+    line = _Noise("")  # Dummy.
+    with path.open() as file:
+        for i, text in enumerate(file):
+            new = new_line(text)
+            if not isinstance(line, (_Jump, _Return)):
+                new.predecessors.append(line)
+            lines.append(new)
+            line = new
+    for i, line in enumerate(reversed(lines)):
+        if not isinstance(line, (_Label, _Noise)):
+            break
+    new = new_line("_JIT_CONTINUE:\n")
+    lines.insert(len(lines) - i, new)
+    line = new
+    # Mark hot lines:
+    todo = labels["_JIT_CONTINUE"].predecessors.copy()
+    while todo:
+        line = todo.pop()
+        line.hot = True
+        for predecessor in line.predecessors:
+            if not predecessor.hot:
+                todo.append(predecessor)
+    for pair in itertools.pairwise(
+        filter(lambda line: line.text and not isinstance(line, _Noise), lines)
+    ):
+        match pair:
+            case (_Branch(hot=True) as branch, _Jump(hot=False) as jump):
+                branch.invert(jump)
+                jump.hot = True
+    for pair in itertools.pairwise(lines):
+        match pair:
+            case (_Jump() | _Return(), _):
+                pass
+            case (_Line(hot=True), _Line(hot=False) as cold):
+                cold.hot = True
+    # Reorder blocks:
+    hot = []
+    cold = []
+    for line in lines:
+        if line.hot:
+            hot.append(line)
+        else:
+            cold.append(line)
+    lines = hot + cold
+    # Remove zero-length jumps:
+    again = True
+    while again:
+        again = False
+        for pair in itertools.pairwise(
+            filter(lambda line: line.text and not isinstance(line, _Noise), lines)
+        ):
+            match pair:
+                case (_Jump(target=target) as jump, label) if target is label:
+                    jump.remove()
+                    again = True
+    # Write new assembly:
+    with path.open("w") as file:
+        file.writelines(line.text for line in lines)
+
 
 @dataclasses.dataclass
 class _Target(typing.Generic[_S, _R]):
@@ -118,8 +336,9 @@ def _handle_relocation(
     async def _compile(
         self, opname: str, c: pathlib.Path, tempdir: pathlib.Path
     ) -> _stencils.StencilGroup:
+        s = tempdir / f"{opname}.s"
         o = tempdir / f"{opname}.o"
-        args = [
+        args_s = [
             f"--target={self.triple}",
             "-DPy_BUILD_CORE_MODULE",
             "-D_DEBUG" if self.debug else "-DNDEBUG",
@@ -133,7 +352,8 @@ async def _compile(
             f"-I{CPYTHON / 'Python'}",
             f"-I{CPYTHON / 'Tools' / 'jit'}",
             "-O3",
-            "-c",
+            "-S",
+            # "-c",
             # Shorten full absolute file paths in the generated code (like the
             # __FILE__ macro and assert failure messages) for reproducibility:
             f"-ffile-prefix-map={CPYTHON}=.",
@@ -152,11 +372,20 @@ async def _compile(
             "-fno-stack-protector",
             "-std=c11",
             "-o",
-            f"{o}",
+            f"{s}",
             f"{c}",
             *self.args,
         ]
-        await _llvm.run("clang", args, echo=self.verbose)
+        await _llvm.run("clang", args_s, echo=self.verbose)
+        _apply_asm_transformations(s)
+        args_o = [
+            f"--target={self.triple}",
+            "-c",
+            "-o",
+            f"{o}",
+            f"{s}",
+        ]
+        await _llvm.run("clang", args_o, echo=self.verbose)
         return await self._parse(o)
 
     async def _build_stencils(self) -> dict[str, _stencils.StencilGroup]: