feat(retrieval-benchmark): add support for external pre-prepared CAM retrieval datasets with recall@k metric

- Add just recipes for preparing CIFAR10/100 hash artifacts and running benchmarks
- Add CAM_RETRIEVAL_DATASET env var support in Makefile
- Add load_retrieval_dataset_npz() to load pre-prepared retrieval datasets
- Add label_hits counter and recall@k metric for retrieval evaluation
- Rename macro_recall to retrieval_recall to clarify semantics

.justfile

@@ -81,3 +81,19 @@ cam-test-retrieval-no-noise:
 # Run CAM retrieval benchmark with read noise enabled
 cam-test-retrieval-read-noise:
     just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 WRITE_NOISE_EN=0 READ_NOISE_EN=1 READ_NOISE_RATE_NUM=1 READ_NOISE_RATE_DEN=100 READ_NOISE_BITS=8"
+
+# Prepare CIFAR10 hash artifact for CAM retrieval smoke benchmark
+cam-prepare-retrieval-cifar10 ROWS="512" QUERIES="128":
+    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt"
+
+# Prepare CIFAR100 hash artifact for CAM retrieval benchmark
+cam-prepare-retrieval-cifar100 ROWS="512" QUERIES="128":
+    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt"
+
+# Run CAM retrieval benchmark on a prepared artifact without hardware noise
+cam-test-retrieval-artifact DATASET_PATH NUM_ROWS="4096":
+    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 READ_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"
+
+# Run CAM retrieval benchmark on a prepared artifact with read noise enabled
+cam-test-retrieval-artifact-read-noise DATASET_PATH NUM_ROWS="4096":
+    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 READ_NOISE_EN=1 READ_NOISE_RATE_NUM=1 READ_NOISE_RATE_DEN=100 READ_NOISE_BITS=8 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"

hw/sim/benchmarks/retrieval/Makefile

@@ -7,7 +7,11 @@ COCOTB_TEST_MODULES := benchmarks.retrieval.test_retrieval_benchmark
 VERILOG_SOURCES := $(RTL_CAM_TOP)
 
 TOPK_K ?= 5
+NUM_ROWS ?= 4096
 WRITE_NOISE_EN ?= 0
 READ_NOISE_EN ?= 0
 
+CAM_RETRIEVAL_DATASET ?=
+export CAM_RETRIEVAL_DATASET
+
 include $(SIM_ROOT)/mk/cocotb-common.mk

hw/sim/benchmarks/retrieval/test_retrieval_benchmark.py

@@ -54,14 +54,16 @@ class MetricAccumulator:
     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, exact: bool) -> "MetricAccumulator":
+    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,
         )
 
@@ -69,15 +71,17 @@ class MetricAccumulator:
         if self.count == 0:
             return {
                 "macro_precision": 0.0,
-                "macro_recall": 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,
-            "macro_recall": self.recall_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,
         }
 
 
@@ -96,6 +100,37 @@ def _flip_exact_bits(rng: np.random.Generator, width: int, n_bits: int) -> int:
     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,
@@ -195,8 +230,8 @@ def write_outputs(out_dir: Path, result: dict) -> None:
         "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", "macro_recall", "macro_f1",
+        "num_queries", "k", "macro_precision", "retrieval_recall", "macro_f1",
-        "exact_match_rate", "status",
+        "recall@k", "exact_match_rate", "status",
     ]
     with metrics_csv.open("w", newline="", encoding="utf-8") as f:
         writer = csv.DictWriter(f, fieldnames=fieldnames)
@@ -218,8 +253,9 @@ def write_outputs(out_dir: Path, result: dict) -> None:
                 "num_queries": result["dataset"]["num_queries"],
                 "k": int(k),
                 "macro_precision": metrics["macro_precision"],
-                "macro_recall": metrics["macro_recall"],
+                "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"],
             }
@@ -233,13 +269,13 @@ def write_outputs(out_dir: Path, result: dict) -> None:
         f"- status: `{result['status']}`",
         f"- num_queries: `{result['dataset']['num_queries']}`",
         "",
