[Bench] Make roofline graph lines and title clearer (#6852)

Use the first compute bound data point as the end point for the
bandwidth-bound line so that it can connect with the compute-bound line.
This is visually better espically if we have sparse data points.

Along the way also make the title clear regarding the benchmarked case.

Author: antiagainst

python/triton_kernels/bench/bench_mlp.py

@@ -175,7 +175,8 @@ def roofline_mlp(batch_ranges, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_
     batches = list(chain(*[range(*r) for r in batch_ranges]))
     # collect performance data
     perfs = []
-    print(f"Benchmarking {name} ({x_dtype}x{w_dtype}, TP={TP}, EP={EP})...")
+    bench_case = f"{name} ({x_dtype}x{w_dtype}, TP={TP}, EP={EP})"
+    print(f"Benchmarking {bench_case}...")
     print("===============================================================")
     for batch in batches:
         perfs += [bench_mlp(batch, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_dtype, TP, EP, name)]
@@ -186,7 +187,7 @@ def roofline_mlp(batch_ranges, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_
     fig, ax = plt.subplots(figsize=(7, 5), dpi=120)
     ax.set_xlabel("batch size (toks/expt)")
     ax.set_ylabel("performance  [TFLOP/s]")
-    ax.set_title("roofline")
+    ax.set_title(f"{bench_case} roofline")
     # add a tiny margin so points are not flush with the frame
     xs = [batch * n_expts_act / n_expts_tot for batch in batches]
     perf = [p.tflops for p in perfs]
@@ -200,7 +201,8 @@ def roofline_mlp(batch_ranges, dim1, dim2, n_expts_tot, n_expts_act, x_dtype, w_
     opints = [p.opint for p in perfs]
     knee = bisect_left(opints, max_tflops / max_tbps) - 1
     x_bw, x_comp = xs[:knee], xs[knee:]
-    y_bw = [op * max_tbps for op in opints[:knee]]
+    x_bw = [x_bw[0], x_comp[0]]
+    y_bw = [opints[0] * max_tbps, max_tflops]
     y_comp = [max_tflops] * len(x_comp)
     ax.plot(x_bw, y_bw, "--", label=f"BW-bound  ({max_tbps:.0f} TB/s)")
     ax.plot(x_comp, y_comp, "--", label=f"Compute-bound  ({max_tflops:.0f} TFLOP/s)")