56fd68e7dd
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
393 lines
12 KiB
Python
393 lines
12 KiB
Python
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)
|