bplot.py

import argparse
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pathlib import Path
import textwrap

from convert_to_csv import to_pandas

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--results_file", type=str, help=f"results file generated by benchmarks here")
    parser.add_argument("--notes", type=str, default="", help=f"use to annotate the plot")
    args = parser.parse_args()

    results_file = Path(args.results_file)
    if not results_file.exists():
        raise ValueError(f"can't find {results_file}")

    img_file = results_file.with_suffix('.png')

    df = to_pandas(results_file)

    throughput_col = "throughput" # assumption for now

    # sort out fixed dimensions from the range ones
    fixed_dim = []
    range_cols = []
    for col in df:
        unique_vals = df[col].unique()
        if len(unique_vals) == 1:
            fixed_dim.append(f"{col}={unique_vals[0]}")
        else:
            range_cols.append(col)

    range_cols = list(set(range_cols) - set([throughput_col]))
    # XXX: at the moment assuming that only one dimension is a range, the other are fixed
    if len(range_cols) != 1:
        raise ValueError("Currently supporting plotting for benchmarks with one dimension using range")

    # these go on the xlabel along with the variable dimension
    dim_notes = ", ".join(fixed_dim)

    plt.figure(dpi=500)
    plt.plot(df[range_cols[0]], df[throughput_col])
    plt.xlabel(f"{range_cols[0]} ({dim_notes})")
    plt.ylabel("Throughput \n (TFLOP/s)")
    plt.title("Throughput of GEMMs of Various Sizes")

    # wrap notes - this can now handle several lines of text.
    notes = "\n".join(textwrap.wrap(args.notes, width=60))

    plt.annotate(notes,
                 xy=(0.001, -0.3),
                 xycoords='axes fraction',
                 ha='left',
                 va="center",
                 fontsize=12)

    plt.savefig(img_file, bbox_inches='tight')