diff --git a/Cargo.lock b/Cargo.lock index b1e2f62..4e4483e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -25,12 +25,14 @@ checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" [[package]] name = "bm25x" -version = "0.1.5" +version = "0.2.0" dependencies = [ "bincode", + "cudarc", "memmap2", "rayon", "rust-stemmers", + "rustc-hash", "serde", "tempfile", "unicode-normalization", @@ -67,6 +69,15 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "cudarc" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6468cb7fa330840f3ebcd8df51edc0e7bf5c18df524792ce6004c6821851cdf3" +dependencies = [ + "libloading", +] + [[package]] name = "either" version = "1.15.0" @@ -171,6 +182,16 @@ version = "0.2.183" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -268,6 +289,12 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustix" version = "1.1.4" diff --git a/Cargo.toml b/Cargo.toml index f0a5a89..9b63eb5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,12 @@ serde = { version = "1", features = ["derive"] } rust-stemmers = "1.2" unicode-normalization = "0.1" rayon = "1.10" +rustc-hash = "2" +cudarc = { version = "0.19", default-features = false, features = ["cuda-version-from-build-system", "driver", "dynamic-loading", "nvrtc"], optional = true } + +[features] +default = [] +cuda = ["dep:cudarc"] [dev-dependencies] tempfile = "3" diff --git a/Makefile b/Makefile index 5d1e77b..b185111 100644 --- a/Makefile +++ b/Makefile @@ -27,5 +27,11 @@ release: bump-version test: cargo test --lib +test-cuda: + cargo test --lib --features cuda + build: cd python && maturin develop --release + +build-gpu: + cd python && maturin develop --release --features cuda diff --git a/README.md b/README.md index d7f5fe1..1e2bbbb 100644 --- a/README.md +++ b/README.md @@ -157,7 +157,7 @@ fn is_empty(&self) -> bool bm25x benchmarks

