1
- From 9b33874562c9e62abf4a863657c54f0d349b0f67 Mon Sep 17 00:00:00 2001
1
+ From 1605542c1f4f7982fe0c5447090ae96f84f27484 Mon Sep 17 00:00:00 2001
2
2
From: John Mazanec <jmazane@amazon.com>
3
3
Date: Wed, 21 Feb 2024 15:34:15 -0800
4
4
Subject: [PATCH] Enable precomp table to be shared ivfpq
@@ -22,10 +22,10 @@ Signed-off-by: John Mazanec <jmazane@amazon.com>
22
22
create mode 100644 tests/test_ivfpq_share_table.cpp
23
23
24
24
diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
25
- index 100f499c..09508890 100644
25
+ index e9d6eead2..97c8c010d 100644
26
26
--- a/faiss/IndexIVFPQ.cpp
27
27
+++ b/faiss/IndexIVFPQ.cpp
28
- @@ -59 ,6 +59 ,29 @@ IndexIVFPQ::IndexIVFPQ(
28
+ @@ -58 ,6 +58 ,29 @@ IndexIVFPQ::IndexIVFPQ(
29
29
polysemous_training = nullptr;
30
30
do_polysemous_training = false;
31
31
polysemous_ht = 0;
@@ -55,7 +55,7 @@ index 100f499c..09508890 100644
55
55
}
56
56
57
57
/****************************************************************
58
- @@ -464 ,11 +487 ,23 @@ void IndexIVFPQ::precompute_table() {
58
+ @@ -463 ,11 +486 ,23 @@ void IndexIVFPQ::precompute_table() {
59
59
use_precomputed_table,
60
60
quantizer,
61
61
pq,
@@ -80,7 +80,7 @@ index 100f499c..09508890 100644
80
80
namespace {
81
81
82
82
#define TIC t0 = get_cycles()
83
- @@ -648 ,7 +683 ,7 @@ struct QueryTables {
83
+ @@ -647 ,7 +682 ,7 @@ struct QueryTables {
84
84
85
85
fvec_madd(
86
86
pq.M * pq.ksub,
@@ -89,7 +89,7 @@ index 100f499c..09508890 100644
89
89
-2.0,
90
90
sim_table_2,
91
91
sim_table);
92
- @@ -677 ,7 +712 ,7 @@ struct QueryTables {
92
+ @@ -676 ,7 +711 ,7 @@ struct QueryTables {
93
93
k >>= cpq.nbits;
94
94
95
95
// get corresponding table
@@ -98,7 +98,7 @@ index 100f499c..09508890 100644
98
98
(ki * pq.M + cm * Mf) * pq.ksub;
99
99
100
100
if (polysemous_ht == 0) {
101
- @@ -707 ,7 +742 ,7 @@ struct QueryTables {
101
+ @@ -706 ,7 +741 ,7 @@ struct QueryTables {
102
102
dis0 = coarse_dis;
103
103
104
104
const float* s =
@@ -107,7 +107,7 @@ index 100f499c..09508890 100644
107
107
for (int m = 0; m < pq.M; m++) {
108
108
sim_table_ptrs[m] = s;
109
109
s += pq.ksub;
110
- @@ -727 ,7 +762 ,7 @@ struct QueryTables {
110
+ @@ -726 ,7 +761 ,7 @@ struct QueryTables {
111
111
int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
112
112
k >>= cpq.nbits;
113
113
@@ -116,7 +116,7 @@ index 100f499c..09508890 100644
116
116
(ki * pq.M + cm * Mf) * pq.ksub;
117
117
118
118
for (int m = m0; m < m0 + Mf; m++) {
119
- @@ -1344 ,6 +1379 ,8 @@ IndexIVFPQ::IndexIVFPQ() {
119
+ @@ -1343 ,6 +1378 ,8 @@ IndexIVFPQ::IndexIVFPQ() {
120
120
do_polysemous_training = false;
121
121
polysemous_ht = 0;
122
122
polysemous_training = nullptr;
@@ -126,7 +126,7 @@ index 100f499c..09508890 100644
126
126
127
127
struct CodeCmp {
128
128
diff --git a/faiss/IndexIVFPQ.h b/faiss/IndexIVFPQ.h
129
- index d5d21da4..850bbe44 100644
129
+ index 7bf97ec0f..f647e5f87 100644
130
130
--- a/faiss/IndexIVFPQ.h
131
131
+++ b/faiss/IndexIVFPQ.h
132
132
@@ -48,7 +48,8 @@ struct IndexIVFPQ : IndexIVF {
@@ -167,10 +167,10 @@ index d5d21da4..850bbe44 100644
167
167
};
168
168
169
169
diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp
170
- index 2844ae49..895df342 100644
170
+ index 9d1cdfcae..647644e36 100644
171
171
--- a/faiss/IndexIVFPQFastScan.cpp
172
172
+++ b/faiss/IndexIVFPQFastScan.cpp
173
- @@ -46 ,6 +46 ,8 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(
173
+ @@ -42 ,6 +42 ,8 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(
174
174
: IndexIVFFastScan(quantizer, d, nlist, 0, metric), pq(d, M, nbits) {
175
175
by_residual = false; // set to false by default because it's faster
176
176
@@ -179,7 +179,7 @@ index 2844ae49..895df342 100644
179
179
init_fastscan(M, nbits, nlist, metric, bbs);
180
180
}
181
181
182
- @@ -53 ,6 +55 ,17 @@ IndexIVFPQFastScan::IndexIVFPQFastScan() {
182
+ @@ -49 ,6 +51 ,17 @@ IndexIVFPQFastScan::IndexIVFPQFastScan() {
183
183
by_residual = false;
184
184
bbs = 0;
185
185
M2 = 0;
@@ -197,7 +197,7 @@ index 2844ae49..895df342 100644
197
197
}
198
198
199
199
IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
200
- @@ -71 ,13 +84 ,15 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
200
+ @@ -67 ,13 +80 ,15 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
201
201
ntotal = orig.ntotal;
202
202
is_trained = orig.is_trained;
203
203
nprobe = orig.nprobe;
@@ -218,7 +218,7 @@ index 2844ae49..895df342 100644
218
218
}
219
219
220
220
for (size_t i = 0; i < nlist; i++) {
221
- @@ -102 ,6 +117 ,12 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
221
+ @@ -98 ,6 +113 ,12 @@ IndexIVFPQFastScan::IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs)
222
222
orig_invlists = orig.invlists;
223
223
}
224
224
@@ -231,7 +231,7 @@ index 2844ae49..895df342 100644
231
231
/*********************************************************
232
232
* Training
233
233
*********************************************************/
234
- @@ -127 ,11 +148 ,23 @@ void IndexIVFPQFastScan::precompute_table() {
234
+ @@ -123 ,11 +144 ,23 @@ void IndexIVFPQFastScan::precompute_table() {
235
235
use_precomputed_table,
236
236
quantizer,
237
237
pq,
@@ -256,7 +256,7 @@ index 2844ae49..895df342 100644
256
256
/*********************************************************
257
257
* Code management functions
258
258
*********************************************************/
259
- @@ -229 ,7 +262 ,7 @@ void IndexIVFPQFastScan::compute_LUT(
259
+ @@ -225 ,7 +258 ,7 @@ void IndexIVFPQFastScan::compute_LUT(
260
260
if (cij >= 0) {
261
261
fvec_madd_simd(
262
262
dim12,
@@ -266,7 +266,7 @@ index 2844ae49..895df342 100644
266
266
ip_table.get() + i * dim12,
267
267
tab);
268
268
diff --git a/faiss/IndexIVFPQFastScan.h b/faiss/IndexIVFPQFastScan.h
269
- index 00dd2f11..91f35a6e 100644
269
+ index a2cce3266..1e1f0049c 100644
270
270
--- a/faiss/IndexIVFPQFastScan.h
271
271
+++ b/faiss/IndexIVFPQFastScan.h
272
272
@@ -38,7 +38,8 @@ struct IndexIVFPQFastScan : IndexIVFFastScan {
@@ -302,19 +302,19 @@ index 00dd2f11..91f35a6e 100644
302
302
/// same as the regular IVFPQ encoder. The codes are not reorganized by
303
303
/// blocks a that point
304
304
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
305
- index 87ab2020..a859516c 100644
305
+ index ae6cc7878..16c99e04d 100644
306
306
--- a/tests/CMakeLists.txt
307
307
+++ b/tests/CMakeLists.txt
308
308
@@ -38,6 +38,7 @@ set(FAISS_TEST_SRC
309
- test_common_ivf_empty_index.cpp
310
309
test_callback.cpp
311
310
test_utils.cpp
311
+ test_hamming.cpp
312
312
+ test_ivfpq_share_table.cpp
313
313
)
314
314
315
315
add_executable(faiss_test ${FAISS_TEST_SRC})
316
316
diff --git a/tests/test_disable_pq_sdc_tables.cpp b/tests/test_disable_pq_sdc_tables.cpp
317
- index b211a5c4..a27973d5 100644
317
+ index f94aac870..60c59d7ba 100644
318
318
--- a/tests/test_disable_pq_sdc_tables.cpp
319
319
+++ b/tests/test_disable_pq_sdc_tables.cpp
320
320
@@ -15,7 +15,9 @@
@@ -327,10 +327,10 @@ index b211a5c4..a27973d5 100644
327
327
+ }
328
328
329
329
TEST(IO, TestReadHNSWPQ_whenSDCDisabledFlagPassed_thenDisableSDCTable) {
330
- Tempfilename index_filename(&temp_file_mutex, "/tmp/faiss_TestReadHNSWPQ");
330
+ // Create a temp file name with a randomized component for stress runs
331
331
diff --git a/tests/test_ivfpq_share_table.cpp b/tests/test_ivfpq_share_table.cpp
332
332
new file mode 100644
333
- index 00000000..f827315d
333
+ index 000000000..f827315d8
334
334
--- /dev/null
335
335
+++ b/tests/test_ivfpq_share_table.cpp
336
336
@@ -0,0 +1,173 @@
0 commit comments