extract_metrics.py

#!/usr/bin/env python3
"""
Extract a curated set of per-kernel metrics from every .ncu-rep file in
profiles/ncu/ and write a tidy CSV at profiles/reports/ncu_metrics.csv.

Also extracts kernel time-summary from each .nsys-rep into
profiles/reports/nsys_kernel_summary.csv.
"""
import argparse
import csv
import glob
import os
import re
import subprocess
import sys
from collections import defaultdict

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

_argp = argparse.ArgumentParser(add_help=False)
_argp.add_argument("--base", default=os.path.join(ROOT, "profiles"),
                   help="Base profile dir (default: profiles/). "
                        "Use profiles/a4000 or profiles/3050_laptop to "
                        "process per-GPU subtrees.")
_args, _ = _argp.parse_known_args()
BASE     = os.path.abspath(_args.base)
NCU_DIR  = os.path.join(BASE, "ncu")
NSYS_DIR = os.path.join(BASE, "nsys")
OUT_DIR  = os.path.join(BASE, "reports")
os.makedirs(OUT_DIR, exist_ok=True)

# Map noisy mangled / fully-qualified kernel names to short labels we use
# in the report. Cuts cuFFT internals to a single bucket per family.
def short_kernel(name):
    n = name
    # Our hand-written kernels — drop namespace + signature
    for k in ("scatter_ones_kernel",
              "sparse_row_dft_naive",
              "sparse_row_dft_tile",
              "sparse_row_dft_iter_u16",
              "sparse_row_dft_iter",
              "sparse_row_dft_table",
              "build_twiddle_table",
              "transpose_kernel",
              "tile_stats",
              "stats_pass1"):
        if k in n:
            return k
    # cuFFT kernel families (case-insensitive, matches lowercase _fft variants
    # like regular_bluestein_fft, multi_bluestein_fft, prime_fft_factor, ...)
    cufft_families = [
        "regular_bluestein_fft",
        "multi_bluestein_fft",
        "regular_fft_factor",
        "prime_fft_factor",
        "vector_fft",
        "smem_fft",
        "scale_conjugate",
        "generate_chirp_signal",
        "convert",
        "spRealComplex",
        "spVector",
    ]
    for fam in cufft_families:
        if fam in n:
            return f"cufft::{fam}"
    # Last-resort generic match
    m = re.search(r"(vector_fft|tilegen|c2r_kernel|r2c_kernel|fft)",
                  n, re.IGNORECASE)
    if m:
        return f"cufft::{m.group(1).lower()}"
    return n.split("(")[0][:60]


def parse_ver_mtx(filename):
    """profiles/ncu/v3_benzene.ncu-rep → ('v3', 'benzene')"""
    base = os.path.basename(filename)
    base = base.replace(".ncu-rep", "").replace(".nsys-rep", "")
    parts = base.split("_", 1)
    if len(parts) != 2:
        return None, None
    return parts[0], parts[1]


# ---------- ncu ----------

# Metrics we want, by ncu metric name. With --print-units base:
#   Duration → nsecond
#   throughput percentages → %
#   cache throughputs → %
#   occupancy → %
#   cycles → cycle
WANTED_METRICS = [
    # Section "GPU Speed Of Light Throughput"
    ("Duration", "ns"),
    ("Compute (SM) Throughput", "%"),
    ("Memory Throughput", "%"),
    ("DRAM Throughput", "%"),
    ("L1/TEX Cache Throughput", "%"),
    ("L2 Cache Throughput", "%"),
    ("Elapsed Cycles", ""),
    # Section "Launch Statistics"
    ("Block Size", ""),
    ("Grid Size", ""),
    ("Registers Per Thread", ""),
    ("Static Shared Memory Per Block", ""),
    # Section "Occupancy"
    ("Achieved Occupancy", "%"),
    ("Theoretical Occupancy", "%"),
    # Section "Memory Workload Analysis"
    ("L1/TEX Hit Rate", "%"),
    ("L2 Hit Rate", "%"),
    # Section "Compute Workload Analysis"
    ("Executed Ipc Active", "inst/cycle"),
    # Section "Scheduler Statistics"
    ("Issued Warp Per Scheduler", "warp"),
    # Section "Warp State Statistics"
    ("Warp Cycles Per Issued Instruction", ""),
    # Pipe utilizations
    ("FMA", ""),                   # FMA pipe util
    ("ALU", ""),                   # generic ALU
    ("FP64", ""),
    ("Tensor (FP)", ""),
]

