ncu_capture.py

#!/usr/bin/env python3
"""Run a calibration microbenchmark under Nsight Compute and merge metrics into CSV output."""
from __future__ import annotations

import argparse
import csv
import io
import math
import subprocess
import sys
from pathlib import Path
from collections import defaultdict
from typing import Dict, List, Optional, Sequence, Tuple

# Metrics requested from Nsight Compute; keep order stable for parsing.
NCU_METRICS = (
    "dram__bytes.sum.per_second",
    "dram__bytes.sum",
    "dram__sectors_read.sum",
    "dram__sectors_write.sum",
    "dram__throughput.avg.pct_of_peak_sustained_elapsed",
    "lts__t_bytes.sum.per_second",
    "lts__throughput.avg.pct_of_peak_sustained_elapsed",
    "lts__t_requests",
    "lts__t_requests_lookup_miss",
    "lts__t_requests_op_atom.sum.per_second",
    "lts__t_requests_op_atom_dot_cas.sum.per_second",
    "lts__t_requests_op_atom_dot_alu.sum.per_second",
    "lts__t_sectors_op_atom.sum.per_second",
    "lts__t_sector_op_atom_hit_rate",
    "l1tex__m_xbar2l1tex_read_bytes_mem_global_op_atom.sum.per_second",
    "l1tex__m_l1tex2xbar_write_bytes_mem_global_op_atom.sum.per_second",
    "l1tex__t_sectors_pipe_lsu_mem_global_op_atom.sum",
    "l1tex__t_sectors_pipe_lsu_mem_global_op_red.sum",
    "sm__throughput.avg.pct_of_peak_sustained_elapsed",
    "smsp__inst_executed.sum",
)

HEADER = (
    "test",
    "mode",
    "bytes",
    "elements",
    "blocks",
    "threads",
    "iters",
    "stride",
    "time_ms",
    "eff_GBs",
    "L2_bytes",
    "L2_MB",
    "dram_gbps",
    "dram_bytes_sum",
    "dram_sectors_read",
    "dram_sectors_write",
    "l2_requests",
    "l2_lookup_miss",
    "dram_pct_of_peak",
    "l2_gbps",
    "l2_pct_of_peak",
    "l2_miss_rate_pct",
    "avg_dram_tx_bytes",
    "sm_pct_of_peak",
    "sm_inst_total",
    "global_atomic_transactions",
    "atomic_bytes_ps",
    "atomic_hit_rate",
    "atomic_cas_ratio",
    "atomic_ratio",
    "gpu_name",
    "l2_mb",
)

EXPECTED_KERNEL_SUBSTRINGS = {
    "dram_test": "dram_triad",
    "l2_sweep": "l2_sweep_kernel",
    "histogram": "histogram_kernel",
}


def parse_microbenchmark_output(text: str) -> Tuple[Optional[str], List[Dict[str, str]]]:
    """Extract the microbenchmark CSV block (GPU line + rows)."""
    gpu_line = None
    table_lines: List[str] = []
    in_table = False
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        if line.startswith("GPU:"):
            gpu_line = line
            continue
        if line.startswith("test,mode,"):
            in_table = True
            table_lines = [line]
            continue
        if in_table:
            if line.startswith("==PROF=="):
                in_table = False
                continue
            if "," not in line:
                # Reached the end of the CSV rows emitted by the benchmark.
                in_table = False
                continue
            table_lines.append(line)
    if not table_lines:
        raise RuntimeError("Failed to find microbenchmark CSV output in Nsight Compute stdout")
    reader = csv.DictReader(io.StringIO("\n".join(table_lines)))
    rows = [dict(row) for row in reader]
    if not rows:
        raise RuntimeError("Microbenchmark CSV block did not contain data rows")
    return gpu_line, rows


def parse_csv_line(line: str) -> List[str]:
    """Parse a single CSV line while respecting quoting."""
    return next(csv.reader([line]))


