diff --git a/similarities/bench.py b/similarities/bench.py index 5314a9b..a378da3 100644 --- a/similarities/bench.py +++ b/similarities/bench.py @@ -27,10 +27,14 @@ - Edit Distance baselines: rapidfuzz, python-Levenshtein, jellyfish, editdistance, nltk, edlib, polyleven - StringZilla cross-product: szs.LevenshteinDistances, szs.LevenshteinDistancesUTF8, szs.NeedlemanWunschScores, szs.SmithWatermanScores +- Bounded membership (`within_k`): szs.LevenshteinWithinK cross-product vs cutoff baselines + (polyleven with bound, rapidfuzz with score_cutoff), reported in cmp/s, not CUPS, as bounded + kernels never fill the full DP matrix - BioPython: PairwiseAligner baseline (unary match/mismatch scoring) - cuDF: GPU-accelerated edit distance (optional) -Environment variables (identical to bench.rs / the C++ harness): +Common environment variables mirror bench.rs / the C++ harness; the bounded-only controls are +specific to this Python benchmark: - STRINGWARS_DATASET: Path to the input dataset file - STRINGWARS_TOKENS: Tokenization mode ('lines', 'words', 'file') - STRINGWARS_MAX_TOKENS: Limit on the number of tokens loaded @@ -39,6 +43,14 @@ - STRINGWARS_WARMUP: Uncounted warm-up budget per variant (seconds) - STRINGWARS_SEED: Seed for the token shuffle - STRINGWARS_FILTER: Regex selecting which benchmark variants run +- STRINGWARS_WITHIN_K: Comma-separated bounds for the `within_k` category (default: "1,2,4") +- STRINGWARS_WITHIN_CANDIDATES: Candidate construction for `within_k`: 'random' (default, + reject-heavy natural mix), 'sparse' (one guaranteed diagonal match per query), 'dense' (all + pairs guaranteed within the bound), or 'all' + +Every `within_k` (category, bound) pair is verified before benchmarking: the szs boolean matrix on +a fixed slice is compared cell-by-cell against a byte-level rapidfuzz oracle, and a single mismatch +aborts the run — throughput of a wrong matrix is worthless. A CPU core counts as one core; a GPU streaming multiprocessor (SM) counts as one core. The per-device pair budget is ``STRINGWARS_BATCH_PER_CORE * cores``, and the square cross-product side @@ -80,6 +92,7 @@ # Edit-distance baselines (each optional so a missing wheel skips just its row). try: from rapidfuzz.distance import Levenshtein as rapidfuzz_levenshtein + from rapidfuzz.process import cdist as rapidfuzz_cdist RAPIDFUZZ_AVAILABLE = True except ImportError: @@ -202,13 +215,15 @@ def measure_crossproduct( total_bytes: int, warmup_seconds: float, time_limit_seconds: float, + report: str = "cups", ) -> None: """Run `compute` (one full cross-product per call) for a wall-time budget and report CUPS. The matrix/output buffer lives inside `compute`'s closure and is reused across iterations, so no allocation happens in the measured loop. After an uncounted warm-up the kernel is cycled until the deadline; throughput is the TRUE aggregate cell count divided by elapsed time. Mirrors the - Rust `measure_throughput` cross-product blocks. + Rust `measure_throughput` cross-product blocks. `report` selects the `report_stats` unit; for + bounded membership kernels pass "comparisons" with `total_cells` holding the pair count. """ # Uncounted warm-up so caches and CPU frequency settle. if warmup_seconds > 0: @@ -233,7 +248,7 @@ def measure_crossproduct( elapsed_seconds = (now_nanoseconds() - start_nanoseconds) / 1e9 processed_cells = total_cells * iterations processed_bytes = total_bytes * iterations - report_stats(name, "cups", elapsed_seconds, processed_cells, processed_bytes) + report_stats(name, report, elapsed_seconds, processed_cells, processed_bytes) def measure_pairwise_baseline( @@ -248,6 +263,7 @@ def measure_pairwise_baseline( candidate_byte_lengths: np.ndarray, warmup_seconds: float, time_limit_seconds: float, + report: str = "cups", ) -> None: """Benchmark a one-pair-at-a-time baseline along the cross-product diagonal. @@ -255,6 +271,8 @@ def measure_pairwise_baseline( measured exactly as bench.rs measures them: one ``(query[index], candidate[index])`` pair per call, cycling the diagonal ``index in [0, side)``. Each call's true cells (length product) are accumulated, so the reported CUPS are directly comparable to the StringZilla matrix engines. + With ``report="comparisons"`` each call counts as one pair instead — the right denominator for + bounded membership baselines that never fill the full DP matrix. """ if side <= 0: print(f"{name}: No pairs to process") @@ -283,7 +301,10 @@ def measure_pairwise_baseline( try: while True: checksum += int(scalar_function(queries[pair_index], candidates[pair_index])) - processed_cells += int(query_lengths[pair_index]) * int(candidate_lengths[pair_index]) + if report == "cups": + processed_cells += int(query_lengths[pair_index]) * int(candidate_lengths[pair_index]) + else: + processed_cells += 1 processed_bytes += int(query_byte_lengths[pair_index]) + int(candidate_byte_lengths[pair_index]) pair_index = (pair_index + 1) % side countdown -= 1 @@ -301,7 +322,7 @@ def measure_pairwise_baseline( return elapsed_seconds = (now_nanoseconds() - start_nanoseconds) / 1e9 - report_stats(name, "cups", elapsed_seconds, processed_cells, processed_bytes) + report_stats(name, report, elapsed_seconds, processed_cells, processed_bytes) print(f" {name} checksum={checksum}", file=sys.stderr) @@ -622,6 +643,413 @@ def compute(): print(f" {name} checksum={checksum}", file=sys.stderr) +def parse_within_k_bounds() -> list[int]: + """Bounds for the `within_k` category from STRINGWARS_WITHIN_K (default "1,2,4").""" + raw = os.environ.get("STRINGWARS_WITHIN_K", "1,2,4") + bounds = sorted({int(part.strip()) for part in raw.split(",") if part.strip()}) + return bounds or [2] + + +def mutate_token(token: str, edits: int, rng: random.Random) -> str: + """Copy of `token` with `edits` random unit-cost mutations (substitution / insertion / deletion). + + The result is guaranteed to differ from `token` and to be within `edits` of it, so pairs built + this way always satisfy the `within_k` bound. In the cross-product workload this guarantees + one accepted diagonal pair per query; it does not make the full matrix accept-heavy. + """ + alphabet = "abcdefghijklmnopqrstuvwxyz" + mutated = list(token) + for _ in range(max(1, edits)): + operation = rng.choice("sid") if mutated else "i" + position = rng.randrange(len(mutated)) if mutated else 0 + if operation == "s": + mutated[position] = rng.choice(alphabet) + elif operation == "i": + mutated.insert(position, rng.choice(alphabet)) + else: + del mutated[position] + result = "".join(mutated) + return result if result != token else token + "x" + + +def within_k_candidate_tokens( + tokens: Sequence, + side: int, + bound: int, + candidate_mode: str, + rng: random.Random, +) -> tuple[list, list]: + """Disjoint (queries, candidates) token lists for one `within_k` variant. + + Random mode takes the disjoint slices ``[0, side)`` and ``[side, 2 * side)`` — the reject-heavy + natural mix. Sparse mode derives every candidate from its diagonal query with `bound` edits, + so diagonal pairs accept while off-diagonal pairs stay random. Dense mode repeats one query and + mutates every candidate from it, guaranteeing that every cell accepts. + """ + query_tokens = list(tokens[0:side]) + if candidate_mode == "sparse": + candidate_tokens = [mutate_token(query_tokens[index], bound, rng) for index in range(side)] + elif candidate_mode == "dense": + query_tokens = [query_tokens[0]] * side + candidate_tokens = [mutate_token(query_tokens[0], bound, rng) for _ in range(side)] + else: + candidate_tokens = list(tokens[side : 2 * side]) + return query_tokens, candidate_tokens + + +def within_k_candidate_bytes( + candidate_tokens: list, byte_lengths: np.ndarray, side: int, candidate_mode: str +) -> np.ndarray: + """Per-candidate UTF-8 byte counts, recomputed when candidates were mutated.""" + if candidate_mode == "random": + return byte_lengths[side : 2 * side] + return np.fromiter((len(token.encode("utf-8")) for token in candidate_tokens), dtype=np.int64, count=side) + + +def benchmark_stringzillas_within( + tokens: Sequence, + device_variants: list[DeviceVariant], + category: str, + engine_name: str, + bound: int, + byte_lengths: np.ndarray, + candidate_mode: str, + seed: int, + warmup_seconds: float, + time_limit_seconds: float, + filter_pattern: re.Pattern | None, +) -> None: + """Cross-product benchmark for the bounded szs.LevenshteinWithinK membership engine. + + Mirrors `benchmark_stringzillas_distances`, but the engine is constructed with `bound`, the + output matrix is boolean, and throughput is reported in cmp/s (pairs per second): a bounded + kernel never fills the full DP matrix, so CUPS would misstate the work. Elements are the true + pair count ``side * side`` per cross-product. + """ + for variant in device_variants: + full_name = f"{engine_name}_k{bound}[{candidate_mode}]{variant.label}<{variant.side}^2,reuse>" + allocating_name = f"{engine_name}_k{bound}[{candidate_mode}]{variant.label}<{variant.side}^2,alloc>" + run_reuse = should_run(f"{category}/{full_name}", filter_pattern) + run_allocating = should_run(f"{category}/{allocating_name}", filter_pattern) + if not run_reuse and not run_allocating: + continue + + side = variant.side + rng = random.Random(f"{seed}:{category}:{bound}:{side}") + query_tokens, candidate_tokens = within_k_candidate_tokens(tokens, side, bound, candidate_mode, rng) + queries = sz.Strs(query_tokens) + candidates = sz.Strs(candidate_tokens) + + try: + engine = szs.LevenshteinWithinK(bound=bound, capabilities=variant.scope) + except Exception as creation_error: + print(f"{full_name}: SKIPPED ({creation_error})") + continue + + if not _crossproduct_supported(engine, queries, candidates): + print(f"{full_name}: SKIPPED (installed stringzillas lacks the queries x candidates cross-product API)") + continue + + total_pairs = side * side + candidate_bytes = within_k_candidate_bytes(candidate_tokens, byte_lengths, side, candidate_mode) + query_bytes = np.fromiter((len(token.encode("utf-8")) for token in query_tokens), dtype=np.int64, count=side) + total_bytes = int(query_bytes.sum()) + int(candidate_bytes.sum()) + matrix = np.zeros((side, side), dtype=np.bool_) + + def compute(engine=engine, queries=queries, candidates=candidates, scope=variant.scope, matrix=matrix): + engine(queries, candidates, scope, out=matrix) + + # Mirror the distance engines: a backend that declines surfaces as an exception we SKIP on. + try: + compute() + except Exception as compute_error: + print(f"{full_name}: SKIPPED ({compute_error})") + continue + + if run_reuse: + measure_crossproduct( + full_name, compute, total_pairs, total_bytes, warmup_seconds, time_limit_seconds, report="comparisons" + ) + + # RapidFuzz's cdist API allocates its result on every call. Publish an allocation-inclusive + # StringZilla row as the direct comparison and keep the reuse row as a separately labelled + # steady-state measurement for callers that provide an output buffer. + if run_allocating: + + def compute_allocating(engine=engine, queries=queries, candidates=candidates, scope=variant.scope): + return engine(queries, candidates, scope) + + measure_crossproduct( + allocating_name, + compute_allocating, + total_pairs, + total_bytes, + warmup_seconds, + time_limit_seconds, + report="comparisons", + ) + + +def benchmark_within_k_baselines( + tokens: Sequence, + device_variants: list[DeviceVariant], + bound: int, + byte_lengths: np.ndarray, + candidate_mode: str, + seed: int, + category: str, + warmup_seconds: float, + time_limit_seconds: float, + filter_pattern: re.Pattern | None, +) -> None: + """Cutoff-capable baselines for `within_k`: polyleven with a bound, rapidfuzz with score_cutoff. + + Both return `bound + 1` when the distance exceeds the bound, so the scalar function folds the + result to a boolean membership count and throughput is reported in cmp/s, matching the + StringZilla membership engine. + """ + baseline_side = device_variants[0].side + rng = random.Random(f"{seed}:{category}:{bound}:{baseline_side}") + query_tokens, candidate_tokens = within_k_candidate_tokens(tokens, baseline_side, bound, candidate_mode, rng) + candidate_bytes = within_k_candidate_bytes(candidate_tokens, byte_lengths, baseline_side, candidate_mode) + query_bytes = np.fromiter( + (len(token.encode("utf-8")) for token in query_tokens), dtype=np.int64, count=baseline_side + ) + encoded_queries = [token.encode("utf-8") for token in query_tokens] + encoded_candidates = [token.encode("utf-8") for token in candidate_tokens] + + def run(name: str, scalar_function: Callable[[Any, Any], int]): + if not should_run(f"{category}/{name}", filter_pattern): + return + measure_pairwise_baseline( + name, + scalar_function, + encoded_queries, + encoded_candidates, + baseline_side, + query_bytes, + candidate_bytes, + query_bytes, + candidate_bytes, + warmup_seconds, + time_limit_seconds, + report="comparisons", + ) + + if RAPIDFUZZ_AVAILABLE: + + def rapidfuzz_within(first_string: bytes, second_string: bytes) -> int: + return 1 if rapidfuzz_levenshtein.distance(first_string, second_string, score_cutoff=bound) <= bound else 0 + + run(f"rapidfuzz.Levenshtein.distance_k{bound}[{candidate_mode}]<1cpu,scalar,bytes>", rapidfuzz_within) + if POLYLEVEN_AVAILABLE and all(token.isascii() for token in query_tokens + candidate_tokens): + + def polyleven_within(first_string: str, second_string: str) -> int: + return 1 if polyleven.levenshtein(first_string, second_string, bound) <= bound else 0 + + # polyleven accepts text rather than arbitrary byte strings. It is comparable only for + # ASCII inputs, where codepoint and UTF-8 byte semantics coincide. + def run_polyleven(name: str): + if not should_run(f"{category}/{name}", filter_pattern): + return + measure_pairwise_baseline( + name, + polyleven_within, + query_tokens, + candidate_tokens, + baseline_side, + query_bytes, + candidate_bytes, + query_bytes, + candidate_bytes, + warmup_seconds, + time_limit_seconds, + report="comparisons", + ) + + run_polyleven(f"polyleven.levenshtein_k{bound}[{candidate_mode}]") + + # Batched baseline: use the exact same deterministic inputs, dimensions, and CPU scope as each + # StringZilla CPU variant. cdist allocates the result, so compare this row to StringZilla's + # explicitly labelled `alloc` row rather than its caller-buffer `reuse` row. + if RAPIDFUZZ_AVAILABLE: + for variant in device_variants: + if "gpu" in variant.label: + continue + cdist_side = variant.side + cdist_rng = random.Random(f"{seed}:{category}:{bound}:{cdist_side}") + cdist_queries_text, cdist_candidates_text = within_k_candidate_tokens( + tokens, cdist_side, bound, candidate_mode, cdist_rng + ) + cdist_queries = [token.encode("utf-8") for token in cdist_queries_text] + cdist_candidates = [token.encode("utf-8") for token in cdist_candidates_text] + total_pairs = cdist_side * cdist_side + cdist_candidate_bytes = within_k_candidate_bytes( + cdist_candidates_text, byte_lengths, cdist_side, candidate_mode + ) + cdist_query_bytes = sum(len(token) for token in cdist_queries) + total_bytes = cdist_query_bytes + int(cdist_candidate_bytes.sum()) + workers = 1 if variant.label == "<1cpu>" else -1 + raw_name = ( + f"rapidfuzz.process.cdist_k{bound}[{candidate_mode}]" + f"{variant.label}<{cdist_side}^2,uint8-distance>" + ) + membership_name = ( + f"rapidfuzz.process.cdist_k{bound}[{candidate_mode}]" + f"{variant.label}<{cdist_side}^2,bool-membership>" + ) + + def compute_raw(workers=workers, cdist_queries=cdist_queries, cdist_candidates=cdist_candidates): + return rapidfuzz_cdist( + cdist_queries, + cdist_candidates, + scorer=rapidfuzz_levenshtein.distance, + score_cutoff=bound, + dtype=np.uint8, + workers=workers, + ) + + if should_run(f"{category}/{raw_name}", filter_pattern): + try: + compute_raw() + except Exception as compute_error: + print(f"{raw_name}: SKIPPED ({compute_error})") + else: + measure_crossproduct( + raw_name, + compute_raw, + total_pairs, + total_bytes, + warmup_seconds, + time_limit_seconds, + report="comparisons", + ) + + # This is the like-for-like public result: RapidFuzz's batched API returns cutoff + # distances, so include the conversion needed to produce the boolean membership matrix. + if should_run(f"{category}/{membership_name}", filter_pattern): + + def compute_membership(compute_raw=compute_raw, bound=bound): + return compute_raw() <= bound + + try: + compute_membership() + except Exception as compute_error: + print(f"{membership_name}: SKIPPED ({compute_error})") + else: + measure_crossproduct( + membership_name, + compute_membership, + total_pairs, + total_bytes, + warmup_seconds, + time_limit_seconds, + report="comparisons", + ) + + +def verify_within_k( + tokens: Sequence, + bound: int, + candidate_mode: str, + category: str, + seed: int, +) -> None: + """Oracle check: the szs boolean matrix must match a byte-level rapidfuzz reference. + + The engine is byte-metric, so the oracle scores UTF-8 encodings (not codepoints). Runs once + per (category, bound) on a fixed 512 x 512 slice with its own seed, independent of the timed + loops. Any mismatch aborts the run: throughput of a wrong matrix is worthless. + """ + verify_side = min(512, len(tokens) // 2) + if verify_side < 2 or not RAPIDFUZZ_AVAILABLE: + return + rng = random.Random(f"{seed}:verify:{category}:{bound}") + query_tokens, candidate_tokens = within_k_candidate_tokens(tokens, verify_side, bound, candidate_mode, rng) + ours = np.asarray( + szs.LevenshteinWithinK(bound=bound)(sz.Strs(query_tokens), sz.Strs(candidate_tokens)), dtype=np.bool_ + ) + query_bytes = [token.encode("utf-8") for token in query_tokens] + candidate_bytes = [token.encode("utf-8") for token in candidate_tokens] + reference = np.array( + [ + [rapidfuzz_levenshtein.distance(q, c, score_cutoff=bound) <= bound for c in candidate_bytes] + for q in query_bytes + ], + dtype=np.bool_, + ) + mismatches = int(np.count_nonzero(ours != reference)) + print(f" verify within_k k={bound} [{category}]: {ours.size} cells, {mismatches} mismatches", file=sys.stderr) + if mismatches: + raise SystemExit(f"within_k verification FAILED at bound={bound} ({category})") + + +def perform_within_k_benchmarks( + tokens: Sequence, + device_variants: list[DeviceVariant], + byte_lengths: np.ndarray, + warmup_seconds: float, + time_limit_seconds: float, + filter_pattern: re.Pattern | None, + seed: int, +) -> None: + """Bounded-membership group: szs.LevenshteinWithinK vs cutoff baselines, per bound. + + STRINGWARS_WITHIN_K selects the bounds; STRINGWARS_WITHIN_CANDIDATES selects the candidate + construction — 'random' (reject-heavy, category `within_k`), 'sparse' (one guaranteed + diagonal acceptance per query, category `within_k_sparse`), 'dense' (every pair accepted, + category `within_k_dense`), or 'all'. Skipped entirely when the installed stringzillas lacks + the LevenshteinWithinK engine. + """ + if not hasattr(szs, "LevenshteinWithinK"): + print("within_k: SKIPPED (installed stringzillas lacks the LevenshteinWithinK engine)") + return + + bounds = parse_within_k_bounds() + mode = os.environ.get("STRINGWARS_WITHIN_CANDIDATES", "random").strip().lower() + candidate_modes = { + "random": ["random"], + "sparse": ["sparse"], + "dense": ["dense"], + "all": ["random", "sparse", "dense"], + }.get(mode) + if candidate_modes is None: + print(f"within_k: unknown STRINGWARS_WITHIN_CANDIDATES={mode!r}, expected random|sparse|dense|all") + return + + for candidate_mode in candidate_modes: + category = "within_k" if candidate_mode == "random" else f"within_k_{candidate_mode}" + print(f"\n## {category}") + for bound in bounds: + # Oracle first: never benchmark an engine whose matrix disagrees with the reference. + verify_within_k(tokens, bound, candidate_mode, category, seed) + benchmark_within_k_baselines( + tokens, + device_variants, + bound, + byte_lengths, + candidate_mode, + seed, + category, + warmup_seconds, + time_limit_seconds, + filter_pattern, + ) + benchmark_stringzillas_within( + tokens, + device_variants, + category, + "stringzillas.LevenshteinWithinK", + bound, + byte_lengths, + candidate_mode, + seed, + warmup_seconds, + time_limit_seconds, + filter_pattern, + ) + + def benchmark_biopython_baseline( tokens: Sequence, baseline_side: int, @@ -893,6 +1321,17 @@ def main() -> int: args.batch_size, ) + print("\n# within_k") + perform_within_k_benchmarks( + tokens, + device_variants, + byte_lengths, + warmup_seconds, + time_limit_seconds, + filter_pattern, + seed, + ) + if args.bio: print("\n# linear") perform_score_benchmarks(