SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
This commit is contained in:
SikongJueluo
2026-06-04 16:57:53 +08:00
parent e5b764520c
commit b5a40819cc
5 changed files with 283 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
scripts/prepare_cam_retrieval_dataset.py
View File

@@ -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"),

157
scripts/sw_retrieval_benchmark.py
View File

@@ -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)
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} "

5
tests/test_prepare_cam_retrieval_dataset.py
View File

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

23
tests/test_retrieval_benchmark_performance.py
View File

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

126
tests/test_sw_retrieval_benchmark.py
View File

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