Load NCU data for a single method across all batch sizes. Kernel names are shortened and durations aggregated by short name, so multiple raw kernels that map to the same short name are summed.
(base: Path, bsz_dirs: list[Path], method: str)
| 56 | |
| 57 | |
| 58 | def load_method(base: Path, bsz_dirs: list[Path], method: str) -> pd.DataFrame: |
| 59 | """Load NCU data for a single method across all batch sizes. |
| 60 | |
| 61 | Kernel names are shortened and durations aggregated by short name, so |
| 62 | multiple raw kernels that map to the same short name are summed. |
| 63 | """ |
| 64 | fname = METHOD_FILENAMES[method] |
| 65 | frames = [] |
| 66 | for d in bsz_dirs: |
| 67 | bsz = int(d.name[3:]) |
| 68 | path = d / fname |
| 69 | if not path.exists(): |
| 70 | continue |
| 71 | method_df = parse_ncu_csv(path).assign(bsz=bsz) |
| 72 | method_df["kernel"] = method_df["kernel_name"].map(_shorten_kernel) |
| 73 | frames.append(method_df) |
| 74 | df = pd.concat(frames, ignore_index=True) |
| 75 | return df.groupby(["bsz", "kernel"], as_index=False, sort=False).agg( |
| 76 | duration_us=("duration_us", "sum") |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | def plot_breakdown(df: pd.DataFrame, out: Path, method: str) -> None: |