文件 simd_result_handlers.h

namespace faiss

使用多种变体的 k-means 聚类实现。

版权所有 (c) Facebook, Inc. 及其附属公司。

此源代码根据 MIT 许可证获得许可，该许可证位于此源树的根目录中的 LICENSE 文件中。

IDSelector 旨在定义要处理的向量子集（用于删除或作为搜索的子集）

PQ4 SIMD 打包和累积函数

基本内核累积 nq 查询向量，其中 bbs = nb * 2 * 16 向量，并生成一个输出矩阵。 对于 nq * nb <= 4 很有意思，否则寄存器溢出会太大。

这些函数的实现分布在 3 个 cpp 文件中，以减少并行编译时间。 模板被显式实例化。

此文件包含用于计算距离的内核的回调。

在整个库中，向量以 float * 指针的形式提供。 当批量处理（添加/搜索）多个向量时，大多数算法可以优化。 在这种情况下，它们作为矩阵传入。 当 n 个大小为 d 的向量作为 float * x 提供时，向量 i 的分量 j 是

x[ i * d + j ]

其中 0 <= i < n 且 0 <= j < d。 换句话说，矩阵始终是紧凑的。 指定矩阵的大小时，我们称它为 n*d 矩阵，这意味着行主存储。

I/O 函数可以读/写到文件名、文件句柄或抽象介质的对象。

读取函数返回的对象应该使用 delete 释放。 这些对象中的所有引用都归对象所有。

反向列表的定义 + 一些实现该接口的常见类。

由于 IVF（反向文件）索引对于大规模用例非常有用，因此我们将一些与它们相关的功能分组到这个小库中。 大多数函数都适用于 IndexIVF 和嵌入在 IndexPreTransform 中的 IndexIVF。

此文件中实现了 L2 和内积之外的额外度量

实现了一些神经网络层，主要用于支持 QINCo

定义了一些对象，这些对象将转换应用于一组向量。 通常这些是预处理步骤。

变量

FAISS_API bool simd_result_handlers_accept_virtual

struct SIMDResultHandler

由 faiss::SIMDResultHandlerToFloat, faiss::simd_result_handlers::DummyResultHandler, faiss::simd_result_handlers::FixedStorageHandler< NQ, BB >, faiss::simd_result_handlers::StoreResultHandler 继承

公共函数

virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) = 0

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

virtual void set_block_origin(size_t i0, size_t j0) = 0

设置正在计算的子矩阵

inline virtual ~SIMDResultHandler()

公共成员

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

struct SIMDResultHandlerToFloat : public faiss::SIMDResultHandler

由 faiss::simd_result_handlers::ResultHandlerCompare< C, false >, faiss::simd_result_handlers::ResultHandlerCompare< C, with_id_map > 继承

公共函数

inline SIMDResultHandlerToFloat(size_t nq, size_t ntotal)

inline virtual void begin(const float *norms)

inline virtual void end()

virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) = 0

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

virtual void set_block_origin(size_t i0, size_t j0) = 0

设置正在计算的子矩阵

公共成员

size_t nq

size_t ntotal

const idx_t *id_map = nullptr

这些字段主要用于IVF变体（with_id_map=true）

const int *q_map = nullptr

const uint16_t *dbias = nullptr

const float *normalizers = nullptr

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

namespace simd_result_handlers

函数

template<class C, bool W, class Consumer, class ...Types>
void dispatch_SIMDResultHandler_fixedCW(SIMDResultHandler &res, Consumer &consumer, Types... args)

template<class C, class Consumer, class ...Types>
void dispatch_SIMDResultHandler_fixedC(SIMDResultHandler &res, Consumer &consumer, Types... args)

template<class Consumer, class ...Types>
void dispatch_SIMDResultHandler(SIMDResultHandler &res, Consumer &consumer, Types... args)

struct DummyResultHandler : public faiss::SIMDResultHandler

#include <simd_result_handlers.h>

只是计算结果校验和的虚拟结构体（以避免计算被优化掉）

