Major Refactor of Functional Warmup

pyproject.toml

@@ -12,27 +12,27 @@ packages = [
 gen-ckpt = "tidalsim.scripts.gen_ckpt:main"
 gen-cache-state = "tidalsim.scripts.gen_cache_state:main"
 tidalsim = "tidalsim.scripts.tidalsim:main"
-analyze = "tidalsim.scripts.analyze:main"
+goldensim = "tidalsim.scripts.goldensim:main"
 bench-spike-bb-extraction = "tidalsim.scripts.bench_spike_bb_extraction:main"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
 tqdm = "^4.66.1"
-numpy = "^1.26.1"
-scikit-learn = "^1.3.2"
-notebook = "^7.0.6"
-matplotlib = "^3.8.1"
-joblib = "^1.3.2"
-more-itertools = "^10.2.0"
-pandas = "^2.2.0"
-pandera = {version="^0.18.0"}
-pyarrow = "^15.0.0"
+numpy = "^1.26"
+scikit-learn = "^1.4"
+notebook = "^7.1"
+matplotlib = "^3.8"
+joblib = "^1.3"
+more-itertools = "^10.2"
+pandas = "^2.2"
+pandera = {version="^0.18"}
+pyarrow = "^15.0"
 
 [tool.poetry.dev-dependencies]
-pytest = "^8.0.0"
-pyright = "^1.1.350"
-setuptools = "^69.1.0"
-black = "^24.2.0"
+pytest = "^8.1"
+pyright = "^1.1"
+setuptools = "^69.2"
+black = "^24.3"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_extrapolation.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+from pandera.typing import DataFrame
+
+from tidalsim.modeling.extrapolation import (
+    fill_perf_metrics,
+    parse_perf_file,
+    PerfMetrics,
+    pick_intervals_for_rtl_sim,
+)
+from tidalsim.modeling.schemas import ClusteringSchema
+
+
+class TestExtrapolation:
+    def test_parse_perf_file(self, tmp_path: Path) -> None:
+        perf_file_csv = """cycles,instret
+180,100
+140,100
+130,100
+135,100"""
+        perf_file = tmp_path / "perf.csv"
+        with perf_file.open("w") as f:
+            f.write(perf_file_csv)
+        cycles = 180 + 140 + 130 + 135
+        perf = parse_perf_file(perf_file, detailed_warmup_insts=0)
+        expected_perf = PerfMetrics(ipc=400 / (180 + 140 + 130 + 135))
+        assert perf == expected_perf
+
+        perf = parse_perf_file(perf_file, detailed_warmup_insts=100)
+        expected_perf = PerfMetrics(ipc=300 / (140 + 130 + 135))
+        assert perf == expected_perf
+
+        perf = parse_perf_file(
+            perf_file, detailed_warmup_insts=150
+        )  # should round up to the next interval after 150 insts
+        expected_perf = PerfMetrics(ipc=200 / (130 + 135))
+        assert perf == expected_perf
+
+    def test_fill_perf_metrics(self, tmp_path: Path) -> None:
+        clustering_df = DataFrame[ClusteringSchema]({
+            "instret": [100, 100, 100, 100],
+            "inst_count": [100, 200, 300, 400],
+            "inst_start": [0, 100, 200, 300],
+            "embedding": [[0, 1], [0, 1], [1, 2], [3, 4]],
+            "cluster_id": [2, 2, 1, 0],
+            "dist_to_centroid": [0.0, 0.0, 0.0, 0.0],
+            "chosen_for_rtl_sim": [False, True, True, True],
+        })
+
+        # Place some dummy perf data on disk
+        perf_file_cold = """cycles,instret
+70,50
+100,50
+        """
+        perf_file_warm = """cycles,instret
+60,50
+80,50
+        """
+        cluster_dir = tmp_path / "checkpoints" / "0x80000000.100"
+        cluster_dir.mkdir(parents=True)
+        (cluster_dir / "perf_cold.csv").write_text(perf_file_cold)
+        (cluster_dir / "perf_warmup.csv").write_text(perf_file_warm)
+
+        perf_df = fill_perf_metrics(clustering_df, tmp_path, detailed_warmup_insts=0)
+        row_with_data = perf_df.iloc[1]
+        assert row_with_data.est_ipc_cold == (100 / 170)
+        assert row_with_data.est_ipc_warm == (100 / 140)
+
+    def test_pick_intervals_for_rtl_sim(self) -> None:
+        clustering_df = DataFrame[ClusteringSchema]({
+            "instret": [100, 100, 100, 100, 100],
+            "inst_count": [100, 200, 300, 400, 500],
+            "inst_start": [0, 100, 200, 300, 400],
+            "embedding": [[0, 1], [0, 1], [], [1, 2], [3, 4]],
+            "cluster_id": [2, 2, 2, 1, 0],
+            "dist_to_centroid": [1.0, 1.0, 0.0, 0.0, 0.0],
+            "chosen_for_rtl_sim": [False, False, False, False, False],
+        })
+        pick_intervals_for_rtl_sim(clustering_df, 1)
+        assert clustering_df["chosen_for_rtl_sim"].to_list() == [True, False, True, True, True]

tests/test_util_cli.py

@@ -0,0 +1,21 @@
+import pytest
+from pathlib import Path
+
+from tidalsim.util.cli import run_rtl_sim_cmd
+
+
+class TestUtilCli:
+    def test_run_rtl_sim_cmd(self, tmp_path: Path) -> None:
+        cmd = run_rtl_sim_cmd(
+            simulator=(tmp_path / "simulator"),
+            perf_file=(tmp_path / "perf.csv"),
+            perf_sample_period=1000,
+            max_instructions=None,
+            chipyard_root=tmp_path,
+            binary=(tmp_path / "binary"),
+            loadarch=(tmp_path / "loadarch"),
+            suppress_exit=True,
+            checkpoint_dir=tmp_path,
+            timeout_cycles=5000,
+        )
+        print(cmd)

tidalsim/modeling/extrapolation.py

@@ -1,13 +1,80 @@
 from pathlib import Path
-from typing import Tuple, List, cast, Tuple, Optional
-import logging
+from typing import Tuple, Optional, cast
+from dataclasses import dataclass
 
 import numpy as np
 import pandas as pd
 from pandera.typing import DataFrame
 
 from tidalsim.util.pickle import load
-from tidalsim.modeling.schemas import *
+from tidalsim.modeling.schemas import ClusteringSchema, EstimatedPerfSchema, GoldenPerfSchema
+
+
+@dataclass
+class PerfMetrics:
+    ipc: float
+
+
+# Build the performance metrics struct from a csv file dumped from RTL simulation
+# Assume perf metrics are captured without any offset from the start of the interval
+# Detailed warmup is accounted for in this function
+def parse_perf_file(
+    perf_file: Path,
+    detailed_warmup_insts: int,
+) -> PerfMetrics:
+    perf_data = pd.read_csv(perf_file)
+    perf_data["insts_retired_after_interval"] = np.cumsum(perf_data["instret"])
+    perf_data["insts_retired_before_interval"] = (
+        perf_data["insts_retired_after_interval"] - perf_data["instret"]
+    )
+    # Find the first row where more than [detailed_warmup_insts] have elapsed, and only begin tracking IPC from that row onwards
+    start_point = (perf_data["insts_retired_before_interval"] >= detailed_warmup_insts).idxmax()
+    perf_data_visible = perf_data[start_point:]
+    ipc = np.sum(perf_data_visible["instret"]) / np.sum(perf_data_visible["cycles"])
+    return PerfMetrics(ipc=ipc)
+
+
+# Fill each row in the [clustering_df] that was chosen for RTL simulation with performance numbers from RTL sim
+# BUT, DO NOT perform any extrapolation yet (leave the estimated perf blank for any row without RTL sim)
+def fill_perf_metrics(
+    clustering_df: DataFrame[ClusteringSchema],
+    cluster_dir: Path,
+    detailed_warmup_insts: int,
+) -> DataFrame[EstimatedPerfSchema]:
+    # Augment the dataframe with zeroed out estimated perf columns
+    perf_df = cast(
+        DataFrame[EstimatedPerfSchema],
+        clustering_df.assign(
+            est_ipc_cold=np.zeros(len(clustering_df.index)),
+            est_ipc_warm=np.zeros(len(clustering_df.index)),
+        ),
+    )
+    simulated_rows = perf_df.loc[perf_df["chosen_for_rtl_sim"] == True]
+    for index, row in simulated_rows.iterrows():
+        checkpoint_dir = cluster_dir / "checkpoints" / f"0x80000000.{row.inst_start}"
+        cold_perf_file = checkpoint_dir / "perf_cold.csv"
+        warm_perf_file = checkpoint_dir / "perf_warmup.csv"
+        if cold_perf_file.exists():
+            cold_perf = parse_perf_file(cold_perf_file, detailed_warmup_insts)
+            row.est_ipc_cold = cold_perf.ipc
+        if warm_perf_file.exists():
+            warm_perf = parse_perf_file(warm_perf_file, detailed_warmup_insts)
+            row.est_ipc_warm = warm_perf.ipc
+        perf_df.iloc[index] = row
+    return perf_df
+
+
+# Pick which intervals we are going to simulate in RTL.
+# Randomly sample from each cluster.
+# AND pick the closest point to each centroid too.
+def pick_intervals_for_rtl_sim(
+    clustering_df: DataFrame[ClusteringSchema], points_per_cluster: int
+) -> None:
+    grouped_intervals = clustering_df.groupby("cluster_id")
+    chosen_intervals = grouped_intervals.sample(points_per_cluster, random_state=1)
+    clustering_df.loc[chosen_intervals.index, "chosen_for_rtl_sim"] = True
+    c = clustering_df.groupby("cluster_id")["dist_to_centroid"].idxmin()
+    clustering_df.loc[c, "chosen_for_rtl_sim"] = True
 
 
 def analyze_tidalsim_results(
@@ -20,35 +87,21 @@ def analyze_tidalsim_results(
 ) -> Tuple[DataFrame[EstimatedPerfSchema], Optional[DataFrame[GoldenPerfSchema]]]:
     interval_dir = run_dir / f"n_{interval_length}_{'elf' if elf else 'spike'}"
     cluster_dir = interval_dir / f"c_{clusters}"
-
     clustering_df = load(cluster_dir / "clustering_df.pickle")
-    simulated_points = (
-        clustering_df.loc[clustering_df["chosen_for_rtl_sim"] == True]
-        .groupby("cluster_id", as_index=False)
-        .nth(0)
-        .sort_values("cluster_id")
-    )
-    ipcs = []
-    for index, row in simulated_points.iterrows():
-        perf_file = cluster_dir / "checkpoints" / f"0x80000000.{row['inst_start']}" / "perf.csv"
-        perf_data = pd.read_csv(perf_file)
-        perf_data["ipc"] = perf_data["instret"] / perf_data["cycles"]
-        perf_data["inst_count"] = np.cumsum(perf_data["instret"])
-        # Find the first row where more than [detailed_warmup_insts] have elapsed, and only begin tracking IPC from that row onwards
-        # mypy can't infer the type of [start_point] correctly
-        start_point = (perf_data["inst_count"] > detailed_warmup_insts).idxmax()
-        # mypy can't say that perf_data[start_point:] is a legal slice
-        ipc: float = np.nanmean(perf_data[start_point:]["ipc"])  # type: ignore
-        ipcs.append(ipc)
-
-    estimated_perf_df: DataFrame[EstimatedPerfSchema]
+    estimated_perf_df = fill_perf_metrics(clustering_df, cluster_dir, detailed_warmup_insts)
+
     if not interpolate_clusters:
-        # If we don't interpolate, we just use the IPC of the simulated point for that cluster
+        # If we don't interpolate, we just use the average IPC of the simulated points for that cluster
+        ipcs = (
+            estimated_perf_df.sort_values("cluster_id").groupby("cluster_id")["est_ipc_warm"].mean()
+        )
         estimated_perf_df = clustering_df.assign(
-            est_ipc=np.array(ipcs)[clustering_df["cluster_id"]],
+            est_ipc=ipcs[clustering_df["cluster_id"]],
             est_cycles=lambda x: np.round(x["instret"] * np.reciprocal(x["est_ipc"])),
         )
     else:
+        assert False
+        # TODO: fix this block
         # If we do interpolate, we use a weighted (by inverse L2 norm) average of the IPCs of all simulated points
         kmeans_file = cluster_dir / "kmeans_model.pickle"
         kmeans = load(kmeans_file)
@@ -63,7 +116,7 @@ def analyze_tidalsim_results(
         norms = 1 / norms
         weight_vecs = norms / norms.sum(axis=1, keepdims=True)
         # multiply weight vecs by ips to get weighted average
-        est_ipc = weight_vecs @ np.array(ipcs)
+        est_ipc = weight_vecs @ ipcs
         # assign to df
         estimated_perf_df = clustering_df.assign(
             est_ipc=est_ipc,