feat: vectorize CAM retrieval with NumPy and add multi-worker support

- Replace scalar hamming distance with NumPy bitwise_count for batch retrieval
- Add ThreadPoolExecutor-based multi-worker query parallelism
- Improve missing dataset error message with generation command hint
- Increase DEFAULT_MAX_QUERIES from 128 to 8192 for meaningful throughput tests

scripts/prepare_cam_retrieval_dataset.py

@@ -24,6 +24,7 @@ from utils import get_device  # noqa: E402
 
 DEFAULT_COMPRESSOR_PATH = Path("outputs/hash_compressor.pt")
 DEFAULT_OUTPUT_ROOT = Path("outputs/cam_retrieval_benchmark/datasets")
+DEFAULT_MAX_QUERIES = 8192
 
 
 @dataclass(frozen=True)
@@ -207,7 +208,7 @@ def prepare_artifact(
 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),
+    max_queries: int = typer.Option(DEFAULT_MAX_QUERIES, 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"),

scripts/sw_retrieval_benchmark.py

@@ -3,12 +3,14 @@ from __future__ import annotations
 import argparse
 import csv
 import json
+import re
 import sys
 import time
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import Callable, Iterable, Sequence
+from typing import Callable, Iterable
 
 import numpy as np
 
@@ -28,6 +30,8 @@ class RetrievalDataset:
     row_labels: list[int]
     queries: list[int]
     query_labels: list[int]
+    rows_words: np.ndarray
+    queries_words: np.ndarray
     hash_bits: int
     num_classes: int
     positives_per_class: int = 0
@@ -76,6 +80,31 @@ def project_root() -> Path:
     return PROJECT_ROOT
 
 
+def missing_dataset_message(dataset_path: Path) -> str:
+    message = f"retrieval dataset not found: {dataset_path}"
+    match = re.fullmatch(
+        r"(?P<dataset>cifar10|cifar100)_hash(?P<hash_bits>\d+)_rows(?P<num_rows>\d+)_queries(?P<max_queries>\d+)\.npz",
+        dataset_path.name,
+    )
+    if match is None:
+        return message
+
+    groups = match.groupdict()
+    command = (
+        "python scripts/prepare_cam_retrieval_dataset.py "
+        f"--dataset {groups['dataset']} "
+        f"--num-rows {groups['num_rows']} "
+        f"--max-queries {groups['max_queries']} "
+        f"--hash-bits {groups['hash_bits']}"
+    )
+    return (
+        f"{message}\n"
+        "The requested benchmark artifact has not been generated yet. "
+        "Create it first, then rerun this benchmark:\n"
+        f"  {command}"
+    )
+
+
 def words_le_to_int(words: np.ndarray) -> int:
     value = 0
     for idx, word in enumerate(words.tolist()):
@@ -88,7 +117,7 @@ def load_retrieval_dataset_npz(path: str | Path) -> RetrievalDataset:
     if not dataset_path.is_absolute():
         dataset_path = project_root() / dataset_path
     if not dataset_path.exists():
-        raise FileNotFoundError(f"retrieval dataset not found: {dataset_path}")
+        raise FileNotFoundError(missing_dataset_message(dataset_path))
 
     loaded = np.load(dataset_path)
     rows_words = loaded["rows_words"]
@@ -98,6 +127,8 @@ def load_retrieval_dataset_npz(path: str | Path) -> RetrievalDataset:
     if rows_words.shape[1] != queries_words.shape[1]:
         raise ValueError("rows_words and queries_words must use the same word width")
 
+    rows_words = np.asarray(rows_words, dtype=np.uint64)
+    queries_words = np.asarray(queries_words, dtype=np.uint64)
     rows = [words_le_to_int(words) for words in rows_words]
     queries = [words_le_to_int(words) for words in queries_words]
     row_labels = [int(x) for x in loaded["row_labels"].tolist()]
@@ -107,28 +138,85 @@ def load_retrieval_dataset_npz(path: str | Path) -> RetrievalDataset:
         row_labels=row_labels,
         queries=queries,
         query_labels=query_labels,
+        rows_words=rows_words,
+        queries_words=queries_words,
         hash_bits=int(rows_words.shape[1] * 64),
         num_classes=len(set(row_labels)),
     )
 
 
