| """Low-level optimization of textual assembly.""" |
| |
| import dataclasses |
| import enum |
| import pathlib |
| import re |
| import typing |
| |
| # Same as saying "not string.startswith('')": |
| _RE_NEVER_MATCH = re.compile(r"(?!)") |
| # Dictionary mapping branch instructions to their inverted branch instructions. |
| # If a branch cannot be inverted, the value is None: |
| _X86_BRANCH_NAMES = { |
| # https://www.felixcloutier.com/x86/jcc |
| "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", |
| "jrcxz": None, |
| "js": "jns", |
| "jz": "jnz", |
| # https://www.felixcloutier.com/x86/loop:loopcc |
| "loop": None, |
| "loope": None, |
| "loopne": None, |
| "loopnz": None, |
| "loopz": None, |
| } |
| # Update with all of the inverted branches, too: |
| _X86_BRANCH_NAMES |= {v: k for k, v in _X86_BRANCH_NAMES.items() if v} |
| # No custom relocations needed |
| _X86_BRANCHES: dict[str, tuple[str | None, str | None]] = { |
| k: (v, None) for k, v in _X86_BRANCH_NAMES.items() |
| } |
| |
| _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", |
| } |
| # MyPy doesn't understand that a invariant variable can be initialized by a covariant value |
| CUSTOM_AARCH64_BRANCH19: str | None = "CUSTOM_AARCH64_BRANCH19" |
| |
| _AARCH64_SHORT_BRANCHES = { |
| "tbz": "tbnz", |
| "tbnz": "tbz", |
| } |
| |
| # Branches are either b.{cond}, bc.{cond}, cbz, cbnz, tbz or tbnz |
| _AARCH64_BRANCHES: dict[str, tuple[str | None, str | None]] = ( |
| { |
| "b." + cond: (("b." + inverse if inverse else None), CUSTOM_AARCH64_BRANCH19) |
| for (cond, inverse) in _AARCH64_COND_CODES.items() |
| } |
| | { |
| "bc." + cond: (("bc." + inverse if inverse else None), CUSTOM_AARCH64_BRANCH19) |
| for (cond, inverse) in _AARCH64_COND_CODES.items() |
| } |
| | { |
| "cbz": ("cbnz", CUSTOM_AARCH64_BRANCH19), |
| "cbnz": ("cbz", CUSTOM_AARCH64_BRANCH19), |
| } |
| | {cond: (inverse, None) for (cond, inverse) in _AARCH64_SHORT_BRANCHES.items()} |
| ) |
| |
| |
| @enum.unique |
| class InstructionKind(enum.Enum): |
| |
| JUMP = enum.auto() |
| LONG_BRANCH = enum.auto() |
| SHORT_BRANCH = enum.auto() |
| CALL = enum.auto() |
| RETURN = enum.auto() |
| SMALL_CONST_1 = enum.auto() |
| SMALL_CONST_2 = enum.auto() |
| SMALL_CONST_MASK = enum.auto() |
| LARGE_CONST_1 = enum.auto() |
| LARGE_CONST_2 = enum.auto() |
| OTHER = enum.auto() |
| |
| |
| @dataclasses.dataclass |
| class Instruction: |
| kind: InstructionKind |
| name: str |
| text: str |
| register: str | None |
| target: str | None |
| |
| def is_branch(self) -> bool: |
| return self.kind in (InstructionKind.LONG_BRANCH, InstructionKind.SHORT_BRANCH) |
| |
| def update_target(self, target: str) -> "Instruction": |
| assert self.target is not None |
| return Instruction( |
| self.kind, |
| self.name, |
| self.text.replace(self.target, target), |
| self.register, |
| target, |
| ) |
| |
| def update_name_and_target(self, name: str, target: str) -> "Instruction": |
| assert self.target is not None |
| return Instruction( |
| self.kind, |
| name, |
| self.text.replace(self.name, name).replace(self.target, target), |
| self.register, |
| target, |
| ) |
| |
| |
| @dataclasses.dataclass(eq=False) |
| class _Block: |
| label: str | None = None |
| # Non-instruction lines like labels, directives, and comments: |
| noninstructions: list[str] = dataclasses.field(default_factory=list) |
| # Instruction lines: |
| instructions: list[Instruction] = dataclasses.field(default_factory=list) |
| # If this block ends in a jump, where to? |
| target: typing.Self | None = None |
| # The next block in the linked list: |
| link: typing.Self | None = None |
| # Whether control flow can fall through to the linked block above: |
| fallthrough: bool = True |
| # Whether this block can eventually reach the next uop (_JIT_CONTINUE): |
| hot: bool = False |
| |
| def resolve(self) -> typing.Self: |
| """Find the first non-empty block reachable from this one.""" |
| block = self |
| while block.link and not block.instructions: |
| block = block.link |
| return block |
| |
| |
| @dataclasses.dataclass |
| class Optimizer: |
| """Several passes of analysis and optimization for textual assembly.""" |
| |
| path: pathlib.Path |
| _: dataclasses.KW_ONLY |
| # Prefixes used to mangle local labels and symbols: |
| label_prefix: str |
| symbol_prefix: str |
| re_global: re.Pattern[str] |
| frame_pointers: bool |
| # The first block in the linked list: |
| _root: _Block = dataclasses.field(init=False, default_factory=_Block) |
| _labels: dict[str, _Block] = dataclasses.field(init=False, default_factory=dict) |
| # No groups: |
| _re_noninstructions: typing.ClassVar[re.Pattern[str]] = re.compile( |
| r"\s*(?:\.|#|//|;|$)" |
| ) |
| # One group (label): |
| _re_label: typing.ClassVar[re.Pattern[str]] = re.compile( |
| r'\s*(?P<label>[\w."$?@]+):' |
| ) |
| # Override everything that follows in subclasses: |
| _supports_external_relocations = True |
| supports_small_constants = False |
| _branches: typing.ClassVar[dict[str, tuple[str | None, str | None]]] = {} |
| # Short branches are instructions that can branch within a micro-op, |
| # but might not have the reach to branch anywhere within a trace. |
| _short_branches: typing.ClassVar[dict[str, str]] = {} |
| # Two groups (instruction and target): |
| _re_branch: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH |
| # One group (target): |
| _re_call: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH |
| # One group (target): |
| _re_jump: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH |
| # No groups: |
| _re_return: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH |
| text: str = "" |
| globals: set[str] = dataclasses.field(default_factory=set) |
| _re_small_const_1 = _RE_NEVER_MATCH |
| _re_small_const_2 = _RE_NEVER_MATCH |
| _re_small_const_mask = _RE_NEVER_MATCH |
| _re_large_const_1 = _RE_NEVER_MATCH |
| _re_large_const_2 = _RE_NEVER_MATCH |
| const_reloc = "<Not supported>" |
| _frame_pointer_modify: typing.ClassVar[re.Pattern[str]] = _RE_NEVER_MATCH |
| label_index: int = 0 |
| |
| def __post_init__(self) -> None: |
| # Split the code into a linked list of basic blocks. A basic block is an |
| # optional label, followed by zero or more non-instruction lines, |
| # followed by zero or more instruction lines (only the last of which may |
| # be a branch, jump, or return): |
| self.text = self._preprocess(self.path.read_text()) |
| block = self._root |
| for line in self.text.splitlines(): |
| # See if we need to start a new block: |
| if match := self._re_label.match(line): |
| # Label. New block: |
| block.link = block = self._lookup_label(match["label"]) |
| block.noninstructions.append(line) |
| continue |
| if match := self.re_global.match(line): |
| self.globals.add(match["label"]) |
| block.noninstructions.append(line) |
| continue |
| if self._re_noninstructions.match(line): |
| if block.instructions: |
| # Non-instruction lines. New block: |
| block.link = block = _Block() |
| block.noninstructions.append(line) |
| continue |
| if block.target or not block.fallthrough: |
| # Current block ends with a branch, jump, or return. New block: |
| block.link = block = _Block() |
| inst = self._parse_instruction(line) |
| block.instructions.append(inst) |
| if inst.is_branch(): |
| # A block ending in a branch has a target and fallthrough: |
| assert inst.target is not None |
| block.target = self._lookup_label(inst.target) |
| assert block.fallthrough |
| elif inst.kind == InstructionKind.CALL: |
| # A block ending in a call has a target and return point after call: |
| assert inst.target is not None |
| block.target = self._lookup_label(inst.target) |
| assert block.fallthrough |
| elif inst.kind == InstructionKind.JUMP: |
| # A block ending in a jump has a target and no fallthrough: |
| assert inst.target is not None |
| block.target = self._lookup_label(inst.target) |
| block.fallthrough = False |
| elif inst.kind == InstructionKind.RETURN: |
| # A block ending in a return has no target and fallthrough: |
| assert not block.target |
| block.fallthrough = False |
| |
| def _preprocess(self, text: str) -> str: |
| # Override this method to do preprocessing of the textual assembly. |
| # In all cases, replace references to the _JIT_CONTINUE symbol with |
| # references to a local _JIT_CONTINUE label (which we will add later): |
| continue_symbol = rf"\b{re.escape(self.symbol_prefix)}_JIT_CONTINUE\b" |
| continue_label = f"{self.label_prefix}_JIT_CONTINUE" |
| return re.sub(continue_symbol, continue_label, text) |
| |
| def _parse_instruction(self, line: str) -> Instruction: |
| target = None |
| reg = None |
| if match := self._re_branch.match(line): |
| target = match["target"] |
| name = match["instruction"] |
| if name in self._short_branches: |
| kind = InstructionKind.SHORT_BRANCH |
| else: |
| kind = InstructionKind.LONG_BRANCH |
| elif match := self._re_jump.match(line): |
| target = match["target"] |
| name = line[: -len(target)].strip() |
| kind = InstructionKind.JUMP |
| elif match := self._re_call.match(line): |
| target = match["target"] |
| name = line[: -len(target)].strip() |
| kind = InstructionKind.CALL |
| elif match := self._re_return.match(line): |
| name = line |
| kind = InstructionKind.RETURN |
| elif match := self._re_small_const_1.match(line): |
| target = match["value"] |
| name = match["instruction"] |
| reg = match["register"] |
| kind = InstructionKind.SMALL_CONST_1 |
| elif match := self._re_small_const_2.match(line): |
| target = match["value"] |
| name = match["instruction"] |
| reg = match["register"] |
| kind = InstructionKind.SMALL_CONST_2 |
| elif match := self._re_small_const_mask.match(line): |
| target = match["value"] |
| name = match["instruction"] |
| reg = match["register"] |
| if reg.startswith("w"): |
| reg = "x" + reg[1:] |
| kind = InstructionKind.SMALL_CONST_MASK |
| elif match := self._re_large_const_1.match(line): |
| target = match["value"] |
| name = match["instruction"] |
| reg = match["register"] |
| kind = InstructionKind.LARGE_CONST_1 |
| elif match := self._re_large_const_2.match(line): |
| target = match["value"] |
| name = match["instruction"] |
| reg = match["register"] |
| kind = InstructionKind.LARGE_CONST_2 |
| else: |
| name, *_ = line.split(" ") |
| kind = InstructionKind.OTHER |
| return Instruction(kind, name, line, reg, target) |
| |
| def _invert_branch(self, inst: Instruction, target: str) -> Instruction | None: |
| assert inst.is_branch() |
| if inst.kind == InstructionKind.SHORT_BRANCH and self._is_far_target(target): |
| return None |
| inverted_reloc = self._branches.get(inst.name) |
| if inverted_reloc is None: |
| return None |
| inverted = inverted_reloc[0] |
| if not inverted: |
| return None |
| return inst.update_name_and_target(inverted, target) |
| |
| def _lookup_label(self, label: str) -> _Block: |
| if label not in self._labels: |
| self._labels[label] = _Block(label) |
| return self._labels[label] |
| |
| def _is_far_target(self, label: str) -> bool: |
| return not label.startswith(self.label_prefix) |
| |
| def _blocks(self) -> typing.Generator[_Block, None, None]: |
| block: _Block | None = self._root |
| while block: |
| yield block |
| block = block.link |
| |
| def _body(self) -> str: |
| lines = ["#" + line for line in self.text.splitlines()] |
| hot = True |
| for block in self._blocks(): |
| if hot != block.hot: |
| hot = block.hot |
| # Make it easy to tell at a glance where cold code is: |
| lines.append(f"# JIT: {'HOT' if hot else 'COLD'} ".ljust(80, "#")) |
| lines.extend(block.noninstructions) |
| for inst in block.instructions: |
| lines.append(inst.text) |
| return "\n".join(lines) |
| |
| def _predecessors(self, block: _Block) -> typing.Generator[_Block, None, None]: |
| # This is inefficient, but it's never wrong: |
| for pre in self._blocks(): |
| if pre.target is block or pre.fallthrough and pre.link is block: |
| yield pre |
| |
| def _insert_continue_label(self) -> None: |
| # Find the block with the last instruction: |
| for end in reversed(list(self._blocks())): |
| if end.instructions: |
| break |
| # Before: |
| # jmp FOO |
| # After: |
| # jmp FOO |
| # _JIT_CONTINUE: |
| # This lets the assembler encode _JIT_CONTINUE jumps at build time! |
| continuation = self._lookup_label(f"{self.label_prefix}_JIT_CONTINUE") |
| assert continuation.label |
| continuation.noninstructions.append(f"{continuation.label}:") |
| 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 |
| # can eventually reach it: |
| todo = list(self._blocks())[-1:] |
| while todo: |
| block = todo.pop() |
| block.hot = True |
| todo.extend(pre for pre in self._predecessors(block) if not pre.hot) |
| |
| def _invert_hot_branches(self) -> None: |
| for branch in self._blocks(): |
| link = branch.link |
| if link is None: |
| continue |
| jump = link.resolve() |
| # Before: |
| # je HOT |
| # jmp COLD |
| # After: |
| # jne COLD |
| # jmp HOT |
| if ( |
| # block ends with a branch to hot code... |
| branch.target |
| and branch.fallthrough |
| and branch.target.hot |
| # ...followed by a jump to cold code with no other predecessors: |
| and jump.target |
| and not jump.fallthrough |
| and not jump.target.hot |
| and len(jump.instructions) == 1 |
| and list(self._predecessors(jump)) == [branch] |
| ): |
| assert jump.target.label |
| assert branch.target.label |
| inverted = self._invert_branch( |
| branch.instructions[-1], jump.target.label |
| ) |
| # Check to see if the branch can even be inverted: |
| if inverted is None: |
| continue |
| branch.instructions[-1] = inverted |
| jump.instructions[-1] = jump.instructions[-1].update_target( |
| branch.target.label |
| ) |
| branch.target, jump.target = jump.target, branch.target |
| jump.hot = True |
| |
| def _remove_redundant_jumps(self) -> None: |
| # Zero-length jumps can be introduced by _insert_continue_label and |
| # _invert_hot_branches: |
| for block in self._blocks(): |
| target = block.target |
| if target is None: |
| continue |
| target = target.resolve() |
| # Before: |
| # jmp FOO |
| # FOO: |
| # After: |
| # FOO: |
| if block.link and target is block.link.resolve(): |
| block.target = None |
| block.fallthrough = True |
| block.instructions.pop() |
| # Before: |
| # branch FOO: |
| # ... |
| # FOO: |
| # jump BAR |
| # After: |
| # br cond BAR |
| # ... |
| elif ( |
| len(target.instructions) == 1 |
| and target.instructions[0].kind == InstructionKind.JUMP |
| ): |
| assert target.target is not None |
| assert target.target.label is not None |
| if block.instructions[ |
| -1 |
| ].kind == InstructionKind.SHORT_BRANCH and self._is_far_target( |
| target.target.label |
| ): |
| continue |
| block.target = target.target |
| block.instructions[-1] = block.instructions[-1].update_target( |
| target.target.label |
| ) |
| |
| def _find_live_blocks(self) -> set[_Block]: |
| live: set[_Block] = set() |
| # Externally reachable blocks are live |
| todo: set[_Block] = {b for b in self._blocks() if b.label in self.globals} |
| while todo: |
| block = todo.pop() |
| live.add(block) |
| if block.fallthrough: |
| next = block.link |
| if next is not None and next not in live: |
| todo.add(next) |
| next = block.target |
| if next is not None and next not in live: |
| todo.add(next) |
| return live |
| |
| def _remove_unreachable(self) -> None: |
| live = self._find_live_blocks() |
| continuation = self._lookup_label(f"{self.label_prefix}_JIT_CONTINUE") |
| # Keep blocks after continuation as they may contain data and |
| # metadata that the assembler needs |
| prev: _Block | None = None |
| block = self._root |
| while block is not continuation: |
| next = block.link |
| assert next is not None |
| if not block in live and prev: |
| prev.link = next |
| else: |
| prev = block |
| block = next |
| assert prev.link is block |
| |
| def _fixup_external_labels(self) -> None: |
| if self._supports_external_relocations: |
| # Nothing to fix up |
| return |
| for index, block in enumerate(self._blocks()): |
| if block.target and block.fallthrough: |
| branch = block.instructions[-1] |
| if branch.kind == InstructionKind.CALL: |
| continue |
| assert branch.is_branch() |
| target = branch.target |
| assert target is not None |
| reloc = self._branches[branch.name][1] |
| if reloc is not None and self._is_far_target(target): |
| name = target[len(self.symbol_prefix) :] |
| label = f"{self.symbol_prefix}{reloc}_JIT_RELOCATION_{name}_JIT_RELOCATION_{index}:" |
| block.instructions[-1] = Instruction( |
| InstructionKind.OTHER, "", label, None, None |
| ) |
| block.instructions.append(branch.update_target("0")) |
| |
| def _fixup_constants(self) -> None: |
| "Fixup loading of constants. Overridden by OptimizerAArch64" |
| pass |
| |
| def _validate(self) -> None: |
| for block in self._blocks(): |
| if not block.instructions: |
| continue |
| for inst in block.instructions: |
| if self.frame_pointers: |
| assert ( |
| self._frame_pointer_modify.match(inst.text) is None |
| ), "Frame pointer should not be modified" |
| |
| def run(self) -> None: |
| """Run this optimizer.""" |
| self._insert_continue_label() |
| self._mark_hot_blocks() |
| # Removing branches can expose opportunities for more branch removal. |
| # Repeat a few times. 2 would probably do, but it's fast enough with 4. |
| for _ in range(4): |
| self._invert_hot_branches() |
| self._remove_redundant_jumps() |
| self._remove_unreachable() |
| self._fixup_external_labels() |
| self._fixup_constants() |
| self._validate() |
| self.path.write_text(self._body()) |
| |
| |
| class OptimizerAArch64(Optimizer): # pylint: disable = too-few-public-methods |
| """aarch64-pc-windows-msvc/aarch64-apple-darwin/aarch64-unknown-linux-gnu""" |
| |
| _branches = _AARCH64_BRANCHES |
| _short_branches = _AARCH64_SHORT_BRANCHES |
| # Mach-O does not support the 19 bit branch locations needed for branch reordering |
| _supports_external_relocations = False |
| _branch_patterns = [name.replace(".", r"\.") for name in _AARCH64_BRANCHES] |
| _re_branch = re.compile( |
| rf"\s*(?P<instruction>{'|'.join(_branch_patterns)})\s+(.+,\s+)*(?P<target>[\w.]+)" |
| ) |
| |
| # https://developer.arm.com/documentation/ddi0406/b/Application-Level-Architecture/Instruction-Details/Alphabetical-list-of-instructions/BL--BLX--immediate- |
| _re_call = re.compile(r"\s*blx?\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") |
| |
| supports_small_constants = True |
| _re_small_const_1 = re.compile( |
| r"\s*(?P<instruction>adrp)\s+(?P<register>x\d\d?),.*(?P<value>_JIT_OP(ARG|ERAND(0|1))_(16|32)).*" |
| ) |
| _re_small_const_2 = re.compile( |
| r"\s*(?P<instruction>ldr)\s+(?P<register>x\d\d?),.*(?P<value>_JIT_OP(ARG|ERAND(0|1))_(16|32)).*" |
| ) |
| _re_small_const_mask = re.compile( |
| r"\s*(?P<instruction>and)\s+[xw]\d\d?, *(?P<register>[xw]\d\d?).*(?P<value>0xffff)" |
| ) |
| _re_large_const_1 = re.compile( |
| r"\s*(?P<instruction>adrp)\s+(?P<register>x\d\d?),.*:got:(?P<value>[_A-Za-z0-9]+).*" |
| ) |
| _re_large_const_2 = re.compile( |
| r"\s*(?P<instruction>ldr)\s+(?P<register>x\d\d?),.*:got_lo12:(?P<value>[_A-Za-z0-9]+).*" |
| ) |
| const_reloc = "CUSTOM_AARCH64_CONST" |
| _frame_pointer_modify = re.compile(r"\s*stp\s+x29.*") |
| |
| def _make_temp_label(self, note: object = None) -> Instruction: |
| marker = f"jit_temp_{self.label_index}:" |
| if note is not None: |
| marker = f"{marker[:-1]}_{note}:" |
| self.label_index += 1 |
| return Instruction(InstructionKind.OTHER, "", marker, None, None) |
| |
| def _both_registers_same(self, inst: Instruction) -> bool: |
| reg = inst.register |
| assert reg is not None |
| if reg not in inst.text: |
| reg = "w" + reg[1:] |
| return inst.text.count(reg) == 2 |
| |
| def _fixup_small_constant_pair( |
| self, output: list[Instruction], label_index: int, inst: Instruction |
| ) -> str | None: |
| first = output[label_index + 1] |
| reg = first.register |
| if reg is None or inst.register != reg: |
| output.append( |
| Instruction(InstructionKind.OTHER, "", "# registers differ", None, None) |
| ) |
| output.append(inst) |
| return None |
| assert first.target is not None |
| if first.target != inst.target: |
| output.append( |
| Instruction(InstructionKind.OTHER, "", "# targets differ", None, None) |
| ) |
| output.append(inst) |
| return None |
| if not self._both_registers_same(inst): |
| output.append( |
| Instruction( |
| InstructionKind.OTHER, "", "# not same register", None, None |
| ) |
| ) |
| output.append(inst) |
| return None |
| pre, _ = first.text.split(first.name) |
| output[label_index + 1] = Instruction( |
| InstructionKind.OTHER, |
| "movz", |
| f"{pre}movz {reg}, 0", |
| reg, |
| None, |
| ) |
| label_text = f"{self.const_reloc}16a_JIT_RELOCATION_CONST{first.target[:-3]}_JIT_RELOCATION_{self.label_index}:" |
| self.label_index += 1 |
| output[label_index] = Instruction( |
| InstructionKind.OTHER, "", label_text, None, None |
| ) |
| assert first.target.endswith("16") or first.target.endswith("32") |
| if first.target.endswith("32"): |
| label_text = f"{self.const_reloc}16b_JIT_RELOCATION_CONST{first.target[:-3]}_JIT_RELOCATION_{self.label_index}:" |
| self.label_index += 1 |
| output.append( |
| Instruction(InstructionKind.OTHER, "", label_text, None, None) |
| ) |
| pre, _ = inst.text.split(inst.name) |
| output.append( |
| Instruction( |
| InstructionKind.OTHER, |
| "movk", |
| f"{pre}movk {reg}, 0, lsl #16", |
| reg, |
| None, |
| ) |
| ) |
| return reg |
| |
| def may_use_reg(self, inst: Instruction, reg: str | None) -> bool: |
| "Return False if `reg` is not explicitly used by this instruction" |
| if reg is None: |
| return False |
| assert reg.startswith("w") or reg.startswith("x") |
| xreg = f"x{reg[1:]}" |
| wreg = f"w{reg[1:]}" |
| if wreg in inst.text: |
| return True |
| if xreg in inst.text: |
| # Exclude false positives like 0x80 for x8 |
| count = inst.text.count(xreg) |
| number_count = inst.text.count("0" + xreg) |
| return count > number_count |
| return False |
| |
| def _fixup_large_constant_pair( |
| self, output: list[Instruction], label_index: int, inst: Instruction |
| ) -> None: |
| first = output[label_index + 1] |
| reg = first.register |
| if reg is None or inst.register != reg: |
| output.append(inst) |
| return |
| assert first.target is not None |
| if first.target != inst.target: |
| output.append(inst) |
| return |
| label = f"{self.const_reloc}33a_JIT_PAIR_{first.target}_JIT_PAIR_{self.label_index}:" |
| output[label_index] = Instruction(InstructionKind.OTHER, "", label, None, None) |
| label = ( |
| f"{self.const_reloc}33b_JIT_PAIR_{inst.target}_JIT_PAIR_{self.label_index}:" |
| ) |
| self.label_index += 1 |
| output.append(Instruction(InstructionKind.OTHER, "", label, None, None)) |
| output.append(inst) |
| |
| def _fixup_mask(self, output: list[Instruction], inst: Instruction) -> None: |
| if self._both_registers_same(inst): |
| # Nop |
| pass |
| else: |
| output.append(inst) |
| |
| def _fixup_constants(self) -> None: |
| for block in self._blocks(): |
| fixed: list[Instruction] = [] |
| small_const_part: dict[str, int | None] = {} |
| small_const_whole: dict[str, str | None] = {} |
| large_const_part: dict[str, int | None] = {} |
| for inst in block.instructions: |
| if inst.kind == InstructionKind.SMALL_CONST_1: |
| assert inst.register is not None |
| small_const_part[inst.register] = len(fixed) |
| small_const_whole[inst.register] = None |
| large_const_part[inst.register] = None |
| fixed.append(self._make_temp_label(inst.register)) |
| fixed.append(inst) |
| elif inst.kind == InstructionKind.SMALL_CONST_2: |
| assert inst.register is not None |
| index = small_const_part.get(inst.register) |
| small_const_part[inst.register] = None |
| if index is None: |
| fixed.append(inst) |
| continue |
| small_const_whole[inst.register] = self._fixup_small_constant_pair( |
| fixed, index, inst |
| ) |
| small_const_part[inst.register] = None |
| elif inst.kind == InstructionKind.SMALL_CONST_MASK: |
| assert inst.register is not None |
| reg = small_const_whole.get(inst.register) |
| if reg is not None: |
| self._fixup_mask(fixed, inst) |
| else: |
| fixed.append(inst) |
| elif inst.kind == InstructionKind.LARGE_CONST_1: |
| assert inst.register is not None |
| small_const_part[inst.register] = None |
| small_const_whole[inst.register] = None |
| large_const_part[inst.register] = len(fixed) |
| fixed.append(self._make_temp_label()) |
| fixed.append(inst) |
| elif inst.kind == InstructionKind.LARGE_CONST_2: |
| assert inst.register is not None |
| small_const_part[inst.register] = None |
| small_const_whole[inst.register] = None |
| index = large_const_part.get(inst.register) |
| large_const_part[inst.register] = None |
| if index is None: |
| fixed.append(inst) |
| continue |
| self._fixup_large_constant_pair(fixed, index, inst) |
| else: |
| for reg in small_const_part: |
| if self.may_use_reg(inst, reg): |
| small_const_part[reg] = None |
| for reg in small_const_whole: |
| if self.may_use_reg(inst, reg): |
| small_const_whole[reg] = None |
| for reg in small_const_part: |
| if self.may_use_reg(inst, reg): |
| large_const_part[reg] = None |
| fixed.append(inst) |
| block.instructions = fixed |
| |
| |
| class OptimizerX86(Optimizer): # pylint: disable = too-few-public-methods |
| """i686-pc-windows-msvc/x86_64-apple-darwin/x86_64-unknown-linux-gnu""" |
| |
| _branches = _X86_BRANCHES |
| _short_branches = {} |
| _re_branch = re.compile( |
| rf"\s*(?P<instruction>{'|'.join(_X86_BRANCHES)})\s+(?P<target>[\w.]+)" |
| ) |
| # https://www.felixcloutier.com/x86/call |
| _re_call = re.compile(r"\s*callq?\s+(?P<target>[\w.]+)") |
| # https://www.felixcloutier.com/x86/jmp |
| _re_jump = re.compile(r"\s*jmp\s+(?P<target>[\w.]+)") |
| # https://www.felixcloutier.com/x86/ret |
| _re_return = re.compile(r"\s*retq?\b") |
| _frame_pointer_modify = re.compile(r"\s*movq?\s+%(\w+),\s+%rbp.*") |