公共函数

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void set_block_origin(size_t, size_t) final

设置正在计算的子矩阵

inline ~DummyResultHandler()

公共成员

size_t cs = 0

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

struct StoreResultHandler : public faiss::SIMDResultHandler

#include <simd_result_handlers.h>

將結果儲存在一個 nq-by-nb 矩陣中。

j0 是矩陣目前左上角的塊

公共函数

inline StoreResultHandler(uint16_t *data, size_t ld)

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

公共成员

uint16_t *data

size_t ld

size_t i0 = 0

size_t j0 = 0

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<int NQ, int BB>
struct FixedStorageHandler : public faiss::SIMDResultHandler

#include <simd_result_handlers.h>

将结果存储在固定大小的矩阵中。

公共函数

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

template<class OtherResultHandler>
inline void to_other_handler(OtherResultHandler &other) const

inline virtual ~FixedStorageHandler()

公共成员

simd16uint16 dis[NQ][BB]

int i0 = 0

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<class C, bool with_id_map>
struct ResultHandlerCompare : public faiss::SIMDResultHandlerToFloat

#include <simd_result_handlers.h>

Result handler that compares distances to check if they need to be kept

Subclassed by faiss::simd_result_handlers::RangeHandler< C, false >

Public Types

using TI = typename C::TI

公共函数

inline ResultHandlerCompare(size_t nq, size_t ntotal, const IDSelector *sel_in)

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

inline void adjust_with_origin(size_t &q, simd16uint16 &d0, simd16uint16 &d1)

inline int64_t adjust_id(size_t b, size_t j)

inline uint32_t get_lt_mask(uint16_t thr, size_t b, simd16uint16 d0, simd16uint16 d1)

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

inline virtual ~ResultHandlerCompare()

inline virtual void begin(const float *norms)

inline virtual void end()

virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) = 0

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

公共成员

bool disable = false

int64_t i0 = 0

int64_t j0 = 0

const IDSelector *sel

size_t nq

size_t ntotal

const idx_t *id_map = nullptr

这些字段主要用于IVF变体（with_id_map=true）

const int *q_map = nullptr

const uint16_t *dbias = nullptr

const float *normalizers = nullptr

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<class C, bool with_id_map = false>
struct SingleResultHandler : public faiss::simd_result_handlers::ResultHandlerCompare<C, false>

#include <simd_result_handlers.h>

k=1 的特殊版本

Public Types

using T = typename C::T

using TI = typename C::TI

using RHC = ResultHandlerCompare<C, with_id_map>

公共函数

inline SingleResultHandler(size_t nq, size_t ntotal, float *dis, int64_t *ids, const IDSelector *sel_in)

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void end()

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

inline void adjust_with_origin(size_t &q, simd16uint16 &d0, simd16uint16 &d1)

inline int64_t adjust_id(size_t b, size_t j)

inline uint32_t get_lt_mask(uint16_t thr, size_t b, simd16uint16 d0, simd16uint16 d1)

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

inline virtual void begin(const float *norms)

公共成员

std::vector<int16_t> idis

float *dis

int64_t *ids

bool disable

int64_t i0

int64_t j0

const IDSelector *sel

size_t nq

size_t ntotal

const idx_t *id_map = nullptr

这些字段主要用于IVF变体（with_id_map=true）

const int *q_map = nullptr

const uint16_t *dbias = nullptr

const float *normalizers = nullptr

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<class C, bool with_id_map = false>
struct HeapHandler : public faiss::simd_result_handlers::ResultHandlerCompare<C, false>

#include <simd_result_handlers.h>

在最小堆或最大堆中收集结果的结构

Public Types

T 是 C::T 类型的别名。

TI 是 C::TI 类型的别名。

RHC 是 ResultHandlerCompare<C, with_id_map> 类型的别名。

公共函数

内联 HeapHandler 构造函数，接受查询数量 nq、总数量 ntotal、最近邻数量 k、距离数组 dis、ID 数组 ids 和 ID 选择器 sel_in 作为参数。

