tinymind-native-colab-handoff/bundle/evaluation/int6_bridge_imma_eval.py

"""Measured fused two-pass IMMA.SP benchmark for TinyMind INT6 bridge."""

from __future__ import annotations

from datetime import datetime, timezone
import csv
import io
import json
from pathlib import Path
import re
import subprocess
import tempfile
import threading
import time


ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / "kernels" / "int6_sparse_ptx" / "int6_bridge_imma_bench.cu"


def _tool(name: str) -> str:
    for base in [
        "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/bin",
        "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.6/bin",
    ]:
        candidate = Path(base) / name
        if candidate.exists():
            return str(candidate)
    return name


def _run(command: list[str], cwd: Path, timeout: int = 240) -> dict:
    proc = subprocess.run(command, cwd=cwd, capture_output=True, text=True, timeout=timeout, check=False)
    return {
        "command": command,
        "exit_code": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
    }


def _run_msvc(command: list[str], cwd: Path) -> dict:
    vcvars = Path("C:/Program Files (x86)/Microsoft Visual Studio/2022/BuildTools/VC/Auxiliary/Build/vcvars64.bat")
    if not vcvars.exists():
        return _run(command, cwd)
    quoted = " ".join(f'"{part}"' if (" " in part or "\\" in part or ":" in part) else part for part in command)
    script = Path(tempfile.gettempdir()) / "tinymind_int6_bridge_build.bat"
    script.write_text(f'@echo off\r\ncall "{vcvars}" >nul\r\n{quoted}\r\n', encoding="utf-8")
    return _run(["cmd.exe", "/d", "/c", str(script)], cwd)


def _monitor_power(stop: threading.Event, rows: list[dict]) -> None:
    command = [
        "nvidia-smi",
        "--query-gpu=power.draw,temperature.gpu,clocks.current.sm,utilization.gpu",
        "--format=csv,noheader,nounits",
    ]
    while not stop.is_set():
        try:
            proc = subprocess.run(command, capture_output=True, text=True, timeout=3, check=False)
            line = proc.stdout.strip().splitlines()[0]
            parsed = next(csv.reader(io.StringIO(line)))
            rows.append(
                {
                    "power_w": float(parsed[0].strip()),
                    "temp_c": float(parsed[1].strip()),
                    "sm_clock_mhz": float(parsed[2].strip()),
                    "gpu_util_pct": float(parsed[3].strip()),
                }
            )
        except Exception:
            pass
        time.sleep(0.05)


def _avg(rows: list[dict], key: str) -> float:
    vals = [float(row[key]) for row in rows if key in row]
    return sum(vals) / len(vals) if vals else 0.0


def _max(rows: list[dict], key: str) -> float:
    vals = [float(row[key]) for row in rows if key in row]
    return max(vals) if vals else 0.0


def _power_quality(rows: list[dict]) -> dict:
    active_rows = [row for row in rows if float(row.get("gpu_util_pct", 0.0)) >= 50.0]
    active_util = _avg(active_rows, "gpu_util_pct")
    active_power = _avg(active_rows, "power_w")
    active_clock = _avg(active_rows, "sm_clock_mhz")
    return {
        "sample_interval_s": 0.05,
        "sample_count": len(rows),
        "avg_gpu_util_pct": _avg(rows, "gpu_util_pct"),
        "max_gpu_util_pct": _max(rows, "gpu_util_pct"),
        "avg_sm_clock_mhz": _avg(rows, "sm_clock_mhz"),
        "max_sm_clock_mhz": _max(rows, "sm_clock_mhz"),
        "active_sample_count": len(active_rows),
        "active_avg_gpu_util_pct": active_util,
        "active_avg_power_w": active_power,
        "active_avg_sm_clock_mhz": active_clock,
        "active_window_policy": "samples with gpu_util_pct >= 50 to remove process startup/cooldown from sustained-load evidence",
        "sufficient_for_power_claim": len(active_rows) >= 20 and active_util >= 90.0,
        "sufficient_for_world_tfw_claim": False,
    }


