Function parse_ncu_csv

benchmarking/parse_ncu_sweep.py:75–103

Parse a single NCU CSV file into a DataFrame. NCU CSV files contain non-CSV lines (==PROF== messages, Python stdout). Only lines starting with '"' are actual CSV rows. The second CSV row is a units row (e.g. "ns", "block") that must be skipped. Numbers use comma as thousands separat

(path: Path)

def parse_ncu_csv(path: Path) -> pd.DataFrame:
    """Parse a single NCU CSV file into a DataFrame.

    NCU CSV files contain non-CSV lines (==PROF== messages, Python stdout).
    Only lines starting with '"' are actual CSV rows. The second CSV row is a
    units row (e.g. "ns", "block") that must be skipped. Numbers use comma as
    thousands separator. Duration is in nanoseconds.

    Returns a DataFrame with columns: kernel_name, duration_us,
    averaged across ranks for TP>1.
    """
    content = path.read_text()
    csv_lines = [line for line in content.splitlines() if line.startswith('"')]
    assert csv_lines, f"No CSV data found in {path}"
    df = pd.read_csv(
        io.StringIO("\n".join(csv_lines)),
        skiprows=[1],  # skip units row
        thousands=",",
    )
    assert "Kernel Name" in df.columns and "gpu__time_duration.sum" in df.columns, (
        f"Expected columns 'Kernel Name' and 'gpu__time_duration.sum' in {path}, "
        f"got: {list(df.columns)}"
    )
    return (
        df.rename(columns={"Kernel Name": "kernel_name", "gpu__time_duration.sum": "duration_ns"})
        .groupby("kernel_name", as_index=False, sort=False)
        .agg(duration_us=("duration_ns", "mean"))
        .assign(duration_us=lambda d: d["duration_us"] / 1000)
    )


def _print_summary(df: pd.DataFrame) -> None: