Load and combine NCU + Proton data into a DataFrame. Columns: bsz, method, matmul_us, sampling_us, total_us.
(ncu_dir: Path, proton_dir: Path)
| 40 | |
| 41 | |
| 42 | def load_data(ncu_dir: Path, proton_dir: Path) -> pd.DataFrame: |
| 43 | """Load and combine NCU + Proton data into a DataFrame. |
| 44 | |
| 45 | Columns: bsz, method, matmul_us, sampling_us, total_us. |
| 46 | """ |
| 47 | bsz_dirs = sorted(ncu_dir.glob("bsz*"), key=lambda p: int(p.name[3:])) |
| 48 | frames = [] |
| 49 | |
| 50 | for d in bsz_dirs: |
| 51 | bsz = int(d.name[3:]) |
| 52 | |
| 53 | # Load Proton data for FMMS split |
| 54 | proton_trace = parse_chrome_trace(proton_dir / f"bsz{bsz}" / "kernel.chrome_trace") |
| 55 | proton_pcts = trace_phase_pcts(proton_trace) if proton_trace else None |
| 56 | |
| 57 | for fname, label, is_fmms in METHODS: |
| 58 | path = d / fname |
| 59 | if not path.exists(): |
| 60 | continue |
| 61 | kdf = parse_ncu_csv(path) |
| 62 | is_matmul = kdf["kernel_name"].str.contains("gemm|gemv|fused_mm_sample", case=False) |
| 63 | assert is_matmul.iloc[0], ( |
| 64 | f"First kernel in {path} is not a matmul: {kdf['kernel_name'].iloc[0]}" |
| 65 | ) |
| 66 | row = { |
| 67 | "bsz": bsz, |
| 68 | "method": label, |
| 69 | "total_us": kdf["duration_us"].sum(), |
| 70 | } |
| 71 | if is_fmms: |
| 72 | if not proton_pcts: |
| 73 | print(f"WARNING: skipping FMMS at bsz={bsz}: no Proton trace found") |
| 74 | continue |
| 75 | # Split the fused kernel using Proton percentages, then add |
| 76 | # auxiliary kernels (local reduce, TP reduce) to sampling. |
| 77 | fused_us = kdf.loc[is_matmul, "duration_us"].sum() |
| 78 | aux_us = kdf.loc[~is_matmul, "duration_us"].sum() |
| 79 | matmul_frac = proton_pcts["matmul"] / 100 |
| 80 | sampling_frac = proton_pcts["sampling"] / 100 |
| 81 | row["matmul_us"] = fused_us * matmul_frac |
| 82 | row["sampling_us"] = fused_us * sampling_frac + aux_us |
| 83 | else: |
| 84 | row["matmul_us"] = kdf.loc[is_matmul, "duration_us"].sum() |
| 85 | row["sampling_us"] = kdf.loc[~is_matmul, "duration_us"].sum() |
| 86 | |
| 87 | frames.append(row) |
| 88 | |
| 89 | df = pd.DataFrame(frames) |
| 90 | df[["matmul_us", "sampling_us", "total_us"]] = df[ |
| 91 | ["matmul_us", "sampling_us", "total_us"] |
| 92 | ].round(1) |
| 93 | return df |
| 94 | |
| 95 | |
| 96 | def save_csv(df: pd.DataFrame, path: Path) -> None: |
no test coverage detected