def parse_ncu_metrics(text: str, kernel_hint: Optional[str]) -> List[Dict[str, str]]:
    """Extract metrics from Nsight Compute CSV output regardless of version format."""

    # First pass: locate the canonical header (ID, Kernel Name, Metric Name, ...).
    header: Optional[List[str]] = None
    records: List[Dict[str, str]] = []
    for row in csv.reader(io.StringIO(text)):
        if not row:
            continue
        if row[0].startswith("==PROF=="):
            # Profiler log messages; skip them.
            continue
        if header is None:
            if "Kernel Name" in row and "Metric Name" in row and "Metric Value" in row:
                header = row
            continue
        if len(row) != len(header):
            continue
        records.append(dict(zip(header, row)))

    if not records:
        raise RuntimeError("Failed to parse Nsight Compute metric table")

    # Group per kernel launch so we can attach all requested metrics together.
    grouped: Dict[Tuple[str, str], Dict[str, str]] = defaultdict(dict)
    for record in records:
        kernel = record.get("Kernel Name", "")
        if kernel_hint and kernel_hint not in kernel:
            continue
        key = (record.get("ID", ""), kernel)
        metric_name = record.get("Metric Name")
        metric_value = record.get("Metric Value")
        if metric_name:
            grouped[key][metric_name] = metric_value or ""

    if not grouped:
        # If kernel_hint filtered everything out, retry without it to aid debugging.
        for record in records:
            key = (record.get("ID", ""), record.get("Kernel Name", ""))
            metric_name = record.get("Metric Name")
            metric_value = record.get("Metric Value")
            if metric_name:
                grouped[key][metric_name] = metric_value or ""
    if not grouped:
        raise RuntimeError("Nsight Compute output did not contain requested metrics")

    # Sort by launch ID for deterministic pairing with microbenchmark rows.
    def sort_key(item: Tuple[Tuple[str, str], Dict[str, str]]) -> Tuple[int, str]:
        launch_id, kernel_name = item[0]
        try:
            return int(launch_id), kernel_name
        except (TypeError, ValueError):
            return 0, kernel_name

    merged_rows: List[Dict[str, str]] = []
    for (launch_id, kernel_name), metrics in sorted(grouped.items(), key=sort_key):
        row = dict(metrics)
        row["Kernel Name"] = kernel_name
        row["ID"] = launch_id
        merged_rows.append(row)

    return merged_rows


def format_float(value: float) -> str:
    return f"{value:.3f}"


def extract_gpu_name(gpu_line: Optional[str]) -> Tuple[str, str]:
    if not gpu_line:
        return "", ""
    # Expected form: "GPU: <name>, L2 MB=<value>"
    name_part, _, rest = gpu_line.partition(",")
    gpu_name = name_part.replace("GPU:", "", 1).strip()
    l2_value = ""
    if "L2 MB=" in rest:
        _, _, tail = rest.partition("L2 MB=")
        l2_value = tail.strip()
    return gpu_name, l2_value


def to_float_maybe(text: Optional[str]) -> Optional[float]:
    if text is None:
        return None
    stripped = text.strip()
    if not stripped:
        return None
    stripped = stripped.replace("\xa0", "")  # non-breaking spaces
    stripped = stripped.replace("\u202f", "")  # narrow no-break space
    stripped = stripped.replace(" ", "")
    stripped = stripped.replace(",", "")
    lowered = stripped.lower()
    if lowered in {"nan", "inf", "+inf", "-inf", ""}:
        return None
    try:
        value = float(stripped)
        if math.isnan(value) or math.isinf(value):
            return None
        return value
    except ValueError:
        return None


def _find_metric(metrics: Dict[str, str], name: str) -> Optional[str]:
    if name in metrics:
        return metrics[name]
    for key, value in metrics.items():
        if key.startswith(name) or name in key:
            return value
    return None


def _derive_bytes_per_second(
    metric: Dict[str, str],
    micro: Dict[str, str],
    per_second_key: str,
    total_key: str,
) -> Optional[float]:
    per_second = to_float_maybe(_find_metric(metric, per_second_key))
    if per_second is not None:
        return per_second
    total = to_float_maybe(_find_metric(metric, total_key))
    if total is None:
        return None
    time_ms = to_float_maybe(micro.get("time_ms"))
    if time_ms is None or time_ms <= 0.0:
        return None
    return total / (time_ms / 1e3)


def _derive_count(
    metric: Dict[str, str],
    per_second_key: str,
    total_key: str,
    time_ms: Optional[float],
) -> Optional[float]:
    per_second = to_float_maybe(_find_metric(metric, per_second_key))
    if per_second is not None and time_ms is not None and time_ms > 0.0:
        return per_second * (time_ms / 1e3)
    total = to_float_maybe(_find_metric(metric, total_key))
    return total

EPS = 1e-9


