Skip to content

Commit 1cb7601

Browse files
mulugetammeta-codesync[bot]
mulugetam
authored and
meta-codesync[bot]
committed
Eliminate per-code denormalization in uniform SQ distance computation (#5166)
Summary: This PR removes per-code denormalization from the inner loop of L2 and inner product distance computations for uniform scalar quantizers (`QT_8bit_uniform` and `QT_4bit_uniform`), yielding a speedup of up to **1.39x**. For uniform integer scalar quantizers, `vmin` and `vdiff` are scalars shared across all dimensions. The reconstructed value for each component is therefore a function of a per-code decode `n` that depends only on the byte code: ``` x_hat = vmin + vdiff * n ``` This structure lets us factor `vmin` and `vdiff` out of the per-database-vector inner loop entirely, instead of recomputing the transform on every code on every distance evaluation. **L2 distance.** ``` ||q - x_hat||^2 = ||q - (vmin + vdiff * n)||^2 = vdiff^2 * ||(q - vmin) / vdiff - n||^2 ``` We pre-adjust the query once in `set_query()` to `q_adj = (q - vmin) / vdiff` and precompute `scale = vdiff^2`. The hot loop then compares codes directly against `q_adj` in the codec's native decode space, applying `scale` exactly once at the end. **Inner product.** ``` <q, x_hat> = sum_i q[i] * (vmin + vdiff * n[i]) = vmin * sum_i q[i] + vdiff * sum_i q[i] * n[i] = bias + scale * <q, n> ``` We compute `bias = vmin * sum_i q[i]` and `scale = vdiff` once in `set_query()`. The hot loop accumulates the dot product against the raw decode `n` and applies `bias + scale` once at the end. In both cases this removes one FMA (the `vmin + vdiff * n` denormalization) from the inner loop per 8 components per database vector. Beyond the raw FLOP reduction, shortening the dependency chain lets SIMD pipelines better overlap the decode of the next lane with the accumulator update of the previous lane. **The Change** ``` Current This PR -------- --------- set_query(q): store q q_adj = (q - vmin) / vdiff [once] per code (×N): raw = decode(bytes) raw = decode(bytes) x = vmin + raw * vdiff ← gone diff = q_adj - raw diff = q - x accu += diff^2 accu += diff^2 ``` The optimization is gated on a C++20 `requires` check for a new `decode_8_raw()` method, defined only on the uniform `QuantizerTemplate` specializations. All other quantizer types fall through to the original `compute_distance` path unchanged. **Speedup** Below are the results from running `benchs/bench_scalar_quantizer.py` for a `dd` build on SPR, compared to the existing implementation. Similar results were observed for `avx-2` as well. ``` | | QT_4bit_uniform | QT_8bit_uniform | |--------------|-----------------|-----------------| | RS_minmax | 0.99x | 1.05x | | RS_minmax | 1.07x | 1.03x | | RS_minmax | 0.83x | 1.03x | | RS_minmax | 0.96x | 1.05x | | RS_minmax | 0.89x | 1.03x | | RS_minmax | 1.14x | 1.03x | | RS_minmax | 0.99x | 1.05x | | RS_meanstd | 1.28x | 1.09x | | RS_meanstd | 1.10x | 1.01x | | RS_meanstd | 1.14x | 1.06x | | RS_meanstd | 1.18x | 1.08x | | RS_meanstd | 1.11x | 1.05x | | RS_meanstd | 1.18x | 1.06x | | RS_meanstd | 1.16x | 1.08x | | RS_quantiles | 1.39x | 1.07x | | RS_quantiles | 1.08x | 1.01x | | RS_quantiles | 1.21x | 0.99x | | RS_quantiles | 1.25x | 1.10x | | RS_optim | 1.03x | 1.03x | ``` The raw performance results are available here: https://gist.github.com/mulugetam/7db50f89279bb270a1fe336206730d60 Pull Request resolved: #5166 Test Plan: Ran `python benchs/bench_scalar_quantizer.py` on SPR to validate performance results (see summary table). Ran `pytest tests/test_index_accuracy.py` to verify correctness is preserved for uniform SQ types (`QT_8bit_uniform`, `QT_4bit_uniform`). Reviewed By: mdouze Differential Revision: D106148760 Pulled By: mnorris11 fbshipit-source-id: 65769cfe9cf8d2d1ef5de05b6bdecf9dbee96d5c
1 parent 34eb989 commit 1cb7601

4 files changed

Lines changed: 242 additions & 8 deletions

File tree

‎benchs/bench_scalar_quantizer.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
variants = [(name, getattr(faiss.ScalarQuantizer, name))
1919
for name in dir(faiss.ScalarQuantizer)
20-
if name.startswith('QT_')]
20+
if name.startswith('QT_') and name != 'QT_count']
2121

