TheRock/rocm-libraries/projects/rocrand/scripts/config-tuning/select_best_config.py

#!/usr/bin/env -S python3

import argparse
import json
import os
from datetime import datetime

import json5
import pandas
import seaborn as sns
from jinja2 import Environment, PackageLoader
from matplotlib import pyplot as plt
from pandas import DataFrame

env = Environment(
    loader=PackageLoader("select_best_config"), lstrip_blocks=True, trim_blocks=True
)


def load_default_configs_json(in_dir: str) -> dict:
    """
    Reads the config_defaults.json object for rocRAND's generators and returns them as dict.
    """
    with open(os.path.join(in_dir, "config_defaults.json")) as f:
        json_data = json5.load(f)

    return json_data


def load_benchmark_result_json(path: str) -> DataFrame:
    """
    Reads and processes a JSON file produced by rocRAND's tuning benchmark.
    """
    with open(path) as f:
        raw_json_data = json.load(f)

    general_ctx = raw_json_data.get("context", {}).get("general", {})
    arch = general_ctx.get("gpu", {}).get("arch", "unknown")

    rows = []
    for spec in raw_json_data.get("specializations", []):
        meta = spec.get("meta", {})
        cfg = meta.get("cfg", {})

        rows.append(
            {
                "generator": meta.get("engine"),
                "distribution": meta.get("distribution"),
                "value_type": meta.get("type"),
                "block_size": int(cfg.get("threads", 0)),
                "grid_size": int(cfg.get("blocks", 0)),
                "gb_per_s": float(spec.get("bytes_per_second", 0.0))
                / 1000
                / 1000
                / 1000,
                "arch": arch,
            }
        )

    return DataFrame(rows)


def get_best_config_for_arch(benchmark_data: DataFrame, default_configs: dict):
    """
    Calculates the best config for each architecture, that on average provides the highest
    performance across generated types and distributions.
    """

    def config_to_goodness(subf: DataFrame) -> pandas.Series:
        relative_perfs = subf["normalized_perf"].reset_index(drop=True)
        # Averages the performance over the different distribution/value_types
        # to get a single average performance per configuration of the generator
        return relative_perfs.sum() / relative_perfs.count()

    best_configs = {}
    for gen, generator_df in benchmark_data.groupby("generator"):
        best_config_for_generator = {}
        default_block_size = default_configs[gen]["block_size"]
        default_grid_size = default_configs[gen]["grid_size"]

        for arch, arch_df in generator_df.groupby("arch"):

            temp_df = DataFrame()
            # Normalize other configurations wrt default config per distribution/value_type
            for _, perf_df in arch_df.groupby(["distribution", "value_type"]):
                default_config_perf = perf_df.loc[
                    (perf_df["block_size"] == default_block_size)
                    & (perf_df["grid_size"] == default_grid_size)
                ].gb_per_s.iloc[0]
                perf_df["normalized_perf"] = perf_df.gb_per_s / default_config_perf
                temp_df = pandas.concat([temp_df, perf_df])

            arch_df = temp_df
            config_goodness = arch_df.groupby(["block_size", "grid_size"]).apply(
                config_to_goodness
            )
            config_goodness = config_goodness.sort_values(ascending=False)
            config_goodness = config_goodness.to_frame(
                name="normalized_perf"
            ).reset_index()

            # config_goodness is sorted by normalized_perf, so just take the first element
            best_block_size = config_goodness.iloc[0]["block_size"].astype(int)
            best_grid_size = config_goodness.iloc[0]["grid_size"].astype(int)
            best_perf = config_goodness.iloc[0]["normalized_perf"]

            print(
                f"{gen} {arch}: best=({best_block_size}, {best_grid_size}) avg={best_perf}"
            )

            # check for performance regressions and possible alternative configurations
            detailed_performance = arch_df[
                (arch_df["block_size"] == best_block_size)
                & (arch_df["grid_size"] == best_grid_size)
            ]
            if (detailed_performance["normalized_perf"] < 0.98).any():
                print(
                    "WARNING: configuration with best average performance has performance regressions for some distributions:"
                )
                print("Average speedup: ", best_perf)
                print(detailed_performance.to_string())
                # set index to config for easier looping
                config_goodness = config_goodness.set_index(["block_size", "grid_size"])
                other_configs = config_goodness[
                    (config_goodness["normalized_perf"] > 1.0)
                    & (config_goodness["normalized_perf"] != best_perf)
                ]
                if len(other_configs.index) == 0:
                    print(
                        "\nNo other configs available that provide an average speedup over the default config!\n"
                    )
                # check other configs that have an average speedup for regressions
                found_config_without_regressions = False
                for config, row in other_configs.iterrows():
                    block_size, grid_size = config
                    cur_detailed_perf = arch_df[
                        (arch_df["block_size"] == block_size)
                        & (arch_df["grid_size"] == grid_size)
                    ]
                    if (cur_detailed_perf["normalized_perf"] > 1.0).all():
                        print(
                            "Next best config that provides an average speedup over the default config without regressions:"
                        )
                        print("Average speedup: ", row["normalized_perf"])
                        print(cur_detailed_perf.to_string())
                        print("\n")
                        found_config_without_regressions = True
                        break
                if not found_config_without_regressions:
                    print(
                        "No other config provides a speedup over the default config without any regressions for certain distributions!\n"
                    )

            best_config_for_generator[arch] = {
                "block_size": best_block_size,
                "grid_size": best_grid_size,
            }
            generator_df.loc[
                (generator_df["generator"] == gen) & (generator_df["arch"] == arch)
            ] = arch_df
        best_configs[gen] = best_config_for_generator
    return best_configs


