Function main

benchmarking/plot_bsz_sweep_runtime.py:200–243

()

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--n_procs", type=int, default=N_PROCS, help="Number of processes (default: 1)"
    )
    parser.add_argument("--case", default=CASE, help="Benchmark case (default: small)")
    parser.add_argument("--ncu-dir", type=Path, default=None)
    parser.add_argument("--proton-dir", type=Path, default=None)
    parser.add_argument("--out-dir", type=Path, default=None)
    parser.add_argument("--fmt", default="png", help="Output image format (default: png)")
    parser.add_argument(
        "--use-name-flashsampling",
        action="store_true",
        help="Use 'FlashSampling' instead of 'FMMS' in plot labels",
    )
    args = parser.parse_args()
    if args.ncu_dir is None:
        args.ncu_dir = SWEEPS / "ncu-txt" / f"tp{args.n_procs}" / f"case-{args.case}"
    if args.proton_dir is None:
        args.proton_dir = SWEEPS / "proton" / f"tp{args.n_procs}" / f"case-{args.case}"
    if args.out_dir is None:
        args.out_dir = SWEEPS / f"tp{args.n_procs}"

    args.out_dir.mkdir(parents=True, exist_ok=True)
    rows = load_data(args.ncu_dir, args.proton_dir)
    if args.use_name_flashsampling:
        rows["method"] = rows["method"].replace(FLASHSAMPLING_RENAMES)
    save_csv(rows, args.out_dir / "runtime-breakdown.csv")
    plot(
        rows,
        args.out_dir / "sampling-latency",
        args.fmt,
        y_col="sampling_us",
        y_label="Sampling Latency (\u00b5s)",
        y_cap=800,
    )
    plot(
        rows,
        args.out_dir / "matmul-latency",
        args.fmt,
        y_col="matmul_us",
        y_label="Matmul Latency (\u00b5s)",
        style="bar",
    )


if __name__ == "__main__":