-def hamming_similarity_score(query_row: int, stored_row: int, *, width: int) -> int:
+def match_topk_numpy(query_words: np.ndarray, rows_words: np.ndarray, *, width: int, k: int) -> list[int]:
-    if width <= 0:
-        raise ValueError("width must be greater than 0")
-    mask = (1 << width) - 1
-    distance = ((int(query_row) ^ int(stored_row)) & mask).bit_count()
-    return int(width - distance)
-
-
-def match_topk_hamming(query: int, rows: Sequence[int], *, width: int, k: int) -> list[int]:
     if k <= 0:
         raise ValueError("k must be greater than 0")
-    scored = [
+    if width <= 0:
-        (-hamming_similarity_score(query, row, width=width), row_index)
+        raise ValueError("width must be greater than 0")
-        for row_index, row in enumerate(rows)
+    query_words = np.asarray(query_words, dtype=np.uint64)
+    rows_words = np.asarray(rows_words, dtype=np.uint64)
+    if rows_words.ndim != 2:
+        raise ValueError("rows_words must have shape [N, words]")
+    if query_words.shape != (rows_words.shape[1],):
+        raise ValueError("query_words must have shape [words]")
+
+    return _match_topk_numpy_batch(query_words[np.newaxis, :], rows_words, width=width, k=k)[0]
+
+
+def _match_topk_numpy_batch(
+    queries_words: np.ndarray,
+    rows_words: np.ndarray,
+    *,
+    width: int,
+    k: int,
+) -> list[list[int]]:
+    if k <= 0:
+        raise ValueError("k must be greater than 0")
+    if width <= 0:
+        raise ValueError("width must be greater than 0")
+    queries_words = np.asarray(queries_words, dtype=np.uint64)
+    rows_words = np.asarray(rows_words, dtype=np.uint64)
+    if queries_words.ndim != 2:
+        raise ValueError("queries_words must have shape [Q, words]")
+    if rows_words.ndim != 2:
+        raise ValueError("rows_words must have shape [N, words]")
+    if queries_words.shape[1] != rows_words.shape[1]:
+        raise ValueError("queries_words and rows_words must use the same word width")
+
+    xor = np.bitwise_xor(queries_words[:, np.newaxis, :], rows_words[np.newaxis, :, :])
+    distances = np.bitwise_count(xor).sum(axis=2, dtype=np.int64)
+    scores = int(width) - distances
+    row_indices = np.arange(rows_words.shape[0], dtype=np.int64)
+    topk_count = min(k, rows_words.shape[0])
+    return [
+        [int(idx) for idx in np.lexsort((row_indices, -score_row))[:topk_count]]
+        for score_row in scores
     ]
-    scored.sort()
+
-    return [row_index for _, row_index in scored[: min(k, len(scored))]]
+
+def match_all_topk_numpy(
+    queries_words: np.ndarray,
+    rows_words: np.ndarray,
+    *,
+    width: int,
+    k: int,
+    workers: int = 1,
+) -> list[list[int]]:
+    if workers <= 0:
+        raise ValueError("workers must be greater than 0")
+    queries_words = np.asarray(queries_words, dtype=np.uint64)
+    if workers == 1 or len(queries_words) <= 1:
+        return _match_topk_numpy_batch(queries_words, rows_words, width=width, k=k)
+
+    chunks = [
+        chunk
+        for chunk in np.array_split(queries_words, min(workers, len(queries_words)))
+        if len(chunk)
+    ]
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        futures = [
+            executor.submit(_match_topk_numpy_batch, chunk, rows_words, width=width, k=k)
+            for chunk in chunks
+        ]
+        parts = [future.result() for future in futures]
+    return [topk for part in parts for topk in part]
 
 
 def compute_metrics(topk_indices: list[int], row_labels: list[int], query_label: int, k: int) -> tuple[float, float, float]:
@@ -154,6 +242,7 @@ def run_benchmark(
     hash_bits: int = 512,
     topk_values: Iterable[int] = BENCHMARK_KS,
     run_id: str | None = None,
+    workers: int = 1,
     timer_ns: Callable[[], int] = time.perf_counter_ns,
 ) -> dict:
     dataset = load_retrieval_dataset_npz(dataset_path)
@@ -165,6 +254,8 @@ def run_benchmark(
         )
     if len(dataset.queries) != len(dataset.query_labels):
         raise ValueError("queries and query_labels must have the same length")
+    if workers <= 0:
+        raise ValueError("workers must be greater than 0")
 
     ks = _normalized_topk_values(topk_values)
     max_k = max(ks)
@@ -172,10 +263,13 @@ def run_benchmark(
         raise ValueError("topk_values cannot exceed the number of dataset rows")
 
     start_ns = timer_ns()
-    all_topk = [
+    all_topk = match_all_topk_numpy(
-        match_topk_hamming(query, dataset.rows, width=hash_bits, k=max_k)
+        dataset.queries_words,
-        for query in dataset.queries
+        dataset.rows_words,
-    ]
+        width=hash_bits,
+        k=max_k,
+        workers=workers,
+    )
     end_ns = timer_ns()
 
     golden_topk = [
@@ -198,16 +292,18 @@ def run_benchmark(
     num_queries = len(dataset.queries)
     ns_per_query = (elapsed_ns / float(num_queries)) if num_queries else 0.0
     qps = (1_000_000_000.0 / ns_per_query) if ns_per_query > 0 else 0.0
-    resolved_run_id = run_id or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-software-hamming"
+    resolved_run_id = run_id or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-software-numpy"
 
     return {
         "run_id": resolved_run_id,
-        "mode": "software-hamming",
+        "mode": "software-numpy",
         "status": "pass",
         "params": {
             "num_rows": len(dataset.rows),
             "hash_bits": int(hash_bits),
             "topk_k": max_k,
+            "workers": int(workers),
+            "engine": "numpy",
         },
         "dataset": {
             "num_classes": dataset.num_classes,
@@ -231,7 +327,7 @@ def write_outputs(out_dir: Path, result: dict) -> None:
     (out_dir / "metrics.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
 
     fieldnames = [
-        "run_id", "mode", "num_rows", "hash_bits", "topk_k", "num_queries", "k",
+        "run_id", "mode", "num_rows", "hash_bits", "topk_k", "workers", "engine", "num_queries", "k",
         "macro_precision", "retrieval_recall", "macro_f1", "recall@k", "exact_match_rate",
         "total_elapsed_ns", "ns_per_query", "queries_per_second", "status",
     ]
@@ -245,6 +341,8 @@ def write_outputs(out_dir: Path, result: dict) -> None:
                 "num_rows": result["params"]["num_rows"],
                 "hash_bits": result["params"]["hash_bits"],
                 "topk_k": result["params"]["topk_k"],
+                "workers": result["params"].get("workers", 1),
+                "engine": result["params"].get("engine", "numpy"),
                 "num_queries": result["dataset"]["num_queries"],
                 "k": int(k),
                 "macro_precision": metrics["macro_precision"],
@@ -265,6 +363,8 @@ def write_outputs(out_dir: Path, result: dict) -> None:
         f"- mode: `{result['mode']}`",
         f"- status: `{result['status']}`",
         f"- num_queries: `{result['dataset']['num_queries']}`",
+        f"- workers: `{result['params'].get('workers', 1)}`",
+        f"- engine: `{result['params'].get('engine', 'numpy')}`",
         f"- ns_per_query: `{result['performance']['ns_per_query']}`",
         f"- queries_per_second: `{result['performance']['queries_per_second']}`",
         "",
@@ -278,7 +378,7 @@ def write_outputs(out_dir: Path, result: dict) -> None:
         )
     lines.extend([
         "",
-        "说明：软件路径直接对 `.npz` 中的 CAM 行整数执行汉明距离 / XNOR-popcount 检索，不使用软件 CAM 时序仿真。",
+        "说明：软件路径直接对 `.npz` 中的 little-endian uint64 words 使用 NumPy bitwise_count 执行汉明距离 / XNOR-popcount 检索，不使用软件 CAM 时序仿真。",
     ])
     (out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
 
@@ -288,10 +388,11 @@ def output_dir_for(run_id: str, output_root: Path) -> Path:
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Run software Hamming CAM retrieval benchmark.")
+    parser = argparse.ArgumentParser(description="Run NumPy software CAM retrieval benchmark.")
     parser.add_argument("--dataset", required=True, help="Prepared CAM retrieval .npz artifact")
     parser.add_argument("--hash-bits", type=int, default=512, help="Hash width in bits")
     parser.add_argument("--topk-k", type=int, default=5, help="Maximum Top-K to report; reports k=1 and this value")
+    parser.add_argument("--workers", type=int, default=1, help="Number of software query worker threads")
     parser.add_argument("--run-id", default=None, help="Output run id")
     parser.add_argument(
         "--output-root",
@@ -305,12 +406,17 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
     topk_values = (1,) if args.topk_k == 1 else (1, args.topk_k)
+    try:
         result = run_benchmark(
             args.dataset,
             hash_bits=args.hash_bits,
             topk_values=topk_values,
             run_id=args.run_id,
+            workers=args.workers,
         )
+    except FileNotFoundError as exc:
+        print(str(exc), file=sys.stderr)
+        raise SystemExit(2) from None
     out_dir = output_dir_for(result["run_id"], args.output_root)
     write_outputs(out_dir, result)
     print(
@@ -318,6 +424,7 @@ def main() -> None:
         f"run_id={result['run_id']} "
         f"num_rows={result['params']['num_rows']} "
         f"hash_bits={result['params']['hash_bits']} "
+        f"workers={result['params']['workers']} "
         f"num_queries={result['dataset']['num_queries']} "
         f"ns_per_query={result['performance']['ns_per_query']:.3f} "
         f"queries_per_second={result['performance']['queries_per_second']:.3f} "

tests/test_prepare_cam_retrieval_dataset.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import numpy as np
 
 from scripts.prepare_cam_retrieval_dataset import (
+    DEFAULT_MAX_QUERIES,
     dataset_config,
     pack_bits_to_words_le,
     stratified_indices,
@@ -15,6 +16,10 @@ def test_dataset_config_resolves_cifar10_and_cifar100() -> None:
     assert dataset_config("cifar100") == ("uoft-cs/cifar100", "fine_label")
 
 
+def test_default_query_count_is_large_enough_for_software_throughput() -> None:
+    assert DEFAULT_MAX_QUERIES == 8192
+
+
 def test_dataset_config_rejects_unknown_dataset() -> None:
     try:
         dataset_config("mnist")

tests/test_retrieval_benchmark_performance.py

@@ -1,20 +1,27 @@
 from __future__ import annotations
 
 import csv
+import importlib.util
 import sys
 from pathlib import Path
 
 
 HW_SIM_DIR = Path(__file__).resolve().parents[1] / "hw" / "sim"
-if str(HW_SIM_DIR) not in sys.path:
+if str(HW_SIM_DIR) in sys.path:
+    sys.path.remove(str(HW_SIM_DIR))
 sys.path.insert(0, str(HW_SIM_DIR))
+benchmark_path = HW_SIM_DIR / "benchmarks" / "retrieval" / "test_retrieval_benchmark.py"
+spec = importlib.util.spec_from_file_location("hw_retrieval_benchmark", benchmark_path)
+assert spec is not None
+assert spec.loader is not None
+hw_retrieval_benchmark = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = hw_retrieval_benchmark
+spec.loader.exec_module(hw_retrieval_benchmark)
 
-from benchmarks.retrieval.test_retrieval_benchmark import (  # noqa: E402
+QueryTiming = hw_retrieval_benchmark.QueryTiming
-    QueryTiming,
+build_hardware_performance = hw_retrieval_benchmark.build_hardware_performance
-    build_hardware_performance,
+summarize_query_timings = hw_retrieval_benchmark.summarize_query_timings
-    summarize_query_timings,
+write_outputs = hw_retrieval_benchmark.write_outputs
-    write_outputs,
-)
 
 
 def test_summarize_query_timings_uses_query_only_accept_to_last_cycles() -> None:

tests/test_sw_retrieval_benchmark.py

@@ -76,8 +76,10 @@ def test_run_benchmark_reports_quality_and_query_speed(tmp_path):
         timer_ns=lambda: next(timer_values),
     )
 
-    assert result["mode"] == "software-hamming"
+    assert result["mode"] == "software-numpy"
     assert result["status"] == "pass"
+    assert result["params"]["engine"] == "numpy"
+    assert result["params"]["workers"] == 1
     assert result["dataset"]["num_queries"] == 2
     assert result["params"]["num_rows"] == 3
     assert result["params"]["topk_k"] == 2
@@ -88,6 +90,59 @@ def test_run_benchmark_reports_quality_and_query_speed(tmp_path):
     assert result["performance"]["queries_per_second"] == 1_000_000.0
 
 
+def test_numpy_topk_matches_reference_with_tiebreak(tmp_path):
+    bench = load_sw_benchmark()
+    dataset_path = tmp_path / "dataset.npz"
+    np.savez_compressed(
+        dataset_path,
+        rows_words=np.array(
+            [
+                [0b11110000],
+                [0b11110000],
+                [0b11100000],
+                [0b00001111],
+            ],
+            dtype=np.uint64,
+        ),
+        row_labels=np.array([0, 0, 0, 1], dtype=np.int64),
+        queries_words=np.array([[0b11110000]], dtype=np.uint64),
+        query_labels=np.array([0], dtype=np.int64),
+    )
+    dataset = bench.load_retrieval_dataset_npz(dataset_path)
+
+    assert bench.match_topk_numpy(dataset.queries_words[0], dataset.rows_words, width=64, k=3) == [0, 1, 2]
+
+
+def test_numpy_batch_topk_vectorizes_across_queries(monkeypatch):
+    bench = load_sw_benchmark()
+    rows_words = np.array(
+        [
+            [0b11110000],
+            [0b11100000],
+            [0b00001111],
+        ],
+        dtype=np.uint64,
+    )
+    queries_words = np.array(
+        [
+            [0b11110000],
+            [0b00001111],
+        ],
+        dtype=np.uint64,
+    )
+
+    monkeypatch.setattr(
+        bench,
+        "match_topk_numpy",
+        lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("scalar path called")),
+    )
+
+    assert bench._match_topk_numpy_batch(queries_words, rows_words, width=64, k=2) == [
+        [0, 1],
+        [2, 1],
+    ]
+
+
 def test_run_benchmark_exact_match_compares_against_reference(tmp_path, monkeypatch):
     bench = load_sw_benchmark()
     dataset_path = tmp_path / "dataset.npz"
@@ -105,6 +160,35 @@ def test_run_benchmark_exact_match_compares_against_reference(tmp_path, monkeypa
     assert result["metrics"]["1"]["exact_match_rate"] == 0.0
 
 
+def test_run_benchmark_threaded_numpy_matches_single_worker(tmp_path):
+    bench = load_sw_benchmark()
+    dataset_path = tmp_path / "dataset.npz"
+    _write_dataset(dataset_path)
+
+    single = bench.run_benchmark(
+        dataset_path,
+        hash_bits=64,
+        topk_values=(1, 2),
+        run_id="single",
+        workers=1,
+        timer_ns=lambda: 0,
+    )
+    threaded = bench.run_benchmark(
+        dataset_path,
+        hash_bits=64,
+        topk_values=(1, 2),
+        run_id="threaded",
+        workers=2,
+        timer_ns=lambda: 0,
+    )
+
+    assert single["mode"] == "software-numpy"
+    assert threaded["mode"] == "software-numpy"
+    assert single["params"]["engine"] == "numpy"
+    assert threaded["params"]["workers"] == 2
+    assert threaded["metrics"] == single["metrics"]
+
+
 def test_run_benchmark_rejects_hash_bits_that_do_not_match_npz_width(tmp_path):
     bench = load_sw_benchmark()
     dataset_path = tmp_path / "dataset.npz"
@@ -114,13 +198,46 @@ def test_run_benchmark_rejects_hash_bits_that_do_not_match_npz_width(tmp_path):
         bench.run_benchmark(dataset_path, hash_bits=8)
 
 
+def test_cli_missing_dataset_prints_generation_hint_without_traceback(monkeypatch, capsys):
+    bench = load_sw_benchmark()
+    missing_path = Path("outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows512_queries8192.npz")
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "sw_retrieval_benchmark.py",
+            "--dataset",
+            str(missing_path),
+            "--hash-bits",
+            "512",
+            "--topk-k",
+            "5",
+            "--workers",
+            "1",
+        ],
+    )
+
+    with pytest.raises(SystemExit) as exc_info:
+        bench.main()
+
+    assert exc_info.value.code == 2
+    captured = capsys.readouterr()
+    assert "retrieval dataset not found" in captured.err
+    assert "python scripts/prepare_cam_retrieval_dataset.py" in captured.err
+    assert "--dataset cifar100" in captured.err
+    assert "--num-rows 512" in captured.err
+    assert "--max-queries 8192" in captured.err
+    assert "--hash-bits 512" in captured.err
+    assert "Traceback" not in captured.err
+
+
 def test_write_outputs_includes_quality_and_performance_csv(tmp_path):
     bench = load_sw_benchmark()
     result = {
         "run_id": "unit-test",
-        "mode": "software-hamming",
+        "mode": "software-numpy",
         "status": "pass",
-        "params": {"num_rows": 3, "hash_bits": 64, "topk_k": 2},
+        "params": {"num_rows": 3, "hash_bits": 64, "topk_k": 2, "workers": 2, "engine": "numpy"},
         "dataset": {"num_classes": 2, "num_queries": 2},
         "metrics": {
             "1": {
@@ -149,4 +266,7 @@ def test_write_outputs_includes_quality_and_performance_csv(tmp_path):
     assert "queries_per_second" in metrics_csv
     assert "1000000.0" in metrics_csv
     assert "Software CAM Retrieval Benchmark Summary" in summary_md
+    assert "software-numpy" in metrics_csv
+    assert "workers" in metrics_csv
+    assert "workers: `2`" in summary_md
     assert "queries_per_second" in summary_md