/
githubmirror
/
cpython
Обзор
Документация
Войти
/
githubmirror
/
cpython
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
Tools/jit/_optimizers.py
788 строк
30 KB
Mark Shannon
GH-146128: Fix AArch64 multi-instruction constants and relocations (GH-148598)
16 апр 2026, 17:33
Не верифицирован
16 апр 2026, 17:33
cecf564
Код
Авторство
О чём код?
"""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.*")