From b5a40819cca399218d8d4a8ed069bb32f6c3e8f2 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Thu, 4 Jun 2026 16:57:53 +0800 Subject: [PATCH] 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 | 3 +- scripts/sw_retrieval_benchmark.py | 169 ++++++++++++++---- tests/test_prepare_cam_retrieval_dataset.py | 5 + tests/test_retrieval_benchmark_performance.py | 23 ++- tests/test_sw_retrieval_benchmark.py | 126 ++++++++++++- 5 files changed, 283 insertions(+), 43 deletions(-) diff --git a/scripts/prepare_cam_retrieval_dataset.py b/scripts/prepare_cam_retrieval_dataset.py index f86d899..c7d449b 100644 --- a/scripts/prepare_cam_retrieval_dataset.py +++ b/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"), diff --git a/scripts/sw_retrieval_benchmark.py b/scripts/sw_retrieval_benchmark.py index 8666ab0..f1d50c2 100644 --- a/scripts/sw_retrieval_benchmark.py +++ b/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"(?Pcifar10|cifar100)_hash(?P\d+)_rows(?P\d+)_queries(?P\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} " diff --git a/tests/test_prepare_cam_retrieval_dataset.py b/tests/test_prepare_cam_retrieval_dataset.py index c26dcaa..f22529e 100644 --- a/tests/test_prepare_cam_retrieval_dataset.py +++ b/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") diff --git a/tests/test_retrieval_benchmark_performance.py b/tests/test_retrieval_benchmark_performance.py index 5f26f7c..9b9145c 100644 --- a/tests/test_retrieval_benchmark_performance.py +++ b/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: - sys.path.insert(0, str(HW_SIM_DIR)) +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, - build_hardware_performance, - summarize_query_timings, - write_outputs, -) +QueryTiming = hw_retrieval_benchmark.QueryTiming +build_hardware_performance = hw_retrieval_benchmark.build_hardware_performance +summarize_query_timings = hw_retrieval_benchmark.summarize_query_timings +write_outputs = hw_retrieval_benchmark.write_outputs def test_summarize_query_timings_uses_query_only_accept_to_last_cycles() -> None: diff --git a/tests/test_sw_retrieval_benchmark.py b/tests/test_sw_retrieval_benchmark.py index 2983f8b..7099891 100644 --- a/tests/test_sw_retrieval_benchmark.py +++ b/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