2222
quantizer = faiss.IndexFlatL2(d)
2323
# quantizer.add(np.zeros((1, d), dtype='float32'))
@@ -67,7 +67,7 @@
6767
index.sq.rangestat = getattr(faiss.ScalarQuantizer,
6868
rsname)
6969

70-
index.rangestat_arg = val
70+
index.sq.rangestat_arg = val
7171

7272
index.train(xt)
7373
index.add(xb)

‎faiss/impl/scalar_quantizer/sq-avx2.cpp‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,12 @@ struct QuantizerTemplate<
196196
return simd8float32(_mm256_fmadd_ps(
197197
xi, _mm256_set1_ps(this->vdiff), _mm256_set1_ps(this->vmin)));
198198
}
199+
200+
/// Raw codec decode without denormalization
201+
FAISS_ALWAYS_INLINE simd8float32
202+
decode_8_raw(const uint8_t* code, int i) const {
203+
return Codec::decode_8_components(code, i);
204+
}
199205
};
200206

201207
template <class Codec>
@@ -399,6 +405,22 @@ struct SimilarityL2<SIMDLevel::AVX2> {
399405
const __m128 v3 = _mm_add_ps(v1, v2);
400406
return _mm_cvtss_f32(v3);
401407
}
408+
409+
static void adjust_query_for_raw_decode(
410+
const float* x,
411+
float* q_adj,
412+
size_t d,
413+
float vmin,
414+
float vdiff,
415+
float& scale_factor,
416+
float& bias) {
417+
float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f;
418+
for (size_t i = 0; i < d; i++) {
419+
q_adj[i] = (x[i] - vmin) * inv_vdiff;
420+
}
421+
scale_factor = vdiff * vdiff;
422+
bias = 0;
423+
}
402424
};
403425

404426
template <>
@@ -442,6 +464,23 @@ struct SimilarityIP<SIMDLevel::AVX2> {
442464
const __m128 v3 = _mm_add_ps(v1, v2);
443465
return _mm_cvtss_f32(v3);
444466
}
467+
468+
static void adjust_query_for_raw_decode(
469+
const float* x,
470+
float* q_adj,
471+
size_t d,
472+
float vmin,
473+
float vdiff,
474+
float& scale_factor,
475+
float& bias) {
476+
float sum_q = 0;
477+
for (size_t i = 0; i < d; i++) {
478+
q_adj[i] = x[i];
479+
sum_q += x[i];
480+
}
481+
scale_factor = vdiff;
482+
bias = vmin * sum_q;
483+
}
445484
};
446485

