from __future__ import annotations import csv import json import os from dataclasses import dataclass from datetime import datetime from pathlib import Path import cocotb import numpy as np from cocotb.clock import Clock from model.ref_model import ( match_topk, match_topk_from_scores, ) from tests.top.utils import ( dut_hash_bits, dut_lanes, dut_num_rows, get_param, query_topk_once, reset_dut, write_rows, ) MAX_BENCHMARK_QUERIES = 128 DEFAULT_POSITIVES_PER_CLASS = 8 DEFAULT_QUERIES_PER_CLASS = 2 DEFAULT_ROW_FLIP_BITS = 16 DEFAULT_QUERY_FLIP_BITS = 16 DEFAULT_SEED = 20260522 BENCHMARK_KS = (1, 5) @dataclass(frozen=True) class RetrievalDataset: rows: list[int] row_labels: list[int] queries: list[int] query_labels: list[int] num_classes: int positives_per_class: int queries_per_class: int seed: int @dataclass(frozen=True) class MetricAccumulator: precision_sum: float = 0.0 recall_sum: float = 0.0 f1_sum: float = 0.0 exact_matches: int = 0 label_hits: int = 0 count: int = 0 def add(self, precision: float, recall: float, f1: float, label_hit: bool, exact: bool) -> "MetricAccumulator": return MetricAccumulator( precision_sum=self.precision_sum + precision, recall_sum=self.recall_sum + recall, f1_sum=self.f1_sum + f1, exact_matches=self.exact_matches + int(exact), label_hits=self.label_hits + int(label_hit), count=self.count + 1, ) def as_dict(self) -> dict[str, float]: if self.count == 0: return { "macro_precision": 0.0, "retrieval_recall": 0.0, "macro_f1": 0.0, "exact_match_rate": 0.0, "recall@k": 0.0, } return { "macro_precision": self.precision_sum / self.count, "retrieval_recall": self.recall_sum / self.count, "macro_f1": self.f1_sum / self.count, "exact_match_rate": self.exact_matches / self.count, "recall@k": self.label_hits / self.count, } def _project_root() -> Path: return Path(__file__).resolve().parents[4] def _flip_exact_bits(rng: np.random.Generator, width: int, n_bits: int) -> int: n_bits = max(0, min(int(n_bits), int(width))) if n_bits == 0: return 0 positions = rng.choice(width, size=n_bits, replace=False) mask = 0 for pos in positions: mask |= 1 << int(pos) return mask def words_le_to_int(words: np.ndarray) -> int: value = 0 for idx, word in enumerate(words.tolist()): value |= int(word) << (64 * idx) return value def load_retrieval_dataset_npz(path: str | os.PathLike[str]) -> RetrievalDataset: dataset_path = Path(path) if not dataset_path.is_absolute(): dataset_path = _project_root() / dataset_path if not dataset_path.exists(): raise AssertionError(f"CAM_RETRIEVAL_DATASET not found: {dataset_path}") loaded = np.load(dataset_path) rows = [words_le_to_int(words) for words in loaded["rows_words"]] queries = [words_le_to_int(words) for words in loaded["queries_words"]] row_labels = [int(x) for x in loaded["row_labels"].tolist()] query_labels = [int(x) for x in loaded["query_labels"].tolist()] return RetrievalDataset( rows=rows, row_labels=row_labels, queries=queries, query_labels=query_labels, num_classes=len(set(row_labels)), positives_per_class=0, queries_per_class=0, seed=0, ) def make_clustered_dataset( *, num_rows: int, hash_bits: int, positives_per_class: int = DEFAULT_POSITIVES_PER_CLASS, queries_per_class: int = DEFAULT_QUERIES_PER_CLASS, row_flip_bits: int = DEFAULT_ROW_FLIP_BITS, query_flip_bits: int = DEFAULT_QUERY_FLIP_BITS, seed: int = DEFAULT_SEED, ) -> RetrievalDataset: usable_rows = int(num_rows) if usable_rows < 5: raise AssertionError("Retrieval benchmark requires at least 5 CAM rows") positives_per_class = min(positives_per_class, usable_rows) num_classes = max(1, usable_rows // positives_per_class) usable_rows = num_classes * positives_per_class # Cap total queries to keep simulation runtime bounded max_queries = min(MAX_BENCHMARK_QUERIES, num_classes * queries_per_class) if max_queries < num_classes * queries_per_class: queries_per_class = max(1, max_queries // num_classes) rng = np.random.default_rng(seed) mask = (1 << hash_bits) - 1 words = (hash_bits + 63) // 64 rows: list[int] = [] row_labels: list[int] = [] queries: list[int] = [] query_labels: list[int] = [] for class_id in range(num_classes): center = 0 for word in range(words): center |= int(rng.integers(0, 1 << 64, dtype=np.uint64)) << (64 * word) center &= mask for _ in range(positives_per_class): rows.append((center ^ _flip_exact_bits(rng, hash_bits, row_flip_bits)) & mask) row_labels.append(class_id) for _ in range(queries_per_class): queries.append((center ^ _flip_exact_bits(rng, hash_bits, query_flip_bits)) & mask) query_labels.append(class_id) return RetrievalDataset( rows=rows, row_labels=row_labels, queries=queries, query_labels=query_labels, num_classes=num_classes, positives_per_class=positives_per_class, queries_per_class=queries_per_class, seed=seed, ) def compute_metrics(topk_indices: list[int], row_labels: list[int], query_label: int, k: int) -> tuple[float, float, float]: retrieved = topk_indices[:k] relevant = {idx for idx, label in enumerate(row_labels) if label == query_label} tp = len(set(retrieved) & relevant) precision = tp / float(k) recall = tp / float(len(relevant)) if relevant else 0.0 f1 = 0.0 if precision + recall == 0 else (2.0 * precision * recall) / (precision + recall) return precision, recall, f1 def mode_from_params(write_noise_en: int) -> str: if write_noise_en: return "write_noise" return "no_noise" def output_dir_for(mode: str) -> Path: run_id = os.environ.get("CAM_RETRIEVAL_RUN_ID") if not run_id: run_id = f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{mode}" out_dir = _project_root() / "outputs" / "cam_retrieval_benchmark" / run_id out_dir.mkdir(parents=True, exist_ok=True) (out_dir / "logs").mkdir(exist_ok=True) return out_dir def write_outputs(out_dir: Path, result: dict) -> None: metrics_json = out_dir / "metrics.json" metrics_csv = out_dir / "metrics.csv" summary_md = out_dir / "summary.md" metrics_json.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") fieldnames = [ "run_id", "mode", "num_rows", "hash_bits", "lanes", "topk_k", "write_noise_en", "write_noise_rate_num", "write_noise_rate_den", "num_queries", "k", "macro_precision", "retrieval_recall", "macro_f1", "recall@k", "exact_match_rate", "status", ] with metrics_csv.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() for k, metrics in result["metrics"].items(): row = { "run_id": result["run_id"], "mode": result["mode"], "num_rows": result["params"]["num_rows"], "hash_bits": result["params"]["hash_bits"], "lanes": result["params"]["lanes"], "topk_k": result["params"]["topk_k"], "write_noise_en": result["params"]["write_noise_en"], "write_noise_rate_num": result["params"]["write_noise_rate_num"], "write_noise_rate_den": result["params"]["write_noise_rate_den"], "num_queries": result["dataset"]["num_queries"], "k": int(k), "macro_precision": metrics["macro_precision"], "retrieval_recall": metrics["retrieval_recall"], "macro_f1": metrics["macro_f1"], "recall@k": metrics["recall@k"], "exact_match_rate": metrics["exact_match_rate"], "status": result["status"], } writer.writerow(row) lines = [ "# CAM Retrieval Benchmark Summary", "", f"- run_id: `{result['run_id']}`", f"- mode: `{result['mode']}`", f"- status: `{result['status']}`", f"- num_queries: `{result['dataset']['num_queries']}`", "", "| k | macro_precision | retrieval_recall | macro_f1 | recall@k | exact_match_rate |", "|---:|---:|---:|---:|---:|---:|", ] for k, metrics in result["metrics"].items(): lines.append( f"| {k} | {metrics['macro_precision']:.6f} | {metrics['retrieval_recall']:.6f} | " f"{metrics['macro_f1']:.6f} | {metrics['recall@k']:.6f} | {metrics['exact_match_rate']:.6f} |" ) lines.extend([ "", "说明：结果来自 Verilator/Cocotb 仿真，不是 FPGA 板上实测。", ]) summary_md.write_text("\n".join(lines) + "\n", encoding="utf-8") @cocotb.test() async def cam_retrieval_benchmark(dut): cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start()) await reset_dut(dut) num_rows = dut_num_rows(dut) hash_bits = dut_hash_bits(dut) lanes = dut_lanes(dut) write_noise_en = int(get_param(dut, "WRITE_NOISE_EN", 0) or 0) write_noise_rate_num = int(get_param(dut, "WRITE_NOISE_RATE_NUM", 0) or 0) write_noise_rate_den = int(get_param(dut, "WRITE_NOISE_RATE_DEN", 100) or 100) mode = mode_from_params(write_noise_en) if num_rows % lanes != 0: raise AssertionError("Retrieval benchmark requires NUM_ROWS divisible by LANES") dataset_path = os.environ.get("CAM_RETRIEVAL_DATASET") if not dataset_path: raise AssertionError("CAM_RETRIEVAL_DATASET is required; run scripts/prepare_cam_retrieval_dataset.py first") dataset = load_retrieval_dataset_npz(dataset_path) if len(dataset.rows) != num_rows: raise AssertionError(f"artifact row count {len(dataset.rows)} must equal DUT NUM_ROWS {num_rows}") await write_rows(dut, dataset.rows) accumulators = {k: MetricAccumulator() for k in BENCHMARK_KS} for query, query_label in zip(dataset.queries, dataset.query_labels): beats, _, _, _ = await query_topk_once(dut, query) if len(beats) < max(BENCHMARK_KS): raise AssertionError(f"Expected at least {max(BENCHMARK_KS)} Top-K beats, got {len(beats)}") dut_topk = [int(beat[1]) for beat in beats[: max(BENCHMARK_KS)]] golden_topk, _ = match_topk(query, dataset.rows, width=hash_bits, k=max(BENCHMARK_KS)) for k in BENCHMARK_KS: precision, recall, f1 = compute_metrics(dut_topk, dataset.row_labels, query_label, k) exact = dut_topk[:k] == golden_topk[:k] retrieved_labels = [dataset.row_labels[idx] for idx in dut_topk[:k]] label_hit = query_label in retrieved_labels accumulators[k] = accumulators[k].add(precision, recall, f1, label_hit, exact) run_id = os.environ.get("CAM_RETRIEVAL_RUN_ID") or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{mode}" result = { "run_id": run_id, "mode": mode, "status": "pass", "params": { "num_rows": len(dataset.rows), "hash_bits": hash_bits, "lanes": lanes, "topk_k": max(BENCHMARK_KS), "write_noise_en": write_noise_en, "write_noise_rate_num": write_noise_rate_num, "write_noise_rate_den": write_noise_rate_den, }, "dataset": { "num_classes": dataset.num_classes, "positives_per_class": dataset.positives_per_class, "queries_per_class": dataset.queries_per_class, "num_queries": len(dataset.queries), "seed": dataset.seed, }, "metrics": {str(k): accumulators[k].as_dict() for k in BENCHMARK_KS}, } out_dir = output_dir_for(mode) write_outputs(out_dir, result) for k in BENCHMARK_KS: metrics = result["metrics"][str(k)] dut._log.info( "RETRIEVAL_RESULT mode=%s k=%d precision=%.6f retrieval_recall=%.6f f1=%.6f recall_at_k=%.6f exact_match=%.6f output_dir=%s", mode, k, metrics["macro_precision"], metrics["retrieval_recall"], metrics["macro_f1"], metrics["recall@k"], metrics["exact_match_rate"], str(out_dir.relative_to(_project_root())), ) if write_noise_en == 0: assert result["metrics"]["5"]["exact_match_rate"] == 1.0, ( f"Expected perfect exact match with no noise, got " f"{result['metrics']['5']['exact_match_rate']}" ) else: dut._log.info( "Noise enabled (WRITE_NOISE_RATE=%d/%d) — exact_match assertion skipped", write_noise_rate_num, write_noise_rate_den, )