Files
IQ.Pilot/artifacts/package_runtime/tinygrad/runtime/support/compiler_cpu.py
T
2026-08-31 23:04:09 -05:00

40 lines
2.4 KiB
Python

import signal, subprocess, sys
from tinygrad.device import Compiler
from tinygrad.helpers import getenv, capstone_flatdump
from tinygrad.runtime.support.elf import jit_loader
def _block_sigusr2():
# msgq notifies readers via pthread_kill(SIGUSR2) by cached TID; a recycled TID can land
# the signal on this clang child and kill the compile mid-run (openpilot modeld crash)
if sys.platform == "linux": signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGUSR2})
class ClangCompiler(Compiler):
def __init__(self, arch:list[str], cachekey="compile_clang_jit"):
assert len(arch) >= 2, f"invalid arch string: {','.join(arch)!r}, expected '<arch>,<cpu>,[<feats>]' (eg. 'x86_64,znver2')"
self.arch, cpu, *feats = arch
match self.arch:
case "x86_64": self.args = [f"-march={cpu}"] + [f"-mno{f}" if f.startswith("-") else f"-m{f}" for f in feats]
# on arm march means "runs on this arch and superset" instead of "optimize for this arch". x86 march == arm mcpu
# x18 is a reserved platform register. It is clobbered on context switch in macos and is used to store TEB pointer in windows on arm
case "arm64": self.args = ["-ffixed-x18", "-mcpu=" + "+".join([cpu] + ["no"+f[1:] if f.startswith("-") else f for f in feats])]
case "riscv64": self.args = ["-march=" + "_".join(["rv64g" if cpu == "native" else cpu] + feats)]
case _: raise RuntimeError(f"unsupported arch: {self.arch!r}")
super().__init__(f"{cachekey}_{'_'.join(arch)}")
def compile_to_obj(self, src:str) -> bytes:
"""Compile C source to ELF object file (before linking)."""
# -fno-math-errno is required for __builtin_sqrt to become an instruction instead of a function call
return subprocess.check_output([getenv("CC", 'clang'), '-c', '-x', 'c', '-O2', '-fPIC', '-ffreestanding', '-fno-math-errno', '-nostdlib',
'-fno-ident', f'--target={self.arch}-none-unknown-elf', *self.args, '-', '-o', '-'], input=src.encode('utf-8'),
preexec_fn=_block_sigusr2 if sys.platform == "linux" else None)
def compile(self, src:str) -> bytes: return jit_loader(self.compile_to_obj(src))
def disassemble(self, lib:bytes): return capstone_flatdump(lib, self.arch)
class X86Compiler(Compiler):
def __init__(self): super().__init__(None)
def compile(self, src:str) -> bytes: return bytes.fromhex(src)
def disassemble(self, lib:bytes): return capstone_flatdump(lib, "x86_64")