内联虚函数 handle，处理给定查询 q 和块 b 的 SIMD 距离 d0 和 d1。 `final` 关键字表示该函数不能在派生类中被重写。

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

内联虚函数 end，表示处理结束。 override 关键字表示该函数重写了基类中的一个虚函数。

内联虚函数 set_block_origin，设置块的起始位置 i0_in 和 j0_in。 `final` 关键字表示该函数不能在派生类中被重写。

设置正在计算的子矩阵

内联函数 adjust_with_origin，根据原点调整查询索引 q 和 SIMD 距离 d0 和 d1。

内联函数 adjust_id，根据块 b 和索引 j 调整 ID。

内联函数 get_lt_mask，获取小于阈值 thr 的 SIMD 距离 d0 和 d1 的掩码，给定块 b。

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

内联虚函数 begin，表示处理开始，接受范数数组 norms 作为参数。

公共成员

idis 是一个 std::vector<uint16_t> 类型的成员变量。

iids 是一个 std::vector<TI> 类型的成员变量。

float *dis

int64_t *ids

int64_t k

bool disable

int64_t i0

int64_t j0

const IDSelector *sel

size_t nq

size_t ntotal

const idx_t *id_map = nullptr

这些字段主要用于IVF变体（with_id_map=true）

const int *q_map = nullptr

const uint16_t *dbias = nullptr

const float *normalizers = nullptr

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<class C, bool with_id_map = false>
struct ReservoirHandler : public faiss::simd_result_handlers::ResultHandlerCompare<C, false>

#include <simd_result_handlers.h>

Simple top-N implementation using a reservoir.

Results are stored when they are below the threshold until the capacity is reached. Then a partition sort is used to update the threshold. Handler built from several ReservoirTopN (one per query)

Public Types

using T = typename C::T

using TI = typename C::TI

using RHC = ResultHandlerCompare<C, with_id_map>

公共函数

inline ReservoirHandler(size_t nq, size_t ntotal, size_t k, size_t cap, float *dis, int64_t *ids, const IDSelector *sel_in)

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void end() override

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

inline void adjust_with_origin(size_t &q, simd16uint16 &d0, simd16uint16 &d1)

inline int64_t adjust_id(size_t b, size_t j)

inline uint32_t get_lt_mask(uint16_t thr, size_t b, simd16uint16 d0, simd16uint16 d1)

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

inline virtual void begin(const float *norms)

公共成员

size_t capacity

float *dis

int64_t *ids

std::vector<TI> all_ids

AlignedTable<T> all_vals

std::vector<ReservoirTopN<C>> reservoirs

bool disable

int64_t i0

int64_t j0

const IDSelector *sel

size_t nq

size_t ntotal

const idx_t *id_map = nullptr

这些字段主要用于IVF变体（with_id_map=true）

const int *q_map = nullptr

const uint16_t *dbias = nullptr

const float *normalizers = nullptr

bool is_CMax = false

uint8_t sizeof_ids = 0

bool with_fields = false

template<class C, bool with_id_map = false>
struct RangeHandler : public faiss::simd_result_handlers::ResultHandlerCompare<C, false>

#include <simd_result_handlers.h>

范围搜索的结果处理器。 困难在于范围距离必须使用缩放器进行缩放。

Public Types

using T = typename C::T

using TI = typename C::TI

using RHC = ResultHandlerCompare<C, with_id_map>

公共函数

inline RangeHandler(RangeSearchResult &rres, float radius, size_t ntotal, const IDSelector *sel_in)

inline virtual void begin(const float *norms) override

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出32个距离并以两个 simd16uint16 提供时调用。(q, b) 表示它在块中的哪个条目。

inline virtual void end() override

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

inline void adjust_with_origin(size_t &q, simd16uint16 &d0, simd16uint16 &d1)

inline int64_t adjust_id(size_t b, size_t j)

inline uint32_t get_lt_mask(uint16_t thr, size_t b, simd16uint16 d0, simd16uint16 d1)

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