def merge_rows(
    micro_rows: List[Dict[str, str]],
    metric_rows: List[Dict[str, str]],
    gpu_name: str,
    l2_mb_text: str,
    debug: bool = False,
) -> List[Dict[str, str]]:
    merged: List[Dict[str, str]] = []
    if len(metric_rows) < len(micro_rows):
        sys.stderr.write(
            f"Warning: microbenchmark emitted {len(micro_rows)} rows but NCU produced {len(metric_rows)} metrics; truncating to the shorter list.\n"
        )
    for micro, metric in zip(micro_rows, metric_rows):
        time_ms_val = to_float_maybe(micro.get("time_ms"))
        dram_bytes_per_second = _derive_bytes_per_second(
            metric,
            micro,
            "dram__bytes.sum.per_second",
            "dram__bytes.sum",
        )
        dram_bytes_sum = to_float_maybe(_find_metric(metric, "dram__bytes.sum"))
        dram_sectors_read = to_float_maybe(_find_metric(metric, "dram__sectors_read.sum"))
        dram_sectors_write = to_float_maybe(_find_metric(metric, "dram__sectors_write.sum"))
        dram_pct = to_float_maybe(_find_metric(metric, "dram__throughput.avg.pct_of_peak_sustained_elapsed"))
        l2_bytes = _derive_bytes_per_second(
            metric,
            micro,
            "lts__t_bytes.sum.per_second",
            "lts__t_bytes.sum",
        )
        l2_pct = to_float_maybe(_find_metric(metric, "lts__throughput.avg.pct_of_peak_sustained_elapsed"))

        l2_requests = to_float_maybe(_find_metric(metric, "lts__t_requests.sum"))
        l2_lookup_miss = to_float_maybe(_find_metric(metric, "lts__t_requests_lookup_miss.sum"))
        l2_miss_rate = None
        if l2_requests is not None and l2_lookup_miss is not None and l2_requests > 0:
            l2_miss_rate = (l2_lookup_miss / l2_requests) * 100.0

        atomic_atom = _derive_count(
            metric,
            "l1tex__t_sectors_pipe_lsu_mem_global_op_atom.sum.per_second",
            "l1tex__t_sectors_pipe_lsu_mem_global_op_atom.sum",
            time_ms_val,
        )
        atomic_red = _derive_count(
            metric,
            "l1tex__t_sectors_pipe_lsu_mem_global_op_red.sum.per_second",
            "l1tex__t_sectors_pipe_lsu_mem_global_op_red.sum",
            time_ms_val,
        )
        atomic_transactions = (atomic_atom or 0.0) + (atomic_red or 0.0)

        atom_req_ps = to_float_maybe(_find_metric(metric, "lts__t_requests_op_atom.sum.per_second"))
        atom_cas_ps = to_float_maybe(_find_metric(metric, "lts__t_requests_op_atom_dot_cas.sum.per_second"))
        atom_alu_ps = to_float_maybe(_find_metric(metric, "lts__t_requests_op_atom_dot_alu.sum.per_second"))
        atom_sectors_ps = to_float_maybe(_find_metric(metric, "lts__t_sectors_op_atom.sum.per_second"))
        atom_hit_rate = to_float_maybe(_find_metric(metric, "lts__t_sector_op_atom_hit_rate"))
        l1_atom_bytes_read_ps = to_float_maybe(
            _find_metric(metric, "l1tex__m_xbar2l1tex_read_bytes_mem_global_op_atom.sum.per_second")
        )
        l1_atom_bytes_write_ps = to_float_maybe(
            _find_metric(metric, "l1tex__m_l1tex2xbar_write_bytes_mem_global_op_atom.sum.per_second")
        )

        sm_pct = to_float_maybe(_find_metric(metric, "sm__throughput.avg.pct_of_peak_sustained_elapsed"))
        sm_inst = to_float_maybe(_find_metric(metric, "smsp__inst_executed.sum"))

        micro_dram = to_float_maybe(micro.get("eff_GBs"))
        micro_l2 = None
        micro_l2_bytes = to_float_maybe(micro.get("L2_bytes"))
        if micro_l2_bytes is not None and time_ms_val is not None and time_ms_val > 0.0:
            micro_l2 = (micro_l2_bytes / (time_ms_val / 1e3)) / 1e9

        if dram_bytes_per_second is None and l2_bytes is None:
            if micro_dram is None and micro_l2 is None:
                message = "Warning: Nsight metrics missing and no fallback values available; row omitted.\n"
                if debug:
                    sys.stderr.write(message.rstrip("\n") + f" {metric}\n")
                else:
                    sys.stderr.write(message)
                continue
        row = dict(micro)
        dram_gbps = None
        if dram_bytes_per_second is not None:
            dram_gbps = dram_bytes_per_second / 1e9
        elif micro_dram is not None:
            dram_gbps = micro_dram
        row["dram_gbps"] = format_float(dram_gbps) if dram_gbps is not None else ""
        row["dram_bytes_sum"] = format_float(dram_bytes_sum) if dram_bytes_sum is not None else ""
        row["dram_sectors_read"] = format_float(dram_sectors_read) if dram_sectors_read is not None else ""
        row["dram_sectors_write"] = format_float(dram_sectors_write) if dram_sectors_write is not None else ""
        row["l2_requests"] = format_float(l2_requests) if l2_requests is not None else ""
        row["l2_lookup_miss"] = format_float(l2_lookup_miss) if l2_lookup_miss is not None else ""
        row["dram_pct_of_peak"] = format_float(dram_pct) if dram_pct is not None else ""

        l2_gbps = None
        if l2_bytes is not None:
            l2_gbps = l2_bytes / 1e9
        elif micro_l2 is not None:
            l2_gbps = micro_l2
        row["l2_gbps"] = format_float(l2_gbps) if l2_gbps is not None else ""
        row["l2_pct_of_peak"] = format_float(l2_pct) if l2_pct is not None else ""
        row["l2_miss_rate_pct"] = format_float(l2_miss_rate) if l2_miss_rate is not None else ""
        total_dram_tx = (dram_sectors_read or 0.0) + (dram_sectors_write or 0.0)
        avg_tx_bytes = (dram_bytes_sum or 0.0) / total_dram_tx if total_dram_tx > 0.0 else None
        row["avg_dram_tx_bytes"] = format_float(avg_tx_bytes) if avg_tx_bytes is not None else ""
        row["sm_pct_of_peak"] = format_float(sm_pct) if sm_pct is not None else ""
        row["sm_inst_total"] = format_float(sm_inst) if sm_inst is not None else ""
        row["global_atomic_transactions"] = (
            format_float(atomic_transactions) if atomic_transactions > 0.0 else ""
        )
        atomic_bytes_ps = None
        if atom_sectors_ps is not None:
            atomic_bytes_ps = 32.0 * atom_sectors_ps
        elif l1_atom_bytes_read_ps is not None or l1_atom_bytes_write_ps is not None:
            atomic_bytes_ps = (l1_atom_bytes_read_ps or 0.0) + (l1_atom_bytes_write_ps or 0.0)
        if atomic_bytes_ps is not None:
            row["atomic_bytes_ps"] = format_float(atomic_bytes_ps)
        else:
            row["atomic_bytes_ps"] = ""
        if atom_hit_rate is not None:
            row["atomic_hit_rate"] = format_float(atom_hit_rate)
        else:
            row["atomic_hit_rate"] = ""
        cas_ratio = None
        if atom_cas_ps is not None and atom_req_ps is not None:
            cas_ratio = atom_cas_ps / max(atom_req_ps, EPS)
        row["atomic_cas_ratio"] = format_float(cas_ratio) if cas_ratio is not None else ""
        atomic_ratio = None
        if atom_req_ps is not None:
            req_total_ps = None
            if l2_requests is not None and time_ms_val is not None and time_ms_val > 0.0:
                req_total_ps = l2_requests / (time_ms_val / 1e3)
            atomic_ratio = atom_req_ps / max(req_total_ps or atom_req_ps, EPS)
        row["atomic_ratio"] = format_float(atomic_ratio) if atomic_ratio is not None else ""
        row["gpu_name"] = gpu_name
        row["l2_mb"] = l2_mb_text or row.get("L2_MB", "")
        merged.append(row)
    if not merged:
        raise RuntimeError("No rows remained after merging microbenchmark and NCU metrics")
    return merged


