perf(py): zero-copy contiguous input + GIL release for pq_encode_batch - #108
Draft
konjoinfinity wants to merge 3 commits into
Draft
perf(py): zero-copy contiguous input + GIL release for pq_encode_batch#108konjoinfinity wants to merge 3 commits into
konjoinfinity wants to merge 3 commits into
Conversation
pq_encode_batch always copied its [N, D] vectors input with .iter().copied().collect(), even when the numpy array was already C-contiguous. Now borrows directly via .as_slice() in that case (mirroring quantize_int8_batch's established idiom) and only copies for the rare non-contiguous case. Also releases the GIL around the pq_encode_into call (already rayon-parallel internally, but previously held the GIL), matching the same fix already shipped elsewhere. centroids stays a direct copy -- it's orders of magnitude smaller than vectors and PQCodebook must own its buffer regardless. Measured (Rust-level copy microbenchmark; no pytest/numpy available in this container for an end-to-end Python benchmark): the avoided copy was ~220-230ms of a ~935ms pq_encode_into call at [200000, 768], M=96, K=256 (~20-25% of total call time). Also reconciles OPTIMIZATION_OPPORTUNITIES.md's FFI/marshalling section: items 7 and 8 (search_batch_arrays, vectorized query normalization) turned out to already be shipped; corrected to prevent re-chasing them. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
Split the varr.as_slice().unwrap_or_else(...) call onto two lines to match rustfmt's line-width preference, flagged by the gates job's repo:fmt-check (now correctly matching content post kiban v1.1.5, unlike the earlier false-positive net-new bug this session worked around). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
The previous fmt fix reverted the as_slice().unwrap_or_else() split based on a manual line-width miscount -- rustfmt genuinely wants that chain split (confirmed via a direct rustfmt --check run against this file), so restore it. Also multi-line the adjacent PQCodebook struct literal, which rustfmt's struct_lit_width heuristic (not the general line-width limit) wants split regardless of the ~90-char line fitting under 100. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EhXzwfPHcfRj7NdSaSeasB
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
OPTIMIZATION_OPPORTUNITIES.mditem 10 flaggedpq_encode_batchas always copying its[N, D]vectorsinput with.iter().copied().collect(), even when the numpy array was already C-contiguous (the common case). Verified this was still true and fixed it: now borrows directly via.as_slice()when contiguous, mirroring thequantize_int8_batchidiom already established elsewhere in this file, and only falls back to an owned copy for the (rare) non-contiguous case.pq_encode_intocall — it's already rayon-parallel internally but previously ran with the GIL held, unlikequantize_int8_batch/trainwhich already release it.centroids([M, K, sub_dim]) is left as a direct copy: it's orders of magnitude smaller thanvectorsfor realisticN, andPQCodebookmust own its centroid buffer regardless, so there's no copy to avoid there.search_batch_arrays, vectorized query normalization) had also already shipped but were still listed as open — corrected those too so a future pass doesn't re-chase them.Measured (Rust-level microbenchmark of the copy step in isolation — this container has no
pytest/numpyto drive an end-to-end Python benchmark): copying a[200,000, 768]f32 array (614 MB) took a stable ~220–230ms across repeated trials, vs. ~935ms for the actualpq_encode_intowork at the same scale (M=96, K=256) — the removed copy is ~20–25% of total call time at this scale, and a larger fraction for smaller batches where the fixed copy cost dominates more.Type of change
Checklist
cargo testpasses locally (238vectro_libtests;pq_encode_into's numerics are untouched, only input ownership changed)cargo clippy -- -D warningsreports no errors (verified against both the project's baseline flags and the full CI gate flags:-D warnings -D clippy::pedantic -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic -D clippy::todo -D clippy::dbg_macro)cargo fmt --checkpasses on the files touched by this PRpytest python/tests/— not run (nopytest/numpyinstalled in this container); no Python source changed, only the Rust FFI binding's internalsruff check python/— not applicable, no Python changedpq_encode_batchfix + doc reconciliation for the same section)Related issues
Continuation of the plan following #99 / #100 / #104 / #105 / #106 / #107.
Generated by Claude Code