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
This commit is contained in:
2026-05-22 21:06:51 +08:00
parent e1bed00cc4
commit 1ff9a5f18b
5 changed files with 373 additions and 15 deletions

View File

@@ -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

View File

@@ -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",
"exact_match_rate", "status",
"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)
@@ -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"{metrics['macro_f1']:.6f} | {metrics['exact_match_rate']:.6f} |"
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([
"",
@@ -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",
mode, k, metrics["macro_precision"], metrics["macro_recall"],
metrics["macro_f1"], metrics["exact_match_rate"],
"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())),
)