Stacked bar chart: one bar per batch size, one segment per kernel.
(df: pd.DataFrame, out: Path, method: str)
| 78 | |
| 79 | |
| 80 | def plot_breakdown(df: pd.DataFrame, out: Path, method: str) -> None: |
| 81 | """Stacked bar chart: one bar per batch size, one segment per kernel.""" |
| 82 | batch_sizes = sorted(df["bsz"].unique()) |
| 83 | # Preserve first-seen order for consistent stacking |
| 84 | kernels = df.drop_duplicates("kernel", keep="first")["kernel"].tolist() |
| 85 | palette = sns.color_palette("tab10", n_colors=len(kernels)) |
| 86 | |
| 87 | pivot = ( |
| 88 | df.assign(duration_ms=df["duration_us"] / 1000) |
| 89 | .pivot(index="bsz", columns="kernel", values="duration_ms") |
| 90 | .reindex(columns=kernels) |
| 91 | .fillna(0) |
| 92 | ) |
| 93 | |
| 94 | sns.set_context("talk") |
| 95 | fig, ax = plt.subplots() |
| 96 | x = range(len(batch_sizes)) |
| 97 | bar_width = 0.6 |
| 98 | |
| 99 | bottom = pd.Series(0.0, index=pivot.index) |
| 100 | for i, kernel in enumerate(kernels): |
| 101 | values = pivot[kernel] |
| 102 | bars = ax.bar( |
| 103 | x, |
| 104 | values, |
| 105 | bottom=bottom, |
| 106 | width=bar_width, |
| 107 | label=kernel, |
| 108 | color=palette[i], |
| 109 | ) |
| 110 | for rect, v in zip(bars.patches, values): |
| 111 | _label_bar(ax, rect, f"{v:.2f}") |
| 112 | bottom = bottom + values |
| 113 | |
| 114 | ax.set_xticks(list(x)) |
| 115 | ax.set_xticklabels([str(b) for b in batch_sizes]) |
| 116 | ax.set_xlabel("Batch Size") |
| 117 | ax.set_ylabel("Time (ms)") |
| 118 | ax.grid(alpha=0.5, axis="y") |
| 119 | |
| 120 | ncol = 2 if len(kernels) > 4 else 1 |
| 121 | ax.legend( |
| 122 | title="Kernel", |
| 123 | loc="lower center", |
| 124 | bbox_to_anchor=(0.5, 1.02), |
| 125 | ncol=ncol, |
| 126 | fontsize=9, |
| 127 | framealpha=0.9, |
| 128 | ) |
| 129 | |
| 130 | fig.savefig(out, dpi=300, bbox_inches="tight") |
| 131 | print(f"Saved {out}") |
| 132 | plt.close(fig) |
| 133 | |
| 134 | |
| 135 | def _shorten_kernel(name: str) -> str: |