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

convert simdlib in distances_simd.cpp (#4884)

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

Templatize fvec_sub, fvec_add, and compute_PQ_dis_tables_dsub2 on SIMDLevel.

Move these 256-bit simdlib-based functions from distances_simd.cpp into a new header simd_impl/distances_simdlib256.h, templatized on THE_SIMDLEVEL. Add with_simd_level_256bit() to simd_dispatch.h that maps AVX512->AVX2 and ARM_SVE->ARM_NEON using simd256_level_selector. Dispatch wrappers use with_simd_level_256bit with lambdas.

Reviewed By: algoriddle

Differential Revision: D95570445

fbshipit-source-id: 557b6546e9526c40a44f7165584e96a9e39a3b89
This commit is contained in:
Matthijs Douze
2026-03-09 07:09:33 -07:00
committed by meta-codesync[bot]
parent 8d8268c8ae
commit bb298a2323
9 changed files with 308 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+1
View File
@@ -322,6 +322,7 @@ set(FAISS_HEADERS
utils/hamming_distance/avx2-inl.h
utils/hamming_distance/avx512-inl.h
utils/simd_impl/distances_autovec-inl.h
utils/simd_impl/distances_simdlib256.h
utils/simd_impl/distances_sse-inl.h
)
faiss/impl/simd_dispatch.h
+17
View File
@@ -136,4 +136,21 @@ inline auto with_simd_level(LambdaType&& action) {
DISPATCH_SIMDLevel(action.template operator());
}
/**
* Like with_simd_level, but maps to the 256-bit SIMD equivalent:
* AVX512, AVX512_SPR -> AVX2
* ARM_SVE -> ARM_NEON
* AVX2, ARM_NEON, NONE -> unchanged
*
* Use for functions implemented with simd8float32 (256-bit) operations
* that don't have dedicated AVX512 or SVE implementations.
*/
template <typename LambdaType>
inline auto with_simd_level_256bit(LambdaType&& action) {
return with_simd_level([&]<SIMDLevel level>() {
constexpr SIMDLevel level256 = simd256_level_selector<level>::value;
return action.template operator()<level256>();
});
}
} // namespace faiss
faiss/utils/distances.cpp
+24
View File
@@ -172,6 +172,30 @@ int fvec_madd_and_argmin(
return fvec_madd_and_argmin_dispatch(n, a, bf, b, c);
}
void fvec_sub(size_t d, const float* a, const float* b, float* c) {
fvec_sub_dispatch(d, a, b, c);
}
void fvec_add(size_t d, const float* a, const float* b, float* c) {
fvec_add_dispatch(d, a, b, c);
}
void fvec_add(size_t d, const float* a, float b, float* c) {
fvec_add_scalar_dispatch(d, a, b, c);
}
void compute_PQ_dis_tables_dsub2(
size_t d,
size_t ksub,
const float* all_centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables) {
compute_PQ_dis_tables_dsub2_dispatch(
d, ksub, all_centroids, nx, x, is_inner_product, dis_tables);
}
/***************************************************************************
* Matrix/vector ops
***************************************************************************/
faiss/utils/distances.h
+19
View File
@@ -261,6 +261,9 @@ void inner_product_to_L2sqr(
*/
void fvec_add(size_t d, const float* a, const float* b, float* c);
template <SIMDLevel>
void fvec_add(size_t d, const float* a, const float* b, float* c);
/** compute c := a + b for a, c vectors and b a scalar
*
* c and a can overlap
@@ -270,6 +273,9 @@ void fvec_add(size_t d, const float* a, const float* b, float* c);
*/
void fvec_add(size_t d, const float* a, float b, float* c);
template <SIMDLevel>
void fvec_add(size_t d, const float* a, float b, float* c);
/** compute c := a - b for vectors
*
* c and a can overlap, c and b can overlap
@@ -280,6 +286,9 @@ void fvec_add(size_t d, const float* a, float b, float* c);
*/
void fvec_sub(size_t d, const float* a, const float* b, float* c);
template <SIMDLevel>
void fvec_sub(size_t d, const float* a, const float* b, float* c);
/***************************************************************************
* Compute a subset of distances
***************************************************************************/
@@ -542,6 +551,16 @@ void compute_PQ_dis_tables_dsub2(
bool is_inner_product,
float* dis_tables);
template <SIMDLevel>
void compute_PQ_dis_tables_dsub2(
size_t d,
size_t ksub,
const float* centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables);
/***************************************************************************
* Templatized versions of distance functions
***************************************************************************/
faiss/utils/distances_dispatch.h
+41
View File
@@ -167,4 +167,45 @@ inline int fvec_madd_and_argmin_dispatch(
DISPATCH_SIMDLevel(fvec_madd_and_argmin, n, a, bf, b, c);
}
inline void fvec_sub_dispatch(
size_t d,
const float* a,
const float* b,
float* c) {
with_simd_level_256bit(
[&]<SIMDLevel level>() { fvec_sub<level>(d, a, b, c); });
}
inline void fvec_add_dispatch(
size_t d,
const float* a,
const float* b,
float* c) {
with_simd_level_256bit(
[&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
}
inline void fvec_add_scalar_dispatch(
size_t d,
const float* a,
float b,
float* c) {
with_simd_level_256bit(
[&]<SIMDLevel level>() { fvec_add<level>(d, a, b, c); });
}
inline void compute_PQ_dis_tables_dsub2_dispatch(
size_t d,
size_t ksub,
const float* centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables) {
with_simd_level_256bit([&]<SIMDLevel level>() {
compute_PQ_dis_tables_dsub2<level>(
d, ksub, centroids, nx, x, is_inner_product, dis_tables);
});
}
} // namespace faiss
faiss/utils/distances_simd.cpp
+4 -174
View File
@@ -22,6 +22,10 @@
// 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>
namespace faiss {
/*******
@@ -168,177 +172,3 @@ int fvec_madd_and_argmin<SIMDLevel::NONE>(
}
} // namespace faiss
namespace faiss {
/***************************************************************************
* PQ tables computations
***************************************************************************/
namespace {
/// compute the IP for dsub = 2 for 8 centroids and 4 sub-vectors at a time
template <bool is_inner_product>
void pq2_8cents_table(
const simd8float32 centroids[8],
const simd8float32 x,
float* out,
size_t ldo,
size_t nout = 4) {
simd8float32 ips[4];
for (int i = 0; i < 4; i++) {
simd8float32 p1, p2;
if (is_inner_product) {
p1 = x * centroids[2 * i];
p2 = x * centroids[2 * i + 1];
} else {
p1 = (x - centroids[2 * i]);
p1 = p1 * p1;
p2 = (x - centroids[2 * i + 1]);
p2 = p2 * p2;
}
ips[i] = hadd(p1, p2);
}
simd8float32 ip02a = geteven(ips[0], ips[1]);
simd8float32 ip02b = geteven(ips[2], ips[3]);
simd8float32 ip0 = getlow128(ip02a, ip02b);
simd8float32 ip2 = gethigh128(ip02a, ip02b);
simd8float32 ip13a = getodd(ips[0], ips[1]);
simd8float32 ip13b = getodd(ips[2], ips[3]);
simd8float32 ip1 = getlow128(ip13a, ip13b);
simd8float32 ip3 = gethigh128(ip13a, ip13b);
switch (nout) {
case 4:
ip3.storeu(out + 3 * ldo);
[[fallthrough]];
case 3:
ip2.storeu(out + 2 * ldo);
[[fallthrough]];
case 2:
ip1.storeu(out + 1 * ldo);
[[fallthrough]];
case 1:
ip0.storeu(out);
}
}
simd8float32 load_simd8float32_partial(const float* x, int n) {
ALIGNED(32) float tmp[8] = {0, 0, 0, 0, 0, 0, 0, 0};
float* wp = tmp;
for (int i = 0; i < n; i++) {
*wp++ = *x++;
}
return simd8float32(tmp);
}
} // anonymous namespace
void compute_PQ_dis_tables_dsub2(
size_t d,
size_t ksub,
const float* all_centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables) {
size_t M = d / 2;
FAISS_THROW_IF_NOT(ksub % 8 == 0);
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 centroids[8];
for (int k = 0; k < 8; k++) {
ALIGNED(32) float centroid[8];
size_t wp = 0;
size_t rp = (m0 * ksub + k + k0) * 2;
for (int m = m0; m < m1; m++) {
centroid[wp++] = all_centroids[rp];
centroid[wp++] = all_centroids[rp + 1];
rp += 2 * ksub;
}
centroids[k] = simd8float32(centroid);
}
for (size_t i = 0; i < nx; i++) {
simd8float32 xi;
if (m1 == m0 + 4) {
xi.loadu(x + i * d + m0 * 2);
} else {
xi = load_simd8float32_partial(
x + i * d + m0 * 2, 2 * (m1 - m0));
}
if (is_inner_product) {
pq2_8cents_table<true>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,
ksub,
m1 - m0);
} else {
pq2_8cents_table<false>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,
ksub,
m1 - m0);
}
}
}
}
}
/*********************************************************
* Vector to vector functions
*********************************************************/
void fvec_sub(size_t d, const float* a, const float* b, float* c) {
size_t i;
for (i = 0; i + 7 < d; i += 8) {
simd8float32 ci, ai, bi;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai - bi;
ci.storeu(c + i);
}
// finish non-multiple of 8 remainder
for (; i < d; i++) {
c[i] = a[i] - b[i];
}
}
void fvec_add(size_t d, const float* a, const float* b, float* c) {
size_t i;
for (i = 0; i + 7 < d; i += 8) {
simd8float32 ci, ai, bi;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai + bi;
ci.storeu(c + i);
}
// finish non-multiple of 8 remainder
for (; i < d; i++) {
c[i] = a[i] + b[i];
}
}
void fvec_add(size_t d, const float* a, float b, float* c) {
size_t i;
simd8float32 bv(b);
for (i = 0; i + 7 < d; i += 8) {
simd8float32 ci, ai;
ai.loadu(a + i);
ci = ai + bv;
ci.storeu(c + i);
}
// finish non-multiple of 8 remainder
for (; i < d; i++) {
c[i] = a[i] + b;
}
}
} // namespace faiss
faiss/utils/simd_impl/distances_aarch64.cpp
+3
View File
@@ -15,6 +15,9 @@
#define AUTOVEC_LEVEL SIMDLevel::ARM_NEON
#include <faiss/utils/simd_impl/distances_autovec-inl.h>
#define THE_SIMDLEVEL SIMDLevel::ARM_NEON
#include <faiss/utils/simd_impl/distances_simdlib256.h>
namespace faiss {
template <>
faiss/utils/simd_impl/distances_avx2.cpp
+4
View File
@@ -13,6 +13,10 @@
// 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>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
#include <faiss/utils/simd_impl/distances_sse-inl.h>
// NOLINTNEXTLINE(facebook-hte-InlineHeader)
faiss/utils/simd_impl/distances_simdlib256.h
+195
View File
@@ -0,0 +1,195 @@
/*
* 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 <algorithm>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/platform_macros.h>
#include <faiss/impl/simdlib/simdlib_dispatch.h>
#include <faiss/utils/distances.h>
namespace faiss {
template <>
void fvec_sub<THE_SIMDLEVEL>(
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;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai - bi;
ci.storeu(c + i);
}
for (; i < d; i++) {
c[i] = a[i] - b[i];
}
}
template <>
void fvec_add<THE_SIMDLEVEL>(
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;
ai.loadu(a + i);
bi.loadu(b + i);
ci = ai + bi;
ci.storeu(c + i);
}
for (; i < d; i++) {
c[i] = a[i] + b[i];
}
}
template <>
void fvec_add<THE_SIMDLEVEL>(size_t d, const float* a, float b, float* c) {
size_t i;
simd8float32_tpl<THE_SIMDLEVEL> bv(b);
for (i = 0; i + 7 < d; i += 8) {
simd8float32_tpl<THE_SIMDLEVEL> ci, ai;
ai.loadu(a + i);
ci = ai + bv;
ci.storeu(c + i);
}
for (; i < d; i++) {
c[i] = a[i] + b;
}
}
/***************************************************************************
* PQ tables computations
***************************************************************************/
namespace {
/// compute the IP for dsub = 2 for 8 centroids and 4 sub-vectors at a time
template <SIMDLevel SL, bool is_inner_product>
void pq2_8cents_table(
const simd8float32_tpl<SL> centroids[8],
const simd8float32_tpl<SL> x,
float* out,
size_t ldo,
size_t nout = 4) {
simd8float32_tpl<SL> ips[4];
for (int i = 0; i < 4; i++) {
simd8float32_tpl<SL> p1, p2;
if (is_inner_product) {
p1 = x * centroids[2 * i];
p2 = x * centroids[2 * i + 1];
} else {
p1 = (x - centroids[2 * i]);
p1 = p1 * p1;
p2 = (x - centroids[2 * i + 1]);
p2 = p2 * p2;
}
ips[i] = hadd(p1, p2);
}
simd8float32_tpl<SL> ip02a = geteven(ips[0], ips[1]);
simd8float32_tpl<SL> ip02b = geteven(ips[2], ips[3]);
simd8float32_tpl<SL> ip0 = getlow128(ip02a, ip02b);
simd8float32_tpl<SL> ip2 = gethigh128(ip02a, ip02b);
simd8float32_tpl<SL> ip13a = getodd(ips[0], ips[1]);
simd8float32_tpl<SL> ip13b = getodd(ips[2], ips[3]);
simd8float32_tpl<SL> ip1 = getlow128(ip13a, ip13b);
simd8float32_tpl<SL> ip3 = gethigh128(ip13a, ip13b);
switch (nout) {
case 4:
ip3.storeu(out + 3 * ldo);
[[fallthrough]];
case 3:
ip2.storeu(out + 2 * ldo);
[[fallthrough]];
case 2:
ip1.storeu(out + 1 * ldo);
[[fallthrough]];
case 1:
ip0.storeu(out);
}
}
template <SIMDLevel SL>
simd8float32_tpl<SL> load_simd8float32_partial(const float* x, int n) {
ALIGNED(32) float tmp[8] = {0, 0, 0, 0, 0, 0, 0, 0};
float* wp = tmp;
for (int i = 0; i < n; i++) {
*wp++ = *x++;
}
return simd8float32_tpl<SL>(tmp);
}
} // anonymous namespace
template <>
void compute_PQ_dis_tables_dsub2<THE_SIMDLEVEL>(
size_t d,
size_t ksub,
const float* all_centroids,
size_t nx,
const float* x,
bool is_inner_product,
float* dis_tables) {
size_t M = d / 2;
FAISS_THROW_IF_NOT(ksub % 8 == 0);
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];
for (int k = 0; k < 8; k++) {
ALIGNED(32) float centroid[8];
size_t wp = 0;
size_t rp = (m0 * ksub + k + k0) * 2;
for (int m = m0; m < m1; m++) {
centroid[wp++] = all_centroids[rp];
centroid[wp++] = all_centroids[rp + 1];
rp += 2 * ksub;
}
centroids[k] = simd8float32_tpl<THE_SIMDLEVEL>(centroid);
}
for (size_t i = 0; i < nx; i++) {
simd8float32_tpl<THE_SIMDLEVEL> xi;
if (m1 == m0 + 4) {
xi.loadu(x + i * d + m0 * 2);
} else {
xi = load_simd8float32_partial<THE_SIMDLEVEL>(
x + i * d + m0 * 2, 2 * (m1 - m0));
}
if (is_inner_product) {
pq2_8cents_table<THE_SIMDLEVEL, true>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,
ksub,
m1 - m0);
} else {
pq2_8cents_table<THE_SIMDLEVEL, false>(
centroids,
xi,
dis_tables + (i * M + m0) * ksub + k0,
ksub,
m1 - m0);
}
}
}
}
}
} // namespace faiss