[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/benchmarks/scripts/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+*.pyo
+*.pyd

cccl_upstream/benchmarks/scripts/analyze.py

@@ -0,0 +1,909 @@
+#!/usr/bin/env python3
+
+import argparse
+import functools
+import itertools
+import json
+import math
+import os
+import re
+
+import cccl
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from scipy.stats import mannwhitneyu
+from scipy.stats.mstats import hdquantiles
+
+pd.options.display.max_colwidth = 100
+
+default_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
+color_cycle = itertools.cycle(default_colors)
+color_map = {}
+
+precision = 0.01
+sensitivity = 0.5
+
+
+def get_bench_columns():
+    return ["variant", "elapsed", "center", "samples", "bw"]
+
+
+def get_extended_bench_columns():
+    return get_bench_columns() + ["speedup", "base_samples"]
+
+
+def compute_speedup(df):
+    bench_columns = get_bench_columns()
+    workload_columns = [col for col in df.columns if col not in bench_columns]
+    base_df = (
+        df[df["variant"] == "base"]
+        .drop(columns=["variant"])
+        .rename(columns={"center": "base_center", "samples": "base_samples"})
+    )
+    base_df.drop(columns=["elapsed", "bw"], inplace=True)
+
+    merged_df = df.merge(
+        base_df, on=[col for col in df.columns if col in workload_columns]
+    )
+    merged_df["speedup"] = merged_df["base_center"] / merged_df["center"]
+    merged_df = merged_df.drop(columns=["base_center"])
+    return merged_df
+
+
+def get_ct_axes(df):
+    ct_axes = []
+    for col in df.columns:
+        if "{ct}" in col:
+            ct_axes.append(col)
+
+    return ct_axes
+
+
+def get_rt_axes(df):
+    rt_axes = []
+    excluded_columns = get_ct_axes(df) + get_extended_bench_columns()
+
+    for col in df.columns:
+        if col not in excluded_columns:
+            rt_axes.append(col)
+
+    return rt_axes
+
+
+def ct_space(df):
+    ct_axes = get_ct_axes(df)
+
+    unique_ct_combinations = []
+    for _, row in df[ct_axes].drop_duplicates().iterrows():
+        unique_ct_combinations.append({})
+        for col in ct_axes:
+            unique_ct_combinations[-1][col] = row[col]
+
+    return unique_ct_combinations
+
+
+def extract_case(df, ct_point):
+    tuning_df_loc = None
+
+    for ct_axis in ct_point:
+        if tuning_df_loc is None:
+            tuning_df_loc = df[ct_axis] == ct_point[ct_axis]
+        else:
+            tuning_df_loc = tuning_df_loc & (df[ct_axis] == ct_point[ct_axis])
+
+    tuning_df = df.loc[tuning_df_loc].copy()
+    for ct_axis in ct_point:
+        tuning_df.drop(columns=[ct_axis], inplace=True)
+
+    return tuning_df
+
+
+def extract_rt_axes_values(df):
+    rt_axes = get_rt_axes(df)
+    rt_axes_values = {}
+
+    for rt_axis in rt_axes:
+        rt_axes_values[rt_axis] = list(df[rt_axis].unique())
+
+    return rt_axes_values
+
+
+def extract_rt_space(df):
+    rt_axes = get_rt_axes(df)
+    rt_axes_values = []
+    for rt_axis in rt_axes:
+        values = df[rt_axis].unique()
+        rt_axes_values.append(["{}={}".format(rt_axis, v) for v in values])
+    return list(itertools.product(*rt_axes_values))
+
+
+def filter_variants(df, group):
+    rt_axes = get_rt_axes(df)
+    unique_combinations = set(df[rt_axes].drop_duplicates().itertuples(index=False))
+    group_combinations = set(group[rt_axes].drop_duplicates().itertuples(index=False))
+    has_all_combinations = group_combinations == unique_combinations
+    return has_all_combinations
+
+
+def extract_complete_variants(df):
+    return df.groupby("variant").filter(functools.partial(filter_variants, df))
+
+
+def compute_workload_score(rt_axes_values, rt_axes_ids, weights, row):
+    rt_workload = []
+    for rt_axis in rt_axes_values:
+        rt_workload.append("{}={}".format(rt_axis, row[rt_axis]))
+
+    weight = cccl.bench.get_workload_weight(
+        rt_workload, rt_axes_values, rt_axes_ids, weights
+    )
+    return row["speedup"] * weight
+
+
+def compute_variant_score(rt_axes_values, rt_axes_ids, weight_matrix, group):
+    workload_score_closure = functools.partial(
+        compute_workload_score, rt_axes_values, rt_axes_ids, weight_matrix
+    )
+    score_sum = group.apply(workload_score_closure, axis=1).sum()
+    return score_sum
+
+
+def extract_scores(dfs):
+    rt_axes_values = {}
+    for subbench in dfs:
+        rt_axes_values[subbench] = extract_rt_axes_values(dfs[subbench])
+
+    rt_axes_ids = cccl.bench.compute_axes_ids(rt_axes_values)
+    weights = cccl.bench.compute_weight_matrices(rt_axes_values, rt_axes_ids)
+
+    score_dfs = []
+    for subbench in dfs:
+        score_closure = functools.partial(
+            compute_variant_score,
+            rt_axes_values[subbench],
+            rt_axes_ids[subbench],
+            weights[subbench],
+        )
+        grouped = dfs[subbench].groupby("variant")
+        scores = grouped.apply(score_closure, include_groups=False).reset_index()
+        scores.columns = ["variant", "score"]
+        stat = grouped.agg(
+            mins=("speedup", "min"), means=("speedup", "mean"), maxs=("speedup", "max")
+        )
+        scores = pd.merge(scores, stat, on="variant")
+        score_dfs.append(scores)
+    score_df = pd.concat(score_dfs)
+    result = (
+        score_df.groupby("variant")
+        .agg({"score": "sum", "mins": "min", "means": "mean", "maxs": "max"})
+        .reset_index()
+    )
+    return result.sort_values(by=["score"], ascending=False)
+
+
+def distributions_are_different(alpha, row):
+    ref_samples = row["base_samples"]
+    cmp_samples = row["samples"]
+
+    # H0: the distributions are not different
+    # H1: the distribution are different
+    _, p = mannwhitneyu(ref_samples, cmp_samples)
+
+    # Reject H0
+    return p < alpha
+
+
+def remove_matching_distributions(alpha, df):
+    closure = functools.partial(distributions_are_different, alpha)
+    return df[df.apply(closure, axis=1)]
+
+
+def get_filenames_map(arr):
+    if not arr:
+        return []
+
+    prefix = arr[0]
+    for string in arr:
+        while not string.startswith(prefix):
+            prefix = prefix[:-1]
+        if not prefix:
+            break
+
+    return {string: string[len(prefix) :] for string in arr}
+
+
+def is_finite(x):
+    if isinstance(x, float):
+        return x != np.inf and x != -np.inf
+    return True
+
+
+def iterate_case_dfs(args, callable):
+    storages = {}
+    algnames = set()
+    filenames_map = get_filenames_map(args.files)
+    for file in args.files:
+        storage = cccl.bench.SQLiteStorage(file)
+        algnames.update(storage.algnames())
+        storages[filenames_map[file]] = storage
+
+    pattern = re.compile(args.R)
+
+    exact_values = {}
+    if args.args:
+        for value in args.args:
+            name, val = value.split("=")
+            exact_values[name] = val
+
+    for algname in algnames:
+        if not pattern.match(algname):
+            continue
+
+        case_dfs = {}
+        for file in storages:
+            storage = storages[file]
+            for subbench in storage.subbenches(algname):
+                df = storage.alg_to_df(algname, subbench)
+
+                df = df.map(lambda x: x if is_finite(x) else np.nan)
+                df = df.dropna(subset=["center"], how="all")
+
+                for _, row in df[["ctk", "cccl"]].drop_duplicates().iterrows():
+                    ctk_version = row["ctk"]
+                    cccl_version = row["cccl"]
+                    ctk_cub_df = df[
+                        (df["ctk"] == ctk_version) & (df["cccl"] == cccl_version)
+                    ]
+
+                    for gpu in ctk_cub_df["gpu"].unique():
+                        target_df = ctk_cub_df[ctk_cub_df["gpu"] == gpu]
+                        target_df = target_df.drop(columns=["ctk", "cccl", "gpu"])
+                        target_df = compute_speedup(target_df)
+
+                        for key in exact_values:
+                            if key in target_df.columns:
+                                target_df = target_df[
+                                    target_df[key] == exact_values[key]
+                                ]
+
+                        for ct_point in ct_space(target_df):
+                            point_str = ", ".join(
+                                ["{}={}".format(k, ct_point[k]) for k in ct_point]
+                            )
+                            case_df = extract_complete_variants(
+                                extract_case(target_df, ct_point)
+                            )
+                            case_df["variant"] = case_df["variant"].astype(
+                                str
+                            ) + " ({})".format(file)
+                            if point_str not in case_dfs:
+                                case_dfs[point_str] = {}
+                            if subbench not in case_dfs[point_str]:
+                                case_dfs[point_str][subbench] = case_df
+                            else:
+                                case_dfs[point_str][subbench] = pd.concat(
+                                    [case_dfs[point_str][subbench], case_df]
+                                )
+
+        for point_str in case_dfs:
+            callable(algname, point_str, case_dfs[point_str])
+
+
+def case_top(alpha, N, algname, ct_point_name, case_dfs):
+    print("{}[{}]:".format(algname, ct_point_name))
+
+    if alpha < 1.0:
+        for subbench in case_dfs:
+            case_dfs[subbench] = remove_matching_distributions(
+                alpha, case_dfs[subbench]
+            )
+
+    for subbench in case_dfs:
+        case_dfs[subbench] = extract_complete_variants(case_dfs[subbench])
+    with pd.option_context("display.max_rows", None):
+        print(extract_scores(case_dfs).head(N))
+
+
+def top(args):
+    iterate_case_dfs(args, functools.partial(case_top, args.alpha, args.top))
+
+
+def case_coverage(algname, ct_point_name, case_dfs):
+    num_variants = cccl.bench.Config().variant_space_size(algname)
+    min_coverage = 100.0
+    for subbench in case_dfs:
+        num_covered_variants = len(case_dfs[subbench]["variant"].unique())
+        coverage = (num_covered_variants / num_variants) * 100
+        min_coverage = min(min_coverage, coverage)
+    case_str = "{}[{}]".format(algname, ct_point_name)
+    print(
+        "{} coverage: {} / {} ({:.4f}%)".format(
+            case_str, num_covered_variants, num_variants, min_coverage
+        )
+    )
+
+
+def coverage(args):
+    iterate_case_dfs(args, case_coverage)
+
+
+def parallel_coordinates_plot(df, title):
+    # Parallel coordinates plot adaptation of https://stackoverflow.com/a/69411450
+    import matplotlib.cm as cm
+    import matplotlib.patches as patches
+    from matplotlib.path import Path
+
+    # Variables (the first variable must be categoric):
+    my_vars = df.columns.tolist()
+    df_plot = df[my_vars]
+    df_plot = df_plot.dropna()
+    df_plot = df_plot.reset_index(drop=True)
+
+    # Convert to numeric matrix:
+    ym = []
+    dics_vars = []
+    for v, var in enumerate(my_vars):
+        if df_plot[var].dtype.kind not in ["i", "u", "f"]:
+            dic_var = dict([(val, c) for c, val in enumerate(df_plot[var].unique())])
+            dics_vars += [dic_var]
+            ym += [[dic_var[i] for i in df_plot[var].tolist()]]
+        else:
+            ym += [df_plot[var].tolist()]
+    ym = np.array(ym).T
+
+    # Padding:
+    ymins = ym.min(axis=0)
+    ymaxs = ym.max(axis=0)
+    dys = ymaxs - ymins
+    ymins -= dys * 0.05
+    ymaxs += dys * 0.05
+
+    dys = ymaxs - ymins
+
+    # Adjust to the main axis:
+    zs = np.zeros_like(ym)
+    zs[:, 0] = ym[:, 0]
+    zs[:, 1:] = (ym[:, 1:] - ymins[1:]) / dys[1:] * dys[0] + ymins[0]
+
+    # Plot:
+    fig, host_ax = plt.subplots(figsize=(20, 10), tight_layout=True)
+
+    # Make the axes:
+    axes = [host_ax] + [host_ax.twinx() for i in range(ym.shape[1] - 1)]
+    dic_count = 0
+    for i, ax in enumerate(axes):
+        ax.set_ylim(bottom=ymins[i], top=ymaxs[i])
+        ax.spines.top.set_visible(False)
+        ax.spines.bottom.set_visible(False)
+        ax.ticklabel_format(style="plain")
+        if ax != host_ax:
+            ax.spines.left.set_visible(False)
+            ax.yaxis.set_ticks_position("right")
+            ax.spines.right.set_position(("axes", i / (ym.shape[1] - 1)))
+        if df_plot.iloc[:, i].dtype.kind not in ["i", "u", "f"]:
+            dic_var_i = dics_vars[dic_count]
+            ax.set_yticks(range(len(dic_var_i)))
+            if i == 0:
+                ax.set_yticklabels([])
+            else:
+                ax.set_yticklabels([key_val for key_val in dics_vars[dic_count].keys()])
+            dic_count += 1
+    host_ax.set_xlim(left=0, right=ym.shape[1] - 1)
+    host_ax.set_xticks(range(ym.shape[1]))
+    host_ax.set_xticklabels(my_vars, fontsize=14)
+    host_ax.tick_params(axis="x", which="major", pad=7)
+
+    # Color map:
+    colormap = cm.get_cmap("turbo")
+
+    # Normalize speedups:
+    df["speedup_normalized"] = (df["speedup"] - df["speedup"].min()) / (
+        df["speedup"].max() - df["speedup"].min()
+    )
+
+    # Make the curves:
+    host_ax.spines.right.set_visible(False)
+    host_ax.xaxis.tick_top()
+    for j in range(ym.shape[0]):
+        verts = list(
+            zip(
+                [
+                    x
+                    for x in np.linspace(0, len(ym) - 1, len(ym) * 3 - 2, endpoint=True)
+                ],
+                np.repeat(zs[j, :], 3)[1:-1],
+            )
+        )
+        codes = [Path.MOVETO] + [Path.CURVE4 for _ in range(len(verts) - 1)]
+        path = Path(verts, codes)
+        color_first_cat_var = colormap(df.loc[j, "speedup_normalized"])
+        patch = patches.PathPatch(
+            path, facecolor="none", lw=2, alpha=0.05, edgecolor=color_first_cat_var
+        )
+        host_ax.add_patch(patch)
+
+    host_ax.set_title(title)
+    plt.show()
+
+
+def case_coverage_plot(algname, ct_point_name, case_dfs):
+    data_list = []
+
+    for subbench in case_dfs:
+        for _, row_description in case_dfs[subbench].iterrows():
+            variant = row_description["variant"]
+            speedup = row_description["speedup"]
+
+            if variant.startswith("base"):
+                continue
+
+            varname, _ = variant.split(" ")
+            params = varname.split(".")
+            data_dict = {"variant": variant}
+
+            for param in params:
+                print(variant)
+                name, val = param.split("_")
+                data_dict[name] = int(val)
+
+            data_dict["speedup"] = speedup
+            # data_dict['variant'] = variant
+            data_list.append(data_dict)
+
+    df = pd.DataFrame(data_list)
+    parallel_coordinates_plot(df, "{} ({})".format(algname, ct_point_name))
+
+
+def coverage_plot(args):
+    iterate_case_dfs(args, case_coverage_plot)
+
+
+def case_pair_plot(algname, ct_point_name, case_dfs):
+    import seaborn as sns
+
+    data_list = []
+
+    for subbench in case_dfs:
+        for _, row_description in case_dfs[subbench].iterrows():
+            variant = row_description["variant"]
+            speedup = row_description["speedup"]
+
+            if variant.startswith("base"):
+                continue
+
+            varname, _ = variant.split(" ")
+            params = varname.split(".")
+            data_dict = {}
+
+            for param in params:
+                print(variant)
+                name, val = param.split("_")
+                data_dict[name] = int(val)
+
+            data_dict["speedup"] = speedup
+            data_list.append(data_dict)
+
+    df = pd.DataFrame(data_list)
+    sns.pairplot(df, hue="speedup")
+    plt.title("{} ({})".format(algname, ct_point_name))
+    plt.show()
+
+
+def pair_plot(args):
+    iterate_case_dfs(args, case_pair_plot)
+
+
+def qrde_hd(samples):
+    """
+    Computes quantile-respectful density estimation based on the Harrell-Davis
+    quantile estimator. The implementation is based on the following post:
+    https://aakinshin.net/posts/qrde-hd by Andrey Akinshin
+    """
+    min_sample, max_sample = min(samples), max(samples)
+    num_quantiles = math.ceil(1.0 / precision)
+    quantiles = np.linspace(precision, 1 - precision, num_quantiles - 1)
+    hd_quantiles = [min_sample] + list(hdquantiles(samples, quantiles)) + [max_sample]
+    width = [hd_quantiles[idx + 1] - hd_quantiles[idx] for idx in range(num_quantiles)]
+    p = 1.0 / precision
+    height = [1.0 / (p * w) for w in width]
+    return width, height
+
+
+def hd_quantiles(samples):
+    min_sample, max_sample = min(samples), max(samples)
+    num_quantiles = math.ceil(1.0 / precision)
+    quantiles = np.linspace(precision, 1 - precision, num_quantiles - 1)
+    hd_quantiles = [min_sample] + list(hdquantiles(samples, quantiles)) + [max_sample]
+    return hd_quantiles
+
+
+def extract_peaks(pdf):
+    peaks = []
+    for i in range(1, len(pdf) - 1):
+        if pdf[i - 1] < pdf[i] > pdf[i + 1]:
+            peaks.append(i)
+    return peaks
+
+
+def extract_modes(samples):
+    """
+    Extract modes from the given samples based on the lowland algorithm:
+    https://aakinshin.net/posts/lowland-multimodality-detection/ by Andrey Akinshin
+    Implementation is based on the https://github.com/AndreyAkinshin/perfolizer
+    LowlandModalityDetector class.
+    """
+    mode_ids = []
+
+    widths, heights = hd_displot(samples)
+    peak_ids = extract_peaks(heights)
+    bin_area = 1.0 / len(heights)
+
+    x = min(samples)
+    peak_xs = []
+    peak_ys = []
+    bin_lower = [x]
+    for idx in range(len(heights)):
+        if idx in peak_ids:
+            peak_ys.append(heights[idx])
+            peak_xs.append(x + widths[idx] / 2)
+        x += widths[idx]
+        bin_lower.append(x)
+
+    def lowland_between(mode_candidate, left_peak, right_peak):
+        left, right = left_peak, right_peak
+        min_height = min(heights[left_peak], heights[right_peak])
+        while left < right and heights[left] > min_height:
+            left += 1
+        while left < right and heights[right] > min_height:
+            right -= 1
+
+        width = bin_lower[right + 1] - bin_lower[left]
+        total_area = width * min_height
+        total_bin_area = (right - left + 1) * bin_area
+
+        if total_bin_area / total_area < sensitivity:
+            mode_ids.append(mode_candidate)
+            return True
+        return False
+
+    previousPeaks = [peak_ids[0]]
+    for i in range(1, len(peak_ids)):
+        currentPeak = peak_ids[i]
+        while previousPeaks and heights[previousPeaks[-1]] < heights[currentPeak]:
+            if lowland_between(previousPeaks[0], previousPeaks[-1], currentPeak):
+                previousPeaks = []
+            else:
+                previousPeaks.pop()
+
+        if previousPeaks and heights[previousPeaks[-1]] > heights[currentPeak]:
+            if lowland_between(previousPeaks[0], previousPeaks[-1], currentPeak):
+                previousPeaks = []
+
+        previousPeaks.append(currentPeak)
+
+    mode_ids.append(previousPeaks[0])
+    return mode_ids
+
+
+def hd_displot(samples, label, ax):
+    if label not in color_map:
+        color_map[label] = next(color_cycle)
+    color = color_map[label]
+    widths, heights = qrde_hd(samples)
+    mode_ids = extract_modes(samples)
+
+    min_sample, max_sample = min(samples), max(samples)
+
+    xs = [min_sample]
+    ys = [0]
+
+    peak_xs = []
+    peak_ys = []
+
+    x = min(samples)
+    for idx in range(len(widths)):
+        xs.append(x + widths[idx] / 2)
+        ys.append(heights[idx])
+        if idx in mode_ids:
+            peak_ys.append(heights[idx])
+            peak_xs.append(x + widths[idx] / 2)
+        x += widths[idx]
+
+    xs = xs + [max_sample]
+    ys = ys + [0]
+
+    ax.fill_between(xs, ys, 0, alpha=0.4, color=color)
+
+    quartiles_of_interest = [0.25, 0.5, 0.75]
+
+    for quartile in quartiles_of_interest:
+        bin = int(quartile / precision) + 1
+        ax.plot([xs[bin], xs[bin]], [0, ys[bin]], color=color)
+
+    ax.plot(xs, ys, label=label, color=color)
+    ax.plot(peak_xs, peak_ys, "o", color=color)
+    ax.legend()
+
+
+def displot(data, ax):
+    for variant in data:
+        hd_displot(data[variant], variant, ax)
+
+
+def variant_ratio(data, variant, ax):
+    if variant not in color_map:
+        color_map[variant] = next(color_cycle)
+    color = color_map[variant]
+
+    variant_samples = data[variant]
+    base_samples = data["base"]
+
+    variant_widths = hd_quantiles(variant_samples)
+    base_widths = hd_quantiles(base_samples)
+
+    quantiles = []
+    ratios = []
+
+    base_x = min(base_samples)
+    variant_x = min(variant_samples)
+
+    for i in range(1, len(variant_widths) - 1):
+        base_x += base_widths[i] / 2
+        variant_x += variant_widths[i] / 2
+        quantiles.append(i * precision)
+        ratios.append(base_x / variant_x)
+
+    ax.plot(quantiles, ratios, label=variant, color=color)
+    ax.axhline(1, color="red", alpha=0.7)
+    ax.legend()
+    ax.tick_params(axis="both", direction="in", pad=-22)
+
+
+def ratio(data, ax):
+    for variant in data:
+        if variant != "base":
+            variant_ratio(data, variant, ax)
+
+
+def case_variants(pattern, mode, algname, ct_point_name, case_dfs):
+    for subbench in case_dfs:
+        case_df = case_dfs[subbench]
+        title = "{}[{}]:".format(algname + "/" + subbench, ct_point_name)
+        df = case_df[case_df["variant"].str.contains(pattern, regex=True)].reset_index(
+            drop=True
+        )
+        rt_axes = get_rt_axes(df)
+        rt_axes_values = extract_rt_axes_values(df)
+
+        vertical_axis_name = rt_axes[0]
+        if "Elements{io}[pow2]" in rt_axes:
+            vertical_axis_name = "Elements{io}[pow2]"
+        horizontal_axes = rt_axes
+        horizontal_axes.remove(vertical_axis_name)
+        vertical_axis_values = rt_axes_values[vertical_axis_name]
+
+        vertical_axis_ids = {}
+        for idx, val in enumerate(vertical_axis_values):
+            vertical_axis_ids[val] = idx
+
+        def extract_horizontal_space(df):
+            values = []
+            for rt_axis in horizontal_axes:
+                values.append(
+                    ["{}={}".format(rt_axis, v) for v in df[rt_axis].unique()]
+                )
+            return list(itertools.product(*values))
+
+        if len(horizontal_axes) > 0:
+            idx = 0
+            horizontal_axis_ids = {}
+            for point in extract_horizontal_space(df):
+                horizontal_axis_ids[" / ".join(point)] = idx
+                idx = idx + 1
+
+        num_rows = len(vertical_axis_ids)
+        num_cols = max(1, len(extract_horizontal_space(df)))
+
+        if num_rows == 0:
+            return
+
+        fig, axes = plt.subplots(
+            nrows=num_rows, ncols=num_cols, gridspec_kw={"wspace": 0, "hspace": 0}
+        )
+
+        for _, vertical_row_description in (
+            df[[vertical_axis_name]].drop_duplicates().iterrows()
+        ):
+            vertical_val = vertical_row_description[vertical_axis_name]
+            vertical_id = vertical_axis_ids[vertical_val]
+            vertical_name = "{}={}".format(vertical_axis_name, vertical_val)
+
+            vertical_df = df[df[vertical_axis_name] == vertical_val]
+
+            for _, horizontal_row_description in (
+                vertical_df[horizontal_axes].drop_duplicates().iterrows()
+            ):
+                horizontal_df = vertical_df
+
+                for axis in horizontal_axes:
+                    horizontal_df = horizontal_df[
+                        horizontal_df[axis] == horizontal_row_description[axis]
+                    ]
+
+                horizontal_id = 0
+
+                if len(horizontal_axes) > 0:
+                    horizontal_point = []
+                    for rt_axis in horizontal_axes:
+                        horizontal_point.append(
+                            "{}={}".format(rt_axis, horizontal_row_description[rt_axis])
+                        )
+                    horizontal_name = " / ".join(horizontal_point)
+                    horizontal_id = horizontal_axis_ids[horizontal_name]
+                    ax = axes[vertical_id, horizontal_id]
+                else:
+                    ax = axes[vertical_id]
+                    ax.set_ylabel(vertical_name)
+
+                data = {}
+                for _, variant in (
+                    horizontal_df[["variant"]].drop_duplicates().iterrows()
+                ):
+                    variant_name = variant["variant"]
+                    if "base" not in data:
+                        data["base"] = horizontal_df[
+                            horizontal_df["variant"] == variant_name
+                        ].iloc[0]["base_samples"]
+                    data[variant_name] = horizontal_df[
+                        horizontal_df["variant"] == variant_name
+                    ].iloc[0]["samples"]
+
+                if mode == "pdf":
+                    # sns.histplot(data=data, ax=ax, kde=True)
+                    displot(data, ax)
+                else:
+                    ratio(data, ax)
+
+                if len(horizontal_axes) > 0:
+                    ax = axes[vertical_id, horizontal_id]
+                    if vertical_id == (num_rows - 1):
+                        ax.set_xlabel(horizontal_name)
+                    if horizontal_id == 0:
+                        ax.set_ylabel(vertical_name)
+                    else:
+                        ax.set_ylabel("")
+
+        for ax in axes.flat:
+            ax.set_xticklabels([])
+
+        fig.suptitle(title)
+        plt.tight_layout()
+        plt.show()
+
+
+def variants(args, mode):
+    pattern = (
+        re.compile(args.variants_pdf)
+        if mode == "pdf"
+        else re.compile(args.variants_ratio)
+    )
+    iterate_case_dfs(args, functools.partial(case_variants, pattern, mode))
+
+
+def file_exists(value):
+    if not os.path.isfile(value):
+        raise argparse.ArgumentTypeError(f"The file '{value}' does not exist.")
+    return value
+
+
+def case_offload(algname, ct_point_name, case_dfs):
+    for subbench in case_dfs:
+        df = case_dfs[subbench]
+        for rt_point in extract_rt_space(df):
+            point_df = df
+            for rt_kv in rt_point:
+                key, value = rt_kv.split("=")
+                point_df = point_df[point_df[key] == value]
+            point_name = ct_point_name + " " + " ".join(rt_point)
+            point_name = point_name.replace(",", "")
+            bench_name = "{}.{}-{}".format(algname, subbench, point_name)
+            bench_name = bench_name.replace(" ", "___")
+            bench_name = "".join(c if c.isalnum() else "_" for c in bench_name)
+            with open(bench_name + ".json", "w") as f:
+                obj = json.loads(point_df.to_json(orient="records"))
+                json.dump(obj, f, indent=2)
+
+
+def offload(args):
+    iterate_case_dfs(args, case_offload)
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Analyze benchmark results.")
+    parser.add_argument(
+        "-R", type=str, default=".*", help="Regex for benchmarks selection."
+    )
+    parser.add_argument(
+        "--list-benches",
+        action=argparse.BooleanOptionalAction,
+        help="Show available benchmarks.",
+    )
+    parser.add_argument(
+        "--coverage",
+        action=argparse.BooleanOptionalAction,
+        help="Show variant space coverage.",
+    )
+    parser.add_argument(
+        "--coverage-plot",
+        action=argparse.BooleanOptionalAction,
+        help="Plot variant space coverage.",
+    )
+    parser.add_argument(
+        "--pair-plot", action=argparse.BooleanOptionalAction, help="Pair plot."
+    )
+    parser.add_argument(
+        "--top",
+        default=7,
+        type=int,
+        action="store",
+        nargs="?",
+        help="Show top N variants with highest score.",
+    )
+    parser.add_argument(
+        "files", type=file_exists, nargs="+", help="At least one file is required."
+    )
+    parser.add_argument("--alpha", default=1.0, type=float)
+    parser.add_argument("--variants-pdf", type=str, help="Show matching variants data.")
+    parser.add_argument(
+        "--variants-ratio", type=str, help="Show matching variants data."
+    )
+    parser.add_argument(
+        "-a",
+        "--args",
+        action="append",
+        type=str,
+        help="Parameter in the format `Param=Value`.",
+    )
+    parser.add_argument(
+        "-o", "--offload", action=argparse.BooleanOptionalAction, help="Offload samples"
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    if args.list_benches:
+        cccl.bench.list_benches()
+        return
+
+    if args.coverage:
+        coverage(args)
+        return
+
+    if args.coverage_plot:
+        coverage_plot(args)
+        return
+
+    if args.pair_plot:
+        pair_plot(args)
+        return
+
+    if args.variants_pdf:
+        variants(args, "pdf")
+        return
+
+    if args.variants_ratio:
+        variants(args, "ratio")
+        return
+
+    if args.offload:
+        offload(args)
+        return
+
+    top(args)
+
+
+if __name__ == "__main__":
+    main()

cccl_upstream/benchmarks/scripts/cccl/__init__.py

@@ -0,0 +1,3 @@
+from . import bench
+
+__all__ = ["bench"]

cccl_upstream/benchmarks/scripts/cccl/bench/__init__.py

@@ -0,0 +1,6 @@
+from .bench import Bench  # noqa: F401
+from .cmake import CMake  # noqa: F401
+from .config import *  # noqa: F403
+from .score import *  # noqa: F403
+from .search import *  # noqa: F403
+from .storage import *  # noqa: F403

cccl_upstream/benchmarks/scripts/cccl/bench/bench.py

@@ -0,0 +1,814 @@
+import itertools
+import json
+import os
+import signal
+import subprocess
+import time
+
+import fpzip
+import numpy as np
+
+from .cmake import CMake
+from .config import BasePoint, Config
+from .logger import Logger
+from .score import compute_axes_ids, compute_weight_matrices, get_workload_weight
+from .storage import Storage, get_bench_table_name
+
+
+def first_val(my_dict):
+    values = list(my_dict.values())
+    first_value = values[0]
+
+    if not all(value == first_value for value in values):
+        raise ValueError(
+            "All values in the dictionary are not equal. First value: {} All values: {}".format(
+                first_value, values
+            )
+        )
+
+    return first_value
+
+
+class JsonCache:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.bench_cache = {}
+            cls._instance.device_cache = {}
+        return cls._instance
+
+    def get_jsonlist(self, algname, listname):
+        benchmark_bin = os.path.join(".", "bin", algname + ".base")
+        if not os.path.exists(benchmark_bin):
+            raise Exception(f"Benchmark binary not found: {benchmark_bin}")
+        return subprocess.check_output([benchmark_bin, f"--jsonlist-{listname}"])
+
+    def get_bench(self, algname):
+        if algname not in self.bench_cache:
+            result = self.get_jsonlist(algname, "benches")
+            self.bench_cache[algname] = json.loads(result)
+        return self.bench_cache[algname]
+
+    def get_device(self, algname):
+        if algname not in self.device_cache:
+            result = self.get_jsonlist(algname, "devices")
+            data = json.loads(result)
+            if "devices" not in data:
+                raise Exception(
+                    "JSON returned from --jsonlist-devices does not contain 'devices' key"
+                )
+            devices = data["devices"]
+            if len(devices) != 1:
+                raise Exception(
+                    "NVBench doesn't work well with multiple GPUs, use `CUDA_VISIBLE_DEVICES`"
+                )
+
+            self.device_cache[algname] = devices[0]
+
+        return self.device_cache[algname]
+
+
+def json_benches(algname):
+    return JsonCache().get_bench(algname)
+
+
+def create_benches_tables(conn, subbench, bench_axes):
+    with conn:
+        conn.execute("""
+        CREATE TABLE IF NOT EXISTS subbenches (
+            algorithm TEXT NOT NULL,
+            bench TEXT NOT NULL,
+            UNIQUE(algorithm, bench)
+        );
+        """)
+
+        for algorithm_name in bench_axes:
+            axes = bench_axes[algorithm_name]
+            column_names = ", ".join(['"{}"'.format(name) for name in axes])
+            columns = ", ".join(['"{}" TEXT'.format(name) for name in axes])
+
+            conn.execute(
+                """
+            INSERT INTO subbenches (algorithm, bench)
+            VALUES (?, ?)
+            ON CONFLICT DO NOTHING;
+            """,
+                (algorithm_name, subbench),
+            )
+
+            if axes:
+                columns = ", " + columns
+                column_names = ", " + column_names
+                conn.execute(
+                    """
+                CREATE TABLE IF NOT EXISTS "{0}" (
+                    ctk TEXT NOT NULL,
+                    cccl TEXT NOT NULL,
+                    gpu TEXT NOT NULL,
+                    variant TEXT NOT NULL,
+                    elapsed REAL,
+                    center REAL,
+                    bw REAL,
+                    samples BLOB
+                    {1}
+                    , UNIQUE(ctk, cccl, gpu, variant {2})
+                );
+                """.format(
+                        get_bench_table_name(subbench, algorithm_name),
+                        columns,
+                        column_names,
+                    )
+                )
+
+
+def read_json(filename):
+    with open(filename, "r") as f:
+        file_root = json.load(f)
+    return file_root
+
+
+def extract_filename(summary):
+    summary_data = summary["data"]
+    value_data = next(filter(lambda v: v["name"] == "filename", summary_data))
+    assert value_data["type"] == "string"
+    return value_data["value"]
+
+
+def extract_size(summary):
+    summary_data = summary["data"]
+    value_data = next(filter(lambda v: v["name"] == "size", summary_data))
+    assert value_data["type"] == "int64"
+    return int(value_data["value"])
+
+
+def extract_bw(summary):
+    summary_data = summary["data"]
+    value_data = next(filter(lambda v: v["name"] == "value", summary_data))
+    assert value_data["type"] == "float64"
+    return float(value_data["value"])
+
+
+def parse_samples_meta(state):
+    summaries = state["summaries"]
+    if not summaries:
+        return None, None
+
+    summary = next(
+        filter(lambda s: s["tag"] == "nv/json/bin:nv/cold/sample_times", summaries),
+        None,
+    )
+    if not summary:
+        return None, None
+
+    sample_filename = extract_filename(summary)
+    sample_count = extract_size(summary)
+    return sample_count, sample_filename
+
+
+def parse_samples(state):
+    sample_count, samples_filename = parse_samples_meta(state)
+    if not sample_count or not samples_filename:
+        return np.array([], dtype=np.float32)
+
+    with open(samples_filename, "rb") as f:
+        samples = np.fromfile(f, "<f4")
+
+    samples.sort()
+
+    assert sample_count == len(samples)
+    return samples
+
+
+def parse_bw(state):
+    bwutil = next(
+        filter(
+            lambda s: s["tag"] == "nv/cold/bw/global/utilization", state["summaries"]
+        ),
+        None,
+    )
+    if not bwutil:
+        return None
+
+    return extract_bw(bwutil)
+
+
+class SubBenchState:
+    def __init__(self, state, axes_names, axes_values):
+        self.samples = parse_samples(state)
+        self.bw = parse_bw(state)
+
+        self.point = {}
+        for axis in state["axis_values"]:
+            name = axes_names[axis["name"]]
+            value = axes_values[axis["name"]][axis["value"]]
+            self.point[name] = value
+
+    def __repr__(self):
+        return str(self.__dict__)
+
+    def name(self):
+        return " ".join(f"{k}={v}" for k, v in self.point.items())
+
+    def center(self, estimator):
+        return estimator(self.samples)
+
+
+class SubBenchResult:
+    def __init__(self, bench):
+        axes_names = {}
+        axes_values = {}
+        for axis in bench["axes"]:
+            short_name = axis["name"]
+            full_name = get_axis_name(axis)
+            axes_names[short_name] = full_name
+            axes_values[short_name] = {}
+            for value in axis["values"]:
+                if "value" in value:
+                    axes_values[axis["name"]][str(value["value"])] = value[
+                        "input_string"
+                    ]
+                else:
+                    axes_values[axis["name"]][value["input_string"]] = value[
+                        "input_string"
+                    ]
+
+        self.states = []
+        for state in bench["states"]:
+            if not state["is_skipped"]:
+                self.states.append(SubBenchState(state, axes_names, axes_values))
+
+    def __repr__(self):
+        return str(self.__dict__)
+
+    def centers(self, estimator):
+        result = {}
+        for state in self.states:
+            result[state.name()] = state.center(estimator)
+        return result
+
+
+class BenchResult:
+    def __init__(self, json_path, code, elapsed):
+        self.code = code
+        self.elapsed = elapsed
+
+        if json_path:
+            self.subbenches = {}
+            if code == 0:
+                for bench in read_json(json_path)["benchmarks"]:
+                    self.subbenches[bench["name"]] = SubBenchResult(bench)
+
+    def __repr__(self):
+        return str(self.__dict__)
+
+    def centers(self, estimator):
+        result = {}
+        for subbench in self.subbenches:
+            result[subbench] = self.subbenches[subbench].centers(estimator)
+        return result
+
+
+def device_json(algname):
+    return JsonCache().get_device(algname)
+
+
+CCCL_BENCH_GPU_ENV = "CCCL_BENCH_GPU"
+
+
+def get_gpu_name_override():
+    override = os.environ.get(CCCL_BENCH_GPU_ENV)
+    if override is not None and override.strip():
+        return override.strip()
+    return None
+
+
+def get_device_name(device):
+    gpu_name = device["name"]
+    bus_width = device["global_memory_bus_width"]
+    sms = device["number_of_sms"]
+    ecc = "eccon" if device["ecc_state"] else "eccoff"
+    name = "{} ({}, {}, {})".format(gpu_name, bus_width, sms, ecc)
+    return name.replace("NVIDIA ", "")
+
+
+def get_gpu_name(algname):
+    override = get_gpu_name_override()
+    if override is not None:
+        return override
+    return get_device_name(device_json(algname))
+
+
+def is_ct_axis(name):
+    return "{ct}" in name
+
+
+def state_to_rt_workload(bench, state):
+    rt_workload = []
+    for param in state.split(" "):
+        name, value = param.split("=")
+        if is_ct_axis(name):
+            continue
+        rt_workload.append("{}={}".format(name, value))
+    return rt_workload
+
+
+def create_runs_table(conn):
+    with conn:
+        conn.execute("""
+        CREATE TABLE IF NOT EXISTS runs (
+            ctk TEXT NOT NULL,
+            cccl TEXT NOT NULL,
+            bench TEXT NOT NULL,
+            code TEXT NOT NULL,
+            elapsed REAL
+        );
+        """)
+
+
+class RunsCache:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+            create_runs_table(Storage().connection())
+        return cls._instance
+
+    def pull_run(self, bench):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        conn = Storage().connection()
+
+        with conn:
+            query = "SELECT code, elapsed FROM runs WHERE ctk = ? AND cccl = ? AND bench = ?;"
+            result = conn.execute(query, (ctk, cccl, bench.label())).fetchone()
+
+            if result:
+                code, elapsed = result
+                return int(code), float(elapsed)
+
+            return result
+
+    def push_run(self, bench, code, elapsed):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        conn = Storage().connection()
+
+        with conn:
+            conn.execute(
+                "INSERT INTO runs (ctk, cccl, bench, code, elapsed) VALUES (?, ?, ?, ?, ?);",
+                (ctk, cccl, bench.label(), code, elapsed),
+            )
+
+
+class BenchCache:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+            cls._instance.existing_tables = set()
+
+        return cls._instance
+
+    def create_table_if_not_exists(self, conn, bench):
+        bench_base = bench.get_base()
+        alg_name = bench_base.algorithm_name()
+
+        if alg_name not in self.existing_tables:
+            subbench_axes_names = bench_base.axes_names()
+            for subbench in subbench_axes_names:
+                create_benches_tables(
+                    conn, subbench, {alg_name: subbench_axes_names[subbench]}
+                )
+                self.existing_tables.add(alg_name)
+
+    def push_bench_centers(self, bench, result, estimator):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        gpu = get_gpu_name(bench.algname)
+        conn = Storage().connection()
+
+        self.create_table_if_not_exists(conn, bench)
+
+        centers = {}
+        with conn:
+            for subbench in result.subbenches:
+                centers[subbench] = {}
+                for state in result.subbenches[subbench].states:
+                    table_name = get_bench_table_name(subbench, bench.algorithm_name())
+                    columns = ""
+                    placeholders = ""
+                    values = []
+
+                    for name in state.point:
+                        value = state.point[name]
+                        columns = columns + ', "{}"'.format(name)
+                        placeholders = placeholders + ", ?"
+                        values.append(value)
+
+                    values = tuple(values)
+                    samples = fpzip.compress(state.samples)
+                    center = estimator(state.samples)
+                    to_insert = (
+                        ctk,
+                        cccl,
+                        gpu,
+                        bench.variant_name(),
+                        result.elapsed,
+                        center,
+                        state.bw,
+                        samples,
+                    ) + values
+
+                    query = """
+                    INSERT INTO "{0}" (ctk, cccl, gpu, variant, elapsed, center, bw, samples {1})
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ? {2})
+                    ON CONFLICT(ctk, cccl, gpu, variant {1}) DO NOTHING;
+                    """.format(table_name, columns, placeholders)
+
+                    conn.execute(query, to_insert)
+                    centers[subbench][state.name()] = center
+
+        return centers
+
+    def pull_bench_centers(self, bench, ct_workload_point, rt_values):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        gpu = get_gpu_name(bench.algname)
+        conn = Storage().connection()
+
+        self.create_table_if_not_exists(conn, bench)
+
+        centers = {}
+
+        with conn:
+            for subbench in rt_values:
+                centers[subbench] = {}
+                table_name = get_bench_table_name(subbench, bench.algorithm_name())
+
+                for rt_point in values_to_space(rt_values[subbench]):
+                    point_map = {}
+                    point_checks = ""
+                    workload_point = list(ct_workload_point) + list(rt_point)
+                    for axis in workload_point:
+                        name, value = axis.split("=")
+                        point_map[name] = value
+                        point_checks = point_checks + ' AND "{}" = "{}"'.format(
+                            name, value
+                        )
+
+                    query = """
+                    SELECT center FROM "{0}" WHERE ctk = ? AND cccl = ? AND gpu = ? AND variant = ?{1};
+                    """.format(table_name, point_checks)
+
+                    result = conn.execute(
+                        query, (ctk, cccl, gpu, bench.variant_name())
+                    ).fetchone()
+                    if result is None:
+                        return None
+
+                    state_name = " ".join(f"{k}={v}" for k, v in point_map.items())
+                    centers[subbench][state_name] = float(result[0])
+
+        return centers
+
+
+def get_axis_name(axis):
+    name = axis["name"]
+    if axis["flags"]:
+        name = name + "[{}]".format(axis["flags"])
+    return name
+
+
+def speedup(base, variant):
+    # If one of the runs failed, dict is empty
+    if not base or not variant:
+        return {}
+
+    benchmarks = set(base.keys())
+    if benchmarks != set(variant.keys()):
+        raise Exception("Benchmarks do not match.")
+
+    result = {}
+    for bench in benchmarks:
+        base_states = base[bench]
+        variant_states = variant[bench]
+
+        state_names = set(base_states.keys())
+        if state_names != set(variant_states.keys()):
+            raise Exception("States do not match.")
+
+        result[bench] = {}
+        for state in state_names:
+            result[bench][state] = base_states[state] / variant_states[state]
+
+    return result
+
+
+def values_to_space(axes):
+    result = []
+    for axis in axes:
+        result.append(["{}={}".format(axis, value) for value in axes[axis]])
+    return list(itertools.product(*result))
+
+
+class ProcessRunner:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not isinstance(cls._instance, cls):
+            cls._instance = super(ProcessRunner, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+    def __init__(self):
+        self.process = None
+        signal.signal(signal.SIGINT, self.signal_handler)
+        signal.signal(signal.SIGTERM, self.signal_handler)
+
+    def new_process(self, cmd):
+        self.process = subprocess.Popen(
+            cmd,
+            start_new_session=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        return self.process
+
+    def signal_handler(self, signum, frame):
+        self.kill_process()
+        raise SystemExit("search was interrupted")
+
+    def kill_process(self):
+        if self.process is not None:
+            self.process.kill()
+
+
+class Bench:
+    def __init__(self, algorithm_name, variant, ct_workload):
+        self.algname = algorithm_name
+        self.variant = variant
+        self.ct_workload = ct_workload
+
+    def label(self):
+        return self.algname + "." + self.variant.label()
+
+    def variant_name(self):
+        return self.variant.label()
+
+    def algorithm_name(self):
+        return self.algname
+
+    def is_base(self):
+        return self.variant.is_base()
+
+    def get_base(self):
+        return BaseBench(self.algorithm_name())
+
+    def exe_name(self):
+        if self.is_base():
+            return self.algorithm_name() + ".base"
+        return self.algorithm_name() + ".variant"
+
+    def bench_names(self):
+        return [bench["name"] for bench in json_benches(self.algname)["benchmarks"]]
+
+    def axes_names(self):
+        subbench_names = {}
+        for bench in json_benches(self.algname)["benchmarks"]:
+            names = []
+            for axis in bench["axes"]:
+                names.append(get_axis_name(axis))
+
+            subbench_names[bench["name"]] = names
+        return subbench_names
+
+    def axes_values(self, sub_space, ct):
+        subbench_space = {}
+        for bench in json_benches(self.algname)["benchmarks"]:
+            space = {}
+            for axis in bench["axes"]:
+                name = get_axis_name(axis)
+
+                if ct:
+                    if "{ct}" not in name:
+                        continue
+                else:
+                    if "{ct}" in name:
+                        continue
+
+                axis_space = []
+                if name in sub_space:
+                    for value in sub_space[name]:
+                        axis_space.append(value)
+                else:
+                    for value in axis["values"]:
+                        axis_space.append(value["input_string"])
+
+                space[name] = axis_space
+
+            subbench_space[bench["name"]] = space
+        return subbench_space
+
+    def ct_axes_value_descriptions(self):
+        subbench_descriptions = {}
+        for bench in json_benches(self.algname)["benchmarks"]:
+            descriptions = {}
+            for axis in bench["axes"]:
+                name = axis["name"]
+                if "{ct}" not in name:
+                    continue
+                if axis["flags"]:
+                    name = name + "[{}]".format(axis["flags"])
+                descriptions[name] = {}
+                for value in axis["values"]:
+                    descriptions[name][value["input_string"]] = value["description"]
+
+            subbench_descriptions[bench["name"]] = descriptions
+        return first_val(subbench_descriptions)
+
+    def axis_values(self, axis_name):
+        result = json_benches(self.algname)
+
+        if len(result["benchmarks"]) != 1:
+            raise Exception("Executable should contain exactly one benchmark")
+
+        for axis in result["benchmarks"][0]["axes"]:
+            name = axis["name"]
+
+            if axis["flags"]:
+                name = name + "[{}]".format(axis["flags"])
+
+            if name != axis_name:
+                continue
+
+            values = []
+            for value in axis["values"]:
+                values.append(value["input_string"])
+
+            return values
+
+        return []
+
+    def build(self):
+        if not self.is_base():
+            self.get_base().build()
+        build = CMake().build(self)
+        return build.code == 0
+
+    def definitions(self):
+        definitions = self.variant.tuning()
+        definitions = definitions + "\n"
+
+        descriptions = self.ct_axes_value_descriptions()
+        for ct_component in self.ct_workload:
+            ct_axis_name, ct_value = ct_component.split("=")
+            description = descriptions[ct_axis_name][ct_value]
+            ct_axis_name = ct_axis_name.replace("{ct}", "")
+            definitions = definitions + "#define TUNE_{} {}\n".format(
+                ct_axis_name, description
+            )
+
+        return definitions
+
+    def do_run(self, ct_point, rt_values, timeout, is_search=True):
+        logger = Logger()
+
+        try:
+            result_path = "result.json"
+            if os.path.exists(result_path):
+                os.remove(result_path)
+
+            bench_path = os.path.join(".", "bin", self.exe_name())
+            cmd = [bench_path]
+
+            for value in ct_point:
+                cmd.append("-a")
+                cmd.append(value)
+
+            cmd.append("--jsonbin")
+            cmd.append(result_path)
+
+            cmd.append("--stopping-criterion")
+            cmd.append("entropy")
+
+            # NVBench is currently broken for multiple GPUs, use `CUDA_VISIBLE_DEVICES`
+            cmd.append("-d")
+            cmd.append("0")
+
+            for bench in rt_values:
+                cmd.append("-b")
+                cmd.append(bench)
+
+                for axis in rt_values[bench]:
+                    cmd.append("-a")
+                    cmd.append("{}=[{}]".format(axis, ",".join(rt_values[bench][axis])))
+
+            logger.info(
+                "starting benchmark {} with {}: {}".format(
+                    self.label(), ct_point, " ".join(cmd)
+                )
+            )
+
+            begin = time.time()
+            p = ProcessRunner().new_process(cmd)
+            p.wait(timeout=timeout)
+            elapsed = time.time() - begin
+
+            logger.info(
+                "finished benchmark {} with {} ({}) in {:.3f}s".format(
+                    self.label(), ct_point, p.returncode, elapsed
+                )
+            )
+
+            return BenchResult(result_path, p.returncode, elapsed)
+        except subprocess.TimeoutExpired:
+            logger.info(
+                "benchmark {} with {} reached timeout of {:.3f}s".format(
+                    self.label(), ct_point, timeout
+                )
+            )
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+            return BenchResult(None, 42, float("inf"))
+
+    def ct_workload_space(self, sub_space):
+        if not self.build():
+            raise Exception("Unable to build benchmark: " + self.label())
+
+        return values_to_space(first_val(self.axes_values(sub_space, True)))
+
+    def rt_axes_values(self, sub_space):
+        if not self.build():
+            raise Exception("Unable to build benchmark: " + self.label())
+
+        return self.axes_values(sub_space, False)
+
+    def run(self, ct_workload_point, rt_values, estimator, is_search=True):
+        logger = Logger()
+        bench_cache = BenchCache()
+        runs_cache = RunsCache()
+        cached_centers = bench_cache.pull_bench_centers(
+            self, ct_workload_point, rt_values
+        )
+        if cached_centers:
+            logger.info("found benchmark {} in cache".format(self.label()))
+            return cached_centers
+
+        timeout = None
+
+        if not self.is_base():
+            code, elapsed = runs_cache.pull_run(self.get_base())
+            if code != 0:
+                raise Exception("Base bench return code = " + code)
+            timeout = elapsed * 50
+
+        result = self.do_run(ct_workload_point, rt_values, timeout, is_search)
+        runs_cache.push_run(self, result.code, result.elapsed)
+        return bench_cache.push_bench_centers(self, result, estimator)
+
+    def speedup(self, ct_workload_point, rt_values, base_estimator, variant_estimator):
+        if self.is_base():
+            return 1.0
+
+        base = self.get_base()
+        base_center = base.run(ct_workload_point, rt_values, base_estimator)
+        self_center = self.run(ct_workload_point, rt_values, variant_estimator)
+        return speedup(base_center, self_center)
+
+    def score(self, ct_workload, rt_values, base_estimator, variant_estimator):
+        if self.is_base():
+            return 1.0
+
+        speedups = self.speedup(
+            ct_workload, rt_values, base_estimator, variant_estimator
+        )
+
+        if not speedups:
+            return float("-inf")
+
+        rt_axes_ids = compute_axes_ids(rt_values)
+        weight_matrices = compute_weight_matrices(rt_values, rt_axes_ids)
+
+        score = 0
+        for bench in speedups:
+            for state in speedups[bench]:
+                rt_workload = state_to_rt_workload(bench, state)
+                weights = weight_matrices[bench]
+                weight = get_workload_weight(
+                    rt_workload, rt_values[bench], rt_axes_ids[bench], weights
+                )
+                speedup = speedups[bench][state]
+                score = score + weight * speedup
+
+        return score
+
+
+class BaseBench(Bench):
+    def __init__(self, algname):
+        super().__init__(algname, BasePoint(), [])

cccl_upstream/benchmarks/scripts/cccl/bench/build.py

@@ -0,0 +1,7 @@
+class Build:
+    def __init__(self, code, elapsed):
+        self.code = code
+        self.elapsed = elapsed
+
+    def __repr__(self):
+        return "Build(code = {}, elapsed = {:.4f}s)".format(self.code, self.elapsed)

cccl_upstream/benchmarks/scripts/cccl/bench/cmake.py

@@ -0,0 +1,138 @@
+import os
+import signal
+import subprocess
+import time
+
+from .build import Build
+from .config import Config
+from .logger import Logger
+from .storage import Storage
+
+
+def create_builds_table(conn):
+    with conn:
+        conn.execute("""
+        CREATE TABLE IF NOT EXISTS builds (
+            ctk TEXT NOT NULL,
+            cccl TEXT NOT NULL,
+            bench TEXT NOT NULL,
+            code TEXT NOT NULL,
+            elapsed REAL
+        );
+        """)
+
+
+class CMakeCache:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+            create_builds_table(Storage().connection())
+        return cls._instance
+
+    def pull_build(self, bench):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        conn = Storage().connection()
+
+        with conn:
+            query = "SELECT code, elapsed FROM builds WHERE ctk = ? AND cccl = ? AND bench = ?;"
+            result = conn.execute(query, (ctk, cccl, bench.label())).fetchone()
+
+            if result:
+                code, elapsed = result
+                return Build(int(code), float(elapsed))
+
+            return result
+
+    def push_build(self, bench, build):
+        config = Config()
+        ctk = config.ctk
+        cccl = config.cccl
+        conn = Storage().connection()
+
+        with conn:
+            conn.execute(
+                "INSERT INTO builds (ctk, cccl, bench, code, elapsed) VALUES (?, ?, ?, ?, ?);",
+                (ctk, cccl, bench.label(), build.code, build.elapsed),
+            )
+
+
+class CMake:
+    def __init__(self):
+        pass
+
+    def do_build(self, bench, timeout):
+        logger = Logger()
+
+        try:
+            if not bench.is_base():
+                with open(bench.exe_name() + ".h", "w") as f:
+                    f.writelines(bench.definitions())
+
+            cmd = ["cmake", "--build", ".", "--target", bench.exe_name()]
+            logger.info(
+                "starting build for {}: {}".format(bench.label(), " ".join(cmd))
+            )
+
+            begin = time.time()
+            p = subprocess.Popen(
+                cmd,
+                start_new_session=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            p.wait(timeout=timeout)
+            elapsed = time.time() - begin
+            logger.info(
+                "finished build for {} (exit code: {}) in {:.3f}s".format(
+                    bench.label(), p.returncode, elapsed
+                )
+            )
+
+            return Build(p.returncode, elapsed)
+        except subprocess.TimeoutExpired:
+            logger.info(
+                "build for {} reached timeout of {}s".format(bench.label(), timeout)
+            )
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+            return Build(424242, float("inf"))
+
+    def build(self, bench):
+        logger = Logger()
+        timeout = None
+
+        cache = CMakeCache()
+
+        if bench.is_base():
+            # Only base build can be pulled from cache
+            build = cache.pull_build(bench)
+
+            if build:
+                logger.info("found cached base build for {}".format(bench.label()))
+                if bench.is_base():
+                    if not os.path.exists("bin/{}".format(bench.exe_name())):
+                        self.do_build(bench, None)
+
+                return build
+        else:
+            base_build = self.build(bench.get_base())
+
+            if base_build.code != 0:
+                raise Exception("Base build failed")
+
+            timeout = base_build.elapsed * 10
+
+        build = self.do_build(bench, timeout)
+        cache.push_build(bench, build)
+        return build
+
+    def clean():
+        cmd = ["cmake", "--build", ".", "--target", "clean"]
+        p = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        p.wait()
+
+        if p.returncode != 0:
+            raise Exception("Unable to clean build directory")

cccl_upstream/benchmarks/scripts/cccl/bench/config.py

@@ -0,0 +1,153 @@
+import os
+import random
+import sys
+
+
+def randomized_cartesian_product(list_of_lists):
+    length = 1
+    for lst in list_of_lists:
+        length *= len(lst)
+
+    visited = set()
+    while len(visited) < length:
+        variant = tuple(map(random.choice, list_of_lists))
+        if variant not in visited:
+            visited.add(variant)
+            yield variant
+
+
+class Range:
+    def __init__(self, definition, label, low, high, step):
+        self.definition = definition
+        self.label = label
+        self.low = low
+        self.high = high
+        self.step = step
+
+
+class RangePoint:
+    def __init__(self, definition, label, value):
+        self.definition = definition
+        self.label = label
+        self.value = value
+
+
+class VariantPoint:
+    def __init__(self, range_points):
+        self.range_points = range_points
+
+    def label(self):
+        if self.is_base():
+            return "base"
+        return ".".join(
+            ["{}_{}".format(point.label, point.value) for point in self.range_points]
+        )
+
+    def is_base(self):
+        return len(self.range_points) == 0
+
+    def tuning(self):
+        if self.is_base():
+            return ""
+
+        tuning = "#pragma once\n\n"
+        for point in self.range_points:
+            tuning += "#define {} {}\n".format(point.definition, point.value)
+        return tuning
+
+
+class BasePoint(VariantPoint):
+    def __init__(self):
+        VariantPoint.__init__(self, [])
+
+
+def parse_ranges(columns):
+    ranges = []
+    for column in columns:
+        definition, label_range = column.split("|")
+        label, range = label_range.split("=")
+        start, end, step = [int(x) for x in range.split(":")]
+        ranges.append(Range(definition, label, start, end + 1, step))
+
+    return ranges
+
+
+def parse_meta():
+    if not os.path.isfile("cccl_meta_bench.csv"):
+        print("cccl_meta_bench.csv not found", file=sys.stderr)
+        print(
+            "make sure to run the script from the CUB build directory", file=sys.stderr
+        )
+
+    benchmarks = {}
+    ctk_version = "0.0.0"
+    cccl_revision = "0.0-0-0000"
+    with open("cccl_meta_bench.csv", "r") as f:
+        lines = f.readlines()
+        for line in lines:
+            if "," in line:
+                columns = line.split(",")
+            else:
+                columns = [" ".join(line.split())]
+
+            name = columns[0]
+
+            if name == "ctk_version":
+                ctk_version = columns[1].rstrip()
+            elif name == "cccl_revision":
+                cccl_revision = columns[1].rstrip()
+            else:
+                if len(columns) > 1:
+                    benchmarks[name] = parse_ranges(columns[1:])
+                else:
+                    benchmarks[name] = []
+
+    return ctk_version, cccl_revision, benchmarks
+
+
+class Config:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+            cls._instance.ctk, cls._instance.cccl, cls._instance.benchmarks = (
+                parse_meta()
+            )
+        return cls._instance
+
+    def label_to_variant_point(self, algname, label):
+        if label == "base":
+            return BasePoint()
+
+        label_to_definition = {}
+        for param_space in self.benchmarks[algname]:
+            label_to_definition[param_space.label] = param_space.definition
+
+        points = []
+        for point in label.split("."):
+            label, value = point.split("_")
+            points.append(RangePoint(label_to_definition[label], label, int(value)))
+
+        return VariantPoint(points)
+
+    def variant_space(self, algname):
+        variants = []
+        for param_space in self.benchmarks[algname]:
+            variants.append([])
+            for value in range(param_space.low, param_space.high, param_space.step):
+                variants[-1].append(
+                    RangePoint(param_space.definition, param_space.label, value)
+                )
+
+        return (
+            VariantPoint(points) for points in randomized_cartesian_product(variants)
+        )
+
+    def variant_space_size(self, algname):
+        num_variants = 1
+        for param_space in self.benchmarks[algname]:
+            num_variants = num_variants * len(
+                range(param_space.low, param_space.high, param_space.step)
+            )
+        return num_variants

cccl_upstream/benchmarks/scripts/cccl/bench/logger.py

@@ -0,0 +1,20 @@
+import logging
+
+
+class Logger:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+            logger = logging.getLogger()
+            logger.setLevel(logging.DEBUG)
+            file_handler = logging.FileHandler("cccl_meta_bench.log")
+            file_handler.setFormatter(logging.Formatter("%(asctime)s: %(message)s"))
+            logger.addHandler(file_handler)
+            cls._instance.logger = logger
+
+        return cls._instance
+
+    def info(self, message):
+        self.logger.info(message)

cccl_upstream/benchmarks/scripts/cccl/bench/score.py

@@ -0,0 +1,105 @@
+import math
+
+import numpy as np
+
+
+def importance_function(x):
+    return 1 - math.exp(-x)
+
+
+def x_by_importance(y):
+    return -math.log(1 - y)
+
+
+def compute_weights(num_values):
+    least_importance = 0.6
+    most_importance = 0.999
+
+    assert least_importance < most_importance
+    assert least_importance >= 0 and least_importance < 1
+    assert most_importance > 0 and most_importance < 1
+
+    begin = x_by_importance(least_importance)
+    end = x_by_importance(most_importance)
+
+    rng = end - begin
+    step = rng / num_values
+
+    return np.array([importance_function(begin + x * step) for x in range(num_values)])
+
+
+def io_weights(values):
+    return compute_weights(len(values))
+
+
+def ei_weights(values):
+    return np.ones(len(values))
+
+
+def compute_axes_ids(rt_axes_values):
+    result = {}
+    for bench in rt_axes_values:
+        rt_axes_ids = {}
+
+        axis_id = 0
+        for rt_axis in rt_axes_values[bench]:
+            rt_axes_ids[rt_axis] = axis_id
+            axis_id = axis_id + 1
+        result[bench] = rt_axes_ids
+    return result
+
+
+def compute_raw_weight_matrix(rt_axes_values, rt_axes_ids):
+    rt_axes_weights = {}
+
+    first_rt_axis = True
+    first_rt_axis_name = None
+    for rt_axis, values in rt_axes_values.items():
+        if first_rt_axis:
+            first_rt_axis_name = rt_axis
+            first_rt_axis = False
+        if "{io}" in rt_axis:
+            rt_axes_weights[rt_axis] = io_weights(values)
+        else:
+            rt_axes_weights[rt_axis] = ei_weights(values)
+
+    num_rt_axes = len(rt_axes_ids)
+    for rt_axis in rt_axes_weights:
+        shape = [1] * num_rt_axes
+        shape[rt_axes_ids[rt_axis]] = -1
+        rt_axes_weights[rt_axis] = rt_axes_weights[rt_axis].reshape(*shape)
+
+    weights_matrix = rt_axes_weights[first_rt_axis_name]
+    for rt_axis in rt_axes_weights:
+        if rt_axis == first_rt_axis_name:
+            continue
+
+        weights_matrix = weights_matrix * rt_axes_weights[rt_axis]
+
+    return weights_matrix
+
+
+def compute_weight_matrices(rt_axes_values, rt_axes_ids):
+    matrices = {}
+    aggregate = 0.0
+    for bench in rt_axes_values:
+        matrices[bench] = compute_raw_weight_matrix(
+            rt_axes_values[bench], rt_axes_ids[bench]
+        )
+        aggregate = aggregate + np.sum(matrices[bench])
+    for bench in rt_axes_values:
+        matrices[bench] = matrices[bench] / aggregate
+    return matrices
+
+
+def get_workload_coordinates(rt_workload, rt_axes_values, rt_axes_ids):
+    coordinates = [0] * len(rt_axes_ids)
+    for point in rt_workload:
+        rt_axis, rt_value = point.split("=")
+        coordinates[rt_axes_ids[rt_axis]] = rt_axes_values[rt_axis].index(rt_value)
+    return coordinates
+
+
+def get_workload_weight(rt_workload, rt_axes_values, rt_axes_ids, weights_matrix):
+    coordinates = get_workload_coordinates(rt_workload, rt_axes_values, rt_axes_ids)
+    return weights_matrix[tuple(coordinates)]

cccl_upstream/benchmarks/scripts/cccl/bench/search.py

@@ -0,0 +1,214 @@
+import argparse
+import re
+
+import numpy as np
+
+from .bench import BaseBench, Bench
+from .cmake import CMake
+from .config import Config
+from .storage import Storage
+
+
+def list_benches(algnames):
+    print("### Benchmarks")
+
+    config = Config()
+
+    for algname in algnames:
+        space_size = config.variant_space_size(algname)
+        print("  * `{}`: {} variants: ".format(algname, space_size))
+
+        for param_space in config.benchmarks[algname]:
+            param_name = param_space.label
+            param_rng = (param_space.low, param_space.high, param_space.step)
+            print("    * `{}`: {}".format(param_name, param_rng))
+
+
+def parse_sub_space(args):
+    sub_space = {}
+    for axis in args:
+        name, value = axis.split("=")
+
+        if "[" in value:
+            value = value.replace("[", "").replace("]", "")
+            values = value.split(",")
+        else:
+            values = [value]
+        sub_space[name] = values
+
+    return sub_space
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Runs benchmarks and stores results in a database."
+    )
+    parser.add_argument(
+        "-R", type=str, default=".*", help="Regex for benchmarks selection."
+    )
+    parser.add_argument(
+        "-a",
+        "--args",
+        action="append",
+        type=str,
+        help="Parameter in the format `Param=Value`.",
+    )
+    parser.add_argument(
+        "-l", "--list-benches", action="store_true", help="Show available benchmarks."
+    )
+    parser.add_argument(
+        "--num-shards",
+        type=int,
+        default=1,
+        help="Split benchmarks into NUM_SHARDS pieces and only run one",
+    )
+    parser.add_argument(
+        "--run-shard",
+        type=int,
+        default=0,
+        help="Run benchmark shard RUN_SHARD from NUM_SHARDS pieces",
+    )
+    parser.add_argument("-P0", action="store_true", help="Run P0 benchmarks")
+    return parser.parse_args()
+
+
+def filter_benchmark_space_for_p0(algname, ct_space, rt_values):
+    if algname in [
+        "cub.bench.merge_sort.pairs",
+        "cub.bench.radix_sort.pairs",
+        "cub.bench.select.unique_by_key",
+    ]:
+        ct_space = list(
+            filter(
+                lambda variant: not (
+                    ("OffsetT{ct}=I64" in variant)
+                    or ("KeyT{ct}=I16" in variant)
+                    or ("ValueT{ct}=I16" in variant)
+                    or ("KeyT{ct}=I128" in variant)
+                    or ("ValueT{ct}=I128" in variant)
+                ),
+                ct_space,
+            )
+        )
+
+    if algname == "cub.bench.merge_sort.pairs":
+        for subbench in rt_values:
+            for axis in rt_values[subbench]:
+                if axis == "Entropy":
+                    rt_values[subbench][axis] = ["1.000"]
+
+    return ct_space, rt_values
+
+
+def run_benches(algnames, sub_space, seeker, args):
+    for algname in algnames:
+        try:
+            bench = BaseBench(algname)
+            ct_space = bench.ct_workload_space(sub_space)
+            rt_values = bench.rt_axes_values(sub_space)
+            if args.P0:
+                ct_space, rt_values = filter_benchmark_space_for_p0(
+                    algname, ct_space, rt_values
+                )
+            seeker(algname, ct_space, rt_values)
+        except Exception as e:
+            print(
+                "#### ERROR exception occurred while running {}: '{}'".format(
+                    algname, e
+                )
+            )
+
+
+def filter_benchmarks_by_regex(benchmarks, R):
+    pattern = re.compile(R)
+    return list(filter(lambda x: pattern.match(x), benchmarks))
+
+
+def filter_benchmarks(benchmarks, args):
+    if args.run_shard >= args.num_shards:
+        raise ValueError("run-shard must be less than num-shards")
+
+    p0_benchmarks = [
+        "cub.bench.merge_sort.keys",
+        "cub.bench.merge_sort.pairs",
+        "cub.bench.radix_sort.keys",
+        "cub.bench.radix_sort.pairs",
+        "cub.bench.reduce.by_key",
+        "cub.bench.reduce.custom",
+        "cub.bench.reduce.sum",
+        "cub.bench.scan.exclusive.deterministic",
+        "cub.bench.scan.exclusive.sum",
+        "cub.bench.select.flagged",
+        "cub.bench.select.if",
+        "cub.bench.select.unique",
+        "cub.bench.select.unique_by_key",
+        "cub.bench.transform.babelstream",
+        "cub.bench.transform.fill",
+    ]
+
+    algnames = filter_benchmarks_by_regex(benchmarks.keys(), args.R)
+    if args.P0:
+        algnames = [name for name in p0_benchmarks if name in algnames]
+    algnames.sort()
+
+    if args.num_shards > 1:
+        algnames = np.array_split(algnames, args.num_shards)[args.run_shard].tolist()
+        return algnames
+
+    return algnames
+
+
+def search(seeker):
+    args = parse_arguments()
+
+    if not Storage().exists():
+        CMake().clean()
+
+    config = Config()
+    print(" ctk: ", config.ctk)
+    print("cccl: ", config.cccl)
+
+    workload_sub_space = {}
+
+    if args.args:
+        workload_sub_space = parse_sub_space(args.args)
+
+    algnames = filter_benchmarks(config.benchmarks, args)
+    if args.list_benches:
+        list_benches(algnames)
+        return
+
+    run_benches(algnames, workload_sub_space, seeker, args)
+
+
+class MedianCenterEstimator:
+    def __init__(self):
+        pass
+
+    def __call__(self, samples):
+        if len(samples) == 0:
+            return float("inf")
+
+        return float(np.median(samples))
+
+
+class BruteForceSeeker:
+    def __init__(self, base_center_estimator, variant_center_estimator):
+        self.base_center_estimator = base_center_estimator
+        self.variant_center_estimator = variant_center_estimator
+
+    def __call__(self, algname, ct_workload_space, rt_values):
+        variants = Config().variant_space(algname)
+
+        for ct_workload in ct_workload_space:
+            for variant in variants:
+                bench = Bench(algname, variant, list(ct_workload))
+                if bench.build():
+                    score = bench.score(
+                        ct_workload,
+                        rt_values,
+                        self.base_center_estimator,
+                        self.variant_center_estimator,
+                    )
+
+                    print(bench.label(), score)

cccl_upstream/benchmarks/scripts/cccl/bench/storage.py

@@ -0,0 +1,392 @@
+import os
+import sqlite3
+
+import fpzip
+import numpy as np
+import pandas as pd
+
+db_name = "cccl_meta_bench.db"
+
+# PostgreSQL support
+try:
+    import psycopg2
+    import psycopg2.extras
+
+    POSTGRES_AVAILABLE = True
+except ImportError:
+    POSTGRES_AVAILABLE = False
+
+
+def get_postgres_config():
+    """Get PostgreSQL configuration from environment variables."""
+    if not POSTGRES_AVAILABLE:
+        return None
+
+    # Check if all required environment variables are set
+    required_vars = [
+        "CCCL_BENCH_PG_HOST",
+        "CCCL_BENCH_PG_USER",
+        "CCCL_BENCH_PG_DB",
+        "CCCL_BENCH_PG_PASSWORD",
+    ]
+    config = {}
+
+    for var in required_vars:
+        value = os.environ.get(var)
+        if not value:
+            return None  # Fall back to SQLite if any required var is missing
+        config[var] = value
+
+    # Optional port (default to 5432)
+    config["CCCL_BENCH_PG_PORT"] = os.environ.get("CCCL_BENCH_PG_PORT", "5432")
+
+    return config
+
+
+def get_bench_table_name(subbench, algname):
+    return "{}.{}".format(algname, subbench)
+
+
+def blob_to_samples(blob):
+    return np.squeeze(fpzip.decompress(blob))
+
+
+class StorageBase:
+    """Abstract base class for storage backends."""
+
+    def connection(self):
+        raise NotImplementedError
+
+    def exists(self):
+        raise NotImplementedError
+
+    def algnames(self):
+        raise NotImplementedError
+
+    def subbenches(self, algname):
+        raise NotImplementedError
+
+    def alg_to_df(self, algname, subbench):
+        raise NotImplementedError
+
+    def store_df(self, algname, df):
+        raise NotImplementedError
+
+
+class SQLiteStorage(StorageBase):
+    def __init__(self, db_path):
+        self.db_path = db_path
+        self.conn = sqlite3.connect(db_path)
+
+    def connection(self):
+        return self.conn
+
+    def exists(self):
+        return os.path.exists(self.db_path)
+
+    def algnames(self):
+        with self.conn:
+            rows = self.conn.execute(
+                "SELECT DISTINCT algorithm FROM subbenches"
+            ).fetchall()
+            return [row[0] for row in rows]
+
+    def subbenches(self, algname):
+        with self.conn:
+            rows = self.conn.execute(
+                "SELECT DISTINCT bench FROM subbenches WHERE algorithm=?", (algname,)
+            ).fetchall()
+            return [row[0] for row in rows]
+
+    def alg_to_df(self, algname, subbench):
+        table = get_bench_table_name(subbench, algname)
+        with self.conn:
+            df = pd.read_sql_query('SELECT * FROM "{}"'.format(table), self.conn)
+            df["samples"] = df["samples"].apply(blob_to_samples)
+
+        return df
+
+    def store_df(self, algname, df):
+        df["samples"] = df["samples"].apply(fpzip.compress)
+        df.to_sql(algname, self.conn, if_exists="replace", index=False)
+
+
+class PostgreSQLConnectionWrapper:
+    """Wrapper to make psycopg2 connection compatible with sqlite3 interface."""
+
+    def __init__(self, pg_conn):
+        self.pg_conn = pg_conn
+        self.pg_conn.autocommit = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            self.pg_conn.commit()
+        else:
+            self.pg_conn.rollback()
+
+    def execute(self, query, params=None):
+        """Execute query with SQLite-style parameter substitution."""
+        # Convert SQLite-style ? placeholders to PostgreSQL %s
+        if params:
+            query = query.replace("?", "%s")
+
+        # Convert SQLite BLOB type to PostgreSQL BYTEA
+        query = query.replace(" BLOB", " BYTEA")
+
+        # Fix SQLite-style double-quoted string literals to PostgreSQL single quotes
+        # This is a simple approach - in production you'd want a proper SQL parser
+        import re
+
+        # Match patterns like = "value" and convert to = 'value'
+        query = re.sub(r'= "([^"]*)"', r"= '\1'", query)
+
+        # Handle ON CONFLICT DO NOTHING (SQLite) -> ON CONFLICT DO NOTHING (PostgreSQL)
+        # Both databases support this syntax, so no conversion needed
+
+        cur = self.pg_conn.cursor()
+        if params:
+            cur.execute(query, params)
+        else:
+            cur.execute(query)
+        return cur
+
+    def commit(self):
+        self.pg_conn.commit()
+
+    def rollback(self):
+        self.pg_conn.rollback()
+
+    def close(self):
+        self.pg_conn.close()
+
+
+if POSTGRES_AVAILABLE:
+
+    class PostgreSQLStorage(StorageBase):
+        def __init__(self, config):
+            self.config = config
+            self.pg_conn = psycopg2.connect(
+                host=config["CCCL_BENCH_PG_HOST"],
+                port=config["CCCL_BENCH_PG_PORT"],
+                database=config["CCCL_BENCH_PG_DB"],
+                user=config["CCCL_BENCH_PG_USER"],
+                password=config["CCCL_BENCH_PG_PASSWORD"],
+            )
+            self.conn = PostgreSQLConnectionWrapper(self.pg_conn)
+
+        def connection(self):
+            return self.conn
+
+        def exists(self):
+            # For PostgreSQL, check if the subbenches table exists
+            with self.conn:
+                cur = self.conn.execute("""
+                    SELECT EXISTS (
+                        SELECT FROM information_schema.tables
+                        WHERE table_name = 'subbenches'
+                    );
+                """)
+                return cur.fetchone()[0]
+
+        def algnames(self):
+            with self.conn:
+                cur = self.conn.execute("SELECT DISTINCT algorithm FROM subbenches")
+                rows = cur.fetchall()
+                return [row[0] for row in rows]
+
+        def subbenches(self, algname):
+            with self.conn:
+                cur = self.conn.execute(
+                    "SELECT DISTINCT bench FROM subbenches WHERE algorithm=?",
+                    (algname,),
+                )
+                rows = cur.fetchall()
+                return [row[0] for row in rows]
+
+        def alg_to_df(self, algname, subbench):
+            table = get_bench_table_name(subbench, algname)
+            with self.conn:
+                # Use proper quoting for PostgreSQL
+                query = 'SELECT * FROM "{}"'.format(table.replace('"', '""'))
+                df = pd.read_sql_query(query, self.pg_conn)
+                df["samples"] = df["samples"].apply(lambda x: blob_to_samples(bytes(x)))
+            return df
+
+        def store_df(self, algname, df):
+            df["samples"] = df["samples"].apply(fpzip.compress)
+            # For PostgreSQL, we need to use a different approach
+            # as pandas doesn't support direct to_sql with psycopg2
+            # We'll need to implement this separately or use SQLAlchemy
+            raise NotImplementedError(
+                "DataFrame storage for PostgreSQL not yet implemented"
+            )
+else:
+    # Define a dummy class when psycopg2 is not available
+    PostgreSQLStorage = None
+
+
+class DualStorageWrapper:
+    """Wrapper that writes to multiple storage backends."""
+
+    def __init__(self, primary, secondary=None):
+        self.primary = primary
+        self.secondary = secondary
+        self._primary_conn = None
+        self._secondary_conn = None
+
+    def connection(self):
+        # Return a wrapper that forwards operations to both backends
+        if not self._primary_conn:
+            self._primary_conn = DualConnectionWrapper(
+                self.primary.connection(),
+                self.secondary.connection() if self.secondary else None,
+            )
+        return self._primary_conn
+
+    def exists(self):
+        # Check primary storage
+        return self.primary.exists()
+
+    def algnames(self):
+        # Read from primary only
+        return self.primary.algnames()
+
+    def subbenches(self, algname):
+        # Read from primary only
+        return self.primary.subbenches(algname)
+
+    def alg_to_df(self, algname, subbench):
+        # Read from primary only
+        return self.primary.alg_to_df(algname, subbench)
+
+    def store_df(self, algname, df):
+        # Write to both databases
+        self.primary.store_df(algname, df)
+        if self.secondary:
+            try:
+                self.secondary.store_df(algname, df)
+            except Exception as e:
+                print(f"Warning: Failed to write to secondary storage: {e}")
+
+
+class DualCursorWrapper:
+    """Wrapper for cursor results from dual storage."""
+
+    def __init__(self, primary_cursor):
+        self.primary_cursor = primary_cursor
+
+    def fetchone(self):
+        return self.primary_cursor.fetchone()
+
+    def fetchall(self):
+        return self.primary_cursor.fetchall()
+
+
+class DualConnectionWrapper:
+    """Wrapper that forwards connection operations to both backends."""
+
+    def __init__(self, primary_conn, secondary_conn=None):
+        self.primary_conn = primary_conn
+        self.secondary_conn = secondary_conn
+
+    def __enter__(self):
+        # SQLite connections are their own context managers
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Commit or rollback based on exception
+        if exc_type is None:
+            self.commit()
+        else:
+            self.rollback()
+        return False
+
+    def execute(self, query, params=None):
+        # Execute on primary
+        if params:
+            primary_result = self.primary_conn.execute(query, params)
+        else:
+            primary_result = self.primary_conn.execute(query)
+
+        # Also execute on secondary if available
+        if self.secondary_conn:
+            try:
+                if params:
+                    self.secondary_conn.execute(query, params)
+                else:
+                    self.secondary_conn.execute(query)
+            except Exception:
+                # Don't print warnings for every query, too noisy
+                pass
+
+        # Return a wrapper that delegates to the primary result
+        return DualCursorWrapper(primary_result)
+
+    def fetchone(self):
+        # Delegate to primary connection
+        return self.primary_conn.fetchone()
+
+    def fetchall(self):
+        # Delegate to primary connection
+        return self.primary_conn.fetchall()
+
+    def commit(self):
+        self.primary_conn.commit()
+        if self.secondary_conn:
+            try:
+                self.secondary_conn.commit()
+            except Exception as e:
+                print(f"Warning: Failed to commit to secondary storage: {e}")
+
+    def rollback(self):
+        self.primary_conn.rollback()
+        if self.secondary_conn:
+            try:
+                self.secondary_conn.rollback()
+            except Exception as e:
+                print(f"Warning: Failed to rollback secondary storage: {e}")
+
+
+class Storage:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls, *args, **kwargs)
+
+            # Always use SQLite as primary
+            sqlite_storage = SQLiteStorage(db_name)
+
+            # Try to add PostgreSQL as secondary if configured
+            pg_config = get_postgres_config()
+            pg_storage = None
+
+            if pg_config and PostgreSQLStorage is not None:
+                try:
+                    pg_storage = PostgreSQLStorage(pg_config)
+                    print(
+                        "Using dual storage: SQLite (primary) + PostgreSQL (secondary)"
+                    )
+                except Exception as e:
+                    print(f"Failed to connect to PostgreSQL: {e}")
+                    print("Using SQLite only")
+
+            # Create wrapper with SQLite as primary and PostgreSQL as optional secondary
+            cls._instance.base = DualStorageWrapper(sqlite_storage, pg_storage)
+
+        return cls._instance
+
+    def connection(self):
+        return self.base.connection()
+
+    def exists(self):
+        return self.base.exists()
+
+    def algnames(self):
+        return self.base.algnames()
+
+    def alg_to_df(self, algname, subbench):
+        return self.base.alg_to_df(algname, subbench)

cccl_upstream/benchmarks/scripts/compare.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+
+import argparse
+import os
+
+import cccl
+import numpy as np
+import pandas as pd
+from colorama import Fore
+
+
+def get_filenames_map(arr):
+    if not arr:
+        return []
+
+    prefix = arr[0]
+    for string in arr:
+        while not string.startswith(prefix):
+            prefix = prefix[:-1]
+        if not prefix:
+            break
+
+    return {string: string[len(prefix) :] for string in arr}
+
+
+def is_finite(x):
+    if isinstance(x, float):
+        return x != np.inf and x != -np.inf
+    return True
+
+
+def filter_by_problem_size(df):
+    min_elements_pow2 = 28
+    if "Elements{io}[pow2]" in df.columns:
+        df["Elements{io}[pow2]"] = df["Elements{io}[pow2]"].astype(int)
+        df = df[df["Elements{io}[pow2]"] >= min_elements_pow2]
+    return df
+
+
+def filter_by_offset_type(df):
+    if "OffsetT{ct}" in df.columns:
+        df = df[(df["OffsetT{ct}"] == "I32") | (df["OffsetT{ct}"] == "U32")]
+    return df
+
+
+def filter_by_type(df):
+    if "T{ct}" in df:
+        # df = df[df['T{ct}'].str.contains('64')]
+        df = df[~df["T{ct}"].str.contains("C")]
+    elif "KeyT{ct}" in df:
+        # df = df[df['KeyT{ct}'].str.contains('64')]
+        df = df[~df["KeyT{ct}"].str.contains("C")]
+    return df
+
+
+def alg_dfs(file):
+    result = {}
+    storage = cccl.bench.SQLiteStorage(file)
+    for algname in storage.algnames():
+        for subbench in storage.subbenches(algname):
+            df = storage.alg_to_df(algname, subbench)
+            df = df.map(lambda x: x if is_finite(x) else np.nan)
+            df = df.dropna(subset=["center"], how="all")
+            # TODO(bgruber): maybe expose the filters under a -p0, or --short flag
+            # df = filter_by_type(filter_by_offset_type(filter_by_problem_size(df)))
+            df["Noise"] = df["samples"].apply(lambda x: np.std(x) / np.mean(x)) * 100
+            df["Mean"] = df["samples"].apply(lambda x: np.mean(x))
+            df = df.drop(columns=["samples", "center", "bw", "elapsed", "variant"])
+            fused_algname = (
+                algname.removeprefix("cub.bench.").removeprefix("thrust.bench.")
+                + "."
+                + subbench
+            )
+            result[fused_algname] = df
+
+    for algname in result:
+        if result[algname]["cccl"].nunique() != 1:
+            print(f"WARNING: Multiple CCCL versions in one db '{algname}'")
+        result[algname] = result[algname].drop(columns=["cccl"])
+
+    return result
+
+
+def file_exists(value):
+    if not os.path.isfile(value):
+        raise argparse.ArgumentTypeError(f"The file '{value}' does not exist.")
+    return value
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Analyze benchmark results.")
+    parser.add_argument("reference", type=file_exists, help="Reference database file.")
+    parser.add_argument("compare", type=file_exists, help="Comparison database file.")
+    return parser.parse_args()
+
+
+config_count = 0
+pass_count = 0
+faster_count = 0
+slower_count = 0
+
+
+def status(frac_diff, noise_ref, noise_cmp):
+    global config_count
+    global pass_count
+    global faster_count
+    global slower_count
+    config_count += 1
+    min_noise = min(noise_ref, noise_cmp)
+    if abs(frac_diff) <= min_noise:
+        pass_count += 1
+        return Fore.BLUE + "SAME" + Fore.RESET
+    if frac_diff < 0:
+        faster_count += 1
+        return Fore.GREEN + "FAST" + Fore.RESET
+    if frac_diff > 0:
+        slower_count += 1
+        return Fore.RED + "SLOW" + Fore.RESET
+
+
+def compare():
+    args = parse_args()
+    reference_df = alg_dfs(args.reference)
+    compare_df = alg_dfs(args.compare)
+    for alg in sorted(reference_df.keys() & compare_df.keys()):
+        print()
+        print()
+        print(f"# {alg}")
+        # use every column except 'Noise', 'Mean', 'ctk', 'gpu' to match runs between reference and comparison file
+        merge_columns = [
+            col
+            for col in reference_df[alg].columns
+            if col not in ["Noise", "Mean", "ctk", "gpu"]
+        ]
+        df = pd.merge(
+            reference_df[alg],
+            compare_df[alg],
+            on=merge_columns,
+            suffixes=("Ref", "Cmp"),
+        )
+        df["Abs. Diff"] = df["MeanCmp"] - df["MeanRef"]
+        df["Rel. Diff"] = (df["Abs. Diff"] / df["MeanRef"]) * 100
+        df["Status"] = list(
+            map(status, df["Rel. Diff"], df["NoiseRef"], df["NoiseCmp"])
+        )
+        df = df.drop(columns=["ctkRef", "ctkCmp", "gpuRef", "gpuCmp"])
+        print()
+        print(df.to_markdown(index=False))
+
+    print("# Summary\n")
+    print("- Total Matches: %d" % config_count)
+    print("  - Pass    (diff <= min_noise): %d" % pass_count)
+    print("  - Faster  (diff > min_noise):  %d" % faster_count)
+    print("  - Slower  (diff > min_noise):  %d" % slower_count)
+
+
+if __name__ == "__main__":
+    compare()

cccl_upstream/benchmarks/scripts/run.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+
+import math
+import os
+import sys
+
+import cccl.bench
+
+
+def elapsed_time_looks_good(x):
+    if isinstance(x, float):
+        if math.isfinite(x):
+            return True
+    return False
+
+
+def get_largest_problem_size(rt_values):
+    # Small problem sizes do not utilize entire GPU.
+    # Benchmarking small problem sizes in environments where we do not control
+    # distributions comparison, e.g. CI, is not useful because of stability issues.
+    elements = []
+    for element in rt_values:
+        if element.isdigit():
+            elements.append(int(element))
+    return [str(max(elements))]
+
+
+def filter_runtime_workloads_for_ci(rt_values):
+    for subbench in rt_values:
+        for axis in rt_values[subbench]:
+            if axis.startswith("Elements") and axis.endswith("[pow2]"):
+                rt_values[subbench][axis] = get_largest_problem_size(
+                    rt_values[subbench][axis]
+                )
+
+    return rt_values
+
+
+class BaseRunner:
+    def __init__(self):
+        self.estimator = cccl.bench.MedianCenterEstimator()
+
+    def __call__(self, algname, ct_workload_space, rt_values):
+        failure_occured = False
+        rt_values = filter_runtime_workloads_for_ci(rt_values)
+
+        for ct_workload in ct_workload_space:
+            bench = cccl.bench.BaseBench(algname)
+            if bench.build():  # might throw
+                results = bench.run(ct_workload, rt_values, self.estimator, False)
+                for subbench in results:
+                    for point in results[subbench]:
+                        bench_name = "{}.{}-{}".format(
+                            bench.algorithm_name(), subbench, point
+                        )
+                        bench_name = bench_name.replace(" ", "___")
+                        bench_name = "".join(
+                            c if c.isalnum() else "_" for c in bench_name
+                        )
+                        elapsed_time = results[subbench][point]
+                        if elapsed_time_looks_good(elapsed_time):
+                            print(
+                                "&&&& PERF {} {} -sec".format(bench_name, elapsed_time)
+                            )
+            else:
+                failure_occured = True
+                print("&&&& FAILED {}".format(algname))
+
+        if failure_occured:
+            sys.exit(1)
+
+
+def main():
+    print("&&&& RUNNING bench")
+    os.environ["CUDA_MODULE_LOADING"] = "EAGER"
+    cccl.bench.search(BaseRunner())
+    print("&&&& PASSED bench")
+
+
+if __name__ == "__main__":
+    main()

cccl_upstream/benchmarks/scripts/search.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+import cccl.bench as bench
+
+# TODO:
+# - driver version
+# - host compiler + version
+# - gpu clocks / pm
+# - ecc
+
+
+def main():
+    center_estimator = bench.MedianCenterEstimator()
+    bench.search(bench.BruteForceSeeker(center_estimator, center_estimator))
+
+
+if __name__ == "__main__":
+    main()

cccl_upstream/benchmarks/scripts/sol.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+
+import argparse
+import os
+import re
+
+import cccl
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+
+
+def is_finite(x):
+    if isinstance(x, float):
+        return x != np.inf and x != -np.inf
+    return True
+
+
+def filter_by_problem_size(df):
+    min_elements_pow2 = 28
+    if "Elements{io}[pow2]" in df.columns:
+        df["Elements{io}[pow2]"] = df["Elements{io}[pow2]"].astype(int)
+        df = df[df["Elements{io}[pow2]"] >= min_elements_pow2]
+    return df
+
+
+def filter_by_offset_type(df):
+    if "OffsetT{ct}" in df.columns:
+        filtered = df[
+            (df["OffsetT{ct}"] == "I32") | (df["OffsetT{ct}"] == "U32")
+        ]  # only use 32-bit offset types
+        if not filtered.empty:  # some benchmarks only use a 64-bit offset type
+            df = filtered
+    return df
+
+
+def filter_by_type(df):
+    if "T{ct}" in df:
+        # df = df[df['T{ct}'].str.contains('64')]
+        df = df[~df["T{ct}"].str.contains("C")]
+    elif "KeyT{ct}" in df:
+        # df = df[df['KeyT{ct}'].str.contains('64')]
+        df = df[~df["KeyT{ct}"].str.contains("C")]
+    return df
+
+
+def alg_dfs(files, alg_regex):
+    pattern = re.compile(alg_regex)
+    result = {}
+    for file in files:
+        storage = cccl.bench.SQLiteStorage(file)
+        for algname in storage.algnames():
+            if pattern.match(algname):
+                for subbench in storage.subbenches(algname):
+                    df = storage.alg_to_df(algname, subbench)
+                    df = df.map(lambda x: x if is_finite(x) else np.nan)
+                    df = df.dropna(subset=["center"], how="all")
+                    df = filter_by_type(
+                        filter_by_offset_type(filter_by_problem_size(df))
+                    )
+                    df = df.filter(items=["ctk", "cccl", "gpu", "variant", "bw"])
+                    fused_algname = algname.replace("bench.", "") + "." + subbench
+                    if df.empty:
+                        print(
+                            f"WARNING: Skipped {fused_algname} because no data is present"
+                        )
+                        print(df)
+                        continue
+                    if df["bw"].dropna().empty:
+                        print(
+                            f"WARNING: Skipped {fused_algname} because it does not report bandwidth"
+                        )
+                        continue
+                    df["variant"] = df["variant"].astype(str)
+                    df["bw"] = df["bw"] * 100
+                    if fused_algname in result:
+                        result[fused_algname] = pd.concat([result[fused_algname], df])
+                    else:
+                        result[fused_algname] = df
+                    print(fused_algname)
+    return result
+
+
+def alg_bws(dfs, verbose):
+    medians = None
+    for algname in dfs:
+        df = dfs[algname]
+        df["alg"] = algname
+        if df is None:
+            medians = df
+        else:
+            medians = pd.concat([medians, df])
+    # print more information if it's not unique across all runs or when requested (verbose)
+    medians["hue"] = ""
+    if verbose or medians["cccl"].unique().size > 1:
+        medians["hue"] = medians["hue"] + "CCCL " + medians["cccl"].astype(str) + " "
+    gpuname = (
+        medians["gpu"]
+        if verbose
+        else medians["gpu"].astype(str).map(lambda x: x[: x.find("(") - 1])
+    )
+    medians["hue"] = medians["hue"] + gpuname + " "
+    if medians["variant"].unique().size > 1:
+        variant = (
+            medians["variant"]
+            .astype(str)
+            .map(lambda x: (" " + x if x != "base" else ""))
+        )
+        medians["hue"] = medians["hue"] + variant + " "
+    if verbose or medians["ctk"].unique().size > 1:
+        medians["hue"] = medians["hue"] + "CTK " + medians["ctk"].astype(str)
+    return medians.drop(columns=["ctk", "cccl", "gpu", "variant"])
+
+
+def file_exists(value):
+    if not os.path.isfile(value):
+        raise argparse.ArgumentTypeError(f"The file '{value}' does not exist.")
+    return value
+
+
+def plot_sol(medians, box):
+    if box:
+        ax = sns.boxenplot(data=medians, x="alg", y="bw", hue="hue")
+    else:
+        ax = sns.barplot(
+            data=medians,
+            x="alg",
+            y="bw",
+            hue="hue",
+            errorbar=lambda x: (x.min(), x.max()),
+        )
+        ax.bar_label(ax.containers[0], fmt="%.1f")
+        for container in ax.containers[1:]:
+            labels = [
+                f"{c:.1f}\n({(c / f) * 100:.0f}%)"
+                for f, c in zip(ax.containers[0].datavalues, container.datavalues)
+            ]
+            ax.bar_label(container, labels=labels)
+
+    ax.legend(title=None)
+    ax.set_xlabel("Algorithm")
+    ax.set_ylabel("Bandwidth (%SOL)")
+    ax.set_xticklabels(
+        ax.get_xticklabels(), rotation=30, rotation_mode="anchor", ha="right"
+    )
+    ax.set_ylim([0, 100])
+    plt.show()
+
+
+def print_speedup(medians):
+    m = medians.groupby(["alg", "hue"], sort=False).mean()
+    m["speedup"] = m["bw"] / m.groupby(["alg"])["bw"].transform("first")
+    print("# Speedups:")
+    print()
+    print(m.drop(columns="bw").sort_values(by="speedup", ascending=False).to_markdown())
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Analyze benchmark results.")
+    parser.add_argument(
+        "files", type=file_exists, nargs="+", help="At least one file is required."
+    )
+    parser.add_argument("--box", action="store_true", help="Plot box instead of bar.")
+    parser.add_argument("-v", action="store_true", help="Verbose legend.")
+    parser.add_argument(
+        "-R", type=str, default=".*", help="Regex for benchmarks selection."
+    )
+    return parser.parse_args()
+
+
+def sol():
+    args = parse_args()
+    dfs = alg_dfs(args.files, args.R)
+    if not dfs:
+        print("ERROR: No benchmark data to process (all benchmarks were skipped).")
+        return
+    medians = alg_bws(dfs, args.v)
+    print_speedup(medians)
+    plot_sol(medians, args.box)
+
+
+if __name__ == "__main__":
+    sol()

cccl_upstream/benchmarks/scripts/submit_benchmark_job.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+
+# This script schedules a SLURM job via crun on computelab to run all CCCL benchmarks and produce a benchmark database
+# TODO: set those accordingly
+scratch=/home/scratch."$USER"_sw
+node_selector="cpu.arch=x86_64 and gpu.product_name='*B200*'"
+container_image="rapidsai/devcontainers:26.06-cpp-gcc14-cuda13.2"
+jobtime="4:00:00"
+benchmark_preset="benchmark"
+
+batch_script=$scratch/batch.sh
+cat << BATCH_SCRIPT > "$batch_script"
+#!/usr/bin/env bash
+
+pip install --break-system-packages fpzip pandas scipy
+
+# clone CCCL
+host=\$(hostname)
+cd $scratch
+if [ -d "\$host/cccl" ]; then
+    rm -r \$host/cccl
+fi
+mkdir \$host
+cd \$host
+git clone --depth 1 git@github.com:NVIDIA/cccl.git
+cd cccl
+
+# configure cmake
+mkdir build_perf
+cd build_perf
+cmake .. --preset $benchmark_preset
+
+# run benchmarks
+export CUDA_VISIBLE_DEVICES=0
+export PYTHONPATH=../benchmarks/scripts/
+../benchmarks/scripts/run.py
+
+echo "Benchmark done. Results in $scratch/\$host/cccl/build_perf/cccl_meta_bench.db"
+BATCH_SCRIPT
+chmod +x "$batch_script"
+
+# schedule SLURM job
+echo "Scheduling script $batch_script"
+echo "#################################################################################"
+cat "$batch_script"
+echo "#################################################################################"
+crun -q "$node_selector" -ex -t "$jobtime" -img "$container_image" -b "$batch_script"

cccl_upstream/benchmarks/scripts/verify.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+
+import argparse
+import sys
+
+import cccl.bench
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="Verify tuning variant")
+    parser.add_argument(
+        "--variant", type=str, help="Variant to verify", default=None, required=True
+    )
+
+    variant = parser.parse_known_args()[0].variant
+    sys.argv.remove("--variant={}".format(variant))
+
+    return variant
+
+
+def workload_header(ct_workload_space, rt_workload_space):
+    for ct_workload in ct_workload_space:
+        for rt_workload in rt_workload_space:
+            workload_point = ct_workload + rt_workload
+            return ", ".join([x.split("=")[0] for x in workload_point])
+
+
+def workload_entry(ct_workload, rt_workload):
+    workload_point = ct_workload + rt_workload
+    return ", ".join([x.split("=")[1] for x in workload_point])
+
+
+class VerifySeeker:
+    def __init__(self, variant_label):
+        self.label = variant_label
+        self.estimator = cccl.bench.MedianCenterEstimator()
+
+    def __call__(self, algname, ct_workload_space, rt_workload_space):
+        variant_point = cccl.bench.Config().label_to_variant_point(algname, self.label)
+
+        print(
+            "{}, MinS, MedianS, MaxS".format(
+                workload_header(ct_workload_space, rt_workload_space)
+            )
+        )
+        for ct_workload in ct_workload_space:
+            bench = cccl.bench.Bench(algname, variant_point, list(ct_workload))
+            if bench.build():
+                base = bench.get_base()
+                for rt_workload in rt_workload_space:
+                    workload_point = ct_workload + rt_workload
+                    base_samples, base_elapsed = base.do_run(workload_point, None)
+                    variant_samples, _ = bench.do_run(workload_point, base_elapsed * 10)
+                    min_speedup = min(base_samples) / min(variant_samples)
+                    median_speedup = self.estimator(base_samples) / self.estimator(
+                        variant_samples
+                    )
+                    max_speedup = max(base_samples) / max(variant_samples)
+                    point_str = workload_entry(ct_workload, rt_workload)
+                    print(
+                        "{}, {}, {}, {}".format(
+                            point_str, min_speedup, median_speedup, max_speedup
+                        )
+                    )
+
+
+def main():
+    cccl.bench.search(VerifySeeker(parse_arguments()))
+
+
+if __name__ == "__main__":
+    main()