Guard ScalarQuantizer compute_codes/decode OMP region for small n#5423
Closed
alibeklfc wants to merge 1 commit into
Closed
Guard ScalarQuantizer compute_codes/decode OMP region for small n#5423alibeklfc wants to merge 1 commit into
alibeklfc wants to merge 1 commit into
Conversation
Summary: `ScalarQuantizer::compute_codes` and `ScalarQuantizer::decode` unconditionally open an `#pragma omp parallel for` region regardless of `n`, unlike every sibling quantizer in the same family: `ProductQuantizer::decode` guards with `if (n > 100)`, `AdditiveQuantizer::decode`/`decode_unpacked` guard with `if (n > 100)`/`if (n > 1000)`, and `RaBitQuantizer::compute_codes_core`/ `decode_core` guard with `if (n > 1000)`. `IndexIVFScalarQuantizer::reconstruct_from_offset` calls `sq.decode(sc.get(), recons, 1)` once per vector, so this hot reconstruction path (used by `reconstruct()`, `reconstruct_n()`, and reranking/`IndexRefine` flows on top of IVF+SQ indexes) pays full OMP thread-team fork/join overhead for a single scalar-encode/decode of one vector. `IndexScalarQuantizer::sa_encode`/ `sa_decode` (the public codec API, also reachable from Python) hit the same cost for any single-vector call. Fix: add `if (n > 100)` to both pragmas, matching the threshold already used by `ProductQuantizer::decode` and `AdditiveQuantizer::decode`. No behavior change for n > 100 (parallel region still opens); for small n the loop now runs on the calling thread directly instead of forking an OMP team. Benchmarked on this devserver (56 OMP threads available), calling `ScalarQuantizer::compute_codes`/`decode` with n=1, 20000 reps, best-of-run median: | qtype | d | compute_codes(n=1) before | after | decode(n=1) before | after | |---------|-----|---------------------------|---------|---------------------|--------| | QT_8bit | 128 | 105.7us | 2.7us | 104.5us | 2.4us | | QT_8bit | 768 | 112.7us | 4.1us | 114.8us | 2.6us | | QT_fp16 | 128 | 117.3us | 2.3us | 113.0us | 2.4us | | QT_fp16 | 768 | 108.3us | 2.9us | 109.9us | 2.6us | ~40-50x reduction in per-call latency for the n=1 case across both qtypes and dimensions tested. n > 100 behavior is unaffected (same code path). Differential Revision: D111693838
Contributor
|
@alibeklfc has exported this pull request. If you are a Meta employee, you can view the originating Diff in D111693838. |
Contributor
|
This pull request has been merged in 764b86f. |
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:
ScalarQuantizer::compute_codesandScalarQuantizer::decodeunconditionallyopen an
#pragma omp parallel forregion regardless ofn, unlike everysibling quantizer in the same family:
ProductQuantizer::decodeguards withif (n > 100),AdditiveQuantizer::decode/decode_unpackedguard withif (n > 100)/if (n > 1000), andRaBitQuantizer::compute_codes_core/decode_coreguard withif (n > 1000).IndexIVFScalarQuantizer::reconstruct_from_offsetcallssq.decode(sc.get(), recons, 1)once per vector, so this hot reconstruction path (used byreconstruct(),reconstruct_n(), and reranking/IndexRefineflows on topof IVF+SQ indexes) pays full OMP thread-team fork/join overhead for a single
scalar-encode/decode of one vector.
IndexScalarQuantizer::sa_encode/sa_decode(the public codec API, also reachable from Python) hit the samecost for any single-vector call.
Fix: add
if (n > 100)to both pragmas, matching the threshold already usedby
ProductQuantizer::decodeandAdditiveQuantizer::decode. No behaviorchange for n > 100 (parallel region still opens); for small n the loop now
runs on the calling thread directly instead of forking an OMP team.
Benchmarked on this devserver (56 OMP threads available), calling
ScalarQuantizer::compute_codes/decodewith n=1, 20000 reps, best-of-runmedian:
~40-50x reduction in per-call latency for the n=1 case across both qtypes
and dimensions tested. n > 100 behavior is unaffected (same code path).
Differential Revision: D111693838