Updated with ranking feature.

README.md

@@ -4,7 +4,7 @@ emoji: 🧠
 colorFrom: blue
 colorTo: green
 sdk: gradio
-sdk_version: 5.42.0
+sdk_version: 6.12.0
 app_file: app.py
 pinned: false
 license: cc-by-nc-sa-4.0

app.py

@@ -1,25 +1,33 @@
-import gradio as gr
-import pandas as pd
-import numpy as np
-import joblib
 import os
-from huggingface_hub import snapshot_download
+import json
+import time
 import tempfile
+import numpy as np
+import pandas as pd
+import xgboost as xgb
+import shap
+import matplotlib.pyplot as plt
+import gradio as gr
 
-# -------------------------------------------------
-# Configuration
-# -------------------------------------------------
-HF_REPO_ID = "uralstech/AIDE-Chip-Surrogates"
-MODEL_DIR = "surrogate_models_v2"
+# =========================================================
+# CONFIG
+# =========================================================
+TOP_K = 20
+MODEL_DIR = "surrogate_models_v2_json"
 
-WORKLOAD_ALIAS = {
-    "matrix": "matrix_mul",
-    "matmul": "matrix_mul",
+WORKLOADS = {
+    "crc32", "dijkstra", "fft", "qsort", "sha", "matrix_mul"
 }
 
-TARGETS = ["ipc", "l2_miss_rate"]
+TARGETS = {"ipc", "l2_miss_rate"}
 
-FEATURE_COLS = [
+REQUIRED_INPUT_COLUMNS = {
+    "l1d_size", "l1i_size", "l2_size",
+    "l1d_assoc", "l1i_assoc", "l2_assoc",
+    "workload"
+}
+
+FEATURES = [
     "l1d_size_log2",
     "l1i_size_log2",
     "l2_size_log2",
@@ -28,98 +36,43 @@ FEATURE_COLS = [
     "l2_assoc_log2",
     "l2_l1d_ratio_log2",
     "l1d_sets_log2",
-    "l2_sets_log2",
+    "l2_sets_log2"
 ]
 
-REQUIRED_COLS = [
-    "workload",
-    "l1d_size",
-    "l1i_size",
-    "l2_size",
-    "l1d_assoc",
-    "l1i_assoc",
-    "l2_assoc",
-]
+# =========================================================
+# LOAD MODELS
+# =========================================================
+def load_models():
+    models = {}
+    explainers = {}
 
-# -------------------------------------------------
-# Global model cache
-# -------------------------------------------------
-MODEL_CACHE = {}
-
-# -------------------------------------------------
-# Model Download
-# -------------------------------------------------
-def ensure_models():
-    if not os.path.exists(MODEL_DIR):
-        snapshot_download(
-            repo_id=HF_REPO_ID,
-            local_dir=".",
-            allow_patterns="*.pkl",
-        )
+    for w in WORKLOADS:
+        for t in TARGETS:
+            name = f"model_{w}_{t}"
 
-# -------------------------------------------------
-# Utilities
-# -------------------------------------------------
-def resolve_workload(workload: str) -> str:
-    return WORKLOAD_ALIAS.get(workload, workload)
-
-def load_model(workload: str, target: str):
-    try:
-        return MODEL_CACHE[(workload, target)]
-    except KeyError:
-        raise RuntimeError(f"Model not preloaded: {workload}, {target}")
-
-def physical_sanity_check(ipc, miss_rate):
-    out = []
-    if ipc < 0 or ipc > 3.5:
-        out.append(f"IPC={ipc:.3f} out of physical range")
-    if miss_rate < 0 or miss_rate > 1:
-        out.append(f"L2 miss rate={miss_rate:.3f} out of [0,1]")
-    return out
-
-# -------------------------------------------------
-# Preload models (runs once at app start)
-# -------------------------------------------------
-def preload_models():
-    ensure_models()
-
-    workloads = {
-        "crc32",
-        "dijkstra",
-        "fft",
-        "matrix_mul",
-        "qsort",
-        "sha",
-    }
-
-    for workload in workloads:
-        for target in TARGETS:
-            path = os.path.join(
-                MODEL_DIR, f"model_{workload}_{target}.pkl"
-            )
-            payload = joblib.load(path)
-            MODEL_CACHE[(workload, target)] = (
-                payload["model"],
-                payload["log_target"],
-            )
-
-    return "ready"
-
-# -------------------------------------------------
-# Inference Core
-# -------------------------------------------------
-def run_inference(df: pd.DataFrame) -> pd.DataFrame:
-    missing = set(REQUIRED_COLS) - set(df.columns)
-    if missing:
-        raise ValueError(f"Missing required columns: {missing}")
+            model = xgb.XGBRegressor()
+            model.load_model(os.path.join(MODEL_DIR, f"{name}.json"))
+
+            with open(os.path.join(MODEL_DIR, f"{name}_meta.json")) as f:
+                meta = json.load(f)
+
+            models[name] = (model, meta["log_target"])
+            explainers[name] = shap.TreeExplainer(model)
+
+    return models, explainers
+
+
+MODELS, EXPLAINERS = load_models()
+
+# =========================================================
+# FEATURE ENGINEERING
+# =========================================================
+def engineer_features(df):
+    df = df.copy()
 
     for col in [
-        "l1d_size",
-        "l1i_size",
-        "l2_size",
-        "l1d_assoc",
-        "l1i_assoc",
-        "l2_assoc",
+        "l1d_size", "l1i_size", "l2_size",
+        "l1d_assoc", "l1i_assoc", "l2_assoc"
     ]:
         df[f"{col}_log2"] = np.log2(df[col])
 
@@ -127,126 +80,191 @@ def run_inference(df: pd.DataFrame) -> pd.DataFrame:
     df["l1d_sets_log2"] = df["l1d_size_log2"] - df["l1d_assoc_log2"]
     df["l2_sets_log2"] = df["l2_size_log2"] - df["l2_assoc_log2"]
 
+    return df
+
+
+# =========================================================
+# MAIN INFERENCE
+# =========================================================
+def run_inference(file):
+    df = pd.read_csv(file.name)
+
+    missing = REQUIRED_INPUT_COLUMNS - set(df.columns)
+    if missing:
+        raise gr.Error(f"Missing required columns: {missing}")
+
+    start = time.perf_counter()
+
+    df = engineer_features(df)
+
     df["pred_ipc"] = np.nan
     df["pred_l2_miss_rate"] = np.nan
-    df["warnings"] = ""
+    df["was_clipped"] = False
 
-    for idx, row in df.iterrows():
-        workload = resolve_workload(row["workload"])
-        X = row[FEATURE_COLS].values.reshape(1, -1)
+    shap_summary_data = []
 
-        preds = {}
-        warn_msgs = []
+    for workload, idx in df.groupby("workload").groups.items():
+        if workload not in WORKLOADS:
+            raise gr.Error(f"Unknown workload: {workload}")
 
-        for target in TARGETS:
-            model, is_log = load_model(workload, target)
-            pred_raw = model.predict(X)[0]
-            pred = np.expm1(pred_raw) if is_log else pred_raw
-            if target == "l2_miss_rate":
-                pred = np.clip(pred, 0, 1)
-            preds[target] = float(pred)
+        X = df.loc[idx, FEATURES]
 
-        warn_msgs.extend(
-            physical_sanity_check(preds["ipc"], preds["l2_miss_rate"])
-        )
+        model_ipc, log_ipc = MODELS[f"model_{workload}_ipc"]
+        model_miss, log_miss = MODELS[f"model_{workload}_l2_miss_rate"]
 
-        df.at[idx, "pred_ipc"] = preds["ipc"]
-        df.at[idx, "pred_l2_miss_rate"] = preds["l2_miss_rate"]
-        df.at[idx, "warnings"] = "; ".join(warn_msgs)
+        raw_ipc = model_ipc.predict(X)
+        ipc = np.expm1(raw_ipc) if log_ipc else raw_ipc
 
-    return df
+        raw_miss = model_miss.predict(X)
+        miss = np.expm1(raw_miss) if log_miss else raw_miss
 
-# -------------------------------------------------
-# Gradio Wrapper
-# -------------------------------------------------
-def infer_from_csv(file):
-    df = pd.read_csv(file.name)
-    out_df = run_inference(df)
+        clipped = (ipc > 3.5) | (ipc < 0)
 
-    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
-    out_df.to_csv(tmp.name, index=False)
+        ipc = np.clip(ipc, 0, 3.5)
+        miss = np.clip(miss, 0, 1)
 
-    warning_rows = out_df[out_df["warnings"] != ""]
-    warning_text = (
-        f"{len(warning_rows)} rows triggered sanity warnings."
-        if not warning_rows.empty
-        else "No sanity warnings detected."
-    )
+        df.loc[idx, "pred_ipc"] = ipc
+        df.loc[idx, "pred_l2_miss_rate"] = miss
+        df.loc[idx, "was_clipped"] = clipped
 
-    return out_df.head(20), tmp.name, warning_text
+        shap_values = EXPLAINERS[f"model_{workload}_ipc"].shap_values(X)
+        shap_summary_data.append((X, shap_values))
 
-# -------------------------------------------------
-# UI
-# -------------------------------------------------
-with gr.Blocks(title="AIDE Chip Surrogate Inference") as demo:
-    # Loading screen ONLY
-    loading_md = gr.Markdown(
-        "## Downloading surrogate models…\n\nThis may take a while.",
-        visible=True,
+    inference_time = time.perf_counter() - start
+
+    # Ranking
+    df["score"] = df["pred_ipc"] / (1 + df["pred_l2_miss_rate"])
+
+    group_cols = [
+        "l1d_size", "l1i_size", "l2_size",
+        "l1d_assoc", "l1i_assoc", "l2_assoc"
+    ]
+
+    agg = df.groupby(group_cols).agg({
+        "pred_ipc": "mean",
+        "pred_l2_miss_rate": "mean"
+    }).reset_index()
+
+    ipc_std = df.groupby(group_cols)["pred_ipc"].std().reset_index()
+    ipc_std.rename(columns={"pred_ipc": "ipc_std"}, inplace=True)
+
+    agg = agg.merge(ipc_std, on=group_cols)
+
+    agg["score"] = agg["pred_ipc"] / (1 + agg["pred_l2_miss_rate"])
+    agg["robust_score"] = agg["score"] / (1 + agg["ipc_std"])
+
+    agg = agg.sort_values("robust_score", ascending=False)
+
+    topk = agg.head(TOP_K)
+
+    # Save CSVs
+    ranked_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
+    topk_path = tempfile.NamedTemporaryFile(delete=False, suffix=".csv").name
+
+    agg.to_csv(ranked_path, index=False)
+    topk.to_csv(topk_path, index=False)
+
+    # Plot topk robust score
+    fig_rank, ax = plt.subplots(figsize=(10, 5))
+    ax.bar(range(len(topk)), topk["robust_score"])
+    ax.set_title("Top-K Robust Scores")
+    ax.set_xlabel("Rank")
+    ax.set_ylabel("Robust Score")
+    plt.tight_layout()
+
+    # SHAP Global Summary
+    X_all = pd.concat([x for x, _ in shap_summary_data], axis=0)
+    shap_all = np.vstack([s for _, s in shap_summary_data])
+
+    fig_shap = plt.figure(figsize=(10, 6))
+    shap.summary_plot(
+        shap_all,
+        X_all,
+        feature_names=FEATURES,
+        show=False
+    )
+    plt.tight_layout()
+
+    summary = f"""
+### Inference Complete
+
+- **Inference Time:** {inference_time:.3f}s  
+- **Total Pairs:** {len(df):,}  
+- **Unique Configs:** {len(agg):,}  
+- **OOD / Clipped:** {df['was_clipped'].sum():,}
+"""
+
+    return (
+        summary,
+        topk,
+        fig_rank,
+        fig_shap,
+        ranked_path,
+        topk_path
     )
 
-    # Main app (hidden initially)
-    with gr.Column(visible=False) as app_ui:
-        gr.Markdown(
-            """
-            # AIDE Chip Surrogate Inference
-
-            Upload a CSV describing cache configurations and workloads.
-            The app will run surrogate models to predict:
-            - IPC
-            - L2 Miss Rate
-
-            ## Expected CSV Format
-
-            The input CSV **must** contain the following columns:
-
-            **Required columns**
-            - `workload` — one of: `crc32`, `dijkstra`, `fft`, `matrix_mul`, `qsort`, `sha`
-            - `l1d_size` — L1 data cache size (kibibytes, power of two)
-            - `l1i_size` — L1 instruction cache size (kibibytes, power of two)
-            - `l2_size` — L2 cache size (kibibytes, power of two)
-            - `l1d_assoc` — L1D associativity (power of two)
-            - `l1i_assoc` — L1I associativity (power of two)
-            - `l2_assoc` — L2 associativity (power of two)
-
-            **Notes**
-            - All size and associativity values must be positive and powers of two.
-            - One row corresponds to one cache configuration.
-
-            **Example**
-            ```
-            workload,l1d_size,l1i_size,l2_size,l1d_assoc,l1i_assoc,l2_assoc
-            matrix_mul,128,64,1024,16,8,16
-            fft,128,64,2048,16,8,32
-            ```
-            """
+
+# =========================================================
+# UI
+# =========================================================
+custom_css = """
+.gradio-container {
+    max-width: 1400px !important;
+    margin: auto;
+}
+footer {display:none !important;}
+"""
+
+with gr.Blocks(
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="indigo"
+    ),
+    css=custom_css,
+    title="AIDE-Chip Cache Ranker"
+) as demo:
+
+    gr.Markdown("""
+    # 🚀 AIDE-Chip Cache Configuration Ranker
+    ### Surrogate Model Inference + SHAP Explainability
+    Upload workload/config CSV to rank optimal cache configurations.
+    """)
+
+    with gr.Row():
+        file_input = gr.File(
+            label="Upload Input CSV",
+            file_types=[".csv"]
         )
 
-        csv_input = gr.File(label="Input CSV", file_types=[".csv"])
-        run_btn = gr.Button("Run Inference")
+    run_btn = gr.Button("Run Inference", variant="primary", size="lg")
 
-        preview = gr.Dataframe(label="Preview (first 20 rows)")
-        output_csv = gr.File(label="Download Full Output CSV")
-        warnings_box = gr.Textbox(label="Sanity Check Summary")
+    summary_md = gr.Markdown()
 
-        run_btn.click(
-            infer_from_csv,
-            inputs=csv_input,
-            outputs=[preview, output_csv, warnings_box],
-        )
+    with gr.Tabs():
+        with gr.Tab("🏆 Top Configurations"):
+            topk_df = gr.Dataframe()
+
+        with gr.Tab("📈 Ranking Plot"):
+            rank_plot = gr.Plot()
+
+        with gr.Tab("🧠 SHAP Explainability"):
+            shap_plot = gr.Plot()
+
+        with gr.Tab("⬇ Downloads"):
+            ranked_csv = gr.File(label="Download Ranked CSV")
+            topk_csv = gr.File(label="Download Top-K CSV")
 
-    # Startup load hook
-    demo.load(
-        preload_models,
-        inputs=None,
-        outputs=None,
-    ).then(
-        lambda: (
-            gr.update(visible=False),
-            gr.update(visible=True),
-        ),
-        outputs=[loading_md, app_ui],
+    run_btn.click(
+        fn=run_inference,
+        inputs=file_input,
+        outputs=[
+            summary_md,
+            topk_df,
+            rank_plot,
+            shap_plot,
+            ranked_csv,
+            topk_csv
+        ]
     )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()

requirements.txt

@@ -1,4 +1,5 @@
-pandas
+shap
 numpy
-joblib
+pandas
 xgboost
+matplotlib

surrogate_models_v2_json/model_crc32_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_crc32_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}

surrogate_models_v2_json/model_dijkstra_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_dijkstra_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}

surrogate_models_v2_json/model_fft_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_fft_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}

surrogate_models_v2_json/model_matrix_mul_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_matrix_mul_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}

surrogate_models_v2_json/model_qsort_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_qsort_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}

surrogate_models_v2_json/model_sha_ipc_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": false
+}

surrogate_models_v2_json/model_sha_l2_miss_rate_meta.json

@@ -0,0 +1,14 @@
+{
+    "features": [
+        "l1d_size_log2",
+        "l1i_size_log2",
+        "l2_size_log2",
+        "l1d_assoc_log2",
+        "l1i_assoc_log2",
+        "l2_assoc_log2",
+        "l2_l1d_ratio_log2",
+        "l1d_sets_log2",
+        "l2_sets_log2"
+    ],
+    "log_target": true
+}