def _project_requested_metrics(row: Dict[str, str], metrics: List[str]) -> Dict[str, object]:
    """Extract stable solo-scaling metrics from a merged Nsight row."""
    normalized = {
        "time_ms": to_float_maybe(row.get("time_ms")),
        "dram_gbps": to_float_maybe(row.get("dram_gbps")),
        "l2_gbps": to_float_maybe(row.get("l2_gbps")),
        "sm_util_pct": to_float_maybe(row.get("sm_pct_of_peak")),
        "l2_miss_rate_pct": to_float_maybe(row.get("l2_miss_rate_pct")),
        "gpu_name": row.get("gpu_name"),
        "raw": row,
    }
    if metrics:
        projected = {key: normalized.get(key) for key in metrics if key in normalized}
        projected["gpu_name"] = normalized.get("gpu_name")
        projected["raw"] = row
        return projected
    return normalized


def run_ncu_profile(cmd: List[str], metrics: List[str]) -> Dict[str, object]:
    """Run Nsight Compute with given command and metrics, return parsed JSON."""
    ncu_cmd = build_ncu_command("ncu", cmd, timeout=None)
    result = run_ncu(ncu_cmd)
    gpu_line, micro_rows = parse_microbenchmark_output(result.stdout)
    gpu_name, l2_mb_text = extract_gpu_name(gpu_line)
    test_name = micro_rows[0].get("test", "") if micro_rows else ""
    kernel_hint = EXPECTED_KERNEL_SUBSTRINGS.get(test_name)
    metric_rows = parse_ncu_metrics(result.stdout, kernel_hint)
    merged_rows = merge_rows(micro_rows, metric_rows, gpu_name, l2_mb_text)
    if not merged_rows:
        return {}
    merged_row = merged_rows[0]
    return _project_requested_metrics(merged_row, metrics)