447486
/**********************************************************
@@ -454,8 +493,23 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
454493

455494
Quantizer quant;
456495

496+
// Pre-adjusted query buffer for uniform quantizers
497+
std::vector<float> q_adj;
498+
float scale_factor = 0;
499+
float bias = 0;
500+
501+
static constexpr bool has_decode_raw() {
502+
return requires(const Quantizer& q, const uint8_t* c, int i) {
503+
{ q.decode_8_raw(c, i) };
504+
};
505+
}
506+
457507
DCTemplate(size_t d, const std::vector<float>& trained)
458-
: quant(d, trained) {}
508+
: quant(d, trained) {
509+
if constexpr (has_decode_raw()) {
510+
q_adj.resize(d);
511+
}
512+
}
459513

460514
float compute_distance(const float* x, const uint8_t* code) const {
461515
Similarity sim(x);
@@ -484,6 +538,26 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
484538

485539
void set_query(const float* x) final {
486540
q = x;
541+
if constexpr (has_decode_raw()) {
542+
Sim::adjust_query_for_raw_decode(
543+
x,
544+
q_adj.data(),
545+
quant.d,
546+
quant.vmin,
547+
quant.vdiff,
548+
scale_factor,
549+
bias);
550+
}
551+
}
552+
553+
float query_to_code_predecoded(const uint8_t* code) const {
554+
Similarity sim(q_adj.data());
555+
sim.begin_8();
556+
for (size_t i = 0; i < quant.d; i += 8) {
557+
simd8float32 xi = quant.decode_8_raw(code, static_cast<int>(i));
558+
sim.add_8_components(xi);
559+
}
560+
return bias + scale_factor * sim.result_8();
487561
}
488562

489563
float symmetric_dis(idx_t i, idx_t j) override {
@@ -492,7 +566,11 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
492566
}
493567

494568
float query_to_code(const uint8_t* code) const final {
495-
return compute_distance(q, code);
569+
if constexpr (has_decode_raw()) {
570+
return query_to_code_predecoded(code);
571+
} else {
572+
return compute_distance(q, code);
573+
}
496574
}
497575

498576
void query_to_codes_batch_4(

‎faiss/impl/scalar_quantizer/sq-avx512.cpp‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,12 @@ struct QuantizerTemplate<
214214
return simd16float32(_mm512_fmadd_ps(
215215
xi, _mm512_set1_ps(this->vdiff), _mm512_set1_ps(this->vmin)));
216216
}
217+
218+
/// Raw codec decode without denormalization
219+
FAISS_ALWAYS_INLINE simd16float32
220+
decode_16_raw(const uint8_t* code, int i) const {
221+
return Codec::decode_16_components(code, i);
222+
}
217223
};
218224

219225
template <class Codec>
@@ -411,6 +417,22 @@ struct SimilarityL2<SIMDLevel::AVX512> {
411417
FAISS_ALWAYS_INLINE float result_16() {
412418
return horizontal_add(accu16);
413419
}
420+
421+
static void adjust_query_for_raw_decode(
422+
const float* x,
423+
float* q_adj,
424+
size_t d,
425+
float vmin,
426+
float vdiff,
427+
float& scale_factor,
428+
float& bias) {
429+
float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f;
430+
for (size_t i = 0; i < d; i++) {
431+
q_adj[i] = (x[i] - vmin) * inv_vdiff;
432+
}
433+
scale_factor = vdiff * vdiff;
434+
bias = 0;
435+
}
414436
};
415437

416438
template <>
@@ -445,6 +467,23 @@ struct SimilarityIP<SIMDLevel::AVX512> {
445467
FAISS_ALWAYS_INLINE float result_16() {
446468
return horizontal_add(accu16);
447469
}
470+
471+
static void adjust_query_for_raw_decode(
472+
const float* x,
473+
float* q_adj,
474+
size_t d,
475+
float vmin,
476+
float vdiff,
477+
float& scale_factor,
478+
float& bias) {
479+
float sum_q = 0;
480+
for (size_t i = 0; i < d; i++) {
481+
q_adj[i] = x[i];
482+
sum_q += x[i];
483+
}
484+
scale_factor = vdiff;
485+
bias = vmin * sum_q;
486+
}
448487
};
449488

450489
/**********************************************************
@@ -458,8 +497,23 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
458497

459498
Quantizer quant;
460499

500+
// Pre-adjusted query buffer for uniform quantizers
501+
std::vector<float> q_adj;
502+
float scale_factor = 0;
503+
float bias = 0;
504+
505+
static constexpr bool has_decode_raw() {
506+
return requires(const Quantizer& q, const uint8_t* c, int i) {
507+
{ q.decode_16_raw(c, i) };
508+
};
509+
}
510+
461511
DCTemplate(size_t d, const std::vector<float>& trained)
462-
: quant(d, trained) {}
512+
: quant(d, trained) {
513+
if constexpr (has_decode_raw()) {
514+
q_adj.resize(d);
515+
}
516+
}
463517

464518
float compute_distance(const float* x, const uint8_t* code) const {
465519
Similarity sim(x);
@@ -485,6 +539,26 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
485539

486540
void set_query(const float* x) final {
487541
q = x;
542+
if constexpr (has_decode_raw()) {
543+
Sim::adjust_query_for_raw_decode(
544+
x,
545+
q_adj.data(),
546+
quant.d,
547+
quant.vmin,
548+
quant.vdiff,
549+
scale_factor,
550+
bias);
551+
}
552+
}
553+
554+
float query_to_code_predecoded(const uint8_t* code) const {
555+
Similarity sim(q_adj.data());
556+
sim.begin_16();
557+
for (size_t i = 0; i < quant.d; i += 16) {
558+
simd16float32 xi = quant.decode_16_raw(code, i);
559+
sim.add_16_components(xi);
560+
}
561+
return bias + scale_factor * sim.result_16();
488562
}
489563

490564
float symmetric_dis(idx_t i, idx_t j) override {
@@ -493,7 +567,11 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
493567
}
494568

495569
float query_to_code(const uint8_t* code) const final {
496-
return compute_distance(q, code);
570+
if constexpr (has_decode_raw()) {
571+
return query_to_code_predecoded(code);
572+
} else {
573+
return compute_distance(q, code);
574+
}
497575
}
498576

499577
void query_to_codes_batch_4(

‎faiss/impl/scalar_quantizer/sq-neon.cpp‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,12 @@ struct QuantizerTemplate<
180180
xi.data.val[1],
181181
this->vdiff)});
182182
}
183+
184+
/// Raw codec decode without denormalization (for pre-decode opt)
185+
FAISS_ALWAYS_INLINE simd8float32
186+
decode_8_raw(const uint8_t* code, int i) const {
187+
return Codec::decode_8_components(code, i);
188+
}
183189
};
184190

185191
template <class Codec>
@@ -397,6 +403,22 @@ struct SimilarityL2<SIMDLevel::ARM_NEON> {
397403
FAISS_ALWAYS_INLINE float result_8() {
398404
return horizontal_add(accu8);
399405
}
406+
407+
static void adjust_query_for_raw_decode(
408+
const float* x,
409+
float* q_adj,
410+
size_t d,
411+
float vmin,
412+
float vdiff,
413+
float& scale_factor,
414+
float& bias) {
415+
float inv_vdiff = (vdiff != 0) ? 1.0f / vdiff : 0.0f;
416+
for (size_t i = 0; i < d; i++) {
417+
q_adj[i] = (x[i] - vmin) * inv_vdiff;
418+
}
419+
scale_factor = vdiff * vdiff;
420+
bias = 0;
421+
}
400422
};
401423

402424
template <>
@@ -431,6 +453,23 @@ struct SimilarityIP<SIMDLevel::ARM_NEON> {
431453
FAISS_ALWAYS_INLINE float result_8() {
432454
return horizontal_add(accu8);
433455
}
456+
457+
static void adjust_query_for_raw_decode(
458+
const float* x,
459+
float* q_adj,
460+
size_t d,
461+
float vmin,
462+
float vdiff,
463+
float& scale_factor,
464+
float& bias) {
465+
float sum_q = 0;
466+
for (size_t i = 0; i < d; i++) {
467+
q_adj[i] = x[i];
468+
sum_q += x[i];
469+
}
470+
scale_factor = vdiff;
471+
bias = vmin * sum_q;
472+
}
434473
};
435474

436475
/**********************************************************
@@ -444,8 +483,23 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
444483

445484
Quantizer quant;
446485

486+
// Pre-adjusted query buffer for uniform quantizers
487+
std::vector<float> q_adj;
488+
float scale_factor = 0;
489+
float bias = 0;
490+
491+
static constexpr bool has_decode_raw() {
492+
return requires(const Quantizer& q, const uint8_t* c, int i) {
493+
{ q.decode_8_raw(c, i) };
494+
};
495+
}
496+
447497
DCTemplate(size_t d, const std::vector<float>& trained)
448-
: quant(d, trained) {}
498+
: quant(d, trained) {
499+
if constexpr (has_decode_raw()) {
500+
q_adj.resize(d);
501+
}
502+
}
449503

450504
float compute_distance(const float* x, const uint8_t* code) const {
451505
Similarity sim(x);
@@ -471,6 +525,26 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
471525

472526
void set_query(const float* x) final {
473527
q = x;
528+
if constexpr (has_decode_raw()) {
529+
Sim::adjust_query_for_raw_decode(
530+
x,
531+
q_adj.data(),
532+
quant.d,
533+
quant.vmin,
534+
quant.vdiff,
535+
scale_factor,
536+
bias);
537+
}
538+
}
539+
540+
float query_to_code_predecoded(const uint8_t* code) const {
541+
Similarity sim(q_adj.data());
542+
sim.begin_8();
543+
for (size_t i = 0; i < quant.d; i += 8) {
544+
simd8float32 xi = quant.decode_8_raw(code, i);
545+
sim.add_8_components(xi);
546+
}
547+
return bias + scale_factor * sim.result_8();
474548
}
475549

476550
float symmetric_dis(idx_t i, idx_t j) override {
@@ -479,7 +553,11 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
479553
}
480554

481555
float query_to_code(const uint8_t* code) const final {
482-
return compute_distance(q, code);
556+
if constexpr (has_decode_raw()) {
557+
return query_to_code_predecoded(code);
558+
} else {
559+
return compute_distance(q, code);
560+
}
483561
}
484562

485563
void query_to_codes_batch_4(

0 commit comments

Comments
 (0)