公共成员

RangeSearchResult &rres

float radius

std::vector<uint16_t> thresholds

每个查询的结果数量：std::vector<size_t> 类型，存储每个查询返回的结果数量。

查询的起始索引：size_t 类型，指示当前处理的查询在查询集合中的起始索引，默认为 0。

三元组列表：std::vector<Triplet> 类型，存储包含查询索引、数据库索引和距离值的三元组结果。

禁用标志：bool 类型，如果为真，则禁用该结果处理程序。

起始行索引：int64_t 类型，矩阵运算中的起始行索引。

起始列索引：int64_t 类型，矩阵运算中的起始列索引。

ID选择器：指向 IDSelector 对象的常量指针，用于根据ID进行选择。

查询数量：size_t 类型，指示查询的数量。

数据库总大小：size_t 类型，指示数据库中向量的总数。

ID映射：指向 idx_t 类型的常量指针，用于将内部索引映射到外部ID，可以为空。

这些字段主要用于IVF变体（with_id_map=true）

查询映射：指向 int 类型的常量指针，用于映射查询索引，可以为空。

距离偏差：指向 uint16_t 类型的常量指针，用于校正距离值，可以为空。

归一化器：指向 float 类型的常量指针，用于归一化距离值，可以为空。

是否为 CMax：bool 类型，指示是否使用 CMax 操作，默认为 false。

ID的大小：uint8_t 类型，以字节为单位指示ID的大小，默认为 0。

是否包含字段：bool 类型，指示结果是否包含额外的字段，默认为 false。

三元组结构体：定义一个名为 Triplet 的结构体。

公共成员

三元组的查询索引：idx_t 类型，表示查询索引。

idx_t b

uint16_t dis

template<class C, bool with_id_map = false>
struct PartialRangeHandler : public faiss::simd_result_handlers::RangeHandler<C, false>

Public Types

using T = typename C::T

using TI = typename C::TI

using RHC = RangeHandler<C, with_id_map>

公共函数

inline PartialRangeHandler(RangeSearchPartialResult &pres, float radius, size_t ntotal, size_t q0, size_t q1, const IDSelector *sel_in)

inline void shift_n_per_query()

inline virtual void end() override

inline virtual void begin(const float *norms) override

inline virtual void handle(size_t q, size_t b, simd16uint16 d0, simd16uint16 d1) final

当计算出 32 个距离并以两个 simd16uint16 提供时调用。(q, b) 指示它在块中的哪个条目。

inline virtual void set_block_origin(size_t i0_in, size_t j0_in) final

设置正在计算的子矩阵

inline void adjust_with_origin(size_t &q, simd16uint16 &d0, simd16uint16 &d1)

inline int64_t adjust_id(size_t b, size_t j)

inline uint32_t get_lt_mask(uint16_t thr, size_t b, simd16uint16 d0, simd16uint16 d1)

return binary mask of elements below thr in (d0, d1) inverse_test returns elements above

公共成员

RangeSearchPartialResult &pres

size_t q0

std::vector<Triplet> triplets

std::vector<size_t> n_per_query

RangeSearchResult &rres

float radius

std::vector<uint16_t> thresholds

disable：bool 类型，默认值为 false。

i0：int64_t 类型，默认值为 0。

j0：int64_t 类型，默认值为 0。

sel：指向 IDSelector 对象的常量指针。

nq：size_t 类型。

ntotal：size_t 类型。

id_map：指向 idx_t 类型的常量指针，默认值为 nullptr。

这些字段主要用于IVF变体（with_id_map=true）

q_map：指向 int 类型的常量指针，默认值为 nullptr。

dbias：指向 uint16_t 类型的常量指针，默认值为 nullptr。

normalizers：指向 float 类型的常量指针，默认值为 nullptr。

is_CMax：bool 类型，默认值为 false。

sizeof_ids：uint8_t 类型，默认值为 0。

with_fields：bool 类型，默认值为 false。