575 lines
14 KiB
C++
575 lines
14 KiB
C++
/**
|
|
* Copyright (c) Facebook, Inc. and its affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
// -*- c++ -*-
|
|
|
|
/*
|
|
* 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 {
|
|
|
|
|
|
extern size_t hamming_batch_size;
|
|
|
|
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,
|
|
* using a max heap.
|
|
* @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_hc (
|
|
int_maxheap_array_t * ha,
|
|
const uint8_t * a,
|
|
const uint8_t * b,
|
|
size_t nb,
|
|
size_t ncodes,
|
|
int ordered);
|
|
|
|
/* Legacy alias to hammings_knn_hc. */
|
|
void hammings_knn (
|
|
int_maxheap_array_t * ha,
|
|
const uint8_t * a,
|
|
const uint8_t * b,
|
|
size_t nb,
|
|
size_t ncodes,
|
|
int ordered);
|
|
|
|
/** Return the k smallest Hamming distances for a set of binary query vectors,
|
|
* using counting max.
|
|
* @param a queries, size na * ncodes
|
|
* @param b database, size nb * ncodes
|
|
* @param na number of query vectors
|
|
* @param nb number of database vectors
|
|
* @param k number of vectors/distances to return
|
|
* @param ncodes size of the binary codes (bytes)
|
|
* @param distances output distances from each query vector to its k nearest
|
|
* neighbors
|
|
* @param labels output ids of the k nearest neighbors to each query vector
|
|
*/
|
|
void hammings_knn_mc (
|
|
const uint8_t * a,
|
|
const uint8_t * b,
|
|
size_t na,
|
|
size_t nb,
|
|
size_t k,
|
|
size_t ncodes,
|
|
int32_t *distances,
|
|
long *labels);
|
|
|
|
/* 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 () {}
|
|
|
|
HammingComputer4 (const uint8_t *a, int code_size) {
|
|
set (a, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputer8 (const uint8_t *a, int code_size) {
|
|
set (a, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputer16 (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputer20 (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputer32 (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputer64 (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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);
|
|
}
|
|
|
|
};
|
|
|
|
// very inefficient...
|
|
struct HammingComputerDefault {
|
|
const uint8_t *a;
|
|
int n;
|
|
|
|
HammingComputerDefault () {}
|
|
|
|
HammingComputerDefault (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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 () {}
|
|
|
|
HammingComputerM8 (const uint8_t *a8, int code_size) {
|
|
set (a8, code_size);
|
|
}
|
|
|
|
void set (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;
|
|
}
|
|
|
|
};
|
|
|
|
// even more inefficient!
|
|
struct HammingComputerM4 {
|
|
const uint32_t *a;
|
|
int n;
|
|
|
|
HammingComputerM4 () {}
|
|
|
|
HammingComputerM4 (const uint8_t *a4, int code_size) {
|
|
set (a4, code_size);
|
|
}
|
|
|
|
void set (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_hc (
|
|
int_maxheap_array_t * ha,
|
|
const uint8_t * a,
|
|
const uint8_t * b,
|
|
size_t nb,
|
|
size_t code_size,
|
|
int ordered = true);
|
|
|
|
|
|
|
|
/** This class maintains a list of best distances seen so far.
|
|
*
|
|
* Since the distances are in a limited range (0 to nbit), the
|
|
* object maintains one list per possible distance, and fills
|
|
* in only the n-first lists, such that the sum of sizes of the
|
|
* n lists is below k.
|
|
*/
|
|
template<class HammingComputer>
|
|
struct HCounterState {
|
|
int *counters;
|
|
long *ids_per_dis;
|
|
|
|
HammingComputer hc;
|
|
int thres;
|
|
int count_lt;
|
|
int count_eq;
|
|
int k;
|
|
|
|
HCounterState(int *counters, long *ids_per_dis,
|
|
const uint8_t *x, int d, int k)
|
|
: counters(counters),
|
|
ids_per_dis(ids_per_dis),
|
|
hc(x, d / 8),
|
|
thres(d + 1),
|
|
count_lt(0),
|
|
count_eq(0),
|
|
k(k) {}
|
|
|
|
void update_counter(const uint8_t *y, size_t j) {
|
|
int32_t dis = hc.hamming(y);
|
|
|
|
if (dis <= thres) {
|
|
if (dis < thres) {
|
|
ids_per_dis[dis * k + counters[dis]++] = j;
|
|
++count_lt;
|
|
while (count_lt == k && thres > 0) {
|
|
--thres;
|
|
count_eq = counters[thres];
|
|
count_lt -= count_eq;
|
|
}
|
|
} else if (count_eq < k) {
|
|
ids_per_dis[dis * k + count_eq++] = j;
|
|
counters[dis] = count_eq;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
|
|
} // namespace faiss
|
|
|
|
|
|
#endif /* FAISS_hamming_h */
|