forked from facebookresearch/faiss
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hamming.h
446 lines (342 loc) · 11.7 KB
/
hamming.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
/* Copyright 2004-present Facebook. All Rights Reserved.
*
* Hamming distances. The binary vector dimensionality should be a multiple
* of 64, as the elementary operations operate on words. If you really need
* other sizes, just pad with 0s (this is done by function fvecs2bitvecs).
*
* User-defined type hamdis_t is used for distances because at this time
* it is still uncler clear how we will need to balance
* - flexibility in vector size (may need 16- or even 8-bit vectors)
* - memory usage
* - cache-misses when dealing with large volumes of data (fewer bits is better)
*
* hamdis_t should optimally be compatibe with one of the Torch Storage
* (Byte,Short,Long) and therefore should be signed for 2-bytes and 4-bytes.
*/
#ifndef FAISS_hamming_h
#define FAISS_hamming_h
#include <stdint.h>
#include "Heap.h"
/* The Hamming distance type should be exportable to Lua Tensor, which
excludes most unsigned type */
typedef int32_t hamdis_t;
namespace faiss {
inline int popcount64(uint64_t x) {
return __builtin_popcountl(x);
}
/** Compute a set of Hamming distances between na and nb binary vectors
*
* @param a size na * nbytespercode
* @param b size nb * nbytespercode
* @param nbytespercode should be multiple of 8
* @param dis output distances, size na * nb
*/
void hammings (
const uint8_t * a,
const uint8_t * b,
size_t na, size_t nb,
size_t nbytespercode,
hamdis_t * dis);
void bitvec_print (const uint8_t * b, size_t d);
/* Functions for casting vectors of regular types to compact bits.
They assume proper allocation done beforehand, meaning that b
should be be able to receive as many bits as x may produce. */
/* Makes an array of bits from the signs of a float array. The length
of the output array b is rounded up to byte size (allocate
accordingly) */
void fvecs2bitvecs (
const float * x,
uint8_t * b,
size_t d,
size_t n);
void fvec2bitvec (const float * x, uint8_t * b, size_t d);
/** Return the k smallest Hamming distances for a set of binary query vectors
* @param a queries, size ha->nh * ncodes
* @param b database, size nb * ncodes
* @param nb number of database vectors
* @param ncodes size of the binary codes (bytes)
* @param ordered if != 0: order the results by decreasing distance
* (may be bottleneck for k/n > 0.01) */
void hammings_knn (
int_maxheap_array_t * ha,
const uint8_t * a,
const uint8_t * b,
size_t nb,
size_t ncodes,
int ordered);
/* The core function, without initialization and no re-ordering */
void hammings_knn_core (
int_maxheap_array_t * ha,
const uint8_t * a,
const uint8_t * b,
size_t nb,
size_t ncodes);
/* Counting the number of matches or of cross-matches (without returning them)
For use with function that assume pre-allocated memory */
void hamming_count_thres (
const uint8_t * bs1,
const uint8_t * bs2,
size_t n1,
size_t n2,
hamdis_t ht,
size_t ncodes,
size_t * nptr);
/* Return all Hamming distances/index passing a thres. Pre-allocation of output
is required. Use hamming_count_thres to determine the proper size. */
size_t match_hamming_thres (
const uint8_t * bs1,
const uint8_t * bs2,
size_t n1,
size_t n2,
hamdis_t ht,
size_t ncodes,
long * idx,
hamdis_t * dis);
/* Cross-matching in a set of vectors */
void crosshamming_count_thres (
const uint8_t * dbs,
size_t n,
hamdis_t ht,
size_t ncodes,
size_t * nptr);
/* compute the Hamming distances between two codewords of nwords*64 bits */
hamdis_t hamming (
const uint64_t * bs1,
const uint64_t * bs2,
size_t nwords);
/******************************************************************
* The HammingComputer series of classes compares a single code of
* size 4 to 32 to incoming codes. They are intended for use as a
* template class where it would be inefficient to switch on the code
* size in the inner loop. Hopefully the compiler will inline the
* hamming() functions and put the a0, a1, ... in registers.
******************************************************************/
struct HammingComputer4 {
uint32_t a0;
HammingComputer4 (const uint8_t *a, int code_size) {
assert (code_size == 4);
a0 = *(uint32_t *)a;
}
inline int hamming (const uint8_t *b) const {
return popcount64 (*(uint32_t *)b ^ a0);
}
};
struct HammingComputer8 {
uint64_t a0;
HammingComputer8 (const uint8_t *a, int code_size) {
assert (code_size == 8);
a0 = *(uint64_t *)a;
}
inline int hamming (const uint8_t *b) const {
return popcount64 (*(uint64_t *)b ^ a0);
}
};
struct HammingComputer16 {
uint64_t a0, a1;
HammingComputer16 (const uint8_t *a8, int code_size) {
assert (code_size == 16);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return popcount64 (b[0] ^ a0) + popcount64 (b[1] ^ a1);
}
};
// when applied to an array, 1/2 of the 64-bit accesses are unaligned.
// This incurs a penalty of ~10% wrt. fully aligned accesses.
struct HammingComputer20 {
uint64_t a0, a1;
uint32_t a2;
HammingComputer20 (const uint8_t *a8, int code_size) {
assert (code_size == 20);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1]; a2 = a[2];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return popcount64 (b[0] ^ a0) + popcount64 (b[1] ^ a1) +
popcount64 (*(uint32_t*)(b + 2) ^ a2);
}
};
struct HammingComputer32 {
uint64_t a0, a1, a2, a3;
HammingComputer32 (const uint8_t *a8, int code_size) {
assert (code_size == 32);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1]; a2 = a[2]; a3 = a[3];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return popcount64 (b[0] ^ a0) + popcount64 (b[1] ^ a1) +
popcount64 (b[2] ^ a2) + popcount64 (b[3] ^ a3);
}
};
struct HammingComputer64 {
uint64_t a0, a1, a2, a3, a4, a5, a6, a7;
HammingComputer64 (const uint8_t *a8, int code_size) {
assert (code_size == 64);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1]; a2 = a[2]; a3 = a[3];
a4 = a[4]; a5 = a[5]; a6 = a[6]; a7 = a[7];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return popcount64 (b[0] ^ a0) + popcount64 (b[1] ^ a1) +
popcount64 (b[2] ^ a2) + popcount64 (b[3] ^ a3) +
popcount64 (b[4] ^ a4) + popcount64 (b[5] ^ a5) +
popcount64 (b[6] ^ a6) + popcount64 (b[7] ^ a7);
}
};
struct HammingComputerDefault {
const uint8_t *a;
int n;
HammingComputerDefault (const uint8_t *a8, int code_size) {
a = a8;
n = code_size;
}
int hamming (const uint8_t *b8) const {
int accu = 0;
for (int i = 0; i < n; i++)
accu += popcount64 (a[i] ^ b8[i]);
return accu;
}
};
struct HammingComputerM8 {
const uint64_t *a;
int n;
HammingComputerM8 (const uint8_t *a8, int code_size) {
assert (code_size % 8 == 0);
a = (uint64_t *)a8;
n = code_size / 8;
}
int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
int accu = 0;
for (int i = 0; i < n; i++)
accu += popcount64 (a[i] ^ b[i]);
return accu;
}
};
// very inefficient...
struct HammingComputerM4 {
const uint32_t *a;
int n;
HammingComputerM4 (const uint8_t *a4, int code_size) {
assert (code_size % 4 == 0);
a = (uint32_t *)a4;
n = code_size / 4;
}
int hamming (const uint8_t *b8) const {
const uint32_t *b = (uint32_t *)b8;
int accu = 0;
for (int i = 0; i < n; i++)
accu += popcount64 (a[i] ^ b[i]);
return accu;
}
};
/***************************************************************************
* Equivalence with a template class when code size is known at compile time
**************************************************************************/
// default template
template<int CODE_SIZE>
struct HammingComputer: HammingComputerM8 {
HammingComputer (const uint8_t *a, int code_size):
HammingComputerM8(a, code_size) {}
};
#define SPECIALIZED_HC(CODE_SIZE) \
template<> struct HammingComputer<CODE_SIZE>: \
HammingComputer ## CODE_SIZE { \
HammingComputer (const uint8_t *a): \
HammingComputer ## CODE_SIZE(a, CODE_SIZE) {} \
}
SPECIALIZED_HC(4);
SPECIALIZED_HC(8);
SPECIALIZED_HC(16);
SPECIALIZED_HC(20);
SPECIALIZED_HC(32);
SPECIALIZED_HC(64);
#undef SPECIALIZED_HC
/***************************************************************************
* generalized Hamming = number of bytes that are different between
* two codes.
***************************************************************************/
inline int generalized_hamming_64 (uint64_t a) {
a |= a >> 1;
a |= a >> 2;
a |= a >> 4;
a &= 0x0101010101010101UL;
return popcount64 (a);
}
struct GenHammingComputer8 {
uint64_t a0;
GenHammingComputer8 (const uint8_t *a, int code_size) {
assert (code_size == 8);
a0 = *(uint64_t *)a;
}
inline int hamming (const uint8_t *b) const {
return generalized_hamming_64 (*(uint64_t *)b ^ a0);
}
};
struct GenHammingComputer16 {
uint64_t a0, a1;
GenHammingComputer16 (const uint8_t *a8, int code_size) {
assert (code_size == 16);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return generalized_hamming_64 (b[0] ^ a0) +
generalized_hamming_64 (b[1] ^ a1);
}
};
struct GenHammingComputer32 {
uint64_t a0, a1, a2, a3;
GenHammingComputer32 (const uint8_t *a8, int code_size) {
assert (code_size == 32);
const uint64_t *a = (uint64_t *)a8;
a0 = a[0]; a1 = a[1]; a2 = a[2]; a3 = a[3];
}
inline int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
return generalized_hamming_64 (b[0] ^ a0) +
generalized_hamming_64 (b[1] ^ a1) +
generalized_hamming_64 (b[2] ^ a2) +
generalized_hamming_64 (b[3] ^ a3);
}
};
struct GenHammingComputerM8 {
const uint64_t *a;
int n;
GenHammingComputerM8 (const uint8_t *a8, int code_size) {
assert (code_size % 8 == 0);
a = (uint64_t *)a8;
n = code_size / 8;
}
int hamming (const uint8_t *b8) const {
const uint64_t *b = (uint64_t *)b8;
int accu = 0;
for (int i = 0; i < n; i++)
accu += generalized_hamming_64 (a[i] ^ b[i]);
return accu;
}
};
/** generalized Hamming distances (= count number of code bytes that
are the same) */
void generalized_hammings_knn (
int_maxheap_array_t * ha,
const uint8_t * a,
const uint8_t * b,
size_t nb,
size_t code_size,
int ordered = true);
} // namespace faiss
#endif /* FAISS_hamming_h */