facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

move distances implementation into SIMD specific compile units (#4906)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4906

The main objective is to move the IVFFlat scanning code inside SIMD compile units, so that the loop over codes runs inside SIMD code.

Unfortunately, I could not instant
iate all the implementations with with_MetricType, hence the macro to define them.

Reviewed By: algoriddle

Differential Revision: D95936389

fbshipit-source-id: ae0515fcdcd7c880f329014b264b9dd05065f1c6
This commit is contained in:
Matthijs Douze
2026-03-12 02:30:36 -07:00
committed by meta-codesync[bot]
parent 000eec283a
commit d8c1a9785f
13 changed files with 344 additions and 317 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+1 -1
View File
@@ -289,7 +289,6 @@ set(FAISS_HEADERS
utils/WorkerThread.h
utils/distances.h
utils/distances_dispatch.h
utils/extra_distances-inl.h
utils/extra_distances.h
utils/fp16-fp16c.h
utils/fp16-inl.h
@@ -331,6 +330,7 @@ set(FAISS_HEADERS
utils/hamming_distance/avx512-inl.h
utils/simd_impl/distances_autovec-inl.h
utils/simd_impl/distances_simdlib256.h
utils/simd_impl/IVFFlatScanner-inl.h
utils/simd_impl/distances_sse-inl.h
)
faiss/IndexIVFFlat.cpp
+5 -43
View File
@@ -13,7 +13,6 @@
#include <cinttypes>
#include <cstdio>
#include <numeric>
#include <faiss/IndexFlat.h>
@@ -23,6 +22,11 @@
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/expanded_scanners.h>
#include <faiss/utils/extra_distances.h>
#define THE_SIMD_LEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
#include <faiss/utils/utils.h>
namespace faiss {
@@ -145,48 +149,6 @@ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
}
}
namespace {
template <typename VectorDistance>
struct IVFFlatScanner : InvertedListScanner {
VectorDistance vd;
using C = typename VectorDistance::C;
IVFFlatScanner(
const VectorDistance& vd,
bool store_pairs,
const IDSelector* sel)
: InvertedListScanner(store_pairs, sel), vd(vd) {
keep_max = vd.is_similarity;
code_size = vd.d * sizeof(float);
}
const float* xi;
void set_query(const float* query) override {
this->xi = query;
}
void set_list(idx_t list_no, float /* coarse_dis */) override {
this->list_no = list_no;
}
float distance_to_code(const uint8_t* code) const final {
const float* yj = (float*)code;
return vd(xi, yj);
}
// redefining the scan_codes allows to inline the distance_to_code
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const {
return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler);
}
};
} // anonymous namespace
InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
bool store_pairs,
const IDSelector* sel,
faiss/IndexIVFFlat.h
+32
View File
@@ -62,6 +62,38 @@ struct IndexIVFFlat : IndexIVF {
IndexIVFFlat();
};
template <typename VectorDistance>
struct IVFFlatScanner : InvertedListScanner {
VectorDistance vd;
using C = typename VectorDistance::C;
IVFFlatScanner(
const VectorDistance& vd,
bool store_pairs,
const IDSelector* sel)
: InvertedListScanner(store_pairs, sel), vd(vd) {
keep_max = vd.is_similarity;
code_size = vd.d * sizeof(float);
}
const float* xi;
void set_query(const float* query) override {
this->xi = query;
}
void set_list(idx_t list_no, float /* coarse_dis */) override {
this->list_no = list_no;
}
float distance_to_code(const uint8_t* code) const final;
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const;
};
struct IndexIVFFlatDedup : IndexIVFFlat {
/** Maps ids stored in the index to the ids of vectors that are
* the same. When a vector is unique, it does not appear in the
faiss/utils/distances_simd.cpp
+1 -2
View File
@@ -18,11 +18,10 @@
#include <faiss/impl/platform_macros.h>
#include <faiss/impl/simdlib/simdlib_dispatch.h>
#define AUTOVEC_LEVEL SIMDLevel::NONE
#define THE_SIMD_LEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
#define THE_SIMDLEVEL SIMDLevel::NONE
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/distances_simdlib256.h>
faiss/utils/extra_distances-inl.h
-235
View File
@@ -1,235 +0,0 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
/** In this file are the implementations of extra metrics beyond L2
* and inner product */
#include <faiss/MetricType.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/distances.h>
#include <cmath>
#include <type_traits>
namespace faiss {
/***************************************************************************
* VectorDistance base class - contains common data members and type defs
**************************************************************************/
template <MetricType mt>
struct VectorDistanceBase {
size_t d;
float metric_arg;
static constexpr MetricType metric = mt;
static constexpr bool is_similarity = is_similarity_metric(mt);
using C = typename std::conditional<
is_similarity_metric(mt),
CMin<float, int64_t>,
CMax<float, int64_t>>::type;
};
/***************************************************************************
* VectorDistance struct template - specializations for each metric type
**************************************************************************/
template <MetricType mt, SIMDLevel level>
struct VectorDistance : VectorDistanceBase<mt> {
inline float operator()(const float* x, const float* y) const;
};
template <SIMDLevel level>
struct VectorDistance<METRIC_L2, level> : VectorDistanceBase<METRIC_L2> {
inline float operator()(const float* x, const float* y) const {
return fvec_L2sqr<level>(x, y, this->d);
}
};
template <SIMDLevel level>
struct VectorDistance<METRIC_INNER_PRODUCT, level>
: VectorDistanceBase<METRIC_INNER_PRODUCT> {
inline float operator()(const float* x, const float* y) const {
return fvec_inner_product<level>(x, y, this->d);
}
};
template <SIMDLevel level>
struct VectorDistance<METRIC_L1, level> : VectorDistanceBase<METRIC_L1> {
inline float operator()(const float* x, const float* y) const {
return fvec_L1<level>(x, y, this->d);
}
};
template <SIMDLevel level>
struct VectorDistance<METRIC_Linf, level> : VectorDistanceBase<METRIC_Linf> {
inline float operator()(const float* x, const float* y) const {
return fvec_Linf<level>(x, y, this->d);
}
};
template <>
struct VectorDistance<METRIC_Lp, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_Lp> {
inline float operator()(const float* x, const float* y) const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float diff = fabs(x[i] - y[i]);
accu += powf(diff, this->metric_arg);
}
return accu;
}
};
template <>
struct VectorDistance<METRIC_Canberra, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_Canberra> {
inline float operator()(const float* x, const float* y) const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
accu += fabs(xi - yi) / (fabs(xi) + fabs(yi));
}
return accu;
}
};
template <>
struct VectorDistance<METRIC_BrayCurtis, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_BrayCurtis> {
inline float operator()(const float* x, const float* y) const {
float accu_num = 0, accu_den = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
accu_num += fabs(xi - yi);
accu_den += fabs(xi + yi);
}
return accu_num / accu_den;
}
};
template <>
struct VectorDistance<METRIC_JensenShannon, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_JensenShannon> {
inline float operator()(const float* x, const float* y) const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
float mi = 0.5 * (xi + yi);
float kl1 = -xi * log(mi / xi);
float kl2 = -yi * log(mi / yi);
accu += kl1 + kl2;
}
return 0.5 * accu;
}
};
template <>
struct VectorDistance<METRIC_Jaccard, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_Jaccard> {
inline float operator()(const float* x, const float* y) const {
// WARNING: this distance is defined only for positive input vectors.
// Providing vectors with negative values would lead to incorrect
// results.
float accu_num = 0, accu_den = 0;
for (size_t i = 0; i < this->d; i++) {
accu_num += fmin(x[i], y[i]);
accu_den += fmax(x[i], y[i]);
}
return accu_num / accu_den;
}
};
template <>
struct VectorDistance<METRIC_NaNEuclidean, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_NaNEuclidean> {
inline float operator()(const float* x, const float* y) const {
// https://scikit-learn.org/stable/modules/generated/sklearn.metrics.pairwise.nan_euclidean_distances.html
float accu = 0;
size_t present = 0;
for (size_t i = 0; i < this->d; i++) {
if (!std::isnan(x[i]) && !std::isnan(y[i])) {
float diff = x[i] - y[i];
accu += diff * diff;
present++;
}
}
if (present == 0) {
return NAN;
}
return float(this->d) / float(present) * accu;
}
};
template <>
struct VectorDistance<METRIC_GOWER, SIMDLevel::NONE>
: VectorDistanceBase<METRIC_GOWER> {
inline float operator()(const float* x, const float* y) const {
float accu = 0;
size_t valid_dims = 0;
for (size_t i = 0; i < this->d; i++) {
if (std::isnan(x[i]) || std::isnan(y[i])) {
continue;
}
if (x[i] >= 0 && y[i] >= 0) {
if (x[i] > 1 || y[i] > 1) {
return std::numeric_limits<float>::quiet_NaN();
}
accu += fabs(x[i] - y[i]);
} else if (x[i] < 0 && y[i] < 0) {
accu += float(int(x[i] != y[i]));
} else {
return std::numeric_limits<float>::quiet_NaN();
}
valid_dims++;
}
if (valid_dims == 0) {
return std::numeric_limits<float>::quiet_NaN();
}
return accu / valid_dims;
}
};
/***************************************************************************
* Dispatching function that takes a lambda directly.
* The lambda should be templated on VectorDistance, eg.:
*
* auto result = with_VectorDistance(
* metric, metric_arg, [&]<class VD>(VD vd) {
* return vd(x, y);
* });
**************************************************************************/
template <typename LambdaType>
auto with_VectorDistance(
size_t d,
MetricType metric,
float metric_arg,
LambdaType&& action) {
auto dispatch_metric = [&]<MetricType mt>() {
auto call = [&]<SIMDLevel level>() {
VectorDistance<mt, level> vd = {d, metric_arg};
return action(vd);
};
constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
if constexpr (!has_simd) {
return call.template operator()<SIMDLevel::NONE>();
} else {
DISPATCH_SIMDLevel(call.template operator());
}
};
return with_metric_type(metric, dispatch_metric);
}
} // namespace faiss
faiss/utils/extra_distances.h
+64 -6
View File
@@ -10,12 +10,12 @@
/** In this file are the implementations of extra metrics beyond L2
* and inner product */
#include <stdint.h>
#include <cstdint>
#include <faiss/Index.h>
#include <faiss/MetricType.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/Heap.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/utils/ordered_key_value.h>
namespace faiss {
@@ -100,6 +100,64 @@ inline auto with_metric_type(MetricType metric, LambdaType&& action) {
}
#endif // SWIG
} // namespace faiss
#ifndef SWIG
#include <faiss/utils/extra_distances-inl.h>
/***************************************************************************
* VectorDistance base class - contains common data members and type defs
* VectorDistance struct template - specializations for each metric type
**************************************************************************/
template <MetricType mt, SIMDLevel level>
struct VectorDistance {
size_t d;
float metric_arg;
VectorDistance(size_t d, float metric_arg) : d(d), metric_arg(metric_arg) {}
static constexpr MetricType metric = mt;
static constexpr bool is_similarity = is_similarity_metric(mt);
using C = typename std::conditional<
is_similarity_metric(mt),
CMin<float, int64_t>,
CMax<float, int64_t>>::type;
float operator()(const float* x, const float* y) const;
};
/***************************************************************************
* Dispatching function that takes a lambda directly.
* The lambda should be templated on VectorDistance, eg.:
*
* auto result = with_VectorDistance(
* metric, metric_arg, [&]<class VD>(VD vd) {
* return vd(x, y);
* });
**************************************************************************/
template <typename LambdaType>
auto with_VectorDistance(
size_t d,
MetricType metric,
float metric_arg,
LambdaType&& action) {
auto dispatch_metric = [&]<MetricType mt>() {
auto call = [&]<SIMDLevel level>() {
VectorDistance<mt, level> vd = {d, metric_arg};
return action(vd);
};
constexpr bool has_simd = mt == METRIC_INNER_PRODUCT ||
mt == METRIC_L2 || mt == METRIC_L1 || mt == METRIC_Linf;
if constexpr (!has_simd) {
return call.template operator()<SIMDLevel::NONE>();
} else {
DISPATCH_SIMDLevel(call.template operator());
}
};
return with_metric_type(metric, dispatch_metric);
}
#endif // SWIG
} // namespace faiss
faiss/utils/simd_impl/IVFFlatScanner-inl.h
+51
View File
@@ -0,0 +1,51 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <faiss/IndexIVFFlat.h>
#include <faiss/impl/expanded_scanners.h>
#ifndef THE_SIMD_LEVEL
#error "THE_SIMD_LEVEL not defined"
#endif
namespace faiss {
constexpr faiss::SIMDLevel THE_SL = THE_SIMD_LEVEL;
#define DEFINE_IVFFLAT_SCANNER_METHODS(mt) \
template <> \
float IVFFlatScanner<VectorDistance<mt, THE_SL>>::distance_to_code( \
const uint8_t* code) const { \
const float* yj = (float*)code; \
return vd(xi, yj); \
} \
template <> \
size_t IVFFlatScanner<VectorDistance<mt, THE_SL>>::scan_codes( \
size_t list_size, \
const uint8_t* codes, \
const idx_t* ids, \
ResultHandler& handler) const { \
return run_scan_codes_fix_C<C>(*this, list_size, codes, ids, handler); \
}
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_L2)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_INNER_PRODUCT)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_L1)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Linf)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Lp)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Canberra)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_BrayCurtis)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_JensenShannon)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_Jaccard)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_NaNEuclidean)
DEFINE_IVFFLAT_SCANNER_METHODS(METRIC_GOWER)
#undef DEFINE_IVFFLAT_SCANNER_METHODS
} // namespace faiss
faiss/utils/simd_impl/distances_aarch64.cpp
+3 -2
View File
@@ -12,10 +12,11 @@
#include <arm_neon.h>
#include <limits>
#define AUTOVEC_LEVEL SIMDLevel::ARM_NEON
#define THE_SIMD_LEVEL SIMDLevel::ARM_NEON
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
#define THE_SIMDLEVEL SIMDLevel::ARM_NEON
#include <faiss/utils/simd_impl/distances_simdlib256.h>
namespace faiss {
faiss/utils/simd_impl/distances_arm_sve.cpp
+3 -1
View File
@@ -9,8 +9,10 @@
#include <faiss/utils/distances.h>
#define AUTOVEC_LEVEL SIMDLevel::ARM_SVE
#define THE_SIMD_LEVEL SIMDLevel::ARM_SVE
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
namespace faiss {
faiss/utils/simd_impl/distances_autovec-inl.h
+163 -10
View File
@@ -7,14 +7,26 @@
#pragma once
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/extra_distances.h>
#ifndef THE_SIMD_LEVEL
#error "THE_SIMD_LEVEL not defined"
#endif
namespace faiss {
constexpr faiss::SIMDLevel SL = THE_SIMD_LEVEL;
/******************************************************************
* These functions are simple enough that the compile will do a good job
* vectorizing them given the appropriate flags.
******************************************************************/
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
float fvec_norm_L2sqr<AUTOVEC_LEVEL>(const float* x, size_t d) {
float fvec_norm_L2sqr<SL>(const float* x, size_t d) {
// the double in the _ref is suspected to be a typo. Some of the manual
// implementations this replaces used float.
float res = 0;
@@ -29,7 +41,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
float fvec_L2sqr<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
float fvec_L2sqr<SL>(const float* x, const float* y, size_t d) {
size_t i;
float res = 0;
FAISS_PRAGMA_IMPRECISE_LOOP
@@ -43,10 +55,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
float fvec_inner_product<AUTOVEC_LEVEL>(
const float* x,
const float* y,
size_t d) {
float fvec_inner_product<SL>(const float* x, const float* y, size_t d) {
float res = 0.F;
FAISS_PRAGMA_IMPRECISE_LOOP
for (size_t i = 0; i != d; ++i) {
@@ -58,7 +67,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
float fvec_L1<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
float fvec_L1<SL>(const float* x, const float* y, size_t d) {
size_t i;
float res = 0;
FAISS_PRAGMA_IMPRECISE_LOOP
@@ -72,7 +81,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
float fvec_Linf<AUTOVEC_LEVEL>(const float* x, const float* y, size_t d) {
float fvec_Linf<SL>(const float* x, const float* y, size_t d) {
float res = 0;
FAISS_PRAGMA_IMPRECISE_LOOP
for (size_t i = 0; i < d; i++) {
@@ -84,7 +93,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
void fvec_inner_product_batch_4<AUTOVEC_LEVEL>(
void fvec_inner_product_batch_4<SL>(
const float* x,
const float* y0,
const float* y1,
@@ -116,7 +125,7 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
template <>
void fvec_L2sqr_batch_4<AUTOVEC_LEVEL>(
void fvec_L2sqr_batch_4<SL>(
const float* x,
const float* y0,
const float* y1,
@@ -150,4 +159,148 @@ void fvec_L2sqr_batch_4<AUTOVEC_LEVEL>(
}
FAISS_PRAGMA_IMPRECISE_FUNCTION_END
/******************************************************************
* VectorDistance::operator() specializations — defined out-of-class
* so that SIMD compilation units produce externally-linkable symbols.
******************************************************************/
template <>
float VectorDistance<METRIC_L2, SL>::operator()(const float* x, const float* y)
const {
return fvec_L2sqr<SL>(x, y, this->d);
}
template <>
float VectorDistance<METRIC_INNER_PRODUCT, SL>::operator()(
const float* x,
const float* y) const {
return fvec_inner_product<SL>(x, y, this->d);
}
template <>
float VectorDistance<METRIC_L1, SL>::operator()(const float* x, const float* y)
const {
return fvec_L1<SL>(x, y, this->d);
}
template <>
float VectorDistance<METRIC_Linf, SL>::operator()(
const float* x,
const float* y) const {
return fvec_Linf<SL>(x, y, this->d);
}
template <>
float VectorDistance<METRIC_Lp, SL>::operator()(const float* x, const float* y)
const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float diff = fabs(x[i] - y[i]);
accu += powf(diff, this->metric_arg);
}
return accu;
}
template <>
float VectorDistance<METRIC_Canberra, SL>::operator()(
const float* x,
const float* y) const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
accu += fabs(xi - yi) / (fabs(xi) + fabs(yi));
}
return accu;
}
template <>
float VectorDistance<METRIC_BrayCurtis, SL>::operator()(
const float* x,
const float* y) const {
float accu_num = 0, accu_den = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
accu_num += fabs(xi - yi);
accu_den += fabs(xi + yi);
}
return accu_num / accu_den;
}
template <>
float VectorDistance<METRIC_JensenShannon, SL>::operator()(
const float* x,
const float* y) const {
float accu = 0;
for (size_t i = 0; i < this->d; i++) {
float xi = x[i], yi = y[i];
float mi = 0.5 * (xi + yi);
float kl1 = -xi * log(mi / xi);
float kl2 = -yi * log(mi / yi);
accu += kl1 + kl2;
}
return 0.5 * accu;
}
template <>
float VectorDistance<METRIC_Jaccard, SL>::operator()(
const float* x,
const float* y) const {
float accu_num = 0, accu_den = 0;
for (size_t i = 0; i < this->d; i++) {
accu_num += fmin(x[i], y[i]);
accu_den += fmax(x[i], y[i]);
}
return accu_num / accu_den;
}
template <>
float VectorDistance<METRIC_NaNEuclidean, SL>::operator()(
const float* x,
const float* y) const {
float accu = 0;
size_t present = 0;
for (size_t i = 0; i < this->d; i++) {
if (!std::isnan(x[i]) && !std::isnan(y[i])) {
float diff = x[i] - y[i];
accu += diff * diff;
present++;
}
}
if (present == 0) {
return NAN;
}
return float(this->d) / float(present) * accu;
}
template <>
float VectorDistance<METRIC_GOWER, SL>::operator()(
const float* x,
const float* y) const {
float accu = 0;
size_t valid_dims = 0;
for (size_t i = 0; i < this->d; i++) {
if (std::isnan(x[i]) || std::isnan(y[i])) {
continue;
}
if (x[i] >= 0 && y[i] >= 0) {
if (x[i] > 1 || y[i] > 1) {
return std::numeric_limits<float>::quiet_NaN();
}
accu += fabs(x[i] - y[i]);
} else if (x[i] < 0 && y[i] < 0) {
accu += float(int(x[i] != y[i]));
} else {
return std::numeric_limits<float>::quiet_NaN();
}
valid_dims++;
}
if (valid_dims == 0) {
return std::numeric_limits<float>::quiet_NaN();
}
return accu / valid_dims;
}
} // namespace faiss
faiss/utils/simd_impl/distances_avx2.cpp
+3 -2
View File
@@ -9,11 +9,10 @@
#include <immintrin.h>
#define AUTOVEC_LEVEL SIMDLevel::AVX2
#define THE_SIMD_LEVEL SIMDLevel::AVX2
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
#define THE_SIMDLEVEL SIMDLevel::AVX2
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/distances_simdlib256.h>
@@ -21,6 +20,8 @@
#include <faiss/utils/simd_impl/distances_sse-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/transpose/transpose-avx2-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
namespace faiss {
faiss/utils/simd_impl/distances_avx512.cpp
+4 -1
View File
@@ -9,8 +9,11 @@
#include <immintrin.h>
#define AUTOVEC_LEVEL SIMDLevel::AVX512
#define THE_SIMD_LEVEL SIMDLevel::AVX512
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/IVFFlatScanner-inl.h>
#include <faiss/utils/simd_impl/distances_sse-inl.h>
#include <faiss/utils/transpose/transpose-avx512-inl.h>
faiss/utils/simd_impl/distances_simdlib256.h
+14 -14
View File
@@ -17,14 +17,14 @@
namespace faiss {
template <>
void fvec_sub<THE_SIMDLEVEL>(
void fvec_sub<THE_SIMD_LEVEL>(
size_t d,
const float* a,
const float* b,
float* c) {
size_t i;
for (i = 0; i + 7 < d; i += 8) {
simd8float32_tpl<THE_SIMDLEVEL> ci, ai, bi;
simd8float32_tpl<THE_SIMD_LEVEL> ci, ai, bi;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai - bi;
@@ -36,14 +36,14 @@ void fvec_sub<THE_SIMDLEVEL>(
}
template <>
void fvec_add<THE_SIMDLEVEL>(
void fvec_add<THE_SIMD_LEVEL>(
size_t d,
const float* a,
const float* b,
float* c) {
size_t i;
for (i = 0; i + 7 < d; i += 8) {
simd8float32_tpl<THE_SIMDLEVEL> ci, ai, bi;
simd8float32_tpl<THE_SIMD_LEVEL> ci, ai, bi;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai + bi;
@@ -55,11 +55,11 @@ void fvec_add<THE_SIMDLEVEL>(
}
template <>
void fvec_add<THE_SIMDLEVEL>(size_t d, const float* a, float b, float* c) {
void fvec_add<THE_SIMD_LEVEL>(size_t d, const float* a, float b, float* c) {
size_t i;
simd8float32_tpl<THE_SIMDLEVEL> bv(b);
simd8float32_tpl<THE_SIMD_LEVEL> bv(b);
for (i = 0; i + 7 < d; i += 8) {
simd8float32_tpl<THE_SIMDLEVEL> ci, ai;
simd8float32_tpl<THE_SIMD_LEVEL> ci, ai;
ai.loadu(a + i);
ci = ai + bv;
ci.storeu(c + i);
@@ -137,7 +137,7 @@ simd8float32_tpl<SL> load_simd8float32_partial(const float* x, int n) {
} // anonymous namespace
template <>
void compute_PQ_dis_tables_dsub2<THE_SIMDLEVEL>(
void compute_PQ_dis_tables_dsub2<THE_SIMD_LEVEL>(
size_t d,
size_t ksub,
const float* all_centroids,
@@ -151,7 +151,7 @@ void compute_PQ_dis_tables_dsub2<THE_SIMDLEVEL>(
for (size_t m0 = 0; m0 < M; m0 += 4) {
int m1 = std::min(M, m0 + 4);
for (int k0 = 0; k0 < ksub; k0 += 8) {
simd8float32_tpl<THE_SIMDLEVEL> centroids[8];
simd8float32_tpl<THE_SIMD_LEVEL> centroids[8];
for (int k = 0; k < 8; k++) {
ALIGNED(32) float centroid[8];
size_t wp = 0;
@@ -161,26 +161,26 @@ void compute_PQ_dis_tables_dsub2<THE_SIMDLEVEL>(
centroid[wp++] = all_centroids[rp + 1];
rp += 2 * ksub;
}
centroids[k] = simd8float32_tpl<THE_SIMDLEVEL>(centroid);
centroids[k] = simd8float32_tpl<THE_SIMD_LEVEL>(centroid);
}
for (size_t i = 0; i < nx; i++) {
simd8float32_tpl<THE_SIMDLEVEL> xi;
simd8float32_tpl<THE_SIMD_LEVEL> xi;
if (m1 == m0 + 4) {
xi.loadu(x + i * d + m0 * 2);
} else {
xi = load_simd8float32_partial<THE_SIMDLEVEL>(
xi = load_simd8float32_partial<THE_SIMD_LEVEL>(
x + i * d + m0 * 2, 2 * (m1 - m0));
}
if (is_inner_product) {
pq2_8cents_table<THE_SIMDLEVEL, true>(
pq2_8cents_table<THE_SIMD_LEVEL, true>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,
ksub,
m1 - m0);
} else {
pq2_8cents_table<THE_SIMDLEVEL, false>(
pq2_8cents_table<THE_SIMD_LEVEL, false>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,