Mini-Nav/hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py

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 (
    lane_seed_128,
    match_topk,
    match_topk_from_scores,
    score_rows_with_read_noise,
)
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, read_noise_en: int) -> str:
    if write_noise_en and read_noise_en:
        return "write_read_noise"
    if write_noise_en:
        return "write_noise"
    if read_noise_en:
        return "read_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", "read_noise_en", "write_noise_rate_num",
        "write_noise_rate_den", "read_noise_rate_num", "read_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"],
                "read_noise_en": result["params"]["read_noise_en"],
                "write_noise_rate_num": result["params"]["write_noise_rate_num"],
                "write_noise_rate_den": result["params"]["write_noise_rate_den"],
                "read_noise_rate_num": result["params"]["read_noise_rate_num"],
                "read_noise_rate_den": result["params"]["read_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)
    read_noise_en = int(get_param(dut, "READ_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)
    read_noise_rate_num = int(get_param(dut, "READ_NOISE_RATE_NUM", 0) or 0)
    read_noise_rate_den = int(get_param(dut, "READ_NOISE_RATE_DEN", 100) or 100)
    read_noise_bits = int(get_param(dut, "READ_NOISE_BITS", 8) or 8)
    mode = mode_from_params(write_noise_en, read_noise_en)

    if write_noise_en:
        raise AssertionError("First retrieval benchmark version only supports WRITE_NOISE_EN=0")
    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}
    read_lane_states = [lane_seed_128(0x6A09_E667_F3BC_C909, lane) for lane in range(lanes)]

    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)]]

        if read_noise_en:
            scores, read_lane_states = score_rows_with_read_noise(
                query, dataset.rows, lane_states=read_lane_states,
                width=hash_bits, lanes=lanes, noise_bits=read_noise_bits,
                rate_num=read_noise_rate_num, rate_den=read_noise_rate_den,
            )
            golden_topk = match_topk_from_scores(scores, max(BENCHMARK_KS))
        else:
            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,
            "read_noise_en": read_noise_en,
            "write_noise_rate_num": write_noise_rate_num,
            "write_noise_rate_den": write_noise_rate_den,
            "read_noise_rate_num": read_noise_rate_num,
            "read_noise_rate_den": read_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())),
        )

    assert result["metrics"]["5"]["exact_match_rate"] == 1.0