slime/research/experiments/full_search_feasibility/run.py

195 lines
9.9 KiB
Python

#!/usr/bin/env python3
"""Reproducible CPU validation and sequential CUDA experiments (no network)."""
from __future__ import annotations
import argparse
import hashlib
import json
import platform
import random
import statistics
import subprocess
import time
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo
HERE = Path(__file__).resolve().parent
POS = HERE / "positions"
OUT = HERE / "results"
BIN = "/tmp/slime_feasibility_cuda"
def save_positions(name, points):
(POS / f"{name}.txt").write_text("".join(f"{x} {z}\n" for x, z in points))
def prepare():
POS.mkdir(exist_ok=True)
OUT.mkdir(exist_ok=True)
save_positions("known", [(51069, 419124)])
offsets = [(0, 0), (0, 1), (1, 0), (-1, 0), (0, -1), (2, 0), (0, 2),
(2, 2), (-2, -2), (-1, 1), (1, -1), (3, 0), (0, 3), (3, 3),
(-3, 0), (0, -3)]
save_positions("stress", [(51069 + x, 419124 + z) for x, z in offsets] +
[(0, 0), (-1875000, -1875000), (1874995, 1874995),
(2147483643, 2147483643), (-2147483648, -2147483648)])
rng = random.Random(2026101101)
coords = [(rng.randrange(-1875000, 1875000), rng.randrange(-1875000, 1875000))
for _ in range(65536)]
save_positions("random4096", coords[:4096])
save_positions("random65536", coords)
save_positions("tail128", [(51069, 419124)] + coords[:127])
return coords
def run_case(name, positions, task, repeats=5, cells=None):
cmd = [BIN, str(POS / f"{positions}.txt"), task, str(repeats)]
if cells is not None:
cmd.append(str(cells))
start = time.perf_counter()
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
(OUT / f"{name}.jsonl").write_text(result.stdout)
(OUT / f"{name}.stderr").write_text(result.stderr)
if result.returncode:
raise RuntimeError(f"{name} failed: {result.stderr}")
rows = [json.loads(line) for line in result.stdout.splitlines()]
record = {"command": cmd, "wall_ms": (time.perf_counter() - start) * 1000,
"rows": rows}
print(json.dumps({"finished": name, "wall_ms": record["wall_ms"]}), flush=True)
return record
def median(rows, kind, mode):
return statistics.median(r["kernel_ms"] for r in rows
if r["type"] == kind and r["mode"] == mode)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--analyse-only", action="store_true")
ap.add_argument("--complete-only", action="store_true")
args = ap.parse_args()
if args.analyse_only or args.complete_only:
records = json.loads((OUT / "all.json").read_text())
else:
coords = prepare()
subprocess.run(["nvcc", "-O3", "-std=c++17", "-arch=sm_86",
str(HERE / "cuda_bench.cu"), "-o", BIN], check=True)
subprocess.run(["g++", "-O3", "-std=c++17", "-march=native", "-Wall", "-Wextra",
str(HERE / "cpu_verify.cpp"), "-o", "/tmp/slime_feasibility_verify"], check=True)
cpu = subprocess.check_output(["/tmp/slime_feasibility_verify"], text=True)
(OUT / "cpu_verify.json").write_text(cpu)
gpu_query = ["nvidia-smi", "--query-gpu=name,driver_version,temperature.gpu,utilization.gpu", "--format=csv"]
records = {"metadata": {
"date": datetime.now(ZoneInfo("Asia/Hong_Kong")).isoformat(),
"platform": platform.platform(), "cpu_validation": json.loads(cpu),
"gpu_before": subprocess.check_output(gpu_query, text=True),
"source_sha256": {p.name: hashlib.sha256(p.read_bytes()).hexdigest()
for p in [HERE / "cuda_bench.cu", HERE / "interval_solver.hpp",
HERE / "cpu_verify.cpp", HERE / "run.py",
HERE.parent.parent / "slime_solver_core.hpp",
HERE.parents[2] / "include/slime/solver_core.hpp",
HERE.parents[2] / "include/slime/upper_interval.hpp"]}}}
def record(name, positions, task, repeats=5, cells=None):
records[name] = run_case(name, positions, task, repeats, cells)
(OUT / "all.json").write_text(json.dumps(records, indent=2))
record("known", "known", "compare")
record("stress", "stress", "compare")
record("random4096", "random4096", "compare")
record("random65536", "random65536", "fast")
record("roots65536", "random65536", "root")
counts = next(r["counts"] for r in records["random4096"]["rows"] if r["type"] == "root_counts")
classes = [(x % 65536, z % 131072) for (x, z), n in zip(coords, counts) if n][:32]
save_positions("classes32", classes)
lifts, groups = [], []
for x, z in classes:
current = [(xx, zz)
for xx in range(-1875000 + (x + 1875000) % 65536, 1875000, 65536)
for zz in range(-1875000 + (z + 1875000) % 131072, 1875000, 131072)]
groups.append(len(current))
lifts.extend(current)
save_positions("lifts", lifts)
(POS / "lift_groups.json").write_text(json.dumps(groups))
record("classes32", "classes32", "root")
record("lifts", "lifts", "root")
record("tail_one_cell", "known", "tail", cells=1)
record("tail_two_cells", "known", "tail", cells=2)
record("tail_five_cells", "known", "tail", cells=5)
record("tail128", "tail128", "tail")
records["metadata"]["gpu_after"] = subprocess.check_output(gpu_query, text=True)
(OUT / "all.json").write_text(json.dumps(records, indent=2))
if not args.analyse_only:
if args.complete_only:
subprocess.run(["nvcc", "-O3", "-std=c++17", "-arch=sm_86",
str(HERE / "cuda_bench.cu"), "-o", BIN], check=True)
records["complete_stress"] = run_case("complete_stress", "stress", "complete")
records["complete65536"] = run_case("complete65536", "random65536", "complete")
records["metadata"]["complete_source_sha256"] = {
p.name: hashlib.sha256(p.read_bytes()).hexdigest()
for p in [HERE / "cuda_bench.cu", HERE / "interval_solver.hpp", HERE / "run.py",
HERE.parents[2] / "include/slime/solver_core.hpp",
HERE.parents[2] / "include/slime/upper_interval.hpp"]}
(OUT / "all.json").write_text(json.dumps(records, indent=2))
summary = {}
for name in ["known", "stress", "random4096", "random65536"]:
rows = records[name]["rows"]
modes = sorted({r["mode"] for r in rows if r["type"] == "search_timing"})
summary[name] = {"kernel_median_ms": {str(m): median(rows, "search_timing", m) for m in modes}}
summary[name].update(next(r for r in rows if r["type"] == "search_info"))
counts = next(r["counts"] for r in records["random65536"]["rows"] if r["type"] == "root_counts")
n = len(counts)
mean = statistics.mean(counts)
sem = statistics.stdev(counts) / n**0.5
summary["distribution"] = {"positions": n, "roots": sum(counts), "mean": mean,
"mean_normal95": [mean - 1.96 * sem, mean + 1.96 * sem],
"zero_fraction": counts.count(0) / n, "max": max(counts)}
for name in ["roots65536", "classes32", "lifts"]:
summary[name] = {str(m): median(records[name]["rows"], "root_timing", m) for m in [0, 1]}
for name in ["tail_one_cell", "tail_two_cells", "tail_five_cells", "tail128"]:
rows = records[name]["rows"]
summary[name] = {"kernel_median_ms": {str(m): median(rows, "tail_timing", m) for m in [0, 1]},
"matches": next(r["matches"] for r in rows if r["type"] == "tail_timing")}
# All projections are throughput extrapolations, not measured full-world scans.
world = 3750000**2
year = 365.25 * 86400
fast_ms = summary["random65536"]["kernel_median_ms"]["3"]
tail_ms = summary["tail128"]["kernel_median_ms"]["1"]
root_ms = min(summary["roots65536"].values())
summary["projections"] = {
"world_anchors": world,
"lower35_visits": world * mean * 131072,
"first_draw_kernel_years": fast_ms / 1000 / n * world / year,
"root_no_reuse_years": root_ms / 1000 / n * world / year,
"root_class_reuse_hours": root_ms / 1000 / n * (2**33) / 3600,
"canonical_tail_kernel_years": tail_ms / 1000 / 128 * world / year,
"complete_compute_years": (fast_ms / 1000 / n + tail_ms / 1000 / 128) * world / year,
}
for name in ["complete_stress", "complete65536"]:
if name in records:
rows = records[name]["rows"]
summary[name] = {key: statistics.median(r[key] for r in rows if r["type"] == "complete_timing")
for key in ["first_ms", "tail_ms", "compute_ms", "scan_ms"]}
summary[name].update(next(r for r in rows if r["type"] == "complete_info"))
if "complete65536" in summary:
summary["projections"]["complete_same_batch_kernel_years"] = summary["complete65536"]["compute_ms"] / 1000 / n * world / year
summary["projections"]["complete_same_batch_scan_years"] = summary["complete65536"]["scan_ms"] / 1000 / n * world / year
assert summary["complete_stress"]["matches"] == summary["stress"]["matches"]
# Count equality is checked here; the CUDA program additionally compares complete
# normalized root lists for every repeated coordinate class in its input.
base_counts = next(r["counts"] for r in records["classes32"]["rows"] if r["type"] == "root_counts")
lifted_counts = next(r["counts"] for r in records["lifts"]["rows"] if r["type"] == "root_counts")
groups = json.loads((POS / "lift_groups.json").read_text())
assert lifted_counts == [c for c, number in zip(base_counts, groups) for _ in range(number)]
summary["cache_counts_equal"] = True
(OUT / "summary.json").write_text(json.dumps(summary, indent=2))
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()