Skip to content

Commit e7e0a7a

Browse files
limqiyinglimqiying
authored
Merge branch 'main' into avx512-bf16-opt
2 parents eb3b48f + 9d6b2e7 commit e7e0a7a

33 files changed

Lines changed: 1009 additions & 735 deletions

‎faiss/CMakeLists.txt‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -121,9 +121,9 @@ set(FAISS_SRC
121121
impl/kmeans1d.cpp
122122
impl/lattice_Zn.cpp
123123
impl/mapped_io.cpp
124-
impl/pq4_fast_scan.cpp
125-
impl/pq4_fast_scan_search_1.cpp
126-
impl/pq4_fast_scan_search_qbs.cpp
124+
impl/fast_scan/pq4_fast_scan.cpp
125+
impl/fast_scan/pq4_fast_scan_search_1.cpp
126+
impl/fast_scan/pq4_fast_scan_search_qbs.cpp
127127
impl/residual_quantizer_encode_steps.cpp
128128
impl/zerocopy_io.cpp
129129
impl/NNDescent.cpp
@@ -231,8 +231,8 @@ set(FAISS_HEADERS
231231
impl/HNSW.h
232232
impl/LocalSearchQuantizer.h
233233
impl/ProductAdditiveQuantizer.h
234-
impl/LookupTableScaler.h
235-
impl/FastScanDistancePostProcessing.h
234+
impl/fast_scan/LookupTableScaler.h
235+
impl/fast_scan/FastScanDistancePostProcessing.h
236236
impl/maybe_owned_vector.h
237237
impl/NNDescent.h
238238
impl/NSG.h
@@ -266,10 +266,13 @@ set(FAISS_HEADERS
266266
impl/kmeans1d.h
267267
impl/lattice_Zn.h
268268
impl/platform_macros.h
269-
impl/pq4_fast_scan.h
269+
impl/fast_scan/pq4_fast_scan.h
270+
impl/fast_scan/decompose_qbs.h
271+
impl/fast_scan/kernels_simd256.h
272+
impl/fast_scan/kernels_simd512.h
270273
impl/residual_quantizer_encode_steps.h
271274
impl/simd_dispatch.h
272-
impl/simd_result_handlers.h
275+
impl/fast_scan/simd_result_handlers.h
273276
impl/zerocopy_io.h
274277
utils/pq_code_distance.h
275278
impl/pq_code_distance/pq_code_distance-inl.h

‎faiss/IndexAdditiveQuantizerFastScan.cpp‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,10 @@
1111
#include <memory>
1212

1313
#include <faiss/impl/FaissAssert.h>
14-
#include <faiss/impl/FastScanDistancePostProcessing.h>
1514
#include <faiss/impl/LocalSearchQuantizer.h>
16-
#include <faiss/impl/LookupTableScaler.h>
1715
#include <faiss/impl/ResidualQuantizer.h>
18-
#include <faiss/impl/pq4_fast_scan.h>
16+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
17+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
1918
#include <faiss/utils/quantize_lut.h>
2019
#include <faiss/utils/utils.h>
2120

@@ -202,9 +201,8 @@ void IndexAdditiveQuantizerFastScan::search(
202201
return;
203202
}
204203

205-
NormTableScaler scaler(norm_scale);
206204
FastScanDistancePostProcessing context;
207-
context.norm_scaler = &scaler;
205+
context.pq2x4_scale = norm_scale;
208206
if (metric_type == METRIC_L2) {
209207
search_dispatch_implem<true>(n, x, k, distances, labels, context);
210208
} else {

‎faiss/IndexFastScan.cpp‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,11 @@
1313

1414
#include <faiss/impl/CodePacker.h>
1515
#include <faiss/impl/FaissAssert.h>
16-
#include <faiss/impl/FastScanDistancePostProcessing.h>
1716
#include <faiss/impl/IDSelector.h>
18-
#include <faiss/impl/LookupTableScaler.h>
1917
#include <faiss/impl/RaBitQUtils.h>
20-
#include <faiss/impl/pq4_fast_scan.h>
21-
#include <faiss/impl/simd_result_handlers.h>
18+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
19+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
20+
#include <faiss/impl/fast_scan/simd_result_handlers.h>
2221
#include <faiss/utils/hamming.h>
2322
#include <faiss/utils/quantize_lut.h>
2423
#include <faiss/utils/utils.h>
@@ -187,18 +186,18 @@ void estimators_from_tables_generic(
187186
BitstringReader bsr(codes + j * index.code_size, index.code_size);
188187
accu_t dis = 0;
189188
const dis_t* dt = dis_table;
190-
int nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
189+
int nscale = context.pq2x4_scale ? 2 : 0;
191190

192191
for (size_t m = 0; m < index.M - nscale; m++) {
193192
uint64_t c = bsr.read(index.nbits);
194193
dis += dt[c];
195194
dt += index.ksub;
196195
}
197196

198-
if (nscale && context.norm_scaler) {
197+
if (nscale) {
199198
for (size_t m = 0; m < nscale; m++) {
200199
uint64_t c = bsr.read(index.nbits);
201-
dis += context.norm_scaler->scale_one(dt[c]);
200+
dis += dt[c] * context.pq2x4_scale;
202201
dt += index.ksub;
203202
}
204203
}
@@ -545,7 +544,7 @@ void IndexFastScan::search_implem_12(
545544
codes.get(),
546545
LUT.get(),
547546
*handler.get(),
548-
context.norm_scaler,
547+
context.pq2x4_scale,
549548
get_block_stride());
550549
}
551550
if (!(skip & 8)) {
@@ -629,7 +628,7 @@ void IndexFastScan::search_implem_14(
629628
codes.get(),
630629
LUT.get(),
631630
*handler.get(),
632-
context.norm_scaler,
631+
context.pq2x4_scale,
633632
get_block_stride());
634633
}
635634
if (!(skip & 8)) {

‎faiss/IndexFastScan.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,12 @@
88
#pragma once
99

1010
#include <faiss/Index.h>
11-
#include <faiss/impl/FastScanDistancePostProcessing.h>
11+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
1212
#include <faiss/utils/AlignedTable.h>
1313

1414
namespace faiss {
1515

1616
struct CodePacker;
17-
struct NormTableScaler;
1817
struct IDSelector;
1918
struct SIMDResultHandlerToFloat;
2019

‎faiss/IndexIVFAdditiveQuantizerFastScan.cpp‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,8 @@
1414

1515
#include <faiss/impl/AuxIndexStructures.h>
1616
#include <faiss/impl/FaissAssert.h>
17-
#include <faiss/impl/FastScanDistancePostProcessing.h>
18-
#include <faiss/impl/LookupTableScaler.h>
19-
#include <faiss/impl/pq4_fast_scan.h>
17+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
18+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
2019
#include <faiss/impl/simd_dispatch.h>
2120
#include <faiss/invlists/BlockInvertedLists.h>
2221
#include <faiss/utils/distances.h>
@@ -317,9 +316,8 @@ void IndexIVFAdditiveQuantizerFastScan::search(
317316
return;
318317
}
319318

320-
NormTableScaler scaler(norm_scale);
321319
FastScanDistancePostProcessing context;
322-
context.norm_scaler = &scaler;
320+
context.pq2x4_scale = norm_scale;
323321
IndexIVFFastScan::CoarseQuantized cq{nprobe};
324322
search_dispatch_implem(n, x, k, distances, labels, cq, context);
325323
}

‎faiss/IndexIVFAdditiveQuantizerFastScan.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
#include <faiss/IndexIVFAdditiveQuantizer.h>
1313
#include <faiss/IndexIVFFastScan.h>
1414
#include <faiss/impl/AdditiveQuantizer.h>
15-
#include <faiss/impl/FastScanDistancePostProcessing.h>
1615
#include <faiss/impl/ProductAdditiveQuantizer.h>
16+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
1717
#include <faiss/utils/AlignedTable.h>
1818

1919
namespace faiss {

‎faiss/IndexIVFFastScan.cpp‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,10 @@
1717
#include <faiss/IndexIVFPQ.h>
1818
#include <faiss/impl/AuxIndexStructures.h>
1919
#include <faiss/impl/FaissAssert.h>
20-
#include <faiss/impl/FastScanDistancePostProcessing.h>
21-
#include <faiss/impl/LookupTableScaler.h>
2220
#include <faiss/impl/RaBitQUtils.h>
23-
#include <faiss/impl/pq4_fast_scan.h>
24-
#include <faiss/impl/simd_result_handlers.h>
21+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
22+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
23+
#include <faiss/impl/fast_scan/simd_result_handlers.h>
2524
#include <faiss/invlists/BlockInvertedLists.h>
2625
#include <faiss/utils/hamming.h>
2726
#include <faiss/utils/quantize_lut.h>
@@ -239,7 +238,7 @@ void estimators_from_tables_generic(
239238
int64_t* heap_ids,
240239
const FastScanDistancePostProcessing& context) {
241240
using accu_t = typename C::T;
242-
size_t nscale = context.norm_scaler ? context.norm_scaler->nscale : 0;
241+
size_t nscale = context.pq2x4_scale ? 2 : 0;
243242
for (size_t j = 0; j < ncodes; ++j) {
244243
BitstringReader bsr(codes + j * index.code_size, index.code_size);
245244
accu_t dis = bias;
@@ -251,10 +250,10 @@ void estimators_from_tables_generic(
251250
dt += index.ksub;
252251
}
253252

254-
if (context.norm_scaler) {
253+
if (nscale) {
255254
for (size_t m = 0; m < nscale; m++) {
256255
uint64_t c = bsr.read(index.nbits);
257-
dis += context.norm_scaler->scale_one(dt[c]);
256+
dis += dt[c] * context.pq2x4_scale;
258257
dt += index.ksub;
259258
}
260259
}
@@ -1031,7 +1030,7 @@ void IndexIVFFastScan::search_implem_10(
10311030
codes.get(),
10321031
LUT,
10331032
handler,
1034-
context.norm_scaler,
1033+
context.pq2x4_scale,
10351034
get_block_stride());
10361035

10371036
ndis += ls;
@@ -1183,7 +1182,7 @@ void IndexIVFFastScan::search_implem_12(
11831182
codes.get(),
11841183
LUT.get(),
11851184
handler,
1186-
context.norm_scaler,
1185+
context.pq2x4_scale,
11871186
get_block_stride());
11881187
// prepare for next loop
11891188
i0 = i1;
@@ -1407,7 +1406,7 @@ void IndexIVFFastScan::search_implem_14(
14071406
codes.get(),
14081407
LUT.get(),
14091408
*handler.get(),
1410-
context.norm_scaler,
1409+
context.pq2x4_scale,
14111410
get_block_stride());
14121411
}
14131412

‎faiss/IndexIVFFastScan.h‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,11 @@
88
#pragma once
99

1010
#include <faiss/IndexIVF.h>
11-
#include <faiss/impl/FastScanDistancePostProcessing.h>
11+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
1212
#include <faiss/utils/AlignedTable.h>
1313

1414
namespace faiss {
1515

16-
struct NormTableScaler;
1716
struct SIMDResultHandlerToFloat;
1817
struct Quantizer;
1918

‎faiss/IndexIVFPQFastScan.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@
2222

2323
#include <faiss/invlists/BlockInvertedLists.h>
2424

25-
#include <faiss/impl/pq4_fast_scan.h>
26-
#include <faiss/impl/simd_result_handlers.h>
25+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
26+
#include <faiss/impl/fast_scan/simd_result_handlers.h>
2727

2828
namespace faiss {
2929

@@ -419,7 +419,7 @@ struct IVFPQFastScanScanner : InvertedListScanner {
419419
codes,
420420
LUT,
421421
*handler,
422-
nullptr,
422+
0,
423423
index.get_block_stride());
424424

425425
// The handler is for the results of this iteration.

‎faiss/IndexIVFRaBitQFastScan.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@
1414

1515
#include <faiss/impl/CodePackerRaBitQ.h>
1616
#include <faiss/impl/FaissAssert.h>
17-
#include <faiss/impl/FastScanDistancePostProcessing.h>
1817
#include <faiss/impl/RaBitQUtils.h>
1918
#include <faiss/impl/RaBitQuantizerMultiBit.h>
20-
#include <faiss/impl/pq4_fast_scan.h>
21-
#include <faiss/impl/simd_result_handlers.h>
19+
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
20+
#include <faiss/impl/fast_scan/pq4_fast_scan.h>
21+
#include <faiss/impl/fast_scan/simd_result_handlers.h>
2222
#include <faiss/invlists/BlockInvertedLists.h>
2323
#include <faiss/utils/distances.h>
2424
#include <faiss/utils/utils.h>
@@ -938,7 +938,7 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
938938
codes,
939939
LUT,
940940
*handler,
941-
nullptr,
941+
0,
942942
index.get_block_stride());
943943

944944
// Combine results across iterations

0 commit comments

Comments
 (0)