Print all summary sections from the combined sweep DataFrame.
(df: pd.DataFrame)
| 104 | |
| 105 | |
| 106 | def _print_summary(df: pd.DataFrame) -> None: |
| 107 | """Print all summary sections from the combined sweep DataFrame.""" |
| 108 | labels = df["method"].unique().tolist() |
| 109 | baselines = [m for m in labels if m != FUSED] |
| 110 | has_fused = FUSED in labels |
| 111 | |
| 112 | # Tag each kernel as matmul or sampling |
| 113 | matmul_patterns = ("gemm", "gemv", "fused_mm_sample") |
| 114 | df = df.assign( |
| 115 | is_matmul=df["kernel_name"] |
| 116 | .str.lower() |
| 117 | .apply(lambda n: any(p in n for p in matmul_patterns)) |
| 118 | ) |
| 119 | # The first kernel per (bsz, method) must be the matmul |
| 120 | first_kernels = df.groupby(["bsz", "method"]).first() |
| 121 | bad = first_kernels.query("not is_matmul") |
| 122 | assert bad.empty, ( |
| 123 | f"First kernel is not a matmul (expected one of {matmul_patterns} in name):\n" |
| 124 | + bad[["kernel_name"]].to_string() |
| 125 | ) |
| 126 | |
| 127 | # Per-method totals: bsz x method -> total_us, matmul_us, sampling_us |
| 128 | totals = ( |
| 129 | df.groupby(["bsz", "method"]) |
| 130 | .apply( |
| 131 | lambda g: pd.Series( |
| 132 | { |
| 133 | "total_us": g["duration_us"].sum(), |
| 134 | "matmul_us": g.query("is_matmul")["duration_us"].sum(), |
| 135 | "sampling_us": g.query("not is_matmul")["duration_us"].sum(), |
| 136 | } |
| 137 | ), |
| 138 | include_groups=False, |
| 139 | ) |
| 140 | .reset_index() |
| 141 | ) |
| 142 | |
| 143 | # ── Section 1: Total time per method ── |
| 144 | df_total = totals.pivot(index="bsz", columns="method", values="total_us").round(1) |
| 145 | df_total = df_total.reindex(columns=labels) |
| 146 | df_total.index.name = "N" |
| 147 | print("1. TOTAL TIME (us)") |
| 148 | print() |
| 149 | print(df_total.to_markdown()) |
| 150 | |
| 151 | # ── Section 2: Speedup vs fused-triton ── |
| 152 | if has_fused and baselines: |
| 153 | fused_totals = totals.query("method == @FUSED")[["bsz", "total_us"]].rename( |
| 154 | columns={"total_us": "fused_us"} |
| 155 | ) |
| 156 | speedup = ( |
| 157 | totals.query("method != @FUSED") |
| 158 | .merge(fused_totals, on="bsz") |
| 159 | .assign(speedup=lambda d: d["total_us"] / d["fused_us"]) |
| 160 | ) |
| 161 | df_speedup = ( |
| 162 | speedup.pivot(index="bsz", columns="method", values="speedup") |
| 163 | .reindex(columns=baselines) |