def _parse_bench(text: str) -> list[dict]:
    rows = []
    in_realdata = False
    for line in text.splitlines():
        if line.startswith("realdata_pass,"):
            in_realdata = True
            continue
        if line.startswith("pass,"):
            in_realdata = False
            continue
        if re.match(r"^\d+,", line):
            p = line.split(",")
            rows.append(
                {
                    "mode": "real_data" if in_realdata else "compute_peak",
                    "pass": int(p[0]),
                    "ms": float(p[1]),
                    "logical_int6_tops": float(p[2]),
                    "hardware_imma_tops": float(p[3]),
                    "checksum": int(p[4]),
                }
            )
    return rows


def _parse_realdata_stream_bytes(text: str) -> int:
    match = re.search(r"realdata_stream_bytes,(\d+)", text)
    return int(match.group(1)) if match else 0


def _parse_sass_summary(text: str) -> dict:
    lines = text.splitlines()
    return {
        "imma_sp_count": sum(1 for line in lines if "IMMA.SP" in line or "MMA.SP" in line),
        "imma_count": sum(1 for line in lines if "IMMA" in line),
        "shift_or_funnel_count": sum(1 for line in lines if "SHF" in line or "SHL" in line),
        "logic_count": sum(1 for line in lines if "LOP3" in line or "XOR" in line),
        "integer_add_count": sum(1 for line in lines if "IADD" in line or "LEA" in line),
        "raw_excerpt": text[-4000:],
    }


def _parse_ptxas_info(text: str) -> dict:
    combined = text or ""
    registers = None
    spill_stores = None
    spill_loads = None
    cmem_bytes = []
    for line in combined.splitlines():
        reg_match = re.search(r"Used\s+(\d+)\s+registers", line)
        if reg_match:
            registers = int(reg_match.group(1))
        spill_match = re.search(r"(\d+)\s+bytes\s+spill\s+stores,\s+(\d+)\s+bytes\s+spill\s+loads", line)
        if spill_match:
            spill_stores = int(spill_match.group(1))
            spill_loads = int(spill_match.group(2))
        cmem_match = re.search(r"(\d+)\s+bytes\s+cmem\[(\d+)\]", line)
        if cmem_match:
            cmem_bytes.append({"bytes": int(cmem_match.group(1)), "bank": int(cmem_match.group(2))})
    return {
        "registers_per_thread": registers,
        "spill_stores_bytes": spill_stores,
        "spill_loads_bytes": spill_loads,
        "spilling_observed": bool((spill_stores or 0) > 0 or (spill_loads or 0) > 0),
        "constant_memory": cmem_bytes,
        "raw_excerpt": combined[-4000:],
    }


def _trim(result: dict) -> dict:
    return {
        "command": result["command"],
        "exit_code": result["exit_code"],
        "stdout_tail": result.get("stdout", "")[-4000:],
        "stderr_tail": result.get("stderr", "")[-4000:],
    }


