File size: 3,823 Bytes
15c2580 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""GEAK task runner — Triton kernel tasks."""
import argparse
import json
import os
import re
import subprocess
import sys
TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.chdir(TASK_DIR)
BUILD_DIR = os.path.join(TASK_DIR, "build")
KERNEL_DIR = os.path.join(TASK_DIR, "kernel")
HARNESS = os.path.join(KERNEL_DIR, "test_harness.py")
def _run_harness(argv):
return subprocess.run(
[sys.executable, HARNESS] + argv,
cwd=KERNEL_DIR,
capture_output=True,
text=True,
timeout=7200,
)
def run_compile():
try:
p = subprocess.run(
[
sys.executable,
"-c",
"import sys; sys.path.insert(0, 'kernel'); import host, kernel_jit",
],
cwd=TASK_DIR,
capture_output=True,
text=True,
timeout=600,
)
if p.returncode != 0:
return False, (p.stderr or "") + "\n" + (p.stdout or "")
return True, None
except Exception as e:
return False, str(e)
def run_correctness():
p = _run_harness(["--correctness"])
out = (p.stdout or "") + (p.stderr or "")
ok = p.returncode == 0 and "CORRECTNESS_OVERALL: PASS" in out
if not ok:
tail = out[-8000:] if len(out) > 8000 else out
return False, tail or f"exit {p.returncode}"
return True, None
def run_performance():
p = _run_harness(["--benchmark"])
out = (p.stdout or "") + (p.stderr or "")
rows = []
for m in re.finditer(r"CASE=(\S+)\s+GEAK_RESULT_LATENCY_MS=([0-9.eE+-]+)", out):
rows.append(
{
"test_case_id": m.group(1),
"execution_time_ms": float(m.group(2)),
"params": {},
}
)
if not rows:
m = re.search(r"GEAK_RESULT_LATENCY_MS=([0-9.eE+-]+)", out)
if m:
rows.append(
{
"test_case_id": "aggregate",
"execution_time_ms": float(m.group(1)),
"params": {},
}
)
ok = p.returncode == 0 and rows and all(r.get("execution_time_ms", -1) >= 0 for r in rows)
if not ok:
return [
{
"test_case_id": "_error",
"execution_time_ms": -1.0,
"params": {"stderr": out[-4000:]},
}
]
return rows
def main():
ap = argparse.ArgumentParser()
ap.add_argument("mode", choices=["compile", "correctness", "performance"])
args = ap.parse_args()
os.makedirs(BUILD_DIR, exist_ok=True)
if args.mode == "compile":
ok, err = run_compile()
with open(os.path.join(BUILD_DIR, "compile_report.json"), "w") as f:
json.dump({"status": "ok" if ok else "fail", "error": err}, f)
print(f"Compilation: {'PASS' if ok else 'FAIL'}")
if err:
print("Error:", err)
sys.exit(0 if ok else 1)
if args.mode == "correctness":
ok, err = run_correctness()
with open(os.path.join(BUILD_DIR, "correctness_report.json"), "w") as f:
json.dump({"status": "ok" if ok else "fail", "error": err}, f)
print(f"Correctness: {'PASS' if ok else 'FAIL'}")
if err:
print("Error:", err)
sys.exit(0 if ok else 1)
cases = run_performance()
with open(os.path.join(BUILD_DIR, "performance_report.json"), "w") as f:
json.dump({"test_cases": cases}, f, indent=2)
for c in cases:
ms = c.get("execution_time_ms", -1)
if ms >= 0:
print(f"Performance: {ms:.4f} ms ({c['test_case_id']})")
bad = any(c.get("execution_time_ms", -1) < 0 for c in cases)
sys.exit(1 if bad else 0)
if __name__ == "__main__":
main()
|