lean4-htt/tests/bench/build/leanir_wrapper.py
Sebastian Ullrich c6a89cc716
feat: experimental option to move non-meta compilation out of lean build step (#10291)
The ultimate goal of this work is to turn production of `.ir` files into
separate build step so that it does not block non-`meta` imports and can
be skipped entirely when not needed. This PR implements the main logic
of this new `leanir` compiler executable and runs it after `lean` inside
the same Lake build step but leaves its use disabled behind a
`compiler.postponeCompile` flag until further Lake adjustments move it
to a separate facet so that its use can be actually beneficial.

---------

Co-authored-by: Joscha <joscha@plugh.de>
2026-03-20 10:39:39 +00:00

101 lines
2.7 KiB
Python
Executable file

#!/usr/bin/env python3
import argparse
import json
import os
import re
import subprocess
import sys
from collections import Counter
from pathlib import Path
# Global paths
TEST_DIR = Path(os.environ["TEST_DIR"])
WRAPPER_OUT = Path(os.environ["WRAPPER_OUT"])
WRAPPER_PREFIX = Path(os.environ["WRAPPER_PREFIX"])
# Other config
BENCHMARK = "build"
sys.path.append(str(TEST_DIR))
import measure # noqa: E402
def save_measurement(metric: str, value: float, unit: str | None = None) -> None:
data = {"metric": metric, "value": value}
if unit is not None:
data["unit"] = unit
with open(WRAPPER_OUT, "a") as f:
f.write(f"{json.dumps(data)}\n")
def run(*command: str) -> None:
result = subprocess.run(command)
if result.returncode != 0:
sys.exit(result.returncode)
def run_capture(*command: str) -> tuple[str, str]:
result = subprocess.run(command, capture_output=True, encoding="utf-8")
if result.returncode != 0:
print(result.stdout, end="", file=sys.stdout)
print(result.stderr, end="", file=sys.stderr)
sys.exit(result.returncode)
return result.stdout, result.stderr
def get_module(setup: Path) -> str:
with open(setup) as f:
return json.load(f)["name"]
def count_bytes(module: str, path: Path, suffix: str) -> None:
try:
bytes = path.with_suffix(suffix).stat().st_size
except FileNotFoundError:
return
save_measurement(f"{BENCHMARK}/module/ir/{module}//bytes {suffix}", bytes, "B")
def run_leanir(module: str) -> None:
stdout, stderr = measure.main(
cmd=[
"leanir",
*sys.argv[1:],
*("-Dprofiler=true", "-Dprofiler.threshold=9999999"),
],
output=WRAPPER_OUT,
topics=[f"{BENCHMARK}/module/ir/{module}"],
metrics={"instructions", "cycles"},
append=True,
capture=True,
)
# Output of `leanir -Dprofiler=true`
# See timeit.cpp for the time format
for line in stderr.splitlines():
if match := re.fullmatch(r"\t(.*) ([\d.]+)(m?s)", line):
name = match.group(1)
seconds = float(match.group(2))
if match.group(3) == "ms":
seconds = seconds / 1000
save_measurement(f"{BENCHMARK}/profile/ir/{name}//wall-clock", seconds, "s")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("setup", type=Path)
parser.add_argument("ir", type=Path)
parser.add_argument("c", type=Path)
args, _ = parser.parse_known_args()
module = get_module(args.setup)
run_leanir(module)
count_bytes(module, args.ir, ".ir")
count_bytes(module, args.c, ".c")
if __name__ == "__main__":
main()