def plot_benchmark_data_for_all_arches(benchmark_data: DataFrame, out_path: str):
    """
    Plots comparative figures of the benchmark results for each architecture.
    """

    def plot_benchmark_results(
        benchmark_data: DataFrame, arch: str, gen: str, out_path: str
    ):
        def draw_heatmap(*args, **kwargs):
            data: DataFrame = kwargs.pop("data")
            d = data.pivot_table(index=args[1], columns=args[0], values=args[2])
            sns.heatmap(d, **kwargs)

        grid = sns.FacetGrid(benchmark_data, col="benchmark", col_wrap=3)
        grid.map_dataframe(
            draw_heatmap,
            "grid_size_",
            "block_size_",
            "gb_per_s",
            cbar=False,
            square=True,
            annot=True,
            fmt=".2f",
            cmap="coolwarm",
            linewidth=1,
        )
        fig = plt.gcf()
        fig.suptitle(f"rocRAND benchmarks ({gen} on {arch}, GB/s)")
        num_grid_sizes = benchmark_data["grid_size"].unique().astype(bool).sum()
        fig.set_size_inches(16 / 6 * num_grid_sizes, 18)
        fig.savefig(out_path)
        plt.close()

    benchmark_data["block_size_"] = (
        benchmark_data["block_size"].astype(str).str.pad(4, "left", " ")
    )
    benchmark_data["grid_size_"] = (
        benchmark_data["grid_size"].astype(str).str.pad(4, "left", " ")
    )
    benchmark_data["benchmark"] = (
        benchmark_data["value_type"].astype(str)
        + ","
        + benchmark_data["distribution"].astype(str)
    )
    base_path, full_filename = os.path.split(out_path)
    filename, ext = os.path.splitext(full_filename)
    for arch in benchmark_data["arch"].unique():
        arch_benchmark_data = benchmark_data[benchmark_data["arch"] == arch]
        for gen in arch_benchmark_data["generator"].unique():
            gen_benchmark_data = arch_benchmark_data[
                arch_benchmark_data["generator"] == gen
            ]
            plot_benchmark_results(
                gen_benchmark_data,
                arch,
                gen,
                os.path.join(base_path, f"{filename}_{arch}_{gen}{ext}"),
            )


def generate_config_files(out_dir: str, best_data: dict):
    def generator_type(generator_name: str) -> str:
        if "sobol" in generator_name.lower():
            return "quasi"
        return "pseudo"

    year = datetime.now().year
    algorithm_template = env.get_template("config_template")
    for generator, best_configs in best_data.items():
        with open(os.path.join(out_dir, f"{generator}_config.hpp"), "w") as text_file:
            text_file.write(
                algorithm_template.render(
                    generator=generator,
                    generator_type=generator_type(generator),
                    configs=best_configs,
                    year=year,
                )
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("json", nargs="+")
    parser.add_argument(
        "--plot-out",
        help="Output directory for plots. Plots will only be generated, when this argument is provided. Generating these can take quite some while.",
    )
    parser.add_argument("--out-dir")
    parser.add_argument(
        "--default-config-dir",
        default=os.path.dirname(os.path.realpath(__file__)),
        help="Path to directory of 'config_defaults.json'",
    )
    args = parser.parse_args()

    default_configs = load_default_configs_json(args.default_config_dir)
    benchmark = pandas.concat(
        [load_benchmark_result_json(j) for j in args.json], ignore_index=True
    )
    best_configs = get_best_config_for_arch(benchmark, default_configs)
    if args.plot_out:
        plot_benchmark_data_for_all_arches(benchmark, args.plot_out)
    if args.out_dir:
        generate_config_files(args.out_dir, best_configs)