def build_int6_bridge_imma_eval(
    out_dir: str | Path,
    blocks: int = 160,
    threads: int = 256,
    iterations: int = 20000,
    passes: int = 5,
    min_duration_s: float = 0.0,
    mode: str = "both",
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    exe = out / "int6_bridge_imma_bench.exe"
    nvcc = _tool("nvcc.exe")
    cuobjdump = _tool("cuobjdump.exe")
    build = _run_msvc([nvcc, "-arch=sm_86", "-O3", "-std=c++17", "-Xptxas=-v", str(SOURCE), "-o", str(exe)], ROOT)

    power_rows: list[dict] = []
    stop = threading.Event()
    monitor = threading.Thread(target=_monitor_power, args=(stop, power_rows), daemon=True)
    run_iterations = int(iterations)
    run_passes = int(passes)
    if min_duration_s > 0:
        run_passes = max(run_passes, 8)
        run_iterations = max(run_iterations, 50_000)
    run = {"command": [str(exe)], "exit_code": -1, "stdout": "", "stderr": "build failed"}
    if build["exit_code"] == 0:
        monitor.start()
        try:
            started = time.time()
            run = _run([str(exe), str(blocks), str(threads), str(run_iterations), str(run_passes), mode], ROOT, timeout=900)
            if min_duration_s > 0 and time.time() - started < min_duration_s:
                scale = max(2, int(min_duration_s / max(time.time() - started, 0.1)) + 1)
                run_iterations = min(run_iterations * scale, 5_000_000)
                started = time.time()
                run = _run([str(exe), str(blocks), str(threads), str(run_iterations), str(run_passes), mode], ROOT, timeout=1800)
        finally:
            stop.set()
            monitor.join(timeout=5)

    sass = _run([cuobjdump, "--dump-sass", str(exe)], ROOT) if build["exit_code"] == 0 else {"command": [], "exit_code": -1, "stdout": "", "stderr": "build failed"}
    bench_rows = _parse_bench(run.get("stdout", ""))
    compute_rows = [row for row in bench_rows if row.get("mode") == "compute_peak"]
    realdata_rows = [row for row in bench_rows if row.get("mode") == "real_data"]
    logical = _avg(compute_rows, "logical_int6_tops")
    hardware = _avg(compute_rows, "hardware_imma_tops")
    real_logical = _avg(realdata_rows, "logical_int6_tops")
    real_hardware = _avg(realdata_rows, "hardware_imma_tops")
    power = _avg(power_rows, "power_w")
    ptxas = _parse_ptxas_info(build.get("stdout", "") + "\n" + build.get("stderr", ""))
    sass_text = sass.get("stdout", "")
    sass_summary = _parse_sass_summary(sass_text)
    stream_bytes = _parse_realdata_stream_bytes(run.get("stdout", ""))
    power_quality = _power_quality(power_rows)
    report = {
        "schema_version": "tinymind-int6-bridge-imma-eval-v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "kernel": "fused_two_pass_imma_sp_int6_bridge",
        "equation": "q6 = 4*q_hi + q_lo; fused accumulator = 4*IMMA.SP(hi) + IMMA.SP(lo)",
        "source": str(SOURCE),
        "exe": str(exe),
        "params": {
            "blocks": blocks,
            "threads": threads,
            "requested_iterations": iterations,
            "requested_passes": passes,
            "run_iterations": run_iterations,
            "run_passes": run_passes,
            "min_duration_s": min_duration_s,
            "mode": mode,
        },
        "build": _trim(build),
        "ptxas": ptxas,
        "run": _trim(run),
        "sass": {
            "contains_imma_sp": bool(re.search(r"IMMA\.SP|IMMA|MMA\.SP", sass_text)),
            "dump_exit_code": sass.get("exit_code"),
            "summary": sass_summary,
        },
        "bench_rows": bench_rows,
        "real_data_stream": {
            "measured": bool(realdata_rows),
            "stream_bytes": stream_bytes,
            "note": "real_data mode loads hi/lo/b/metadata streams from global memory before issuing the two-pass IMMA.SP bridge.",
        },
        "power_samples": {
            "count": len(power_rows),
            "avg_power_w": power,
            "max_power_w": _max(power_rows, "power_w"),
            "avg_temp_c": _avg(power_rows, "temp_c"),
            "max_temp_c": _max(power_rows, "temp_c"),
            "avg_sm_clock_mhz": _avg(power_rows, "sm_clock_mhz"),
            "max_sm_clock_mhz": _max(power_rows, "sm_clock_mhz"),
            "avg_gpu_util_pct": _avg(power_rows, "gpu_util_pct"),
            "max_gpu_util_pct": _max(power_rows, "gpu_util_pct"),
            "quality_gate": power_quality,
        },
        "metrics": {
            "compute_peak": {
                "avg_logical_int6_tops": logical,
                "avg_hardware_imma_tops": hardware,
                "avg_logical_int6_tops_per_watt": logical / power if power > 0 else 0.0,
                "avg_hardware_imma_tops_per_watt": hardware / power if power > 0 else 0.0,
            },
            "real_data": {
                "avg_logical_int6_tops": real_logical,
                "avg_hardware_imma_tops": real_hardware,
                "avg_logical_int6_tops_per_watt": real_logical / power if power > 0 else 0.0,
                "avg_hardware_imma_tops_per_watt": real_hardware / power if power > 0 else 0.0,
                "slowdown_vs_compute_peak": logical / real_logical if real_logical > 0 else 0.0,
            },
        },
        "claim_gate": {
            "fused_two_pass_kernel_measured": build["exit_code"] == 0 and run["exit_code"] == 0 and bool(bench_rows),
            "imma_sp_sass_observed": bool(re.search(r"IMMA\.SP|IMMA|MMA\.SP", sass_text)),
            "register_spilling_observed": ptxas["spilling_observed"],
            "register_pressure_safe": ptxas["registers_per_thread"] is not None and ptxas["registers_per_thread"] < 128 and not ptxas["spilling_observed"],
            "real_data_movement_measured": bool(realdata_rows),
            "power_measurement_representative": power_quality["sufficient_for_power_claim"],
            "sass_instruction_summary_available": sass_summary["imma_count"] > 0,
            "int6_bottleneck_removed": False,
            "world_highest_tfw_claim_allowed": False,
        },
    }
    report["claim_gate"]["int6_bottleneck_removed"] = (
        report["claim_gate"]["fused_two_pass_kernel_measured"]
        and report["claim_gate"]["imma_sp_sass_observed"]
        and metrics_above_reference(report)
    )
    json_path = out / "int6_bridge_imma_eval_report.json"
    md_path = out / "int6_bridge_imma_eval_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def metrics_above_reference(report: dict) -> bool:
    compute = float(report["metrics"]["compute_peak"]["avg_logical_int6_tops"])
    real = float(report["metrics"]["real_data"]["avg_logical_int6_tops"])
    return max(compute, real) > 10.0


def _markdown(report: dict) -> str:
    compute = report["metrics"]["compute_peak"]
    real = report["metrics"]["real_data"]
    p = report["power_samples"]
    ptxas = report.get("ptxas", {})
    sass = report.get("sass", {}).get("summary", {})
    return "\n".join(
        [
            "# TinyMind INT6 Bridge IMMA Eval",
            "",
            f"- Compute peak logical INT6 TOPS: {compute['avg_logical_int6_tops']:.6f}",
            f"- Compute peak hardware IMMA TOPS: {compute['avg_hardware_imma_tops']:.6f}",
            f"- Real-data logical INT6 TOPS: {real['avg_logical_int6_tops']:.6f}",
            f"- Real-data hardware IMMA TOPS: {real['avg_hardware_imma_tops']:.6f}",
            f"- Real-data slowdown vs compute peak: {real['slowdown_vs_compute_peak']:.3f}x",
            f"- Avg power W: {p['avg_power_w']:.2f}",
            f"- Real-data logical INT6 TOPS/W: {real['avg_logical_int6_tops_per_watt']:.6f}",
            f"- Power representative: {p['quality_gate']['sufficient_for_world_tfw_claim']}",
            f"- Registers/thread: {ptxas.get('registers_per_thread')}",
            f"- Register spilling observed: {ptxas.get('spilling_observed')}",
            f"- SASS IMMA count: {sass.get('imma_count')}",
            f"- IMMA.SP SASS observed: {report['claim_gate']['imma_sp_sass_observed']}",
            f"- Real data movement measured: {report['claim_gate']['real_data_movement_measured']}",
            f"- INT6 bottleneck removed: {report['claim_gate']['int6_bottleneck_removed']}",
            "",
        ]
    )