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/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:
-    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]:
+def match_topk_numpy(query_words: np.ndarray, rows_words: np.ndarray, *, width: int, k: int) -> list[int]:
     if k <= 0:
         raise ValueError("k must be greater than 0")
-    scored = [
-        (-hamming_similarity_score(query, row, width=width), row_index)
-        for row_index, row in enumerate(rows)
+    if width <= 0:
+        raise ValueError("width must be greater than 0")
+    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 = [
-        match_topk_hamming(query, dataset.rows, width=hash_bits, k=max_k)
-        for query in dataset.queries
-    ]
+    all_topk = match_all_topk_numpy(
+        dataset.queries_words,
+        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)
-    result = run_benchmark(
-        args.dataset,
-        hash_bits=args.hash_bits,
-        topk_values=topk_values,
-        run_id=args.run_id,
-    )
+    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} "