def write_csv(path: Path, rows: List[Dict[str, str]], gpu_line: Optional[str], append: bool) -> None:
    ensure_parent(path)
    should_append = append and path.exists()
    mode = "a" if should_append else "w"
    with path.open(mode, newline="") as fh:
        if not should_append:
            if gpu_line:
                fh.write(f"{gpu_line}\n")
            writer = csv.DictWriter(fh, fieldnames=HEADER)
            writer.writeheader()
        else:
            writer = csv.DictWriter(fh, fieldnames=HEADER)
        for row in rows:
            writer.writerow({key: row.get(key, "") for key in HEADER})


def run_ncu(command: Sequence[str]) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(
        command,
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if result.returncode != 0:
        sys.stderr.write(
            f"Warning: Nsight Compute exited with code {result.returncode}; attempting to parse output anyway.\n"
        )
    return result


def ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def build_ncu_command(ncu_path: str, user_cmd: List[str], timeout: Optional[int]) -> List[str]:
    command = [
        ncu_path,
        "--csv",
        "--metrics",
        ",".join(NCU_METRICS),
        "--target-processes",
        "all",
    ]
    if timeout is not None:
        # Some NCU builds (newer releases) accept --timeout; older ones reject it, so gate on user request.
        command.extend(["--timeout", str(timeout)])
    command.extend(["--", *user_cmd])
    return command


def main() -> None:
    parser = argparse.ArgumentParser(description="Capture Nsight Compute metrics for calibration runs")
    parser.add_argument("--out", type=Path, required=True, help="Destination CSV path for merged metrics")
    parser.add_argument(
        "--ncu-path",
        default="ncu",
        help="Path to the Nsight Compute CLI executable (default: ncu)",
    )
    parser.add_argument("--append", action="store_true", help="Append to existing CSV instead of overwriting")
    parser.add_argument(
        "--timeout",
        type=int,
        help="Optional Nsight Compute timeout value; omit to skip passing --timeout",
    )
    parser.add_argument("--debug", action="store_true", help="Print raw Nsight metrics when parsing issues arise")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to profile (must follow --)")
    args = parser.parse_args()

    if not args.command or args.command[0] != "--":
        parser.error("Expected command after --")
    user_cmd = args.command[1:]
    if not user_cmd:
        parser.error("No executable provided for profiling")

    ncu_cmd = build_ncu_command(args.ncu_path, user_cmd, args.timeout)
    result = run_ncu(ncu_cmd)

    gpu_line, micro_rows = parse_microbenchmark_output(result.stdout)
    gpu_name, l2_mb_text = extract_gpu_name(gpu_line)
    test_name = micro_rows[0].get("test", "")
    kernel_hint = EXPECTED_KERNEL_SUBSTRINGS.get(test_name)
    metric_rows = parse_ncu_metrics(result.stdout, kernel_hint)
    if args.debug and metric_rows:
        sys.stderr.write(f"NCU metric sample: {metric_rows[0]}\n")
    merged_rows = merge_rows(micro_rows, metric_rows, gpu_name, l2_mb_text, debug=args.debug)
    write_csv(args.out, merged_rows, gpu_line, append=args.append)
    print(f"Wrote {len(merged_rows)} rows to {args.out}")


if __name__ == "__main__":
    main()