-        "| k | macro_precision | macro_recall | macro_f1 | exact_match_rate |",
+        "| 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['macro_recall']:.6f} | "
+            f"| {k} | {metrics['macro_precision']:.6f} | {metrics['retrieval_recall']:.6f} | "
-            f"{metrics['macro_f1']:.6f} | {metrics['exact_match_rate']:.6f} |"
+            f"{metrics['macro_f1']:.6f} | {metrics['recall@k']:.6f} | {metrics['exact_match_rate']:.6f} |"
         )
     lines.extend([
         "",
@@ -270,7 +306,12 @@ async def cam_retrieval_benchmark(dut):
     if num_rows % lanes != 0:
         raise AssertionError("Retrieval benchmark requires NUM_ROWS divisible by LANES")
 
-    dataset = make_clustered_dataset(num_rows=num_rows, hash_bits=hash_bits)
+    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}
@@ -296,7 +337,9 @@ async def cam_retrieval_benchmark(dut):
         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]
-            accumulators[k] = accumulators[k].add(precision, recall, f1, exact)
+            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 = {
@@ -331,9 +374,9 @@ async def cam_retrieval_benchmark(dut):
     for k in BENCHMARK_KS:
         metrics = result["metrics"][str(k)]
         dut._log.info(
-            "RETRIEVAL_RESULT mode=%s k=%d precision=%.6f recall=%.6f f1=%.6f exact_match=%.6f output_dir=%s",
+            "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["macro_recall"],
+            mode, k, metrics["macro_precision"], metrics["retrieval_recall"],
-            metrics["macro_f1"], metrics["exact_match_rate"],
+            metrics["macro_f1"], metrics["recall@k"], metrics["exact_match_rate"],
             str(out_dir.relative_to(_project_root())),
         )
 

scripts/prepare_cam_retrieval_dataset.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import json
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import numpy as np
+import torch
+import typer
+from datasets import load_dataset
+from rich.progress import track
+from torch.utils.data import DataLoader
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+MINI_NAV_ROOT = PROJECT_ROOT / "mini-nav"
+if str(MINI_NAV_ROOT) not in sys.path:
+    sys.path.insert(0, str(MINI_NAV_ROOT))
+
+from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor  # noqa: E402
+from utils import get_device  # noqa: E402
+
+DEFAULT_COMPRESSOR_PATH = Path("outputs/hash_compressor.pt")
+DEFAULT_OUTPUT_ROOT = Path("outputs/cam_retrieval_benchmark/datasets")
+
+
+@dataclass(frozen=True)
+class PreparedArtifact:
+    npz_path: Path
+    metadata_path: Path
+    num_rows: int
+    num_queries: int
+
+
+def dataset_config(dataset: str) -> tuple[str, str]:
+    if dataset == "cifar10":
+        return "uoft-cs/cifar10", "label"
+    if dataset == "cifar100":
+        return "uoft-cs/cifar100", "fine_label"
+    raise ValueError("dataset must be cifar10 or cifar100")
+
+
+def stratified_indices(labels: list[int], *, total: int, seed: int) -> list[int]:
+    if total <= 0:
+        raise ValueError("total must be greater than 0")
+    if total > len(labels):
+        raise ValueError("total cannot exceed number of labels")
+
+    rng = np.random.default_rng(seed)
+    by_label: dict[int, list[int]] = {}
+    for idx, label in enumerate(labels):
+        by_label.setdefault(int(label), []).append(idx)
+
+    label_ids = sorted(by_label)
+    base = total // len(label_ids)
+    remainder = total % len(label_ids)
+    selected: list[int] = []
+
+    for offset, label in enumerate(label_ids):
+        want = base + (1 if offset < remainder else 0)
+        candidates = list(by_label[label])
+        if want > len(candidates):
+            raise ValueError(f"not enough examples for label {label}: need {want}, have {len(candidates)}")
+        chosen = rng.choice(candidates, size=want, replace=False).tolist()
+        selected.extend(int(i) for i in chosen)
+
+    selected.sort()
+    return selected
+
+
+def pack_bits_to_words_le(bits: np.ndarray, *, hash_bits: int) -> np.ndarray:
+    if bits.ndim != 2:
+        raise ValueError("bits must have shape [N, hash_bits]")
+    if bits.shape[1] != hash_bits:
+        raise ValueError(f"expected {hash_bits} bits, got {bits.shape[1]}")
+    if hash_bits % 64 != 0:
+        raise ValueError("hash_bits must be divisible by 64")
+
+    words = np.zeros((bits.shape[0], hash_bits // 64), dtype=np.uint64)
+    bits_u8 = bits.astype(np.uint8, copy=False)
+    for word_idx in range(hash_bits // 64):
+        word = np.zeros(bits.shape[0], dtype=np.uint64)
+        for bit in range(64):
+            word |= bits_u8[:, word_idx * 64 + bit].astype(np.uint64) << np.uint64(bit)
+        words[:, word_idx] = word
+    return words
+
+
+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
+
+
+@torch.no_grad()
+def encode_dataset_bits(
+    dataset,
+    *,
+    img_column: str,
+    batch_size: int,
+    dino_model: str,
+    compressor_path: Path,
+    hash_bits: int,
+) -> np.ndarray:
+    if not compressor_path.exists():
+        raise FileNotFoundError(f"missing HashCompressor weights: {compressor_path}")
+
+    device = get_device()
+    processor, dino = load_dino_model(dino_model)
+    compressor = load_hash_compressor(
+        input_dim=get_dino_dim(dino_model),
+        hash_bits=hash_bits,
+        compressor_path=str(compressor_path),
+    ).to(device)
+    compressor.eval()
+
+    loader = DataLoader(dataset.with_format("torch"), batch_size=batch_size, shuffle=False, num_workers=4)
+    all_bits: list[np.ndarray] = []
+    for batch in track(loader, description="Encoding CIFAR hash bits"):
+        images = batch[img_column]
+        inputs = processor(images=images, return_tensors="pt").to(device)
+        tokens = dino(**inputs).last_hidden_state
+        bits = compressor.encode(tokens).detach().cpu().numpy().astype(np.uint8)
+        all_bits.append(bits)
+    return np.concatenate(all_bits, axis=0)
+
+
+def prepare_artifact(
+    *,
+    dataset_name: Literal["cifar10", "cifar100"],
+    num_rows: int,
+    max_queries: int,
+    compressor_path: Path,
+    output_root: Path,
+    dino_model: str = "facebook/dinov2-large",
+    hash_bits: int = 512,
+    batch_size: int = 64,
+    seed: int = 20260522,
+) -> PreparedArtifact:
+    hf_id, label_column = dataset_config(dataset_name)
+    loaded = load_dataset(hf_id)
+    train = loaded["train"]
+    test = loaded["test"] if "test" in loaded else loaded["validation"]
+
+    row_indices = stratified_indices([int(x) for x in train[label_column]], total=num_rows, seed=seed)
+    query_indices = stratified_indices([int(x) for x in test[label_column]], total=max_queries, seed=seed + 1)
+
+    row_subset = train.select(row_indices)
+    query_subset = test.select(query_indices)
+
+    row_bits = encode_dataset_bits(
+        row_subset,
+        img_column="img",
+        batch_size=batch_size,
+        dino_model=dino_model,
+        compressor_path=compressor_path,
+        hash_bits=hash_bits,
+    )
+    query_bits = encode_dataset_bits(
+        query_subset,
+        img_column="img",
+        batch_size=batch_size,
+        dino_model=dino_model,
+        compressor_path=compressor_path,
+        hash_bits=hash_bits,
+    )
+
+    rows_words = pack_bits_to_words_le(row_bits, hash_bits=hash_bits)
+    queries_words = pack_bits_to_words_le(query_bits, hash_bits=hash_bits)
+    row_labels = np.array([int(x) for x in row_subset[label_column]], dtype=np.int64)
+    query_labels = np.array([int(x) for x in query_subset[label_column]], dtype=np.int64)
+
+    output_root.mkdir(parents=True, exist_ok=True)
+    stem = f"{dataset_name}_hash{hash_bits}_rows{num_rows}_queries{max_queries}"
+    npz_path = output_root / f"{stem}.npz"
+    metadata_path = output_root / f"{stem}.json"
+
+    np.savez_compressed(
+        npz_path,
+        rows_words=rows_words,
+        row_labels=row_labels,
+        queries_words=queries_words,
+        query_labels=query_labels,
+    )
+    metadata = {
+        "dataset": dataset_name,
+        "hf_id": hf_id,
+        "label_column": label_column,
+        "compressor_path": str(compressor_path),
+        "dino_model": dino_model,
+        "hash_bits": hash_bits,
+        "bit_source": "HashCompressor.encode(tokens)",
+        "bit_values": "{0,1} int32",
+        "num_rows": int(num_rows),
+        "num_queries": int(max_queries),
+        "rows_source_split": "train",
+        "queries_source_split": "test",
+        "seed": int(seed),
+    }
+    metadata_path.write_text(json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    return PreparedArtifact(npz_path=npz_path, metadata_path=metadata_path, num_rows=num_rows, num_queries=max_queries)
+
+
+def main(
+    dataset: Literal["cifar10", "cifar100"] = typer.Option("cifar100"),
+    num_rows: int = typer.Option(512, min=5),
+    max_queries: int = typer.Option(128, min=1),
+    compressor_path: Path = typer.Option(DEFAULT_COMPRESSOR_PATH),
+    output_root: Path = typer.Option(DEFAULT_OUTPUT_ROOT),
+    dino_model: str = typer.Option("facebook/dinov2-large"),
+    hash_bits: int = typer.Option(512),
+    batch_size: int = typer.Option(64, min=1),
+    seed: int = typer.Option(20260522),
+) -> None:
+    artifact = prepare_artifact(
+        dataset_name=dataset,
+        num_rows=num_rows,
+        max_queries=max_queries,
+        compressor_path=compressor_path,
+        output_root=output_root,
+        dino_model=dino_model,
+        hash_bits=hash_bits,
+        batch_size=batch_size,
+        seed=seed,
+    )
+    print(f"CAM_RETRIEVAL_DATASET={artifact.npz_path}")
+    print(f"CAM_RETRIEVAL_METADATA={artifact.metadata_path}")
+
+
+if __name__ == "__main__":
+    typer.run(main)

tests/test_prepare_cam_retrieval_dataset.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import numpy as np
+
+from scripts.prepare_cam_retrieval_dataset import (
+    dataset_config,
+    pack_bits_to_words_le,
+    stratified_indices,
+    words_le_to_int,
+)
+
+
+def test_dataset_config_resolves_cifar10_and_cifar100() -> None:
+    assert dataset_config("cifar10") == ("uoft-cs/cifar10", "label")
+    assert dataset_config("cifar100") == ("uoft-cs/cifar100", "fine_label")
+
+
+def test_dataset_config_rejects_unknown_dataset() -> None:
+    try:
+        dataset_config("mnist")
+    except ValueError as exc:
+        assert "dataset must be cifar10 or cifar100" in str(exc)
+    else:
+        raise AssertionError("expected ValueError")
+
+
+def test_stratified_indices_balances_labels_and_is_deterministic() -> None:
+    labels = [0, 0, 0, 1, 1, 1, 2, 2, 2]
+
+    first = stratified_indices(labels, total=6, seed=123)
+    second = stratified_indices(labels, total=6, seed=123)
+
+    assert first == second
+    selected_labels = [labels[i] for i in first]
+    assert selected_labels.count(0) == 2
+    assert selected_labels.count(1) == 2
+    assert selected_labels.count(2) == 2
+
+
+def test_stratified_indices_fills_remainder() -> None:
+    labels = [0, 0, 0, 1, 1, 1]
+
+    indices = stratified_indices(labels, total=5, seed=7)
+
+    assert len(indices) == 5
+    assert len(set(indices)) == 5
+
+
+def test_pack_bits_to_words_le_roundtrip() -> None:
+    bits = np.zeros((2, 128), dtype=np.uint8)
+    bits[0, 0] = 1
+    bits[0, 65] = 1
+    bits[1, 63] = 1
+    bits[1, 127] = 1
+
+    words = pack_bits_to_words_le(bits, hash_bits=128)
+
+    assert words.dtype == np.uint64
+    assert words.shape == (2, 2)
+    assert words_le_to_int(words[0]) == (1 << 0) | (1 << 65)
+    assert words_le_to_int(words[1]) == (1 << 63) | (1 << 127)