-> [BEIR](https://github.com/beir-cellar/beir) datasets. Same or better NDCG@10 than bm25s, **3.5-6x faster** indexing, **1-4x faster** search. +> [BEIR](https://github.com/beir-cellar/beir) datasets (log scale). Same or better NDCG@10 than bm25s. CPU: **3.5-6x faster** indexing, **1-4x faster** search. GPU (H100): up to **13x faster** indexing, **214x faster** search on MS MARCO 8.8M docs. GPU search has per-query kernel launch overhead, making it best suited for **batch querying** on **large datasets**. --- diff --git a/assets/benchmarks.png b/assets/benchmarks.png index 6095e24..6cdbd2c 100644 Binary files a/assets/benchmarks.png and b/assets/benchmarks.png differ diff --git a/benchmarks/bench_cuda.py b/benchmarks/bench_cuda.py new file mode 100644 index 0000000..ded31e2 --- /dev/null +++ b/benchmarks/bench_cuda.py @@ -0,0 +1,122 @@ +""" +Benchmark bm25x CUDA-accelerated indexing vs CPU (rayon) indexing. + +Uses BEIR SciFact dataset for realistic evaluation. +""" + +import time + +import bm25x +from beir import util +from beir.datasets.data_loader import GenericDataLoader + +# ───────────────────────── load dataset ────────────────────────── + +print("Downloading SciFact dataset...") +url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip" +data_path = util.download_and_unzip(url, "benchmarks/data") + +corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="test") + +corpus_ids = list(corpus.keys()) +corpus_texts = [ + (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "") + for did in corpus_ids +] +test_qids = sorted(qrels.keys()) +test_queries = [queries[qid] for qid in test_qids if qid in queries] + +gpu_available = bm25x.is_gpu_available() +print(f"GPU available: {gpu_available}") +print(f"Corpus: {len(corpus_texts)} docs | Queries: {len(test_queries)}") +print("=" * 65) + +# ───────────────────────── CPU benchmark ───────────────────────── + +print("\n--- CPU indexing (rayon) ---") + +# Warmup +warmup = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) +warmup.add(corpus_texts[:100]) +del warmup + +cpu_times = [] +for trial in range(5): + idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + t0 = time.perf_counter() + idx.add(corpus_texts) + t = time.perf_counter() - t0 + cpu_times.append(t) + if trial < 4: + del idx + +cpu_index = idx # keep last one for correctness check +best_cpu = min(cpu_times) +dps_cpu = len(corpus_texts) / best_cpu +print(f" Times: {[f'{t:.4f}s' for t in cpu_times]}") +print(f" Best: {best_cpu:.4f}s ({dps_cpu:,.0f} d/s)") + +# Search sanity check +cpu_results = [cpu_index.search(q, 10) for q in test_queries[:10]] +print( + f" Search check: first 10 queries OK ({sum(len(r) for r in cpu_results)} results)" +) + +# ───────────────────────── GPU benchmark ───────────────────────── + +if gpu_available: + print("\n--- GPU indexing (CUDA) ---") + + # Warmup GPU (first call compiles NVRTC kernels) + warmup = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + warmup.add(corpus_texts[:100]) + del warmup + + gpu_times = [] + for trial in range(5): + idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + t0 = time.perf_counter() + idx.add(corpus_texts) + t = time.perf_counter() - t0 + gpu_times.append(t) + if trial < 4: + del idx + + gpu_index = idx + best_gpu = min(gpu_times) + dps_gpu = len(corpus_texts) / best_gpu + print(f" Times: {[f'{t:.4f}s' for t in gpu_times]}") + print(f" Best: {best_gpu:.4f}s ({dps_gpu:,.0f} d/s)") + + # Correctness: compare search results CPU vs GPU + print("\n--- Correctness verification ---") + mismatches = 0 + for q in test_queries[:50]: + cpu_r = cpu_index.search(q, 10) + gpu_r = gpu_index.search(q, 10) + cpu_ids = [r[0] for r in cpu_r] + gpu_ids = [r[0] for r in gpu_r] + if cpu_ids != gpu_ids: + mismatches += 1 + print(f" Tested 50 queries: {mismatches} mismatches") + if mismatches == 0: + print(" PASS: GPU and CPU produce identical results") + else: + print(" WARN: Some results differ (may be due to tie-breaking)") + + # Summary + print("\n" + "=" * 65) + speedup = best_cpu / best_gpu + print(f"CPU: {best_cpu:.4f}s ({dps_cpu:,.0f} d/s)") + print(f"GPU: {best_gpu:.4f}s ({dps_gpu:,.0f} d/s)") + print(f"Speedup: {speedup:.2f}x") +else: + print( + "\nNOTE: GPU not available. Build with CUDA support to compare:\n" + " make build-gpu\n" + " # or: cd python && maturin develop --release --features cuda" + ) + print("\n" + "=" * 65) + print(f"CPU only: {best_cpu:.4f}s ({dps_cpu:,.0f} d/s)") + +print("Done.") diff --git a/benchmarks/bench_msmarco_cuda.py b/benchmarks/bench_msmarco_cuda.py new file mode 100644 index 0000000..edbb585 --- /dev/null +++ b/benchmarks/bench_msmarco_cuda.py @@ -0,0 +1,134 @@ +""" +Benchmark bm25x indexing on MS MARCO passage corpus (~8.8M passages). +Compares CPU (rayon) vs GPU (CUDA) indexing throughput. + +Usage: + CUDA_VISIBLE_DEVICES=3 python benchmarks/bench_msmarco_cuda.py +""" + +import gc +import os +import time + +import bm25x + +# ───────────────────────── load MS MARCO ───────────────────────── + + +def load_msmarco(max_docs=None): + """Load MS MARCO passage corpus via BEIR.""" + from beir import util + from beir.datasets.data_loader import GenericDataLoader + + dataset = "msmarco" + data_dir = "benchmarks/data" + data_path = os.path.join(data_dir, dataset) + + if not os.path.exists(data_path): + print("Downloading MS MARCO dataset (this may take a while)...") + url = f"https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{dataset}.zip" + data_path = util.download_and_unzip(url, data_dir) + + print("Loading corpus...") + corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="dev") + + corpus_ids = list(corpus.keys()) + if max_docs: + corpus_ids = corpus_ids[:max_docs] + + corpus_texts = [ + (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "") + for did in corpus_ids + ] + + test_qids = sorted(qrels.keys())[:200] + test_queries = [queries[qid] for qid in test_qids if qid in queries] + + return corpus_texts, test_queries + + +def bench_indexing(corpus_texts, label, trials=3): + """Benchmark indexing and return (best_time, index).""" + # Warmup + w = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + w.add(corpus_texts[:500]) + del w + gc.collect() + + times = [] + for trial in range(trials): + idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + t0 = time.perf_counter() + idx.add(corpus_texts) + t = time.perf_counter() - t0 + times.append(t) + print( + f" [{label}] trial {trial + 1}: {t:.3f}s ({len(corpus_texts) / t:,.0f} d/s)" + ) + if trial < trials - 1: + del idx + gc.collect() + + best = min(times) + dps = len(corpus_texts) / best + print(f" [{label}] best: {best:.3f}s ({dps:,.0f} d/s)") + return best, idx + + +# ───────────────────────── main ────────────────────────────────── + + +def main(): + gpu_available = bm25x.is_gpu_available() + print(f"GPU available: {gpu_available}") + + # Load corpus + corpus_texts, test_queries = load_msmarco() + n = len(corpus_texts) + print(f"Corpus: {n:,} docs | Queries: {len(test_queries)}") + print("=" * 65) + + # ── CPU baseline ── + print(f"\n--- CPU indexing (rayon) on {n:,} docs ---") + cpu_time, cpu_index = bench_indexing(corpus_texts, "CPU", trials=3) + + # Quick search sanity check + t0 = time.perf_counter() + for q in test_queries: + cpu_index.search(q, 10) + t_search = time.perf_counter() - t0 + print( + f" [CPU] search {len(test_queries)} queries: {t_search:.3f}s ({len(test_queries) / t_search:,.0f} q/s)" + ) + + if not gpu_available: + print("\nNo GPU — done.") + return + + # ── GPU ── + del cpu_index + gc.collect() + + print(f"\n--- GPU indexing (CUDA) on {n:,} docs ---") + gpu_time, gpu_index = bench_indexing(corpus_texts, "GPU", trials=3) + + # Search check + t0 = time.perf_counter() + for q in test_queries: + gpu_index.search(q, 10) + t_search_gpu = time.perf_counter() - t0 + print( + f" [GPU] search {len(test_queries)} queries: {t_search_gpu:.3f}s ({len(test_queries) / t_search_gpu:,.0f} q/s)" + ) + + # ── Summary ── + print("\n" + "=" * 65) + speedup = cpu_time / gpu_time + print(f"{'':>20} {'Time':>10} {'Throughput':>15}") + print(f"{'CPU (rayon)':>20} {cpu_time:>9.3f}s {n / cpu_time:>14,.0f} d/s") + print(f"{'GPU (CUDA)':>20} {gpu_time:>9.3f}s {n / gpu_time:>14,.0f} d/s") + print(f"{'Speedup':>20} {speedup:>9.2f}x") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_msmarco_profiled.py b/benchmarks/bench_msmarco_profiled.py new file mode 100644 index 0000000..1d1733d --- /dev/null +++ b/benchmarks/bench_msmarco_profiled.py @@ -0,0 +1,66 @@ +""" +Profiled MS MARCO benchmark — shows per-phase timing breakdown. +Usage: CUDA_VISIBLE_DEVICES=3 python benchmarks/bench_msmarco_profiled.py +""" + +import gc +import os +import time + +import bm25x + + +def load_msmarco(): + from beir import util + from beir.datasets.data_loader import GenericDataLoader + + data_path = os.path.join("benchmarks/data", "msmarco") + if not os.path.exists(data_path): + url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/msmarco.zip" + data_path = util.download_and_unzip(url, "benchmarks/data") + corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="dev") + corpus_ids = list(corpus.keys()) + corpus_texts = [ + (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "") + for did in corpus_ids + ] + return corpus_texts + + +corpus = load_msmarco() +n = len(corpus) +print(f"GPU available: {bm25x.is_gpu_available()}") +print(f"Corpus: {n:,} docs") +print("=" * 65) + +# Warmup +w = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) +w.add(corpus[:1000]) +del w +gc.collect() + +# CPU trial +print("\n--- CPU ---") +idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) +t0 = time.perf_counter() +idx.add(corpus) +t_cpu = time.perf_counter() - t0 +print(f" Total: {t_cpu:.3f}s ({n / t_cpu:,.0f} d/s)") +del idx +gc.collect() + +# GPU trial (if available) +if bm25x.is_gpu_available(): + # GPU warmup (kernel compilation) + w = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + w.add(corpus[:1000]) + del w + gc.collect() + + print("\n--- GPU ---") + idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + t0 = time.perf_counter() + idx.add(corpus) + t_gpu = time.perf_counter() - t0 + print(f" Total: {t_gpu:.3f}s ({n / t_gpu:,.0f} d/s)") + print(f"\n Speedup: {t_cpu / t_gpu:.2f}x") diff --git a/benchmarks/bench_search_cuda.py b/benchmarks/bench_search_cuda.py new file mode 100644 index 0000000..4114d6b --- /dev/null +++ b/benchmarks/bench_search_cuda.py @@ -0,0 +1,122 @@ +""" +Benchmark bm25x search throughput on MS MARCO. +Measures single-query latency and batch QPS. + +Usage: CUDA_VISIBLE_DEVICES=3 BM25X_PROFILE=1 python benchmarks/bench_search_cuda.py +""" + +import os +import statistics +import time + +import bm25x + +# ───────────────────────── load & index ────────────────────────── + + +def load_and_index(): + from beir import util + from beir.datasets.data_loader import GenericDataLoader + + data_path = os.path.join("benchmarks/data", "msmarco") + if not os.path.exists(data_path): + url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/msmarco.zip" + data_path = util.download_and_unzip(url, "benchmarks/data") + + print("Loading corpus...") + corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="dev") + + corpus_ids = list(corpus.keys()) + corpus_texts = [ + (corpus[did].get("title") or "") + " " + (corpus[did].get("text") or "") + for did in corpus_ids + ] + + test_qids = sorted(qrels.keys()) + test_queries = [queries[qid] for qid in test_qids if qid in queries] + + print(f"Corpus: {len(corpus_texts):,} docs | Queries: {len(test_queries)}") + + # Index + print("Indexing...") + idx = bm25x.BM25(method="lucene", k1=1.5, b=0.75, use_stopwords=True) + t0 = time.perf_counter() + idx.add(corpus_texts) + t_index = time.perf_counter() - t0 + print(f" Indexed in {t_index:.1f}s ({len(corpus_texts) / t_index:,.0f} d/s)") + + return idx, test_queries, corpus_texts + + +def bench_search(idx, queries, label, k=10, warmup=50, trials=3): + """Benchmark search: measure per-query latency and batch throughput.""" + n = len(queries) + + # Warmup + for q in queries[:warmup]: + idx.search(q, k) + + # Per-query latency (sample 200 queries) + sample = queries[:200] + latencies = [] + for q in sample: + t0 = time.perf_counter() + idx.search(q, k) + latencies.append((time.perf_counter() - t0) * 1000) # ms + + p50 = statistics.median(latencies) + p95 = sorted(latencies)[int(0.95 * len(latencies))] + p99 = sorted(latencies)[int(0.99 * len(latencies))] + mean_lat = statistics.mean(latencies) + + # Batch throughput + batch_times = [] + for trial in range(trials): + t0 = time.perf_counter() + for q in queries: + idx.search(q, k) + batch_times.append(time.perf_counter() - t0) + + best = min(batch_times) + qps = n / best + + print(f"\n--- {label} (k={k}, {n} queries, {idx.__len__():,} docs) ---") + print( + f" Latency: mean={mean_lat:.2f}ms p50={p50:.2f}ms p95={p95:.2f}ms p99={p99:.2f}ms" + ) + print(f" Batch: {[f'{t:.3f}s' for t in batch_times]}") + print(f" Best: {best:.3f}s ({qps:,.0f} q/s)") + return { + "mean_lat_ms": mean_lat, + "p50_ms": p50, + "p95_ms": p95, + "p99_ms": p99, + "qps": qps, + "best_time": best, + } + + +# ───────────────────────── main ────────────────────────────────── + + +def main(): + print(f"GPU available: {bm25x.is_gpu_available()}") + idx, queries, corpus = load_and_index() + + # Search benchmarks at different k values + results = {} + for k in [10, 100, 1000]: + r = bench_search(idx, queries, f"Search top-{k}", k=k) + results[k] = r + + # Summary + print("\n" + "=" * 65) + print(f"{'k':>6} {'Mean lat':>10} {'p50':>8} {'p95':>8} {'QPS':>10}") + for k, r in results.items(): + print( + f"{k:>6} {r['mean_lat_ms']:>9.2f}ms {r['p50_ms']:>7.2f}ms {r['p95_ms']:>7.2f}ms {r['qps']:>9,.0f}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/generate_charts.py b/benchmarks/generate_charts.py index 5512431..f1483f3 100644 --- a/benchmarks/generate_charts.py +++ b/benchmarks/generate_charts.py @@ -1,4 +1,4 @@ -"""Generate benchmark bar charts for README.""" +"""Generate benchmark bar charts for README — bm25s vs bm25x vs bm25x GPU.""" import matplotlib.pyplot as plt import numpy as np @@ -22,8 +22,11 @@ BM25S_COLOR = "#bbb" BM25X_COLOR = "#2d8cf0" +BM25X_GPU_COLOR = "#22c55e" -# ── Benchmark data (BEIR datasets, measured on Apple M3) ──────────── +# ── Benchmark data (BEIR datasets) ───────────────────────────────── +# CPU: measured on Apple M3 +# GPU: measured on NVIDIA H100, all datasets datasets = [ "NFCorpus\n3.6k docs", "SciFact\n5k docs", @@ -32,17 +35,20 @@ "MS MARCO\n8.8M docs", ] -# NDCG@10 +# NDCG@10 — GPU identical to CPU (verified on all datasets) ndcg_bm25s = [0.3064, 0.6617, 0.1538, 0.2326, 0.2124] ndcg_bm25x = [0.3287, 0.6904, 0.1600, 0.2514, 0.2186] +ndcg_bm25x_gpu = [0.3287, 0.6904, 0.1600, 0.2514, 0.2240] # Index throughput (docs/s) index_tput_bm25s = [13_658, 15_138, 17_567, 23_698, 23_395] index_tput_bm25x = [70_621, 77_644, 94_390, 144_060, 82_910] +index_tput_bm25x_gpu = [32_551, 57_570, 68_360, 205_115, 295_458] # H100 (warm CUDA) # Search throughput (queries/s) search_tput_bm25s = [36_992, 18_969, 6_543, 2_431, 16] search_tput_bm25x = [128_245, 25_992, 7_032, 4_760, 65] +search_tput_bm25x_gpu = [6_940, 5_682, 6_197, 5_935, 3_430] # H100 GPU search (warm) def _fmt_tput(v): @@ -51,11 +57,12 @@ def _fmt_tput(v): return f"{v:,.0f}" -def grouped_bar( +def grouped_bar_3( ax, labels, vals_s, vals_x, + vals_gpu, ylabel, title, fmt="{:.4f}", @@ -63,13 +70,14 @@ def grouped_bar( log_scale=False, abs_s=None, abs_x=None, + abs_gpu=None, abs_unit="", ): x = np.arange(len(labels)) - w = 0.35 + w = 0.25 bars_s = ax.bar( - x - w / 2, + x - w, vals_s, w, label="bm25s", @@ -79,7 +87,7 @@ def grouped_bar( zorder=3, ) bars_x = ax.bar( - x + w / 2, + x, vals_x, w, label="bm25x", @@ -88,6 +96,16 @@ def grouped_bar( linewidth=0.5, zorder=3, ) + bars_gpu = ax.bar( + x + w, + vals_gpu, + w, + label="bm25x GPU", + color=BM25X_GPU_COLOR, + edgecolor="#ffffff", + linewidth=0.5, + zorder=3, + ) if log_scale: ax.set_yscale("log") @@ -99,115 +117,117 @@ def grouped_bar( ax.yaxis.grid(True, linestyle="--", linewidth=0.5) ax.set_axisbelow(True) - # Value labels on bars - for i, (bar_s, bar_x, vs, vx) in enumerate(zip(bars_s, bars_x, vals_s, vals_x)): - label_s = fmt.format(vs) - label_x = fmt.format(vx) - ax.text( - bar_s.get_x() + bar_s.get_width() / 2, - bar_s.get_height(), - label_s, - ha="center", - va="bottom", - fontsize=7.5, - color="#888", - ) - ax.text( - bar_x.get_x() + bar_x.get_width() / 2, - bar_x.get_height(), - label_x, - ha="center", - va="bottom", - fontsize=7.5, - color="#2d8cf0", - fontweight="bold", - ) - - # Absolute throughput inside bars - if abs_s is not None and abs_x is not None: - txt_s = f"{_fmt_tput(abs_s[i])} {abs_unit}" - txt_x = f"{_fmt_tput(abs_x[i])} {abs_unit}" + # Value labels on top of bars + for i, (bar_s, bar_x, bar_g, vs, vx, vg) in enumerate( + zip(bars_s, bars_x, bars_gpu, vals_s, vals_x, vals_gpu) + ): + for bar, v, color, fw in [ + (bar_s, vs, "#888", "normal"), + (bar_x, vx, "#2d8cf0", "bold"), + (bar_g, vg, "#16a34a", "bold"), + ]: + label = fmt.format(v) + y_pos = bar.get_height() ax.text( - bar_s.get_x() + bar_s.get_width() / 2, - bar_s.get_height() / 2, - txt_s, + bar.get_x() + bar.get_width() / 2, + y_pos, + label, ha="center", - va="center", + va="bottom", fontsize=6.5, - color="#555", - rotation=90, - ) - ax.text( - bar_x.get_x() + bar_x.get_width() / 2, - bar_x.get_height() / 2, - txt_x, - ha="center", - va="center", - fontsize=6.5, - color="#fff", - fontweight="bold", - rotation=90, + color=color, + fontweight=fw, ) + # Absolute throughput inside bars + if abs_s is not None and abs_x is not None and abs_gpu is not None: + for bar, abs_v, color in [ + (bar_s, abs_s[i], "#555"), + (bar_x, abs_x[i], "#fff"), + (bar_g, abs_gpu[i], "#fff"), + ]: + txt = f"{_fmt_tput(abs_v)} {abs_unit}" + h = bar.get_height() + if h > 0.3: + ax.text( + bar.get_x() + bar.get_width() / 2, + h**0.5 if log_scale else h / 2, + txt, + ha="center", + va="center", + fontsize=5.5, + color=color, + fontweight="bold", + rotation=90, + ) + if show_legend: ax.legend( loc="upper left", - fontsize=9, + fontsize=8, facecolor="#fff", edgecolor="#ccc", labelcolor="#333", ) - # Remove top/right spines ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) # ── Compute speedups (bm25s = 1x baseline) ───────────────────────── -index_speedup = [x / s for s, x in zip(index_tput_bm25s, index_tput_bm25x)] -search_speedup = [x / s for s, x in zip(search_tput_bm25s, search_tput_bm25x)] baseline = [1.0] * len(datasets) +index_speedup_x = [x / s for s, x in zip(index_tput_bm25s, index_tput_bm25x)] +index_speedup_gpu = [x / s for s, x in zip(index_tput_bm25s, index_tput_bm25x_gpu)] +search_speedup_x = [x / s for s, x in zip(search_tput_bm25s, search_tput_bm25x)] +search_speedup_gpu = [x / s for s, x in zip(search_tput_bm25s, search_tput_bm25x_gpu)] # ── Generate figure ───────────────────────────────────────────────── -fig, axes = plt.subplots(1, 3, figsize=(18, 5.5)) +fig, axes = plt.subplots(1, 3, figsize=(19, 5.5)) fig.subplots_adjust(wspace=0.3, left=0.05, right=0.97, top=0.88, bottom=0.15) -grouped_bar( +grouped_bar_3( axes[0], datasets, ndcg_bm25s, ndcg_bm25x, + ndcg_bm25x_gpu, "NDCG@10", "Retrieval Quality (NDCG@10)", fmt="{:.4f}", show_legend=True, ) -grouped_bar( +grouped_bar_3( axes[1], datasets, baseline, - index_speedup, - "speedup vs bm25s", + index_speedup_x, + index_speedup_gpu, + "speedup vs bm25s (log)", "Indexing Speedup", fmt="{:.1f}x", show_legend=False, + log_scale=True, abs_s=index_tput_bm25s, abs_x=index_tput_bm25x, + abs_gpu=index_tput_bm25x_gpu, abs_unit="d/s", ) -grouped_bar( +grouped_bar_3( axes[2], datasets, baseline, - search_speedup, - "speedup vs bm25s", + search_speedup_x, + search_speedup_gpu, + "speedup vs bm25s (log)", "Search Speedup", fmt="{:.1f}x", show_legend=False, + log_scale=True, abs_s=search_tput_bm25s, abs_x=search_tput_bm25x, + abs_gpu=search_tput_bm25x_gpu, abs_unit="q/s", ) diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..4a6b6b4 --- /dev/null +++ b/build.rs @@ -0,0 +1,42 @@ +fn main() { + #[cfg(feature = "cuda")] + { + // Add standard CUDA library search paths for the linker. + let candidates = [ + "/usr/local/cuda/lib64/stubs", + "/usr/local/cuda/lib64", + "/usr/lib/x86_64-linux-gnu", + ]; + for path in &candidates { + if std::path::Path::new(path).exists() { + println!("cargo:rustc-link-search=native={}", path); + } + } + + // Support custom CUDA_PATH + if let Ok(cuda_path) = std::env::var("CUDA_PATH") { + let lib = format!("{}/lib64/stubs", cuda_path); + if std::path::Path::new(&lib).exists() { + println!("cargo:rustc-link-search=native={}", lib); + } + let lib64 = format!("{}/lib64", cuda_path); + if std::path::Path::new(&lib64).exists() { + println!("cargo:rustc-link-search=native={}", lib64); + } + } + + // Glob for versioned cuda installs: /usr/local/cuda-12*/lib64/stubs + if let Ok(entries) = std::fs::read_dir("/usr/local") { + for entry in entries.flatten() { + let name = entry.file_name(); + let name = name.to_string_lossy(); + if name.starts_with("cuda-") { + let stubs = format!("/usr/local/{}/lib64/stubs", name); + if std::path::Path::new(&stubs).exists() { + println!("cargo:rustc-link-search=native={}", stubs); + } + } + } + } + } +} diff --git a/python/Cargo.lock b/python/Cargo.lock index 1343075..38990f5 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -19,19 +19,21 @@ dependencies = [ [[package]] name = "bm25x" -version = "0.1.5" +version = "0.2.0" dependencies = [ "bincode", + "cudarc", "memmap2", "rayon", "rust-stemmers", + "rustc-hash", "serde", "unicode-normalization", ] [[package]] name = "bm25x-python" -version = "0.1.5" +version = "0.2.0" dependencies = [ "bm25x", "pyo3", @@ -68,6 +70,15 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "cudarc" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6468cb7fa330840f3ebcd8df51edc0e7bf5c18df524792ce6004c6821851cdf3" +dependencies = [ + "libloading", +] + [[package]] name = "either" version = "1.15.0" @@ -95,6 +106,16 @@ version = "0.2.183" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link", +] + [[package]] name = "memmap2" version = "0.9.10" @@ -236,6 +257,12 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustversion" version = "1.0.22" @@ -324,3 +351,9 @@ name = "unindent" version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" diff --git a/python/Cargo.toml b/python/Cargo.toml index eaf8251..f5d79bf 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -12,3 +12,7 @@ crate-type = ["cdylib"] [dependencies] pyo3 = { version = "0.24", features = ["extension-module"] } bm25x-core = { path = "..", package = "bm25x" } + +[features] +default = [] +cuda = ["bm25x-core/cuda"] diff --git a/python/pyproject.toml b/python/pyproject.toml index b74bdf3..f2669db 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -7,5 +7,8 @@ name = "bm25x" version = "0.2.0" requires-python = ">=3.9" +[project.optional-dependencies] +gpu = [] + [tool.maturin] features = ["pyo3/extension-module"] diff --git a/python/src/lib.rs b/python/src/lib.rs index 57b6dfd..7af8ff3 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,8 +1,67 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::PyList; use bm25x_core::{Method, TokenizerMode}; +/// Wrapper for a raw (*const u8, usize) string pointer pair. +/// These are extracted under the GIL, then used after GIL release. +/// Safety: the Python list must stay alive (keeping the strings alive) +/// for the entire duration these are used. +struct RawStr { + ptr: *const u8, + len: usize, +} + +// Raw pointers are not Send by default, but this is safe because: +// 1. The Python strings are immutable and reference-counted +// 2. The list object keeps all strings alive +// 3. We only read from these pointers, never write +unsafe impl Send for RawStr {} +unsafe impl Sync for RawStr {} + +/// Extract raw UTF-8 pointers from a Python list of strings using +/// direct C-API calls. ~60-200ns per item vs ~2500ns for safe PyO3. +/// +/// Safety: all items must be `str`. The returned pointers are valid +/// as long as the Python list (and its string elements) are alive. +fn extract_raw_ptrs(list: &Bound<'_, PyList>) -> PyResult> { + let len = list.len(); + let mut result = Vec::with_capacity(len); + unsafe { + let list_ptr = list.as_ptr(); + for i in 0..len as pyo3::ffi::Py_ssize_t { + let item = pyo3::ffi::PyList_GET_ITEM(list_ptr, i); + let mut size: pyo3::ffi::Py_ssize_t = 0; + let data = pyo3::ffi::PyUnicode_AsUTF8AndSize(item, &mut size); + if data.is_null() { + // Clear Python error state and return our own error + pyo3::ffi::PyErr_Clear(); + return Err(PyValueError::new_err(format!( + "element {} is not a string", + i + ))); + } + result.push(RawStr { + ptr: data as *const u8, + len: size as usize, + }); + } + } + Ok(result) +} + +/// Reconstruct &str slices from raw pointers. +/// Safety: pointers must be valid (Python strings still alive). +#[inline] +fn raw_to_strs(ptrs: &[RawStr]) -> Vec<&str> { + ptrs.iter() + .map(|r| unsafe { + std::str::from_utf8_unchecked(std::slice::from_raw_parts(r.ptr, r.len)) + }) + .collect() +} + fn parse_method(method: &str) -> PyResult { match method.to_lowercase().as_str() { "lucene" => Ok(Method::Lucene), @@ -34,6 +93,8 @@ fn io_err(e: std::io::Error) -> PyErr { #[pyclass(name = "BM25")] struct PyBM25 { inner: bm25x_core::BM25, + #[cfg(feature = "cuda")] + gpu_search_index: Option, } #[pymethods] @@ -61,22 +122,105 @@ impl PyBM25 { } None => bm25x_core::BM25::with_tokenizer(m, k1, b, delta, tok, use_stopwords), }; - Ok(PyBM25 { inner }) + Ok(PyBM25 { + inner, + #[cfg(feature = "cuda")] + gpu_search_index: None, + }) + } + + /// Upload the index to GPU for fast search. Call once after adding documents. + /// Subsequent search() calls will automatically use the GPU. + #[cfg(feature = "cuda")] + fn upload_to_gpu(&mut self) -> PyResult<()> { + self.gpu_search_index = Some( + self.inner + .to_gpu_search_index() + .map_err(PyValueError::new_err)?, + ); + Ok(()) } /// Add documents to the index. Returns list of assigned indices. - fn add(&mut self, documents: Vec) -> PyResult> { - let refs: Vec<&str> = documents.iter().map(|s| s.as_str()).collect(); - self.inner.add(&refs).map_err(io_err) + /// + /// Uses unsafe FFI for ~12x faster string extraction, then releases + /// the GIL for the entire Rust processing phase. + fn add(&mut self, py: Python<'_>, documents: &Bound<'_, PyList>) -> PyResult> { + // Phase 1: Extract raw UTF-8 pointers via C-API (~2s for 8.8M docs) + let ptrs = extract_raw_ptrs(documents)?; + + // Phase 2: Release GIL, reconstruct &str, run Rust indexing (~11s, GIL-free) + let inner = &mut self.inner; + py.allow_threads(|| { + let refs = raw_to_strs(&ptrs); + inner.add(&refs).map_err(io_err) + }) + } + + /// Add documents from a newline-delimited bytes blob. + /// + /// This is the fastest path: zero-copy buffer extraction (~0.1s for 2.6GB), + /// then GIL-free Rust processing. Use from Python: + /// + /// ```python + /// index.add_bytes(b"\n".join(s.encode() for s in corpus)) + /// # or if corpus is already List[str]: + /// index.add_bytes("\n".join(corpus).encode("utf-8")) + /// ``` + fn add_bytes(&mut self, py: Python<'_>, data: &[u8]) -> PyResult> { + let t0 = std::time::Instant::now(); + let inner = &mut self.inner; + let result = py.allow_threads(|| { + let t_split = std::time::Instant::now(); + let docs: Vec<&str> = data + .split(|&b| b == b'\n') + .map(|chunk| { + std::str::from_utf8(chunk).map_err(|e| { + std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()) + }) + }) + .collect::, _>>()?; + let split_time = t_split.elapsed(); + let t_add = std::time::Instant::now(); + let r = inner.add(&docs); + let add_time = t_add.elapsed(); + if std::env::var("BM25X_PROFILE").is_ok() { + eprintln!( + "[add_bytes] split={:.3}s add={:.3}s docs={}", + split_time.as_secs_f64(), + add_time.as_secs_f64(), + docs.len() + ); + } + r + }); + if std::env::var("BM25X_PROFILE").is_ok() { + eprintln!("[add_bytes] total={:.3}s", t0.elapsed().as_secs_f64()); + } + result.map_err(io_err) } /// Search the index. Returns list of (index, score) tuples. /// If `subset` is provided, only those document IDs are scored (pre-filtering). + /// If to_gpu() was called, uses GPU-accelerated search. #[pyo3(signature = (query, k, subset=None))] - fn search(&self, query: &str, k: usize, subset: Option>) -> Vec<(usize, f32)> { + fn search(&mut self, query: &str, k: usize, subset: Option>) -> Vec<(usize, f32)> { let results = match subset { Some(ids) => self.inner.search_filtered(query, k, &ids), - None => self.inner.search(query, k), + None => { + #[cfg(feature = "cuda")] + { + if let Some(ref mut gpu_idx) = self.gpu_search_index { + return self + .inner + .search_gpu(gpu_idx, query, k) + .into_iter() + .map(|r| (r.index, r.score)) + .collect(); + } + } + self.inner.search(query, k) + } }; results.into_iter().map(|r| (r.index, r.score)).collect() } @@ -101,7 +245,11 @@ impl PyBM25 { #[pyo3(signature = (index, mmap=false))] fn load(index: &str, mmap: bool) -> PyResult { let inner = bm25x_core::BM25::load(index, mmap).map_err(io_err)?; - Ok(PyBM25 { inner }) + Ok(PyBM25 { + inner, + #[cfg(feature = "cuda")] + gpu_search_index: None, + }) } /// Score a query against a list of documents. @@ -132,8 +280,15 @@ impl PyBM25 { } } +/// Returns True if bm25x was compiled with CUDA support and a GPU is available. +#[pyfunction] +fn is_gpu_available() -> bool { + bm25x_core::is_gpu_available() +} + #[pymodule] fn bm25x(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_function(wrap_pyfunction!(is_gpu_available, m)?)?; Ok(()) } diff --git a/src/cuda.rs b/src/cuda.rs new file mode 100644 index 0000000..00af156 --- /dev/null +++ b/src/cuda.rs @@ -0,0 +1,1068 @@ +//! CUDA-accelerated posting list construction for BM25 indexing. +//! +//! The pipeline: CPU tokenization (rayon) → parallel vocab → parallel counting sort +//! (CPU or GPU) → posting lists. The counting sort is stable (preserves doc_id order), +//! so no per-term sort is needed. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; + +use rustc_hash::{FxHashMap, FxHashSet}; + +use cudarc::driver::{ + CudaContext as CudarcContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use cudarc::nvrtc::compile_ptx; +use rayon::prelude::*; + +// --------------------------------------------------------------------------- +// CUDA kernel source +// --------------------------------------------------------------------------- + +const CUDA_KERNELS: &str = r#" +extern "C" __global__ void compute_histogram( + const unsigned int* __restrict__ term_ids, + unsigned int* __restrict__ histogram, + int n) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + atomicAdd(&histogram[term_ids[idx]], 1); + } +} + +extern "C" __global__ void scatter_postings( + const unsigned int* __restrict__ term_ids, + const unsigned int* __restrict__ doc_ids, + const unsigned int* __restrict__ tfs, + const unsigned long long* __restrict__ offsets, + unsigned int* __restrict__ counters, + unsigned int* __restrict__ out_doc_ids, + unsigned int* __restrict__ out_tfs, + int n) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + unsigned int term = term_ids[idx]; + unsigned int pos = (unsigned int)offsets[term] + atomicAdd(&counters[term], 1); + out_doc_ids[pos] = doc_ids[idx]; + out_tfs[pos] = tfs[idx]; + } +} + +extern "C" __global__ void zero_u32(unsigned int* buf, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) buf[idx] = 0; +} + +extern "C" __global__ void zero_f32(float* buf, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) buf[idx] = 0.0f; +} + +// BM25 Lucene scoring kernel: each thread scores one posting entry +// scores[doc_id] += idf * (tf / (k1 * (1 - b + b * dl/avgdl) + tf)) +extern "C" __global__ void bm25_score_lucene( + const unsigned int* __restrict__ posting_doc_ids, + const unsigned int* __restrict__ posting_tfs, + const unsigned int* __restrict__ doc_lengths, + float* __restrict__ scores, + float idf, float k1, float b, float avgdl, + long long offset, + int num_entries) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_entries) { + unsigned int doc_id = posting_doc_ids[offset + idx]; + float tf = (float)posting_tfs[offset + idx]; + float dl = (float)doc_lengths[doc_id]; + float norm = 1.0f - b + b * dl / avgdl; + float tfc = tf / (k1 * norm + tf); + atomicAdd(&scores[doc_id], idf * tfc); + } +} + +// Fused multi-term BM25 scoring: ALL query terms in ONE kernel launch. +// and a prefix-sum array to map global thread idx → term +extern "C" __global__ void bm25_score_fused_v2( + const unsigned int* __restrict__ posting_doc_ids, + const unsigned int* __restrict__ posting_tfs, + const unsigned int* __restrict__ doc_lengths, + float* __restrict__ scores, + const long long* __restrict__ flat_offsets, // [num_terms] offset in flat posting arrays + const long long* __restrict__ virtual_starts, // [num_terms] prefix sum of counts + const float* __restrict__ idfs, // [num_terms] IDF values + float k1, float b, float avgdl, + int num_terms, + int total_entries) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total_entries) return; + + // Binary search: find which term this thread belongs to + int lo = 0, hi = num_terms - 1; + while (lo < hi) { + int mid = (lo + hi + 1) / 2; + if (virtual_starts[mid] <= (long long)idx) lo = mid; + else hi = mid - 1; + } + + int local_pos = idx - (int)virtual_starts[lo]; + long long flat_pos = flat_offsets[lo] + (long long)local_pos; + + unsigned int doc_id = posting_doc_ids[flat_pos]; + float tf = (float)posting_tfs[flat_pos]; + float dl = (float)doc_lengths[doc_id]; + float norm = 1.0f - b + b * dl / avgdl; + float tfc = tf / (k1 * norm + tf); + atomicAdd(&scores[doc_id], idfs[lo] * tfc); +} + +// Top-k extraction: each thread checks one doc, writes to output if score > 0 +// max_results prevents out-of-bounds writes +extern "C" __global__ void collect_nonzero( + const float* __restrict__ scores, + unsigned int* __restrict__ out_doc_ids, + float* __restrict__ out_scores, + unsigned int* __restrict__ count, + int n, + int max_results) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n && scores[idx] > 0.0f) { + unsigned int pos = atomicAdd(count, 1); + if (pos < (unsigned int)max_results) { + out_doc_ids[pos] = (unsigned int)idx; + out_scores[pos] = scores[idx]; + } + } +} + +// ── Correct GPU top-k: per-thread private heaps, then block merge ── +// +// Each thread maintains a private top-k (insertion sort in registers). +// After scanning, threads write their private results to shared memory. +// Thread 0 does a final merge over all 256*k candidates → block top-k. +// No races: each thread's private top-k is exact, merge is sequential. + +#define TOPK_K 10 + +extern "C" __global__ void topk_per_block( + const float* __restrict__ scores, + unsigned int* __restrict__ out_doc_ids, // [num_blocks * TOPK_K] + float* __restrict__ out_scores, // [num_blocks * TOPK_K] + int n, + int k) +{ + int tid = threadIdx.x; + int block_size = blockDim.x; + + // Phase 1: each thread scans its chunk, maintains private top-k in registers + float my_scores[TOPK_K]; + unsigned int my_ids[TOPK_K]; + for (int i = 0; i < TOPK_K; i++) { + my_scores[i] = -1.0f; + my_ids[i] = 0xFFFFFFFF; + } + float my_min = -1.0f; + int my_min_idx = 0; + + // Grid-stride: this block handles a contiguous chunk of the scores array + int chunk_size = (n + gridDim.x - 1) / gridDim.x; + int chunk_start = blockIdx.x * chunk_size; + int chunk_end = min(chunk_start + chunk_size, n); + + int actual_k = min(k, TOPK_K); + + for (int i = chunk_start + tid; i < chunk_end; i += block_size) { + float val = scores[i]; + if (val > my_min) { + // Replace minimum in private top-k + my_scores[my_min_idx] = val; + my_ids[my_min_idx] = (unsigned int)i; + // Recompute minimum + my_min = my_scores[0]; + my_min_idx = 0; + for (int j = 1; j < actual_k; j++) { + if (my_scores[j] < my_min) { + my_min = my_scores[j]; + my_min_idx = j; + } + } + } + } + + // Phase 2: write private results to shared memory + // Shared memory: 256 threads * TOPK_K entries * 8 bytes = 20KB (fits easily) + __shared__ float s_scores[256 * TOPK_K]; + __shared__ unsigned int s_ids[256 * TOPK_K]; + + for (int i = 0; i < actual_k; i++) { + s_scores[tid * TOPK_K + i] = my_scores[i]; + s_ids[tid * TOPK_K + i] = my_ids[i]; + } + __syncthreads(); + + // Phase 3: thread 0 merges all 256*k candidates into block's top-k + if (tid == 0) { + float final_scores[TOPK_K]; + unsigned int final_ids[TOPK_K]; + for (int i = 0; i < actual_k; i++) { + final_scores[i] = -1.0f; + final_ids[i] = 0xFFFFFFFF; + } + float fmin = -1.0f; + int fmin_idx = 0; + + int total_candidates = block_size * actual_k; + for (int i = 0; i < total_candidates; i++) { + float val = s_scores[i]; + if (val > fmin) { + final_scores[fmin_idx] = val; + final_ids[fmin_idx] = s_ids[i]; + fmin = final_scores[0]; + fmin_idx = 0; + for (int j = 1; j < actual_k; j++) { + if (final_scores[j] < fmin) { + fmin = final_scores[j]; + fmin_idx = j; + } + } + } + } + + // Write block's top-k to global memory + int block_offset = blockIdx.x * actual_k; + for (int i = 0; i < actual_k; i++) { + out_doc_ids[block_offset + i] = final_ids[i]; + out_scores[block_offset + i] = final_scores[i]; + } + } +} +"#; + +// --------------------------------------------------------------------------- +// Global CUDA context (lazy, panic-safe) +// --------------------------------------------------------------------------- + +static CUDA_BROKEN: AtomicBool = AtomicBool::new(false); +static GLOBAL_CONTEXT: OnceLock>>> = OnceLock::new(); + +pub struct CudaIndexer { + pub stream: Arc, + histogram_fn: CudaFunction, + scatter_fn: CudaFunction, + zero_fn: CudaFunction, + // Search kernels + zero_f32_fn: CudaFunction, + #[allow(dead_code)] + bm25_score_fn: CudaFunction, + bm25_score_fused_fn: CudaFunction, + #[allow(dead_code)] + collect_nonzero_fn: CudaFunction, + topk_per_block_fn: CudaFunction, +} + +pub fn is_cuda_available() -> bool { + !CUDA_BROKEN.load(Ordering::Relaxed) && get_global_context().is_some() +} + +pub fn mark_cuda_broken() { + CUDA_BROKEN.store(true, Ordering::Relaxed); +} + +pub fn get_global_context() -> Option> { + if CUDA_BROKEN.load(Ordering::Relaxed) { + return None; + } + let mutex = GLOBAL_CONTEXT.get_or_init(|| Mutex::new(None)); + let mut guard = mutex.lock().ok()?; + if let Some(ref ctx) = *guard { + return Some(Arc::clone(ctx)); + } + match CudaIndexer::new(0) { + Ok(ctx) => { + let arc = Arc::new(ctx); + *guard = Some(Arc::clone(&arc)); + Some(arc) + } + Err(e) => { + eprintln!("[bm25x] CUDA init failed: {}, using CPU", e); + CUDA_BROKEN.store(true, Ordering::Relaxed); + None + } + } +} + +#[allow(dead_code)] +pub struct GpuScatterResult { + pub out_doc_ids: Vec, + pub out_tfs: Vec, + pub histogram: Vec, + pub offsets: Vec, +} + +impl CudaIndexer { + fn new(device_id: usize) -> Result { + let device = CudarcContext::new(device_id).map_err(|e| format!("device: {:?}", e))?; + let stream = device.default_stream(); + + let ptx = + compile_ptx(CUDA_KERNELS).map_err(|e| format!("NVRTC compile failed: {:?}", e))?; + let module = device + .load_module(ptx) + .map_err(|e| format!("module load: {:?}", e))?; + + let histogram_fn = module + .load_function("compute_histogram") + .map_err(|e| format!("load compute_histogram: {:?}", e))?; + let scatter_fn = module + .load_function("scatter_postings") + .map_err(|e| format!("load scatter_postings: {:?}", e))?; + let zero_fn = module + .load_function("zero_u32") + .map_err(|e| format!("load zero_u32: {:?}", e))?; + let zero_f32_fn = module + .load_function("zero_f32") + .map_err(|e| format!("load zero_f32: {:?}", e))?; + let bm25_score_fn = module + .load_function("bm25_score_lucene") + .map_err(|e| format!("load bm25_score_lucene: {:?}", e))?; + let collect_nonzero_fn = module + .load_function("collect_nonzero") + .map_err(|e| format!("load collect_nonzero: {:?}", e))?; + let bm25_score_fused_fn = module + .load_function("bm25_score_fused_v2") + .map_err(|e| format!("load bm25_score_fused_v2: {:?}", e))?; + + let topk_per_block_fn = module + .load_function("topk_per_block") + .map_err(|e| format!("load topk_per_block: {:?}", e))?; + + Ok(CudaIndexer { + stream, + histogram_fn, + scatter_fn, + zero_fn, + zero_f32_fn, + bm25_score_fn, + bm25_score_fused_fn, + collect_nonzero_fn, + topk_per_block_fn, + }) + } + + /// GPU scatter: histogram + atomic scatter on device. + #[allow(dead_code)] + pub fn gpu_scatter( + &self, + term_ids: &[u32], + doc_ids: &[u32], + tfs: &[u32], + vocab_size: usize, + ) -> Result { + let n = term_ids.len(); + if n == 0 { + return Ok(GpuScatterResult { + out_doc_ids: vec![], + out_tfs: vec![], + histogram: vec![0; vocab_size], + offsets: vec![0; vocab_size], + }); + } + + let block = 256u32; + let grid_n = (n as u32).div_ceil(block); + let grid_v = (vocab_size as u32).div_ceil(block); + let n_i32 = n as i32; + let v_i32 = vocab_size as i32; + + let d_term_ids: CudaSlice = self + .stream + .clone_htod(term_ids) + .map_err(|e| format!("{:?}", e))?; + let d_doc_ids: CudaSlice = self + .stream + .clone_htod(doc_ids) + .map_err(|e| format!("{:?}", e))?; + let d_tfs: CudaSlice = self + .stream + .clone_htod(tfs) + .map_err(|e| format!("{:?}", e))?; + + let mut d_histogram: CudaSlice = self + .stream + .alloc_zeros(vocab_size) + .map_err(|e| format!("{:?}", e))?; + + unsafe { + self.stream + .launch_builder(&self.zero_fn) + .arg(&mut d_histogram) + .arg(&v_i32) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid_v, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("{:?}", e))?; + self.stream + .launch_builder(&self.histogram_fn) + .arg(&d_term_ids) + .arg(&mut d_histogram) + .arg(&n_i32) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid_n, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("{:?}", e))?; + } + + let histogram: Vec = self + .stream + .clone_dtoh(&d_histogram) + .map_err(|e| format!("{:?}", e))?; + let mut offsets: Vec = Vec::with_capacity(vocab_size); + let mut running: u64 = 0; + for &count in &histogram { + offsets.push(running); + running += count as u64; + } + + let d_offsets: CudaSlice = self + .stream + .clone_htod(&offsets) + .map_err(|e| format!("{:?}", e))?; + let mut d_counters: CudaSlice = self + .stream + .alloc_zeros(vocab_size) + .map_err(|e| format!("{:?}", e))?; + + unsafe { + self.stream + .launch_builder(&self.zero_fn) + .arg(&mut d_counters) + .arg(&v_i32) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid_v, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("{:?}", e))?; + } + + let mut d_out_doc_ids: CudaSlice = + self.stream.alloc_zeros(n).map_err(|e| format!("{:?}", e))?; + let mut d_out_tfs: CudaSlice = + self.stream.alloc_zeros(n).map_err(|e| format!("{:?}", e))?; + + unsafe { + self.stream + .launch_builder(&self.scatter_fn) + .arg(&d_term_ids) + .arg(&d_doc_ids) + .arg(&d_tfs) + .arg(&d_offsets) + .arg(&mut d_counters) + .arg(&mut d_out_doc_ids) + .arg(&mut d_out_tfs) + .arg(&n_i32) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid_n, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("{:?}", e))?; + } + + let out_doc_ids = self + .stream + .clone_dtoh(&d_out_doc_ids) + .map_err(|e| format!("{:?}", e))?; + let out_tfs = self + .stream + .clone_dtoh(&d_out_tfs) + .map_err(|e| format!("{:?}", e))?; + + Ok(GpuScatterResult { + out_doc_ids, + out_tfs, + histogram, + offsets, + }) + } +} + +// --------------------------------------------------------------------------- +// GPU-resident search index +// --------------------------------------------------------------------------- + +/// Index data uploaded to GPU for fast search. Created once, reused across queries. +pub struct GpuSearchIndex { + /// All posting list doc_ids concatenated + flat_doc_ids: CudaSlice, + /// All posting list TFs concatenated + flat_tfs: CudaSlice, + /// Start offset per term_id into flat arrays + offsets: Vec, + /// Entry count per term_id + counts: Vec, + /// Document lengths on GPU + doc_lengths: CudaSlice, + /// Score accumulator (reused across queries) + scores: CudaSlice, + /// Buffer for non-zero doc_ids (for top-k) — reserved for future GPU top-k + #[allow(dead_code)] + result_doc_ids: CudaSlice, + /// Buffer for non-zero scores (for top-k) + #[allow(dead_code)] + result_scores: CudaSlice, + /// Atomic counter for collect_nonzero + #[allow(dead_code)] + result_count: CudaSlice, + /// Top-k per-block output buffers + topk_block_doc_ids: CudaSlice, + topk_block_scores: CudaSlice, + topk_num_blocks: u32, + /// Number of documents + num_docs: u32, +} + +impl GpuSearchIndex { + /// Upload a BM25 index to GPU memory. + pub fn from_index( + ctx: &CudaIndexer, + postings: &[Vec<(u32, u32)>], + doc_lengths: &[u32], + num_docs: u32, + ) -> Result { + // Flatten posting lists into contiguous arrays + let total_entries: usize = postings.iter().map(|p| p.len()).sum(); + let mut flat_doc_ids = Vec::with_capacity(total_entries); + let mut flat_tfs = Vec::with_capacity(total_entries); + let mut offsets = Vec::with_capacity(postings.len()); + let mut counts = Vec::with_capacity(postings.len()); + let mut offset = 0u64; + + for plist in postings { + offsets.push(offset); + counts.push(plist.len() as u32); + for &(doc_id, tf) in plist { + flat_doc_ids.push(doc_id); + flat_tfs.push(tf); + } + offset += plist.len() as u64; + } + + let stream = &ctx.stream; + + // Upload to GPU + let d_flat_doc_ids = stream + .clone_htod(&flat_doc_ids) + .map_err(|e| format!("{:?}", e))?; + let d_flat_tfs = stream + .clone_htod(&flat_tfs) + .map_err(|e| format!("{:?}", e))?; + let d_doc_lengths = stream + .clone_htod(doc_lengths) + .map_err(|e| format!("{:?}", e))?; + + // Allocate score buffer + result buffers + let d_scores: CudaSlice = stream + .alloc_zeros(num_docs as usize) + .map_err(|e| format!("{:?}", e))?; + // For common English queries on 8.8M docs, up to 5M docs may be scored + let max_results = num_docs as usize; + let d_result_doc_ids: CudaSlice = stream + .alloc_zeros(max_results) + .map_err(|e| format!("{:?}", e))?; + let d_result_scores: CudaSlice = stream + .alloc_zeros(max_results) + .map_err(|e| format!("{:?}", e))?; + let d_result_count: CudaSlice = + stream.alloc_zeros(1).map_err(|e| format!("{:?}", e))?; + + // Top-k buffers: use 256 blocks, each producing up to 1024 top-k candidates + let topk_num_blocks = 256u32; + let topk_buf_size = topk_num_blocks as usize * 1024; // max k=1024 + let d_topk_doc_ids: CudaSlice = stream + .alloc_zeros(topk_buf_size) + .map_err(|e| format!("{:?}", e))?; + let d_topk_scores: CudaSlice = stream + .alloc_zeros(topk_buf_size) + .map_err(|e| format!("{:?}", e))?; + + Ok(GpuSearchIndex { + flat_doc_ids: d_flat_doc_ids, + flat_tfs: d_flat_tfs, + offsets, + counts, + doc_lengths: d_doc_lengths, + scores: d_scores, + result_doc_ids: d_result_doc_ids, + result_scores: d_result_scores, + result_count: d_result_count, + topk_block_doc_ids: d_topk_doc_ids, + topk_block_scores: d_topk_scores, + topk_num_blocks, + num_docs, + }) + } + + /// Search on GPU. Returns (doc_id, score) pairs sorted by descending score. + /// + /// Uses a fused kernel that scores ALL query terms in a single launch, + /// eliminating per-term kernel launch overhead (~0.1ms each). + pub fn search( + &mut self, + ctx: &CudaIndexer, + query_term_ids: &[(u32, f32)], // (term_id, idf) + k1: f32, + b: f32, + avgdl: f32, + k: usize, + ) -> Result, String> { + let block = 256u32; + let grid_docs = (self.num_docs).div_ceil(block); + let n_docs = self.num_docs as i32; + + // 1. Zero scores array + unsafe { + ctx.stream + .launch_builder(&ctx.zero_f32_fn) + .arg(&mut self.scores) + .arg(&n_docs) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid_docs, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("zero scores: {:?}", e))?; + } + + // 2. Build fused kernel params: flat_offsets, virtual_starts, idfs + let mut flat_offsets: Vec = Vec::new(); + let mut virtual_starts: Vec = Vec::new(); + let mut idfs: Vec = Vec::new(); + let mut total_entries: i64 = 0; + + for &(term_id, idf_val) in query_term_ids { + let tid = term_id as usize; + if tid >= self.counts.len() || self.counts[tid] == 0 { + continue; + } + flat_offsets.push(self.offsets[tid] as i64); + virtual_starts.push(total_entries); + idfs.push(idf_val); + total_entries += self.counts[tid] as i64; + } + + if total_entries == 0 { + return Ok(Vec::new()); + } + + let num_terms = flat_offsets.len() as i32; + let total = total_entries as i32; + let grid = (total as u32).div_ceil(block); + + // Upload small per-query arrays (typically <100 bytes) + let d_flat_offsets = ctx + .stream + .clone_htod(&flat_offsets) + .map_err(|e| format!("{:?}", e))?; + let d_virtual_starts = ctx + .stream + .clone_htod(&virtual_starts) + .map_err(|e| format!("{:?}", e))?; + let d_idfs = ctx + .stream + .clone_htod(&idfs) + .map_err(|e| format!("{:?}", e))?; + + // Single fused kernel launch for ALL query terms + unsafe { + ctx.stream + .launch_builder(&ctx.bm25_score_fused_fn) + .arg(&self.flat_doc_ids) + .arg(&self.flat_tfs) + .arg(&self.doc_lengths) + .arg(&mut self.scores) + .arg(&d_flat_offsets) + .arg(&d_virtual_starts) + .arg(&d_idfs) + .arg(&k1) + .arg(&b) + .arg(&avgdl) + .arg(&num_terms) + .arg(&total) + .launch(LaunchConfig { + block_dim: (block, 1, 1), + grid_dim: (grid, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("bm25_score_fused: {:?}", e))?; + } + + // 3. GPU top-k: each block finds local top-k, download block results, merge on CPU + let k_clamped = k.min(1024) as i32; // kernel supports max k=1024 + let topk_blocks = self.topk_num_blocks; + let topk_block_size = 256u32; // threads per block + + unsafe { + ctx.stream + .launch_builder(&ctx.topk_per_block_fn) + .arg(&self.scores) + .arg(&mut self.topk_block_doc_ids) + .arg(&mut self.topk_block_scores) + .arg(&n_docs) + .arg(&k_clamped) + .launch(LaunchConfig { + block_dim: (topk_block_size, 1, 1), + grid_dim: (topk_blocks, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| format!("topk_per_block: {:?}", e))?; + } + + // 4. Download only num_blocks * k results (small: 256 * k * 8 bytes) + let download_count = topk_blocks as usize * k_clamped as usize; + let block_doc_ids: Vec = ctx + .stream + .clone_dtoh(&self.topk_block_doc_ids.slice(..download_count)) + .map_err(|e| format!("{:?}", e))?; + let block_scores: Vec = ctx + .stream + .clone_dtoh(&self.topk_block_scores.slice(..download_count)) + .map_err(|e| format!("{:?}", e))?; + + // 5. CPU merge: top-k across all block results (small: 256*k entries) + let mut results: Vec<(u32, f32)> = block_doc_ids + .into_iter() + .zip(block_scores) + .filter(|&(_, s)| s > 0.0) + .collect(); + + if results.len() > k { + results.select_nth_unstable_by(k, |a, b| { + b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal) + }); + results.truncate(k); + } + results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + Ok(results) + } +} + +// --------------------------------------------------------------------------- +// Stable parallel counting sort on CPU — preserves doc_id order, no sort needed +// --------------------------------------------------------------------------- + +/// Parallel counting sort: groups (term_id, doc_id, tf) entries by term_id. +/// Stable: entries within each term preserve their original order (= doc_id order). +/// Returns (out_doc_ids, out_tfs, histogram, offsets). +fn cpu_counting_sort( + flat_term_ids: &[u32], + flat_doc_ids: &[u32], + flat_tfs: &[u32], + vocab_size: usize, + pool: &rayon::ThreadPool, +) -> (Vec, Vec, Vec, Vec) { + let n = flat_term_ids.len(); + if n == 0 { + return (vec![], vec![], vec![0; vocab_size], vec![0; vocab_size]); + } + + let num_threads = pool.current_num_threads().max(1); + let chunk_size = n.div_ceil(num_threads); + + // Step 1: Per-thread histograms (parallel) + let thread_histograms: Vec> = pool.install(|| { + (0..num_threads) + .into_par_iter() + .map(|t| { + let start = t * chunk_size; + let end = (start + chunk_size).min(n); + let mut hist = vec![0u32; vocab_size]; + for i in start..end { + hist[flat_term_ids[i] as usize] += 1; + } + hist + }) + .collect() + }); + + // Step 2: Global histogram + per-thread prefix sums + let mut global_histogram = vec![0u32; vocab_size]; + for hist in &thread_histograms { + for v in 0..vocab_size { + global_histogram[v] += hist[v]; + } + } + + // Global offsets (prefix sum) + let mut global_offsets = vec![0u64; vocab_size]; + let mut running = 0u64; + for v in 0..vocab_size { + global_offsets[v] = running; + running += global_histogram[v] as u64; + } + + // Per-thread offsets: thread t's entries for term v start at + // global_offsets[v] + sum(thread_histograms[0..t][v]) + let mut thread_offsets: Vec> = Vec::with_capacity(num_threads); + let mut cumulative = vec![0u64; vocab_size]; + for hist in thread_histograms.iter().take(num_threads) { + let mut offsets = Vec::with_capacity(vocab_size); + for v in 0..vocab_size { + offsets.push(global_offsets[v] + cumulative[v]); + cumulative[v] += hist[v] as u64; + } + thread_offsets.push(offsets); + } + + // Step 3: Parallel scatter (stable — each thread writes in order) + let out_doc_ids = vec![0u32; n]; + let out_tfs = vec![0u32; n]; + + pool.install(|| { + // Each thread independently scatters its chunk + (0..num_threads).into_par_iter().for_each(|t| { + let start = t * chunk_size; + let end = (start + chunk_size).min(n); + // Thread-local copy of offsets (will be incremented) + let mut local_offsets = thread_offsets[t].clone(); + + for i in start..end { + let term = flat_term_ids[i] as usize; + let pos = local_offsets[term] as usize; + local_offsets[term] += 1; + // Safety: each thread writes to non-overlapping positions + // (guaranteed by the prefix sum computation) + unsafe { + let ptr_d = out_doc_ids.as_ptr() as *mut u32; + let ptr_t = out_tfs.as_ptr() as *mut u32; + *ptr_d.add(pos) = flat_doc_ids[i]; + *ptr_t.add(pos) = flat_tfs[i]; + } + } + }); + }); + + (out_doc_ids, out_tfs, global_histogram, global_offsets) +} + +// --------------------------------------------------------------------------- +// Main entry point +// --------------------------------------------------------------------------- + +pub fn add_documents_cuda( + _ctx: &CudaIndexer, + tokenizer: &crate::tokenizer::Tokenizer, + documents: &[&str], + existing_vocab: &HashMap, + base_doc_id: u32, +) -> Result { + let t_total = std::time::Instant::now(); + // Use more threads than default capped_pool — tokenization is embarrassingly parallel + let n_threads = std::thread::available_parallelism() + .map(|p| p.get()) + .unwrap_or(12); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(n_threads) + .build() + .expect("failed to build rayon pool"); + let num_threads = pool.current_num_threads().max(1); + + // ── Phase 1: Parallel tokenize + collect thread-local vocabs ── + let t0 = std::time::Instant::now(); + let chunk_size = documents.len().div_ceil(num_threads); + + struct ChunkResult { + tokenized: Vec<(u32, FxHashMap)>, + unique_tokens: FxHashSet, + total_tokens: u64, + } + + let chunks: Vec = pool.install(|| { + documents + .par_chunks(chunk_size.max(1)) + .map(|chunk| { + let mut tokenized = Vec::with_capacity(chunk.len()); + let mut unique_tokens: FxHashSet = FxHashSet::default(); + let mut total_tokens = 0u64; + let mut stem_cache: FxHashMap = FxHashMap::default(); + let mut tf_map: FxHashMap = FxHashMap::default(); + let mut buf: Vec = Vec::with_capacity(512); + + for doc in chunk { + // Fused tokenize + TF count: single pass, zero intermediate allocs + let doc_len = + tokenizer.tokenize_and_count(doc, &mut stem_cache, &mut tf_map, &mut buf); + total_tokens += doc_len as u64; + for key in tf_map.keys() { + if !unique_tokens.contains(key.as_str()) { + unique_tokens.insert(key.clone()); + } + } + // Move tf_map out, create fresh one reusing allocated memory + let finished_map = std::mem::take(&mut tf_map); + tokenized.push((doc_len, finished_map)); + } + + ChunkResult { + tokenized, + unique_tokens, + total_tokens, + } + }) + .collect() + }); + let t_tokenize = t0.elapsed(); + + // ── Phase 2: Merge thread-local vocabs (sequential but fast — only unique tokens) ── + let t0 = std::time::Instant::now(); + let mut vocab = existing_vocab.clone(); + let mut next_id = vocab.len() as u32; + + for chunk in &chunks { + for token in &chunk.unique_tokens { + if !vocab.contains_key(token.as_str()) { + vocab.insert(token.clone(), next_id); + next_id += 1; + } + } + } + let vocab_size = vocab.len(); + let t_vocab = t0.elapsed(); + + // ── Phase 3: Parallel flatten ── + let t0 = std::time::Instant::now(); + + // Compute doc_lengths and total_tokens from chunks + let mut doc_lengths: Vec = Vec::with_capacity(documents.len()); + let mut total_tokens: u64 = 0; + let mut chunk_doc_starts: Vec = Vec::with_capacity(chunks.len()); + let mut total_entries = 0usize; + // Also compute per-doc offsets into flat arrays + let mut doc_flat_offsets: Vec = Vec::with_capacity(documents.len()); + + for chunk in &chunks { + chunk_doc_starts.push(doc_lengths.len()); + total_tokens += chunk.total_tokens; + for (doc_len, tf_map) in &chunk.tokenized { + doc_lengths.push(*doc_len); + doc_flat_offsets.push(total_entries); + total_entries += tf_map.len(); + } + } + + // Parallel fill flat arrays + let flat_term_ids = vec![0u32; total_entries]; + let flat_doc_ids = vec![0u32; total_entries]; + let flat_tfs = vec![0u32; total_entries]; + + pool.install(|| { + chunks + .par_iter() + .enumerate() + .for_each(|(chunk_idx, chunk)| { + let doc_start = chunk_doc_starts[chunk_idx]; + for (local_i, (_, tf_map)) in chunk.tokenized.iter().enumerate() { + let doc_idx = doc_start + local_i; + let doc_id = base_doc_id + doc_idx as u32; + let offset = doc_flat_offsets[doc_idx]; + for (j, (token, &tf)) in tf_map.iter().enumerate() { + let term_id = vocab[token.as_str()]; + unsafe { + let ptr_t = flat_term_ids.as_ptr() as *mut u32; + let ptr_d = flat_doc_ids.as_ptr() as *mut u32; + let ptr_f = flat_tfs.as_ptr() as *mut u32; + *ptr_t.add(offset + j) = term_id; + *ptr_d.add(offset + j) = doc_id; + *ptr_f.add(offset + j) = tf; + } + } + } + }); + }); + let t_flatten = t0.elapsed(); + + // ── Phase 4: Stable counting sort (CPU) — preserves doc_id order ── + let t0 = std::time::Instant::now(); + let (out_doc_ids, out_tfs, histogram, offsets) = + cpu_counting_sort(&flat_term_ids, &flat_doc_ids, &flat_tfs, vocab_size, &pool); + let t_sort = t0.elapsed(); + + // ── Phase 5: Build posting Vec> from sorted flat output (parallel) ── + let t0 = std::time::Instant::now(); + let postings: Vec> = pool.install(|| { + (0..vocab_size) + .into_par_iter() + .map(|term_id| { + let start = offsets[term_id] as usize; + let count = histogram[term_id] as usize; + if count == 0 { + return Vec::new(); + } + let mut plist = Vec::with_capacity(count); + for j in start..start + count { + plist.push((out_doc_ids[j], out_tfs[j])); + } + // No sort needed — counting sort is stable, doc_ids already in order + plist + }) + .collect() + }); + let t_postings = t0.elapsed(); + + // ── Phase 6: Parallel drop of intermediate data (avoids single-threaded dealloc) ── + let t0 = std::time::Instant::now(); + // Drop chunks (contains millions of FxHashMap + FxHashSet) + // in parallel to avoid single-threaded deallocation bottleneck (~8s for 8.8M docs) + pool.install(|| { + chunks.into_par_iter().for_each(drop); + }); + // Drop flat arrays in parallel + pool.install(|| { + rayon::join( + || drop(flat_term_ids), + || rayon::join(|| drop(flat_doc_ids), || drop(flat_tfs)), + ); + }); + // Drop counting sort output + pool.install(|| { + rayon::join(|| drop(out_doc_ids), || drop(out_tfs)); + }); + let t_drop = t0.elapsed(); + + let t_total_elapsed = t_total.elapsed(); + if std::env::var("BM25X_PROFILE").is_ok() { + eprintln!( + "[bm25x-cuda] {:.3}s total | tok={:.3}s voc={:.3}s flat={:.3}s sort={:.3}s post={:.3}s drop={:.3}s | {:.0} d/s", + t_total_elapsed.as_secs_f64(), + t_tokenize.as_secs_f64(), + t_vocab.as_secs_f64(), + t_flatten.as_secs_f64(), + t_sort.as_secs_f64(), + t_postings.as_secs_f64(), + t_drop.as_secs_f64(), + documents.len() as f64 / t_total_elapsed.as_secs_f64(), + ); + } + + Ok(AddResult { + postings, + doc_lengths, + vocab, + total_tokens, + }) +} + +pub struct AddResult { + pub postings: Vec>, + pub doc_lengths: Vec, + pub vocab: HashMap, + pub total_tokens: u64, +} diff --git a/src/index.rs b/src/index.rs index bffb1e5..6a03b57 100644 --- a/src/index.rs +++ b/src/index.rs @@ -10,7 +10,7 @@ use crate::storage::MmapData; use crate::tokenizer::{Tokenizer, TokenizerMode}; /// Cap rayon parallelism at 12 threads max. -fn capped_pool() -> rayon::ThreadPool { +pub(crate) fn capped_pool() -> rayon::ThreadPool { let n = std::thread::available_parallelism() .map(|p| p.get().min(12)) .unwrap_or(4); @@ -160,12 +160,43 @@ impl BM25 { } /// Add documents to the index. Returns the assigned document indices. + /// + /// When compiled with the `cuda` feature and a GPU is available, posting list + /// construction is offloaded to CUDA for faster indexing. Tokenization always + /// runs on CPU (rayon). Falls back to CPU automatically on GPU failure. pub fn add(&mut self, documents: &[&str]) -> io::Result> { if self.mmap_data.is_some() { self.materialize_mmap(); } + // Try GPU path + #[cfg(feature = "cuda")] + { + if let Some(ctx) = crate::cuda::get_global_context() { + match crate::cuda::add_documents_cuda( + &ctx, + &self.tokenizer, + documents, + &self.vocab, + self.num_docs, + ) { + Ok(result) => return self.apply_cuda_result(documents.len(), result), + Err(e) => { + eprintln!("[bm25x] CUDA indexing failed: {}, falling back to CPU", e); + } + } + } + } + + self.add_cpu(documents) + } + + /// CPU-only add implementation (original algorithm). + fn add_cpu(&mut self, documents: &[&str]) -> io::Result> { + let t_total = std::time::Instant::now(); + // Phase 1: tokenize + compute TF maps in parallel + let t0 = std::time::Instant::now(); let pool = capped_pool(); let tokenized: Vec<(u32, HashMap)> = pool.install(|| { documents @@ -181,8 +212,10 @@ impl BM25 { }) .collect() }); + let t_tokenize = t0.elapsed(); // Phase 2: merge into index sequentially (cheap — just HashMap lookups + Vec pushes) + let t0 = std::time::Instant::now(); let base_id = self.num_docs; let mut ids = Vec::with_capacity(documents.len()); @@ -199,8 +232,63 @@ impl BM25 { self.total_tokens += doc_len as u64; ids.push(doc_id as usize); } + let t_merge = t0.elapsed(); self.num_docs = base_id + documents.len() as u32; + + let t_total_elapsed = t_total.elapsed(); + if std::env::var("BM25X_PROFILE").is_ok() { + eprintln!( + "[bm25x-cpu] {:.3}s total | tokenize={:.3}s merge={:.3}s | {:.0} d/s", + t_total_elapsed.as_secs_f64(), + t_tokenize.as_secs_f64(), + t_merge.as_secs_f64(), + documents.len() as f64 / t_total_elapsed.as_secs_f64(), + ); + } + + self.auto_save()?; + Ok(ids) + } + + /// Apply the result from CUDA-accelerated indexing into the BM25 struct. + #[cfg(feature = "cuda")] + fn apply_cuda_result( + &mut self, + num_new_docs: usize, + result: crate::cuda::AddResult, + ) -> io::Result> { + let base_id = self.num_docs; + + if self.num_docs == 0 && self.postings.is_empty() { + // Fresh index: move everything directly (O(1), no copying) + self.doc_freqs = result.postings.iter().map(|p| p.len() as u32).collect(); + self.postings = result.postings; + self.vocab = result.vocab; + } else { + // Existing index: merge + let new_vocab_size = result.vocab.len(); + while self.postings.len() < new_vocab_size { + self.postings.push(Vec::new()); + self.doc_freqs.push(0); + } + for (term_id, gpu_plist) in result.postings.into_iter().enumerate() { + if !gpu_plist.is_empty() { + let count = gpu_plist.len() as u32; + self.postings[term_id].extend_from_slice(&gpu_plist); + self.doc_freqs[term_id] += count; + } + } + if result.vocab.len() > self.vocab.len() { + self.vocab = result.vocab; + } + } + + self.doc_lengths.extend_from_slice(&result.doc_lengths); + self.total_tokens += result.total_tokens; + self.num_docs = base_id + num_new_docs as u32; + + let ids: Vec = (base_id as usize..base_id as usize + num_new_docs).collect(); self.auto_save()?; Ok(ids) } @@ -691,6 +779,90 @@ impl BM25 { self.num_docs = num_docs; self.mmap_data = Some(mmap_data); } + + /// Create a GPU search index for fast search. Call once after indexing. + #[cfg(feature = "cuda")] + pub fn to_gpu_search_index(&self) -> Result { + let ctx = + crate::cuda::get_global_context().ok_or_else(|| "CUDA not available".to_string())?; + + // If mmap-backed, need to materialize posting lists + let postings_ref: Vec>; + let postings = if let Some(ref mmap) = self.mmap_data { + postings_ref = (0..self.vocab.len() as u32) + .map(|t| { + let mut entries = Vec::new(); + mmap.for_each_posting(t, &mut |doc_id, tf| entries.push((doc_id, tf))); + entries + }) + .collect(); + &postings_ref + } else { + &self.postings + }; + + let doc_lengths = if let Some(ref mmap) = self.mmap_data { + mmap.all_doc_lengths() + } else { + self.doc_lengths.clone() + }; + + crate::cuda::GpuSearchIndex::from_index(&ctx, postings, &doc_lengths, self.num_docs) + } + + /// Search using GPU-resident index. Much faster than CPU search for large indices. + #[cfg(feature = "cuda")] + pub fn search_gpu( + &self, + gpu_index: &mut crate::cuda::GpuSearchIndex, + query: &str, + k: usize, + ) -> Vec { + if self.num_docs == 0 { + return Vec::new(); + } + + let ctx = match crate::cuda::get_global_context() { + Some(c) => c, + None => return self.search(query, k), + }; + + let query_tokens = self.tokenizer.tokenize_owned(query); + let params = ScoringParams { + k1: self.k1, + b: self.b, + delta: self.delta, + avgdl: self.total_tokens as f32 / self.num_docs as f32, + }; + + // Build (term_id, idf) pairs for query + let mut seen = std::collections::HashSet::new(); + let mut query_terms: Vec<(u32, f32)> = Vec::new(); + for token in &query_tokens { + if let Some(&term_id) = self.vocab.get(token.as_str()) { + if seen.insert(term_id) { + let df = self.doc_freqs.get(term_id as usize).copied().unwrap_or(0); + if df > 0 { + query_terms.push((term_id, scoring::idf(self.method, self.num_docs, df))); + } + } + } + } + + match gpu_index.search(&ctx, &query_terms, params.k1, params.b, params.avgdl, k) { + Ok(results) => results + .into_iter() + .map(|(doc_id, score)| SearchResult { + index: doc_id as usize, + score, + }) + .collect(), + Err(e) => { + eprintln!("[bm25x] GPU search failed: {}, falling back to CPU", e); + self.search(query, k) + } + } + } } impl Default for BM25 { diff --git a/src/lib.rs b/src/lib.rs index 08fb0aa..121727a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,21 @@ pub mod scoring; pub mod storage; pub mod tokenizer; +#[cfg(feature = "cuda")] +pub mod cuda; + pub use index::{SearchResult, BM25}; pub use scoring::{Method, ScoringParams}; pub use tokenizer::TokenizerMode; + +/// Returns true if the library was compiled with CUDA support and a GPU is available. +pub fn is_gpu_available() -> bool { + #[cfg(feature = "cuda")] + { + cuda::is_cuda_available() + } + #[cfg(not(feature = "cuda"))] + { + false + } +} diff --git a/src/tokenizer.rs b/src/tokenizer.rs index 45662ba..590fb82 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -1,4 +1,7 @@ +use std::collections::HashMap; + use rust_stemmers::{Algorithm, Stemmer}; +use rustc_hash::FxHashMap; use unicode_normalization::UnicodeNormalization; /// Tokenizer mode controls text preprocessing. @@ -214,6 +217,205 @@ impl Tokenizer { tokens } + /// Tokenize with a thread-local stem cache for amortized O(1) stemming. + /// + /// The cache maps pre-stem tokens → stemmed form. For large corpora this + /// avoids redundant Snowball stemming of repeated words across documents. + /// ~500K unique English tokens means ~500K stems cached, vs ~500M stem calls. + /// Tokenize with std HashMap stem cache (for non-CUDA paths). + pub fn tokenize_cached( + &self, + text: &str, + stem_cache: &mut HashMap, + ) -> Vec { + self.tokenize_with_cache_impl(text, stem_cache) + } + + /// Tokenize with FxHashMap stem cache (faster hashing, used by CUDA path). + pub fn tokenize_cached_fx( + &self, + text: &str, + stem_cache: &mut FxHashMap, + ) -> Vec { + self.tokenize_with_cache_impl(text, stem_cache) + } + + fn tokenize_with_cache_impl( + &self, + text: &str, + stem_cache: &mut HashMap, + ) -> Vec { + let lowered = text.to_lowercase(); + + let normalized: String = match self.mode { + TokenizerMode::Unicode | TokenizerMode::UnicodeStem => fold_to_ascii(&lowered), + _ => lowered, + }; + + let raw_tokens = split_alphanumeric(&normalized); + + let mut tokens = Vec::with_capacity(raw_tokens.len()); + for token in raw_tokens { + if !self.should_keep(&token) { + continue; + } + let final_token = if let Some(ref stemmer) = self.stemmer { + if let Some(cached) = stem_cache.get(&token) { + cached.clone() + } else { + let stemmed = stemmer.stem(&token).into_owned(); + stem_cache.insert(token, stemmed.clone()); + stemmed + } + } else { + token + }; + if !final_token.is_empty() { + tokens.push(final_token); + } + } + tokens + } + + /// Fused tokenize + TF count: single-pass ASCII fast path, zero intermediate allocations. + /// + /// For ASCII text (99%+ of English corpora), this: + /// - Does in-place lowercase in a reusable buffer (no String alloc) + /// - Skips fold_to_ascii entirely (text is already ASCII) + /// - Splits + stopword-filters + stem-caches + TF-counts in ONE pass + /// - Avoids allocating Vec for intermediate tokens + /// + /// For non-ASCII text, falls back to the standard 3-pass pipeline. + pub fn tokenize_and_count( + &self, + text: &str, + stem_cache: &mut HashMap, + tf_map: &mut HashMap, + buf: &mut Vec, + ) -> u32 { + tf_map.clear(); + + if text.is_ascii() { + self.tokenize_count_ascii(text, stem_cache, tf_map, buf) + } else { + self.tokenize_count_unicode(text, stem_cache, tf_map) + } + } + + /// ASCII fast path: zero-heap-alloc single-pass tokenization. + /// + /// Lowercases on-the-fly into a stack-allocated 128-byte token buffer. + /// Never touches the heap for text processing — only for stem cache + /// misses and new TF map entries. + fn tokenize_count_ascii( + &self, + text: &str, + stem_cache: &mut HashMap, + tf_map: &mut HashMap, + _buf: &mut Vec, + ) -> u32 { + let bytes = text.as_bytes(); + let mut doc_len = 0u32; + // Stack-allocated token buffer — no heap allocation per token. + // 128 bytes covers all English words (longest common word ~30 chars). + let mut token_buf: [u8; 128] = [0; 128]; + let mut token_len: usize = 0; + + for &b in bytes { + if b.is_ascii_alphanumeric() { + if token_len < 128 { + // Branchless ASCII lowercase: 'A'-'Z' → 'a'-'z', others unchanged. + // b | 0x20 works for letters but corrupts digits, so use conditional. + token_buf[token_len] = b.to_ascii_lowercase(); + token_len += 1; + } + // Tokens > 128 chars: silently truncate (never happens in practice) + } else if token_len > 0 { + let token = unsafe { std::str::from_utf8_unchecked(&token_buf[..token_len]) }; + doc_len += self.count_token(token, stem_cache, tf_map); + token_len = 0; + } + } + // Handle last token (no trailing separator) + if token_len > 0 { + let token = unsafe { std::str::from_utf8_unchecked(&token_buf[..token_len]) }; + doc_len += self.count_token(token, stem_cache, tf_map); + } + doc_len + } + + /// Unicode fallback: standard 3-pass pipeline fused with TF counting. + fn tokenize_count_unicode( + &self, + text: &str, + stem_cache: &mut HashMap, + tf_map: &mut HashMap, + ) -> u32 { + let lowered = text.to_lowercase(); + let normalized: String = match self.mode { + TokenizerMode::Unicode | TokenizerMode::UnicodeStem => fold_to_ascii(&lowered), + _ => lowered, + }; + // After fold_to_ascii, result is ASCII — use the fast split+count loop + let bytes = normalized.as_bytes(); + let mut doc_len = 0u32; + let mut start: Option = None; + + for i in 0..bytes.len() { + if bytes[i].is_ascii_alphanumeric() { + if start.is_none() { + start = Some(i); + } + } else if let Some(s) = start { + doc_len += self.count_token(&normalized[s..i], stem_cache, tf_map); + start = None; + } + } + if let Some(s) = start { + doc_len += self.count_token(&normalized[s..], stem_cache, tf_map); + } + doc_len + } + + /// Process a single token: stopword check → stem cache → TF map increment. + /// Returns 1 if token was counted, 0 if filtered. + #[inline] + fn count_token( + &self, + token: &str, + stem_cache: &mut HashMap, + tf_map: &mut HashMap, + ) -> u32 { + if !self.should_keep(token) { + return 0; + } + + if let Some(ref stemmer) = self.stemmer { + if let Some(stem) = stem_cache.get(token) { + // Hot path: stem cached. Try TF map lookup first to avoid clone. + if let Some(count) = tf_map.get_mut(stem.as_str()) { + *count += 1; + } else { + tf_map.insert(stem.clone(), 1); + } + } else { + // Cold path: new token, compute stem + let stemmed = stemmer.stem(token).into_owned(); + if !stemmed.is_empty() { + stem_cache.insert(token.to_string(), stemmed.clone()); + *tf_map.entry(stemmed).or_insert(0) += 1; + } else { + return 0; + } + } + } else if let Some(count) = tf_map.get_mut(token) { + *count += 1; + } else { + tf_map.insert(token.to_string(), 1); + } + 1 + } + #[inline] fn should_keep(&self, token: &str) -> bool { if token.is_empty() {