def load_ncu_csv(path):
    """Read a single .ncu-rep, return list of dicts (one per row).

    --print-units base forces consistent units across rows (ncu's default
    auto-scales into ms/us/ns per row, which makes cross-row comparisons
    misleading). Base units: Duration in ns, throughput in % of peak,
    cycles raw, etc.
    """
    res = subprocess.run(
        ["ncu", "--import", path, "--csv", "--print-units", "base"],
        check=True, capture_output=True, text=True
    )
    rows = list(csv.DictReader(res.stdout.splitlines()))
    return rows


def collect_ncu():
    out_rows = []
    for ncu_path in sorted(glob.glob(os.path.join(NCU_DIR, "*.ncu-rep"))):
        ver, mtx = parse_ver_mtx(ncu_path)
        if ver is None:
            continue
        try:
            rows = load_ncu_csv(ncu_path)
        except subprocess.CalledProcessError as e:
            print(f"[ncu_csv FAIL] {ncu_path}: {e.stderr[:200]}", file=sys.stderr)
            continue
        # Aggregate per (kernel, launch_id) → dict of metric → value
        agg = defaultdict(dict)
        for r in rows:
            kname = short_kernel(r.get("Kernel Name", ""))
            lid = r.get("ID", "")
            key = (kname, lid)
            agg[key]["kernel"] = kname
            agg[key]["launch_id"] = lid
            agg[key]["block_size"] = r.get("Block Size", "")
            agg[key]["grid_size"]  = r.get("Grid Size", "")
            mname = r.get("Metric Name", "")
            mval  = r.get("Metric Value", "").replace(",", "")
            for w, _unit in WANTED_METRICS:
                if mname == w:
                    agg[key][w] = mval
                    break
        for key, d in agg.items():
            d["version"] = ver
            d["matrix"]  = mtx
            out_rows.append(d)
    # Write CSV
    fieldnames = (["version", "matrix", "kernel", "launch_id",
                   "block_size", "grid_size"]
                  + [w for w, _ in WANTED_METRICS])
    out_path = os.path.join(OUT_DIR, "ncu_metrics.csv")
    with open(out_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        w.writeheader()
        for r in out_rows:
            w.writerow(r)
    print(f"[ok] wrote {out_path}  ({len(out_rows)} kernel-launches)")
    return out_path


# ---------- nsys ----------

def collect_nsys():
    """Use `nsys stats --report cuda_gpu_trace` to get per-kernel times."""
    out_path = os.path.join(OUT_DIR, "nsys_kernel_summary.csv")
    rows = []
    for rep in sorted(glob.glob(os.path.join(NSYS_DIR, "*.nsys-rep"))):
        ver, mtx = parse_ver_mtx(rep)
        if ver is None: continue
        try:
            r = subprocess.run(
                ["nsys", "stats", "--report", "cuda_gpu_kern_sum",
                 "--format", "csv", rep],
                check=True, capture_output=True, text=True
            )
        except subprocess.CalledProcessError as e:
            print(f"[nsys_stats FAIL] {rep}: {e.stderr[:200]}", file=sys.stderr)
            continue
        # Output contains some leading lines of metadata; the CSV header begins
        # with "Time (%)".
        out = r.stdout.splitlines()
        # find header line
        try:
            hdr_idx = next(i for i, ln in enumerate(out)
                           if ln.startswith("Time"))
        except StopIteration:
            print(f"[no kern data] {rep}")
            continue
        reader = csv.DictReader(out[hdr_idx:])
        for row in reader:
            rows.append({
                "version":    ver,
                "matrix":     mtx,
                "kernel":     short_kernel(row.get("Name") or row.get("Kernel Name") or ""),
                "time_pct":   row.get("Time (%)", "").strip(),
                "total_ms":   row.get("Total Time (ns)", "0").replace(",","") ,
                "instances":  row.get("Instances", "").strip(),
                "avg_ns":     row.get("Avg (ns)",  "0").replace(",", ""),
                "max_ns":     row.get("Max (ns)",  "0").replace(",", ""),
                "min_ns":     row.get("Min (ns)",  "0").replace(",", ""),
            })
    with open(out_path, "w", newline="") as f:
        if rows:
            w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
    print(f"[ok] wrote {out_path}  ({len(rows)} per-kernel summaries)")
    return out_path


if __name__ == "__main__":
    collect_ncu()
    collect_nsys()