[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/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)