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

ScalarQuantizer: split SIMD specializations into per-SIMD TUs + DD dispatch (#4839)

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

Split the SIMD-gated template specializations out of ScalarQuantizer.cpp
and the shared headers into per-SIMD compilation units and wire up the
Dynamic Dispatch (DD) infrastructure (COMPILE_SIMD_*, with_simd_level).

**What moved where**

- SIMD specializations removed from `codecs.h`, `quantizers.h`,
  `similarities.h`, `distance_computers.h` — these now contain only
  primary templates and scalar (`SIMDLevel::NONE`) specializations.
  (Most use empty primary templates; `quantizers.h` uses an inheriting
  fallback pattern for `QuantizerFP16`, `QuantizerBF16`, etc.)
- SIMD specializations moved into `sq-avx2.cpp` / `sq-avx512.cpp` /
  `sq-neon.cpp`, each guarded by `COMPILE_SIMD_*`.
- `sq-generic.cpp` deleted — the `NONE` level is now instantiated
  directly in `ScalarQuantizer.cpp` via `sq-dispatch.h`.
- `sq-inl.h` renamed to `scanners.h`.

**Dispatch mechanism**

- `sq-dispatch.h` is an X-macro-style header: each per-SIMD `.cpp` file
  `#define`s `THE_LEVEL_TO_DISPATCH` and `#include`s it to stamp out
  explicit template specializations of the selection functions
  (`sq_select_quantizer`, `sq_select_distance_computer`,
  `sq_select_InvertedListScanner`).
- `ScalarQuantizer.cpp` uses `with_simd_level` for runtime dispatch
  and instantiates the `NONE` level via the same `sq-dispatch.h`.
- Each per-SIMD selection function returns `nullptr` when the dimension
  doesn't align, and the caller falls back to `NONE`.
- `sq-neon.cpp` handles both `ARM_NEON` and `ARM_SVE` (SVE forwards
  to NEON — no dedicated SVE SQ implementation yet).

**Build**

- `xplat.bzl`, `CMakeLists.txt` — register new SIMD source files and
  headers.
- Within the SQ module, `COMPILE_SIMD_*` macros gate all SIMD code
  paths. (Compiler-defined macros like `__AVX2__` are still used in
  lower-level shared headers like `simdlib.h` and `fp16.h`.)

Reviewed By: mdouze

Differential Revision: D94375408

fbshipit-source-id: a07c31540242defcc605dd74e07bd25b8c163f43
This commit is contained in:
Gergely Szilvasy
2026-03-03 05:04:31 -08:00
committed by meta-codesync[bot]
parent db9ba35118
commit ccc934f586
14 changed files with 1980 additions and 1699 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+6 -1
View File
@@ -10,13 +10,16 @@
# =============================================================================
set(FAISS_SIMD_AVX2_SRC
impl/pq_code_distance/pq_code_distance-avx2.cpp
impl/scalar_quantizer/sq-avx2.cpp
utils/simd_impl/distances_avx2.cpp
)
set(FAISS_SIMD_AVX512_SRC
impl/pq_code_distance/pq_code_distance-avx512.cpp
impl/scalar_quantizer/sq-avx512.cpp
utils/simd_impl/distances_avx512.cpp
)
set(FAISS_SIMD_NEON_SRC
impl/scalar_quantizer/sq-neon.cpp
utils/simd_impl/distances_aarch64.cpp
)
set(FAISS_SIMD_SVE_SRC
@@ -246,7 +249,9 @@ set(FAISS_HEADERS
impl/scalar_quantizer/codecs.h
impl/scalar_quantizer/distance_computers.h
impl/scalar_quantizer/quantizers.h
impl/scalar_quantizer/scanners.h
impl/scalar_quantizer/similarities.h
impl/scalar_quantizer/sq-dispatch.h
impl/scalar_quantizer/training.h
impl/ThreadedIndex-inl.h
impl/ThreadedIndex.h
@@ -440,7 +445,7 @@ if(FAISS_OPT_LEVEL STREQUAL "dd")
set_source_files_properties(${FAISS_SIMD_AVX512_SRC}
TARGET_DIRECTORY faiss
PROPERTIES COMPILE_OPTIONS
"-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mpopcnt"
"-mavx512f;-mavx512cd;-mavx512vl;-mavx512dq;-mavx512bw;-mfma;-mf16c;-mpopcnt"
)
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
# ARM NEON is always available on aarch64, no special compiler flags needed
faiss/impl/ScalarQuantizer.cpp
+52 -502
View File
@@ -7,233 +7,24 @@
// -*- c++ -*-
#include <cstdio>
#include <cstring>
#include <memory>
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/simdlib.h>
#include <faiss/impl/scalar_quantizer/training.h>
#ifdef __SSE__
#include <immintrin.h>
#endif
#include <faiss/IndexIVF.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/expanded_scanners.h>
#include <faiss/utils/bf16.h>
#include <faiss/utils/fp16.h>
#include <faiss/utils/utils.h>
#include <faiss/impl/simd_dispatch.h>
/*******************************************************************
* ScalarQuantizer implementation
*
* The main source of complexity is to support combinations of 4
* variants without incurring runtime tests or virtual function calls:
*
* - 4 / 6 / 8 bits per code component
* - uniform / non-uniform
* - IP / L2 distance search
* - scalar / SIMD distance computation
*
* The appropriate Quantizer object is returned via select_quantizer
* that hides the template mess.
********************************************************************/
#include <faiss/impl/scalar_quantizer/scanners.h>
#if defined(__AVX512F__) && defined(__F16C__)
#define USE_AVX512_F16C
#elif defined(__AVX2__)
#ifdef __F16C__
#define USE_F16C
#else
#warning \
"Cannot enable AVX optimizations in scalar quantizer if -mf16c is not set as well"
#endif
#endif
#if defined(__aarch64__)
#if defined(__GNUC__) && __GNUC__ < 8
#warning \
"Cannot enable NEON optimizations in scalar quantizer if the compiler is GCC<8"
#else
#define USE_NEON
#endif
#endif
/*******************************************************************
* Codec: converts between values in [0, 1] and an index in a code
* array. The "i" parameter is the vector component index (not byte
* index).
*/
#include <faiss/impl/scalar_quantizer/codecs.h>
/*******************************************************************
* Quantizer: normalizes scalar vector components, then passes them
* through a codec
*******************************************************************/
#include <faiss/impl/scalar_quantizer/quantizers.h>
/*******************************************************************
* Similarity: gets vector components and computes a similarity wrt. a
* query vector stored in the object. The data fields just encapsulate
* an accumulator.
*/
#include <faiss/impl/scalar_quantizer/similarities.h>
/*******************************************************************
* DistanceComputer: combines a similarity and a quantizer to do
* code-to-vector or code-to-code comparisons
*******************************************************************/
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#define THE_LEVEL_TO_DISPATCH SIMDLevel::NONE
#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
namespace faiss {
namespace scalar_quantizer {
using QuantizerType = ScalarQuantizer::QuantizerType;
using RangeStat = ScalarQuantizer::RangeStat;
using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
/*******************************************************************
* select_distance_computer: runtime selection of template
* specialization
*******************************************************************/
template <class Sim>
SQDistanceComputer* select_distance_computer(
QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
constexpr SIMDLevel SL = Sim::simd_level;
switch (qtype) {
case ScalarQuantizer::QT_8bit_uniform:
return new DCTemplate<
QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::UNIFORM,
SL>,
Sim,
SL>(d, trained);
case ScalarQuantizer::QT_4bit_uniform:
return new DCTemplate<
QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::UNIFORM,
SL>,
Sim,
SL>(d, trained);
case ScalarQuantizer::QT_8bit:
return new DCTemplate<
QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Sim,
SL>(d, trained);
case ScalarQuantizer::QT_6bit:
return new DCTemplate<
QuantizerTemplate<
Codec6bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Sim,
SL>(d, trained);
case ScalarQuantizer::QT_4bit:
return new DCTemplate<
QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Sim,
SL>(d, trained);
case ScalarQuantizer::QT_fp16:
return new DCTemplate<QuantizerFP16<SL>, Sim, SL>(d, trained);
case ScalarQuantizer::QT_bf16:
return new DCTemplate<QuantizerBF16<SL>, Sim, SL>(d, trained);
case ScalarQuantizer::QT_8bit_direct:
#if defined(__AVX512F__)
if (d % 32 == 0) {
return new DistanceComputerByte<Sim, SL>(int(d), trained);
} else
#elif defined(__AVX2__)
if (d % 16 == 0) {
return new DistanceComputerByte<Sim, SL>(
static_cast<int>(d), trained);
} else
#endif
{
return new DCTemplate<Quantizer8bitDirect<SL>, Sim, SL>(
d, trained);
}
case ScalarQuantizer::QT_8bit_direct_signed:
return new DCTemplate<Quantizer8bitDirectSigned<SL>, Sim, SL>(
d, trained);
default:
FAISS_THROW_MSG("unknown qtype");
}
}
template <SIMDLevel SL>
ScalarQuantizer::SQuantizer* select_quantizer_1(
QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
switch (qtype) {
case ScalarQuantizer::QT_8bit:
return new QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_6bit:
return new QuantizerTemplate<
Codec6bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_4bit:
return new QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_8bit_uniform:
return new QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_4bit_uniform:
return new QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_fp16:
return new QuantizerFP16<SL>(d, trained);
case ScalarQuantizer::QT_bf16:
return new QuantizerBF16<SL>(d, trained);
case ScalarQuantizer::QT_8bit_direct:
return new Quantizer8bitDirect<SL>(d, trained);
case ScalarQuantizer::QT_8bit_direct_signed:
return new Quantizer8bitDirectSigned<SL>(d, trained);
default:
FAISS_THROW_MSG("unknown qtype");
}
}
} // namespace scalar_quantizer
using namespace scalar_quantizer;
/*******************************************************************
* ScalarQuantizer implementation
********************************************************************/
@@ -277,6 +68,9 @@ void ScalarQuantizer::set_derived_sizes() {
}
void ScalarQuantizer::train(size_t n, const float* x) {
using scalar_quantizer::train_NonUniform;
using scalar_quantizer::train_Uniform;
int bit_per_dim = qtype == QT_4bit_uniform ? 4
: qtype == QT_4bit ? 4
: qtype == QT_6bit ? 6
@@ -319,22 +113,17 @@ void ScalarQuantizer::train(size_t n, const float* x) {
}
ScalarQuantizer::SQuantizer* ScalarQuantizer::select_quantizer() const {
#if defined(USE_AVX512_F16C)
if (d % 16 == 0) {
return select_quantizer_1<SIMDLevel::AVX512>(qtype, d, trained);
} else
#elif defined(USE_F16C)
if (d % 8 == 0) {
return select_quantizer_1<SIMDLevel::AVX2>(qtype, d, trained);
} else
#elif defined(USE_NEON)
if (d % 8 == 0) {
return select_quantizer_1<SIMDLevel::ARM_NEON>(qtype, d, trained);
} else
#endif
{
return select_quantizer_1<SIMDLevel::NONE>(qtype, d, trained);
}
return with_simd_level([&]<SIMDLevel SL>() -> SQuantizer* {
if constexpr (SL != SIMDLevel::NONE) {
auto* q = scalar_quantizer::sq_select_quantizer<SL>(
qtype, d, trained);
if (q) {
return q;
}
}
return scalar_quantizer::sq_select_quantizer<SIMDLevel::NONE>(
qtype, d, trained);
});
}
void ScalarQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
@@ -357,294 +146,55 @@ void ScalarQuantizer::decode(const uint8_t* codes, float* x, size_t n) const {
}
}
SQDistanceComputer* ScalarQuantizer::get_distance_computer(
ScalarQuantizer::SQDistanceComputer* ScalarQuantizer::get_distance_computer(
MetricType metric) const {
FAISS_THROW_IF_NOT(metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT);
#if defined(USE_AVX512_F16C)
if (d % 16 == 0) {
if (metric == METRIC_L2) {
return select_distance_computer<SimilarityL2<SIMDLevel::AVX512>>(
qtype, d, trained);
} else {
return select_distance_computer<SimilarityIP<SIMDLevel::AVX512>>(
qtype, d, trained);
return with_simd_level([&]<SIMDLevel SL>() -> SQDistanceComputer* {
if constexpr (SL != SIMDLevel::NONE) {
auto* dc = scalar_quantizer::sq_select_distance_computer<SL>(
metric, qtype, d, trained);
if (dc) {
return dc;
}
}
} else
#elif defined(USE_F16C)
if (d % 8 == 0) {
if (metric == METRIC_L2) {
return select_distance_computer<SimilarityL2<SIMDLevel::AVX2>>(
qtype, d, trained);
} else {
return select_distance_computer<SimilarityIP<SIMDLevel::AVX2>>(
qtype, d, trained);
}
} else
#elif defined(USE_NEON)
if (d % 8 == 0) {
if (metric == METRIC_L2) {
return select_distance_computer<SimilarityL2<SIMDLevel::ARM_NEON>>(
qtype, d, trained);
} else {
return select_distance_computer<SimilarityIP<SIMDLevel::ARM_NEON>>(
qtype, d, trained);
}
} else
#endif
{
if (metric == METRIC_L2) {
return select_distance_computer<SimilarityL2<SIMDLevel::NONE>>(
qtype, d, trained);
} else {
return select_distance_computer<SimilarityIP<SIMDLevel::NONE>>(
qtype, d, trained);
}
}
return scalar_quantizer::sq_select_distance_computer<SIMDLevel::NONE>(
metric, qtype, d, trained);
});
}
/*******************************************************************
* IndexScalarQuantizer/IndexIVFScalarQuantizer scanner object
*
* It is an InvertedListScanner, but is designed to work with
* IndexScalarQuantizer as well.
********************************************************************/
namespace {
template <class DCClass>
struct IVFSQScannerIP : InvertedListScanner {
DCClass dc;
bool by_residual;
float accu0; /// added to all distances
IVFSQScannerIP(
int d,
const std::vector<float>& trained,
size_t code_size,
bool store_pairs,
const IDSelector* sel,
bool by_residual)
: dc(d, trained), by_residual(by_residual), accu0(0) {
this->store_pairs = store_pairs;
this->sel = sel;
this->code_size = code_size;
this->keep_max = true;
}
void set_query(const float* query) override {
dc.set_query(query);
}
void set_list(idx_t list_no, float coarse_dis) override {
this->list_no = list_no;
accu0 = by_residual ? coarse_dis : 0;
}
float distance_to_code(const uint8_t* code) const final {
return accu0 + dc.query_to_code(code);
}
// 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 override {
return run_scan_codes_fix_C<CMin<float, idx_t>>(
*this, list_size, codes, ids, handler);
}
};
template <class DCClass>
struct IVFSQScannerL2 : InvertedListScanner {
DCClass dc;
bool by_residual;
const Index* quantizer;
const float* x; /// current query
std::vector<float> tmp;
IVFSQScannerL2(
int d,
const std::vector<float>& trained,
size_t code_size,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual)
: dc(d, trained),
by_residual(by_residual),
quantizer(quantizer),
x(nullptr),
tmp(d) {
this->store_pairs = store_pairs;
this->sel = sel;
this->code_size = code_size;
}
void set_query(const float* query) override {
x = query;
if (!quantizer) {
dc.set_query(query);
}
}
void set_list(idx_t list_no, float /*coarse_dis*/) override {
this->list_no = list_no;
if (by_residual) {
// shift of x_in wrt centroid
quantizer->compute_residual(x, tmp.data(), list_no);
dc.set_query(tmp.data());
} else {
dc.set_query(x);
}
}
float distance_to_code(const uint8_t* code) const final {
return dc.query_to_code(code);
}
// 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 override {
return run_scan_codes_fix_C<CMax<float, idx_t>>(
*this, list_size, codes, ids, handler);
}
};
} // anonymous namespace
InvertedListScanner* ScalarQuantizer::select_InvertedListScanner(
MetricType mt,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual) const {
// this maps the runtime selection of the scanner by using a series of
// templatized lambdas
auto scan = [=, this]<class DCClass>() -> InvertedListScanner* {
if constexpr (DCClass::Sim::metric_type == METRIC_L2) {
return new IVFSQScannerL2<DCClass>(
int(d),
trained,
return with_simd_level([&]<SIMDLevel SL>() -> InvertedListScanner* {
if constexpr (SL != SIMDLevel::NONE) {
auto* s = scalar_quantizer::sq_select_InvertedListScanner<SL>(
qtype,
mt,
d,
code_size,
trained,
quantizer,
store_pairs,
sel,
by_residual);
} else if constexpr (
DCClass::Sim::metric_type == METRIC_INNER_PRODUCT) {
return new IVFSQScannerIP<DCClass>(
int(d), trained, code_size, store_pairs, sel, by_residual);
} else {
FAISS_THROW_MSG("unsupported metric type");
if (s) {
return s;
}
}
};
auto select_by_simd_and_metric =
[&,
this]<SIMDLevel SL, class Similarity>() -> InvertedListScanner* {
switch (qtype) {
case QT_8bit_uniform:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::UNIFORM,
SL>,
Similarity,
SL>>();
case QT_4bit_uniform:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::UNIFORM,
SL>,
Similarity,
SL>>();
case QT_8bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec8bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Similarity,
SL>>();
case QT_4bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec4bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Similarity,
SL>>();
case QT_6bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec6bit,
QuantizerTemplateScaling::NON_UNIFORM,
SL>,
Similarity,
SL>>();
case QT_fp16:
return scan.template
operator()<DCTemplate<QuantizerFP16<SL>, Similarity, SL>>();
case QT_bf16:
return scan.template
operator()<DCTemplate<QuantizerBF16<SL>, Similarity, SL>>();
case QT_8bit_direct:
#if defined(__AVX512F__)
if (d % 32 == 0) {
return scan.template
operator()<DistanceComputerByte<Similarity, SL>>();
}
#elif defined(__AVX2__)
if (d % 16 == 0) {
return scan.template
operator()<DistanceComputerByte<Similarity, SL>>();
}
#endif
return scan.template operator()<
DCTemplate<Quantizer8bitDirect<SL>, Similarity, SL>>();
case QT_8bit_direct_signed:
return scan.template operator()<DCTemplate<
Quantizer8bitDirectSigned<SL>,
Similarity,
SL>>();
default:
FAISS_THROW_MSG("unknown qtype");
}
};
auto select_by_simd = [&]<SIMDLevel SL>() {
if (mt == METRIC_L2) {
return select_by_simd_and_metric
.template operator()<SL, SimilarityL2<SL>>();
} else if (mt == METRIC_INNER_PRODUCT) {
return select_by_simd_and_metric
.template operator()<SL, SimilarityIP<SL>>();
}
FAISS_THROW_MSG("unsupported metric type");
};
#if defined(USE_AVX512_F16C)
if (d % 16 == 0) {
return select_by_simd.template operator()<SIMDLevel::AVX512>();
}
#elif defined(USE_F16C)
if (d % 8 == 0) {
return select_by_simd.template operator()<SIMDLevel::AVX2>();
}
#elif defined(USE_NEON)
if (d % 8 == 0) {
return select_by_simd.template operator()<SIMDLevel::ARM_NEON>();
}
#endif
return select_by_simd.template operator()<SIMDLevel::NONE>();
return scalar_quantizer::sq_select_InvertedListScanner<SIMDLevel::NONE>(
qtype,
mt,
d,
code_size,
trained,
quantizer,
store_pairs,
sel,
by_residual);
});
}
} // namespace faiss
faiss/impl/scalar_quantizer/codecs.h
+16 -206
View File
@@ -8,6 +8,7 @@
#pragma once
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/simdlib.h>
namespace faiss {
@@ -20,7 +21,11 @@ namespace scalar_quantizer {
* index).
*/
struct Codec8bit {
template <SIMDLevel SL>
struct Codec8bit {};
template <>
struct Codec8bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE void encode_component(
float x,
uint8_t* code,
@@ -33,46 +38,13 @@ struct Codec8bit {
size_t i) {
return (code[i] + 0.5f) / 255.0f;
}
#if defined(__AVX512F__)
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
const __m128i c16 = _mm_loadu_si128((__m128i*)(code + i));
const __m512i i32 = _mm512_cvtepu8_epi32(c16);
const __m512 f16 = _mm512_cvtepi32_ps(i32);
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 255.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 255.f);
return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
}
#elif defined(__AVX2__)
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
const uint64_t c8 = *(uint64_t*)(code + i);
const __m128i i8 = _mm_set1_epi64x(c8);
const __m256i i32 = _mm256_cvtepu8_epi32(i8);
const __m256 f8 = _mm256_cvtepi32_ps(i32);
const __m256 half_one_255 = _mm256_set1_ps(0.5f / 255.f);
const __m256 one_255 = _mm256_set1_ps(1.f / 255.f);
return simd8float32(_mm256_fmadd_ps(f8, one_255, half_one_255));
}
#endif
#ifdef USE_NEON
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] = decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
#endif
};
struct Codec4bit {
template <SIMDLevel SL>
struct Codec4bit {};
template <>
struct Codec4bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE void encode_component(
float x,
uint8_t* code,
@@ -85,64 +57,13 @@ struct Codec4bit {
size_t i) {
return (((code[i / 2] >> ((i & 1) << 2)) & 0xf) + 0.5f) / 15.0f;
}
#if defined(__AVX512F__)
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
uint64_t c8 = *(uint64_t*)(code + (i >> 1));
uint64_t mask = 0x0f0f0f0f0f0f0f0f;
uint64_t c8ev = c8 & mask;
uint64_t c8od = (c8 >> 4) & mask;
__m128i c16 =
_mm_unpacklo_epi8(_mm_set1_epi64x(c8ev), _mm_set1_epi64x(c8od));
__m256i c8lo = _mm256_cvtepu8_epi32(c16);
__m256i c8hi = _mm256_cvtepu8_epi32(_mm_srli_si128(c16, 8));
__m512i i16 = _mm512_castsi256_si512(c8lo);
i16 = _mm512_inserti32x8(i16, c8hi, 1);
__m512 f16 = _mm512_cvtepi32_ps(i16);
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 15.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 15.f);
return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
}
#elif defined(__AVX2__)
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
uint32_t c4 = *(uint32_t*)(code + (i >> 1));
uint32_t mask = 0x0f0f0f0f;
uint32_t c4ev = c4 & mask;
uint32_t c4od = (c4 >> 4) & mask;
// the 8 lower bytes of c8 contain the values
__m128i c8 =
_mm_unpacklo_epi8(_mm_set1_epi32(c4ev), _mm_set1_epi32(c4od));
__m128i c4lo = _mm_cvtepu8_epi32(c8);
__m128i c4hi = _mm_cvtepu8_epi32(_mm_srli_si128(c8, 4));
__m256i i8 = _mm256_castsi128_si256(c4lo);
i8 = _mm256_insertf128_si256(i8, c4hi, 1);
__m256 f8 = _mm256_cvtepi32_ps(i8);
__m256 half = _mm256_set1_ps(0.5f);
f8 = _mm256_add_ps(f8, half);
__m256 one_255 = _mm256_set1_ps(1.f / 15.f);
return simd8float32(_mm256_mul_ps(f8, one_255));
}
#endif
#ifdef USE_NEON
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] = decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
#endif
};
struct Codec6bit {
template <SIMDLevel SL>
struct Codec6bit {};
template <>
struct Codec6bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE void encode_component(
float x,
uint8_t* code,
@@ -194,117 +115,6 @@ struct Codec6bit {
}
return (bits + 0.5f) / 63.0f;
}
#if defined(__AVX512F__)
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
// pure AVX512 implementation (not necessarily the fastest).
// see:
// https://github.com/zilliztech/knowhere/blob/main/thirdparty/faiss/faiss/impl/ScalarQuantizerCodec_avx512.h
// clang-format off
// 16 components, 16x6 bit=12 bytes
const __m128i bit_6v =
_mm_maskz_loadu_epi8(0b0000111111111111, code + (i >> 2) * 3);
const __m256i bit_6v_256 = _mm256_broadcast_i32x4(bit_6v);
// 00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F
// 00 01 02 03
const __m256i shuffle_mask = _mm256_setr_epi16(
0xFF00, 0x0100, 0x0201, 0xFF02,
0xFF03, 0x0403, 0x0504, 0xFF05,
0xFF06, 0x0706, 0x0807, 0xFF08,
0xFF09, 0x0A09, 0x0B0A, 0xFF0B);
const __m256i shuffled = _mm256_shuffle_epi8(bit_6v_256, shuffle_mask);
// 0: xxxxxxxx xx543210
// 1: xxxx5432 10xxxxxx
// 2: xxxxxx54 3210xxxx
// 3: xxxxxxxx 543210xx
const __m256i shift_right_v = _mm256_setr_epi16(
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U);
__m256i shuffled_shifted = _mm256_srlv_epi16(shuffled, shift_right_v);
// remove unneeded bits
shuffled_shifted =
_mm256_and_si256(shuffled_shifted, _mm256_set1_epi16(0x003F));
// scale
const __m512 f8 =
_mm512_cvtepi32_ps(_mm512_cvtepi16_epi32(shuffled_shifted));
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 63.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 63.f);
return simd16float32(_mm512_fmadd_ps(f8, one_255, half_one_255));
// clang-format on
}
#elif defined(__AVX2__)
/* Load 6 bytes that represent 8 6-bit values, return them as a
* 8*32 bit vector register */
static FAISS_ALWAYS_INLINE __m256i load6(const uint16_t* code16) {
const __m128i perm = _mm_set_epi8(
-1, 5, 5, 4, 4, 3, -1, 3, -1, 2, 2, 1, 1, 0, -1, 0);
const __m256i shifts = _mm256_set_epi32(2, 4, 6, 0, 2, 4, 6, 0);
// load 6 bytes
__m128i c1 =
_mm_set_epi16(0, 0, 0, 0, 0, code16[2], code16[1], code16[0]);
// put in 8 * 32 bits
__m128i c2 = _mm_shuffle_epi8(c1, perm);
__m256i c3 = _mm256_cvtepi16_epi32(c2);
// shift and mask out useless bits
__m256i c4 = _mm256_srlv_epi32(c3, shifts);
__m256i c5 = _mm256_and_si256(_mm256_set1_epi32(63), c4);
return c5;
}
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
// // Faster code for Intel CPUs or AMD Zen3+, just keeping it here
// // for the reference, maybe, it becomes used one day.
// const uint16_t* data16 = (const uint16_t*)(code + (i >> 2) * 3);
// const uint32_t* data32 = (const uint32_t*)data16;
// const uint64_t val = *data32 + ((uint64_t)data16[2] << 32);
// const uint64_t vext = _pdep_u64(val, 0x3F3F3F3F3F3F3F3FULL);
// const __m128i i8 = _mm_set1_epi64x(vext);
// const __m256i i32 = _mm256_cvtepi8_epi32(i8);
// const __m256 f8 = _mm256_cvtepi32_ps(i32);
// const __m256 half_one_255 = _mm256_set1_ps(0.5f / 63.f);
// const __m256 one_255 = _mm256_set1_ps(1.f / 63.f);
// return _mm256_fmadd_ps(f8, one_255, half_one_255);
__m256i i8 = load6((const uint16_t*)(code + (i >> 2) * 3));
__m256 f8 = _mm256_cvtepi32_ps(i8);
// this could also be done with bit manipulations but it is
// not obviously faster
const __m256 half_one_255 = _mm256_set1_ps(0.5f / 63.f);
const __m256 one_255 = _mm256_set1_ps(1.f / 63.f);
return simd8float32(_mm256_fmadd_ps(f8, one_255, half_one_255));
}
#endif
#ifdef USE_NEON
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] = decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
#endif
};
} // namespace scalar_quantizer
faiss/impl/scalar_quantizer/distance_computers.h
+9 -321
View File
@@ -70,161 +70,6 @@ struct DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> : SQDistanceComputer {
}
};
#if defined(USE_AVX512_F16C)
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
: SQDistanceComputer { // Update to handle 16 lanes
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_16();
for (size_t i = 0; i < quant.d; i += 16) {
simd16float32 xi = quant.reconstruct_16_components(code, i);
sim.add_16_components(xi);
}
return sim.result_16();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_16();
for (size_t i = 0; i < quant.d; i += 16) {
simd16float32 x1 = quant.reconstruct_16_components(code1, i);
simd16float32 x2 = quant.reconstruct_16_components(code2, i);
sim.add_16_components_2(x1, x2);
}
return sim.result_16();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
#endif
#if defined(USE_F16C)
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 xi =
quant.reconstruct_8_components(code, static_cast<int>(i));
sim.add_8_components(xi);
}
return sim.result_8();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 x1 =
quant.reconstruct_8_components(code1, static_cast<int>(i));
simd8float32 x2 =
quant.reconstruct_8_components(code2, static_cast<int>(i));
sim.add_8_components_2(x1, x2);
}
return sim.result_8();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
#endif
#ifdef USE_NEON
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
: SQDistanceComputer {
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 xi =
quant.reconstruct_8_components(code, static_cast<int>(i));
sim.add_8_components(xi);
}
return sim.result_8();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 x1 =
quant.reconstruct_8_components(code1, static_cast<int>(i));
simd8float32 x2 =
quant.reconstruct_8_components(code2, static_cast<int>(i));
sim.add_8_components_2(x1, x2);
}
return sim.result_8();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
#endif
/*******************************************************************
* DistanceComputerByte: computes distances in the integer domain
*******************************************************************/
@@ -276,173 +121,16 @@ struct DistanceComputerByte<Similarity, SIMDLevel::NONE> : SQDistanceComputer {
}
};
#if defined(__AVX512F__)
/*******************************************************************
* Selection function
*******************************************************************/
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
: SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
__m512i accu = _mm512_setzero_si512();
for (int i = 0; i < d; i += 32) { // Process 32 bytes at a time
__m512i c1 = _mm512_cvtepu8_epi16(
_mm256_loadu_si256((__m256i*)(code1 + i)));
__m512i c2 = _mm512_cvtepu8_epi16(
_mm256_loadu_si256((__m256i*)(code2 + i)));
__m512i prod32;
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
prod32 = _mm512_madd_epi16(c1, c2);
} else {
__m512i diff = _mm512_sub_epi16(c1, c2);
prod32 = _mm512_madd_epi16(diff, diff);
}
accu = _mm512_add_epi32(accu, prod32);
}
// Horizontally add elements of accu
return _mm512_reduce_add_epi32(accu);
}
void set_query(const float* x) final {
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
#elif defined(__AVX2__)
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
// __m256i accu = _mm256_setzero_ps ();
__m256i accu = _mm256_setzero_si256();
for (int i = 0; i < d; i += 16) {
// load 16 bytes, convert to 16 uint16_t
__m256i c1 = _mm256_cvtepu8_epi16(
_mm_loadu_si128((__m128i*)(code1 + i)));
__m256i c2 = _mm256_cvtepu8_epi16(
_mm_loadu_si128((__m128i*)(code2 + i)));
__m256i prod32;
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
prod32 = _mm256_madd_epi16(c1, c2);
} else {
__m256i diff = _mm256_sub_epi16(c1, c2);
prod32 = _mm256_madd_epi16(diff, diff);
}
accu = _mm256_add_epi32(accu, prod32);
}
__m128i sum = _mm256_extractf128_si256(accu, 0);
sum = _mm_add_epi32(sum, _mm256_extractf128_si256(accu, 1));
sum = _mm_hadd_epi32(sum, sum);
sum = _mm_hadd_epi32(sum, sum);
return _mm_cvtsi128_si32(sum);
}
void set_query(const float* x) final {
/*
for (int i = 0; i < d; i += 8) {
__m256 xi = _mm256_loadu_ps (x + i);
__m256i ci = _mm256_cvtps_epi32(xi);
*/
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
#endif
#ifdef USE_NEON
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
: SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
int accu = 0;
for (int i = 0; i < d; i++) {
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
accu += int(code1[i]) * code2[i];
} else {
int diff = int(code1[i]) - code2[i];
accu += diff * diff;
}
}
return accu;
}
void set_query(const float* x) final {
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
#endif
template <SIMDLevel SL>
SQDistanceComputer* sq_select_distance_computer(
MetricType metric,
ScalarQuantizer::QuantizerType qtype,
size_t d,
const std::vector<float>& trained);
} // namespace scalar_quantizer
} // namespace faiss
faiss/impl/scalar_quantizer/quantizers.h
+29 -395
View File
@@ -8,6 +8,8 @@
#pragma once
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/utils/bf16.h>
#include <faiss/utils/fp16.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/simdlib.h>
@@ -15,6 +17,8 @@ namespace faiss {
namespace scalar_quantizer {
using QuantizerType = ScalarQuantizer::QuantizerType;
/*******************************************************************
* Quantizer: normalizes scalar vector components, then passes them
* through a codec
@@ -67,93 +71,6 @@ struct QuantizerTemplate<
}
};
#if defined(__AVX512F__)
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::AVX512>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m512 xi = Codec::decode_16_components(code, i).f;
return simd16float32(_mm512_fmadd_ps(
xi, _mm512_set1_ps(this->vdiff), _mm512_set1_ps(this->vmin)));
}
};
#elif defined(__AVX2__)
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::AVX2>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m256 xi = Codec::decode_8_components(code, i).f;
return simd8float32(_mm256_fmadd_ps(
xi, _mm256_set1_ps(this->vdiff), _mm256_set1_ps(this->vmin)));
}
};
#endif
#ifdef USE_NEON
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::ARM_NEON>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
float32x4x2_t xi = Codec::decode_8_components(code, i).data;
return simd8float32(
float32x4x2_t{
vfmaq_f32(
vdupq_n_f32(this->vmin),
xi.val[0],
vdupq_n_f32(this->vdiff)),
vfmaq_f32(
vdupq_n_f32(this->vmin),
xi.val[1],
vdupq_n_f32(this->vdiff))});
}
};
#endif
template <class Codec>
struct QuantizerTemplate<
Codec,
@@ -196,101 +113,12 @@ struct QuantizerTemplate<
}
};
#if defined(__AVX512F__)
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::AVX512>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m512 xi = Codec::decode_16_components(code, i).f;
return simd16float32(_mm512_fmadd_ps(
xi,
_mm512_loadu_ps(this->vdiff + i),
_mm512_loadu_ps(this->vmin + i)));
}
};
#elif defined(__AVX2__)
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::AVX2>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m256 xi = Codec::decode_8_components(code, i).f;
return simd8float32(_mm256_fmadd_ps(
xi,
_mm256_loadu_ps(this->vdiff + i),
_mm256_loadu_ps(this->vmin + i)));
}
};
#endif
#ifdef USE_NEON
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::ARM_NEON>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
float32x4x2_t xi = Codec::decode_8_components(code, i).data;
float32x4x2_t vmin_8 = vld1q_f32_x2(this->vmin + i);
float32x4x2_t vdiff_8 = vld1q_f32_x2(this->vdiff + i);
return simd8float32(
float32x4x2_t{
vfmaq_f32(vmin_8.val[0], xi.val[0], vdiff_8.val[0]),
vfmaq_f32(vmin_8.val[1], xi.val[1], vdiff_8.val[1])});
}
};
#endif
/*******************************************************************
* FP16 quantizer
*******************************************************************/
template <SIMDLevel SL>
struct QuantizerFP16 {};
struct QuantizerFP16;
template <>
struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
@@ -317,61 +145,17 @@ struct QuantizerFP16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
}
};
#if defined(USE_AVX512_F16C)
template <>
struct QuantizerFP16<SIMDLevel::AVX512> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m256i codei = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
return simd16float32(_mm512_cvtph_ps(codei));
}
template <SIMDLevel SL>
struct QuantizerFP16 : QuantizerFP16<SIMDLevel::NONE> {
using QuantizerFP16<SIMDLevel::NONE>::QuantizerFP16;
};
#endif
#if defined(USE_F16C)
template <>
struct QuantizerFP16<SIMDLevel::AVX2> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i codei = _mm_loadu_si128((const __m128i*)(code + 2 * i));
return simd8float32(_mm256_cvtph_ps(codei));
}
};
#endif
#ifdef USE_NEON
template <>
struct QuantizerFP16<SIMDLevel::ARM_NEON> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint16x4x2_t codei = vld1_u16_x2((const uint16_t*)(code + 2 * i));
return simd8float32(
{vcvt_f32_f16(vreinterpret_f16_u16(codei.val[0])),
vcvt_f32_f16(vreinterpret_f16_u16(codei.val[1]))});
}
};
#endif
/*******************************************************************
* BF16 quantizer
*******************************************************************/
template <SIMDLevel SL>
struct QuantizerBF16 {};
struct QuantizerBF16;
template <>
struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
@@ -398,64 +182,17 @@ struct QuantizerBF16<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
}
};
#if defined(__AVX512F__)
template <>
struct QuantizerBF16<SIMDLevel::AVX512> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m256i code_256i = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
__m512i code_512i = _mm512_cvtepu16_epi32(code_256i);
code_512i = _mm512_slli_epi32(code_512i, 16);
return simd16float32(_mm512_castsi512_ps(code_512i));
}
template <SIMDLevel SL>
struct QuantizerBF16 : QuantizerBF16<SIMDLevel::NONE> {
using QuantizerBF16<SIMDLevel::NONE>::QuantizerBF16;
};
#elif defined(__AVX2__)
template <>
struct QuantizerBF16<SIMDLevel::AVX2> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i code_128i = _mm_loadu_si128((const __m128i*)(code + 2 * i));
__m256i code_256i = _mm256_cvtepu16_epi32(code_128i);
code_256i = _mm256_slli_epi32(code_256i, 16);
return simd8float32(_mm256_castsi256_ps(code_256i));
}
};
#endif
#ifdef USE_NEON
template <>
struct QuantizerBF16<SIMDLevel::ARM_NEON> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint16x4x2_t codei = vld1_u16_x2((const uint16_t*)(code + 2 * i));
return simd8float32(
{vreinterpretq_f32_u32(
vshlq_n_u32(vmovl_u16(codei.val[0]), 16)),
vreinterpretq_f32_u32(
vshlq_n_u32(vmovl_u16(codei.val[1]), 16))});
}
};
#endif
/*******************************************************************
* 8bit_direct quantizer
*******************************************************************/
template <SIMDLevel SL>
struct Quantizer8bitDirect {};
struct Quantizer8bitDirect;
template <>
struct Quantizer8bitDirect<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
@@ -483,70 +220,17 @@ struct Quantizer8bitDirect<SIMDLevel::NONE> : ScalarQuantizer::SQuantizer {
}
};
#if defined(__AVX512F__)
template <>
struct Quantizer8bitDirect<SIMDLevel::AVX512>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
__m512i y16 = _mm512_cvtepu8_epi32(x16); // 16 * int32
return simd16float32(_mm512_cvtepi32_ps(y16)); // 16 * float32
}
template <SIMDLevel SL>
struct Quantizer8bitDirect : Quantizer8bitDirect<SIMDLevel::NONE> {
using Quantizer8bitDirect<SIMDLevel::NONE>::Quantizer8bitDirect;
};
#elif defined(__AVX2__)
template <>
struct Quantizer8bitDirect<SIMDLevel::AVX2>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i x8 = _mm_loadl_epi64((__m128i*)(code + i)); // 8 * int8
__m256i y8 = _mm256_cvtepu8_epi32(x8); // 8 * int32
return simd8float32(_mm256_cvtepi32_ps(y8)); // 8 * float32
}
};
#endif
#ifdef USE_NEON
template <>
struct Quantizer8bitDirect<SIMDLevel::ARM_NEON>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint8x8_t x8 = vld1_u8((const uint8_t*)(code + i));
uint16x8_t y8 = vmovl_u8(x8);
uint16x4_t y8_0 = vget_low_u16(y8);
uint16x4_t y8_1 = vget_high_u16(y8);
// convert uint16 -> uint32 -> fp32
return simd8float32(
{vcvtq_f32_u32(vmovl_u16(y8_0)),
vcvtq_f32_u32(vmovl_u16(y8_1))});
}
};
#endif
/*******************************************************************
* 8bit_direct_signed quantizer
*******************************************************************/
template <SIMDLevel SL>
struct Quantizer8bitDirectSigned {};
struct Quantizer8bitDirectSigned;
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
@@ -575,71 +259,21 @@ struct Quantizer8bitDirectSigned<SIMDLevel::NONE>
}
};
#if defined(__AVX512F__)
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::AVX512>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
__m512i y16 = _mm512_cvtepu8_epi32(x16); // 16 * int32
__m512i c16 = _mm512_set1_epi32(128);
__m512i z16 = _mm512_sub_epi32(y16, c16); // subtract 128 from all lanes
return simd16float32(_mm512_cvtepi32_ps(z16)); // 16 * float32
}
template <SIMDLevel SL>
struct Quantizer8bitDirectSigned : Quantizer8bitDirectSigned<SIMDLevel::NONE> {
using Quantizer8bitDirectSigned<SIMDLevel::NONE>::Quantizer8bitDirectSigned;
};
#elif defined(__AVX2__)
/*******************************************************************
* Selection function
*******************************************************************/
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::AVX2>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i x8 = _mm_loadl_epi64((__m128i*)(code + i)); // 8 * int8
__m256i y8 = _mm256_cvtepu8_epi32(x8); // 8 * int32
__m256i c8 = _mm256_set1_epi32(128);
__m256i z8 = _mm256_sub_epi32(y8, c8); // subtract 128 from all lanes
return simd8float32(_mm256_cvtepi32_ps(z8)); // 8 * float32
}
};
#endif
#ifdef USE_NEON
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::ARM_NEON>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint8x8_t x8 = vld1_u8((const uint8_t*)(code + i));
uint16x8_t y8 = vmovl_u8(x8); // convert uint8 -> uint16
uint16x4_t y8_0 = vget_low_u16(y8);
uint16x4_t y8_1 = vget_high_u16(y8);
float32x4_t z8_0 = vcvtq_f32_u32(
vmovl_u16(y8_0)); // convert uint16 -> uint32 -> fp32
float32x4_t z8_1 = vcvtq_f32_u32(vmovl_u16(y8_1));
// subtract 128 to convert into signed numbers
return simd8float32(
{vsubq_f32(z8_0, vmovq_n_f32(128.0)),
vsubq_f32(z8_1, vmovq_n_f32(128.0))});
}
};
#endif
// declare for all levels
template <SIMDLevel SL>
ScalarQuantizer::SQuantizer* sq_select_quantizer(
QuantizerType qtype,
size_t d,
const std::vector<float>& trained);
} // namespace scalar_quantizer
faiss/impl/scalar_quantizer/scanners.h
+164
View File
@@ -0,0 +1,164 @@
/*
* 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.
*/
// Private implementation header for per-SIMD scalar quantizer TUs.
// Do not include in public APIs.
#pragma once
#include <faiss/impl/ScalarQuantizer.h>
#include <faiss/utils/simd_levels.h>
#include <faiss/utils/simdlib.h>
#include <faiss/impl/simd_dispatch.h>
#include <faiss/IndexIVF.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/expanded_scanners.h>
#include <faiss/utils/bf16.h>
#include <faiss/utils/fp16.h>
#include <faiss/impl/scalar_quantizer/codecs.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
namespace faiss {
namespace scalar_quantizer {
using QuantizerType = ScalarQuantizer::QuantizerType;
using SQDistanceComputer = ScalarQuantizer::SQDistanceComputer;
/*******************************************************************
* IVFSQScannerIP / IVFSQScannerL2 — moved from anonymous namespace
* in ScalarQuantizer.cpp
*******************************************************************/
template <class DCClass>
struct IVFSQScannerIP : InvertedListScanner {
DCClass dc;
bool by_residual;
float accu0; /// added to all distances
IVFSQScannerIP(
int d,
const std::vector<float>& trained,
size_t code_size,
bool store_pairs,
const IDSelector* sel,
bool by_residual)
: dc(d, trained), by_residual(by_residual), accu0(0) {
this->store_pairs = store_pairs;
this->sel = sel;
this->code_size = code_size;
this->keep_max = true;
}
void set_query(const float* query) override {
dc.set_query(query);
}
void set_list(idx_t list_no, float coarse_dis) override {
this->list_no = list_no;
accu0 = by_residual ? coarse_dis : 0;
}
float distance_to_code(const uint8_t* code) const final {
return accu0 + dc.query_to_code(code);
}
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const override {
return run_scan_codes_fix_C<CMin<float, idx_t>>(
*this, list_size, codes, ids, handler);
}
};
template <class DCClass>
struct IVFSQScannerL2 : InvertedListScanner {
DCClass dc;
bool by_residual;
const Index* quantizer;
const float* x; /// current query
std::vector<float> tmp;
IVFSQScannerL2(
int d,
const std::vector<float>& trained,
size_t code_size,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual)
: dc(d, trained),
by_residual(by_residual),
quantizer(quantizer),
x(nullptr),
tmp(d) {
this->store_pairs = store_pairs;
this->sel = sel;
this->code_size = code_size;
}
void set_query(const float* query) override {
x = query;
if (!quantizer) {
dc.set_query(query);
}
}
void set_list(idx_t list_no, float /*coarse_dis*/) override {
this->list_no = list_no;
if (by_residual) {
quantizer->compute_residual(x, tmp.data(), list_no);
dc.set_query(tmp.data());
} else {
dc.set_query(x);
}
}
float distance_to_code(const uint8_t* code) const final {
return dc.query_to_code(code);
}
size_t scan_codes(
size_t list_size,
const uint8_t* codes,
const idx_t* ids,
ResultHandler& handler) const override {
return run_scan_codes_fix_C<CMax<float, idx_t>>(
*this, list_size, codes, ids, handler);
}
};
/*******************************************************************
* Forward declaration of inverts list scanner
*******************************************************************/
template <SIMDLevel SL>
InvertedListScanner* sq_select_InvertedListScanner(
QuantizerType qtype,
MetricType mt,
size_t d,
size_t code_size,
const std::vector<float>& trained,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual);
} // namespace scalar_quantizer
} // namespace faiss
faiss/impl/scalar_quantizer/similarities.h
-273
View File
@@ -58,143 +58,6 @@ struct SimilarityL2<SIMDLevel::NONE> {
}
};
#if defined(__AVX512F__)
template <>
struct SimilarityL2<SIMDLevel::AVX512> {
static constexpr int simdwidth = 16;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
simd16float32 accu16 = {};
FAISS_ALWAYS_INLINE void begin_16() {
accu16.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
__m512 yiv = _mm512_loadu_ps(yi);
yi += 16;
__m512 tmp = _mm512_sub_ps(yiv, x.f);
accu16 = simd16float32(_mm512_fmadd_ps(tmp, tmp, accu16.f));
}
FAISS_ALWAYS_INLINE void add_16_components_2(
simd16float32 x,
simd16float32 y_2) {
__m512 tmp = _mm512_sub_ps(y_2.f, x.f);
accu16 = simd16float32(_mm512_fmadd_ps(tmp, tmp, accu16.f));
}
FAISS_ALWAYS_INLINE float result_16() {
// performs better than dividing into _mm256 and adding
return _mm512_reduce_add_ps(accu16.f);
}
};
#elif defined(__AVX2__)
template <>
struct SimilarityL2<SIMDLevel::AVX2> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX2;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
simd8float32 accu8 = {};
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
__m256 yiv = _mm256_loadu_ps(yi);
yi += 8;
__m256 tmp = _mm256_sub_ps(yiv, x.f);
accu8 = simd8float32(_mm256_fmadd_ps(tmp, tmp, accu8.f));
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x,
simd8float32 y_2) {
__m256 tmp = _mm256_sub_ps(y_2.f, x.f);
accu8 = simd8float32(_mm256_fmadd_ps(tmp, tmp, accu8.f));
}
FAISS_ALWAYS_INLINE float result_8() {
const __m128 sum = _mm_add_ps(
_mm256_castps256_ps128(accu8.f),
_mm256_extractf128_ps(accu8.f, 1));
const __m128 v0 = _mm_shuffle_ps(sum, sum, _MM_SHUFFLE(0, 0, 3, 2));
const __m128 v1 = _mm_add_ps(sum, v0);
__m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
const __m128 v3 = _mm_add_ps(v1, v2);
return _mm_cvtss_f32(v3);
}
};
#endif
#ifdef USE_NEON
template <>
struct SimilarityL2<SIMDLevel::ARM_NEON> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::ARM_NEON;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y) {}
simd8float32 accu8;
FAISS_ALWAYS_INLINE void begin_8() {
accu8 = simd8float32(
float32x4x2_t{vdupq_n_f32(0.0f), vdupq_n_f32(0.0f)});
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
float32x4x2_t yiv = vld1q_f32_x2(yi);
yi += 8;
float32x4_t sub0 = vsubq_f32(yiv.val[0], x.data.val[0]);
float32x4_t sub1 = vsubq_f32(yiv.val[1], x.data.val[1]);
float32x4_t accu8_0 = vfmaq_f32(accu8.data.val[0], sub0, sub0);
float32x4_t accu8_1 = vfmaq_f32(accu8.data.val[1], sub1, sub1);
accu8 = simd8float32(float32x4x2_t{accu8_0, accu8_1});
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x,
simd8float32 y) {
float32x4_t sub0 = vsubq_f32(y.data.val[0], x.data.val[0]);
float32x4_t sub1 = vsubq_f32(y.data.val[1], x.data.val[1]);
float32x4_t accu8_0 = vfmaq_f32(accu8.data.val[0], sub0, sub0);
float32x4_t accu8_1 = vfmaq_f32(accu8.data.val[1], sub1, sub1);
accu8 = simd8float32(float32x4x2_t{accu8_0, accu8_1});
}
FAISS_ALWAYS_INLINE float result_8() {
float32x4_t sum_0 = vpaddq_f32(accu8.data.val[0], accu8.data.val[0]);
float32x4_t sum_1 = vpaddq_f32(accu8.data.val[1], accu8.data.val[1]);
float32x4_t sum2_0 = vpaddq_f32(sum_0, sum_0);
float32x4_t sum2_1 = vpaddq_f32(sum_1, sum_1);
return vgetq_lane_f32(sum2_0, 0) + vgetq_lane_f32(sum2_1, 0);
}
};
#endif
template <SIMDLevel SL>
struct SimilarityIP {};
@@ -227,141 +90,5 @@ struct SimilarityIP<SIMDLevel::NONE> {
}
};
#if defined(__AVX512F__)
template <>
struct SimilarityIP<SIMDLevel::AVX512> {
static constexpr int simdwidth = 16;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
float accu;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr), accu(0) {}
simd16float32 accu16 = {};
FAISS_ALWAYS_INLINE void begin_16() {
accu16.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
__m512 yiv = _mm512_loadu_ps(yi);
yi += 16;
accu16.f = _mm512_fmadd_ps(yiv, x.f, accu16.f);
}
FAISS_ALWAYS_INLINE void add_16_components_2(
simd16float32 x1,
simd16float32 x2) {
accu16.f = _mm512_fmadd_ps(x1.f, x2.f, accu16.f);
}
FAISS_ALWAYS_INLINE float result_16() {
// performs better than dividing into _mm256 and adding
return _mm512_reduce_add_ps(accu16.f);
}
};
#elif defined(__AVX2__)
template <>
struct SimilarityIP<SIMDLevel::AVX2> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX2;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
float accu;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr), accu(0) {}
simd8float32 accu8 = {};
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
__m256 yiv = _mm256_loadu_ps(yi);
yi += 8;
accu8.f = _mm256_fmadd_ps(yiv, x.f, accu8.f);
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x1,
simd8float32 x2) {
accu8.f = _mm256_fmadd_ps(x1.f, x2.f, accu8.f);
}
FAISS_ALWAYS_INLINE float result_8() {
const __m128 sum = _mm_add_ps(
_mm256_castps256_ps128(accu8.f),
_mm256_extractf128_ps(accu8.f, 1));
const __m128 v0 = _mm_shuffle_ps(sum, sum, _MM_SHUFFLE(0, 0, 3, 2));
const __m128 v1 = _mm_add_ps(sum, v0);
__m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
const __m128 v3 = _mm_add_ps(v1, v2);
return _mm_cvtss_f32(v3);
}
};
#endif
#ifdef USE_NEON
template <>
struct SimilarityIP<SIMDLevel::ARM_NEON> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::ARM_NEON;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr) {}
simd8float32 accu8;
FAISS_ALWAYS_INLINE void begin_8() {
accu8 = simd8float32(
float32x4x2_t{vdupq_n_f32(0.0f), vdupq_n_f32(0.0f)});
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
float32x4x2_t yiv = vld1q_f32_x2(yi);
yi += 8;
float32x4_t accu8_0 =
vfmaq_f32(accu8.data.val[0], yiv.val[0], x.data.val[0]);
float32x4_t accu8_1 =
vfmaq_f32(accu8.data.val[1], yiv.val[1], x.data.val[1]);
accu8 = simd8float32(float32x4x2_t{accu8_0, accu8_1});
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x1,
simd8float32 x2) {
float32x4_t accu8_0 =
vfmaq_f32(accu8.data.val[0], x1.data.val[0], x2.data.val[0]);
float32x4_t accu8_1 =
vfmaq_f32(accu8.data.val[1], x1.data.val[1], x2.data.val[1]);
accu8 = simd8float32(float32x4x2_t{accu8_0, accu8_1});
}
FAISS_ALWAYS_INLINE float result_8() {
float32x4_t sum_0 = vpaddq_f32(accu8.data.val[0], accu8.data.val[0]);
float32x4_t sum_1 = vpaddq_f32(accu8.data.val[1], accu8.data.val[1]);
float32x4_t sum2_0 = vpaddq_f32(sum_0, sum_0);
float32x4_t sum2_1 = vpaddq_f32(sum_1, sum_1);
return vgetq_lane_f32(sum2_0, 0) + vgetq_lane_f32(sum2_1, 0);
}
};
#endif
} // namespace scalar_quantizer
} // namespace faiss
faiss/impl/scalar_quantizer/sq-avx2.cpp
+455
View File
@@ -0,0 +1,455 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_AVX2
#include <faiss/impl/scalar_quantizer/codecs.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
namespace faiss {
namespace scalar_quantizer {
/**********************************************************
* Codecs
**********************************************************/
template <>
struct Codec8bit<SIMDLevel::AVX2> : Codec8bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
const uint64_t c8 = *(uint64_t*)(code + i);
const __m128i i8 = _mm_set1_epi64x(c8);
const __m256i i32 = _mm256_cvtepu8_epi32(i8);
const __m256 f8 = _mm256_cvtepi32_ps(i32);
const __m256 half_one_255 = _mm256_set1_ps(0.5f / 255.f);
const __m256 one_255 = _mm256_set1_ps(1.f / 255.f);
return simd8float32(_mm256_fmadd_ps(f8, one_255, half_one_255));
}
};
template <>
struct Codec4bit<SIMDLevel::AVX2> : Codec4bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
uint32_t c4 = *(uint32_t*)(code + (i >> 1));
uint32_t mask = 0x0f0f0f0f;
uint32_t c4ev = c4 & mask;
uint32_t c4od = (c4 >> 4) & mask;
// the 8 lower bytes of c8 contain the values
__m128i c8 =
_mm_unpacklo_epi8(_mm_set1_epi32(c4ev), _mm_set1_epi32(c4od));
__m128i c4lo = _mm_cvtepu8_epi32(c8);
__m128i c4hi = _mm_cvtepu8_epi32(_mm_srli_si128(c8, 4));
__m256i i8 = _mm256_castsi128_si256(c4lo);
i8 = _mm256_insertf128_si256(i8, c4hi, 1);
__m256 f8 = _mm256_cvtepi32_ps(i8);
__m256 half = _mm256_set1_ps(0.5f);
f8 = _mm256_add_ps(f8, half);
__m256 one_255 = _mm256_set1_ps(1.f / 15.f);
return simd8float32(_mm256_mul_ps(f8, one_255));
}
};
template <>
struct Codec6bit<SIMDLevel::AVX2> : Codec6bit<SIMDLevel::NONE> {
/* Load 6 bytes that represent 8 6-bit values, return them as a
* 8*32 bit vector register */
static FAISS_ALWAYS_INLINE __m256i load6(const uint16_t* code16) {
const __m128i perm = _mm_set_epi8(
-1, 5, 5, 4, 4, 3, -1, 3, -1, 2, 2, 1, 1, 0, -1, 0);
const __m256i shifts = _mm256_set_epi32(2, 4, 6, 0, 2, 4, 6, 0);
// load 6 bytes
__m128i c1 =
_mm_set_epi16(0, 0, 0, 0, 0, code16[2], code16[1], code16[0]);
// put in 8 * 32 bits
__m128i c2 = _mm_shuffle_epi8(c1, perm);
__m256i c3 = _mm256_cvtepi16_epi32(c2);
// shift and mask out useless bits
__m256i c4 = _mm256_srlv_epi32(c3, shifts);
__m256i c5 = _mm256_and_si256(_mm256_set1_epi32(63), c4);
return c5;
}
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
// // Faster code for Intel CPUs or AMD Zen3+, just keeping it here
// // for the reference, maybe, it becomes used one day.
// const uint16_t* data16 = (const uint16_t*)(code + (i >> 2) * 3);
// const uint32_t* data32 = (const uint32_t*)data16;
// const uint64_t val = *data32 + ((uint64_t)data16[2] << 32);
// const uint64_t vext = _pdep_u64(val, 0x3F3F3F3F3F3F3F3FULL);
// const __m128i i8 = _mm_set1_epi64x(vext);
// const __m256i i32 = _mm256_cvtepi8_epi32(i8);
// const __m256 f8 = _mm256_cvtepi32_ps(i32);
// const __m256 half_one_255 = _mm256_set1_ps(0.5f / 63.f);
// const __m256 one_255 = _mm256_set1_ps(1.f / 63.f);
// return _mm256_fmadd_ps(f8, one_255, half_one_255);
__m256i i8 = load6((const uint16_t*)(code + (i >> 2) * 3));
__m256 f8 = _mm256_cvtepi32_ps(i8);
// this could also be done with bit manipulations but it is
// not obviously faster
const __m256 half_one_255 = _mm256_set1_ps(0.5f / 63.f);
const __m256 one_255 = _mm256_set1_ps(1.f / 63.f);
return simd8float32(_mm256_fmadd_ps(f8, one_255, half_one_255));
}
};
/**********************************************************
* Quantizers (uniform and non-uniform)
**********************************************************/
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::AVX2>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m256 xi = Codec::decode_8_components(code, i).f;
return simd8float32(_mm256_fmadd_ps(
xi, _mm256_set1_ps(this->vdiff), _mm256_set1_ps(this->vmin)));
}
};
template <class Codec>
struct QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::AVX2>
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m256 xi = Codec::decode_8_components(code, i).f;
return simd8float32(_mm256_fmadd_ps(
xi,
_mm256_loadu_ps(this->vdiff + i),
_mm256_loadu_ps(this->vmin + i)));
}
};
/**********************************************************
* FP16 Quantizer
**********************************************************/
template <>
struct QuantizerFP16<SIMDLevel::AVX2> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i codei = _mm_loadu_si128((const __m128i*)(code + 2 * i));
return simd8float32(_mm256_cvtph_ps(codei));
}
};
/**********************************************************
* BF16 Quantizer
**********************************************************/
template <>
struct QuantizerBF16<SIMDLevel::AVX2> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i code_128i = _mm_loadu_si128((const __m128i*)(code + 2 * i));
__m256i code_256i = _mm256_cvtepu16_epi32(code_128i);
code_256i = _mm256_slli_epi32(code_256i, 16);
return simd8float32(_mm256_castsi256_ps(code_256i));
}
};
/**********************************************************
* 8bit Direct Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirect<SIMDLevel::AVX2>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i x8 = _mm_loadl_epi64((__m128i*)(code + i)); // 8 * int8
__m256i y8 = _mm256_cvtepu8_epi32(x8); // 8 * int32
return simd8float32(_mm256_cvtepi32_ps(y8)); // 8 * float32
}
};
/**********************************************************
* 8bit Direct Signed Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::AVX2>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
__m128i x8 = _mm_loadl_epi64((__m128i*)(code + i)); // 8 * int8
__m256i y8 = _mm256_cvtepu8_epi32(x8); // 8 * int32
__m256i c8 = _mm256_set1_epi32(128);
__m256i z8 = _mm256_sub_epi32(y8, c8); // subtract 128 from all lanes
return simd8float32(_mm256_cvtepi32_ps(z8)); // 8 * float32
}
};
/**********************************************************
* SimilarityL2 and SimilarityIP
**********************************************************/
template <>
struct SimilarityL2<SIMDLevel::AVX2> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX2;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
simd8float32 accu8 = {};
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
__m256 yiv = _mm256_loadu_ps(yi);
yi += 8;
__m256 tmp = _mm256_sub_ps(yiv, x.f);
accu8 = simd8float32(_mm256_fmadd_ps(tmp, tmp, accu8.f));
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x,
simd8float32 y_2) {
__m256 tmp = _mm256_sub_ps(y_2.f, x.f);
accu8 = simd8float32(_mm256_fmadd_ps(tmp, tmp, accu8.f));
}
FAISS_ALWAYS_INLINE float result_8() {
const __m128 sum = _mm_add_ps(
_mm256_castps256_ps128(accu8.f),
_mm256_extractf128_ps(accu8.f, 1));
const __m128 v0 = _mm_shuffle_ps(sum, sum, _MM_SHUFFLE(0, 0, 3, 2));
const __m128 v1 = _mm_add_ps(sum, v0);
__m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
const __m128 v3 = _mm_add_ps(v1, v2);
return _mm_cvtss_f32(v3);
}
};
template <>
struct SimilarityIP<SIMDLevel::AVX2> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX2;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
float accu;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr), accu(0) {}
simd8float32 accu8 = {};
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
__m256 yiv = _mm256_loadu_ps(yi);
yi += 8;
accu8.f = _mm256_fmadd_ps(yiv, x.f, accu8.f);
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x1,
simd8float32 x2) {
accu8.f = _mm256_fmadd_ps(x1.f, x2.f, accu8.f);
}
FAISS_ALWAYS_INLINE float result_8() {
const __m128 sum = _mm_add_ps(
_mm256_castps256_ps128(accu8.f),
_mm256_extractf128_ps(accu8.f, 1));
const __m128 v0 = _mm_shuffle_ps(sum, sum, _MM_SHUFFLE(0, 0, 3, 2));
const __m128 v1 = _mm_add_ps(sum, v0);
__m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
const __m128 v3 = _mm_add_ps(v1, v2);
return _mm_cvtss_f32(v3);
}
};
/**********************************************************
* Distance computers
**********************************************************/
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 xi =
quant.reconstruct_8_components(code, static_cast<int>(i));
sim.add_8_components(xi);
}
return sim.result_8();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 x1 =
quant.reconstruct_8_components(code1, static_cast<int>(i));
simd8float32 x2 =
quant.reconstruct_8_components(code2, static_cast<int>(i));
sim.add_8_components_2(x1, x2);
}
return sim.result_8();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::AVX2> : SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
// __m256i accu = _mm256_setzero_ps ();
__m256i accu = _mm256_setzero_si256();
for (int i = 0; i < d; i += 16) {
// load 16 bytes, convert to 16 uint16_t
__m256i c1 = _mm256_cvtepu8_epi16(
_mm_loadu_si128((__m128i*)(code1 + i)));
__m256i c2 = _mm256_cvtepu8_epi16(
_mm_loadu_si128((__m128i*)(code2 + i)));
__m256i prod32;
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
prod32 = _mm256_madd_epi16(c1, c2);
} else {
__m256i diff = _mm256_sub_epi16(c1, c2);
prod32 = _mm256_madd_epi16(diff, diff);
}
accu = _mm256_add_epi32(accu, prod32);
}
__m128i sum = _mm256_extractf128_si256(accu, 0);
sum = _mm_add_epi32(sum, _mm256_extractf128_si256(accu, 1));
sum = _mm_hadd_epi32(sum, sum);
sum = _mm_hadd_epi32(sum, sum);
return _mm_cvtsi128_si32(sum);
}
void set_query(const float* x) final {
/*
for (int i = 0; i < d; i += 8) {
__m256 xi = _mm256_loadu_ps (x + i);
__m256i ci = _mm256_cvtps_epi32(xi);
*/
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
} // namespace scalar_quantizer
} // namespace faiss
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX2
#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
#endif // COMPILE_SIMD_AVX2
faiss/impl/scalar_quantizer/sq-avx512.cpp
+430
View File
@@ -0,0 +1,430 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_AVX512
#include <faiss/impl/scalar_quantizer/codecs.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
namespace faiss {
namespace scalar_quantizer {
/**********************************************************
* Codecs
**********************************************************/
template <>
struct Codec8bit<SIMDLevel::AVX512> : Codec8bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
const __m128i c16 = _mm_loadu_si128((__m128i*)(code + i));
const __m512i i32 = _mm512_cvtepu8_epi32(c16);
const __m512 f16 = _mm512_cvtepi32_ps(i32);
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 255.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 255.f);
return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
}
};
template <>
struct Codec4bit<SIMDLevel::AVX512> : Codec4bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
uint64_t c8 = *(uint64_t*)(code + (i >> 1));
uint64_t mask = 0x0f0f0f0f0f0f0f0f;
uint64_t c8ev = c8 & mask;
uint64_t c8od = (c8 >> 4) & mask;
__m128i c16 =
_mm_unpacklo_epi8(_mm_set1_epi64x(c8ev), _mm_set1_epi64x(c8od));
__m256i c8lo = _mm256_cvtepu8_epi32(c16);
__m256i c8hi = _mm256_cvtepu8_epi32(_mm_srli_si128(c16, 8));
__m512i i16 = _mm512_castsi256_si512(c8lo);
i16 = _mm512_inserti32x8(i16, c8hi, 1);
__m512 f16 = _mm512_cvtepi32_ps(i16);
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 15.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 15.f);
return simd16float32(_mm512_fmadd_ps(f16, one_255, half_one_255));
}
};
template <>
struct Codec6bit<SIMDLevel::AVX512> : Codec6bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd16float32
decode_16_components(const uint8_t* code, size_t i) {
// pure AVX512 implementation (not necessarily the fastest).
// see:
// https://github.com/zilliztech/knowhere/blob/main/thirdparty/faiss/faiss/impl/ScalarQuantizerCodec_avx512.h
// clang-format off
// 16 components, 16x6 bit=12 bytes
const __m128i bit_6v =
_mm_maskz_loadu_epi8(0b0000111111111111, code + (i >> 2) * 3);
const __m256i bit_6v_256 = _mm256_broadcast_i32x4(bit_6v);
// 00 01 02 03 04 05 06 07 08 09 0A 0B 0C 0D 0E 0F
// 00 01 02 03
const __m256i shuffle_mask = _mm256_setr_epi16(
0xFF00, 0x0100, 0x0201, 0xFF02,
0xFF03, 0x0403, 0x0504, 0xFF05,
0xFF06, 0x0706, 0x0807, 0xFF08,
0xFF09, 0x0A09, 0x0B0A, 0xFF0B);
const __m256i shuffled = _mm256_shuffle_epi8(bit_6v_256, shuffle_mask);
// 0: xxxxxxxx xx543210
// 1: xxxx5432 10xxxxxx
// 2: xxxxxx54 3210xxxx
// 3: xxxxxxxx 543210xx
const __m256i shift_right_v = _mm256_setr_epi16(
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U,
0x0U, 0x6U, 0x4U, 0x2U);
__m256i shuffled_shifted = _mm256_srlv_epi16(shuffled, shift_right_v);
// remove unneeded bits
shuffled_shifted =
_mm256_and_si256(shuffled_shifted, _mm256_set1_epi16(0x003F));
// scale
const __m512 f8 =
_mm512_cvtepi32_ps(_mm512_cvtepi16_epi32(shuffled_shifted));
const __m512 half_one_255 = _mm512_set1_ps(0.5f / 63.f);
const __m512 one_255 = _mm512_set1_ps(1.f / 63.f);
return simd16float32(_mm512_fmadd_ps(f8, one_255, half_one_255));
// clang-format on
}
};
/**********************************************************
* Quantizers (uniform and non-uniform)
**********************************************************/
template <class Codec>
struct QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::AVX512>
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m512 xi = Codec::decode_16_components(code, i).f;
return simd16float32(_mm512_fmadd_ps(
xi, _mm512_set1_ps(this->vdiff), _mm512_set1_ps(this->vmin)));
}
};
template <class Codec>
struct QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::AVX512>
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m512 xi = Codec::decode_16_components(code, i).f;
return simd16float32(_mm512_fmadd_ps(
xi,
_mm512_loadu_ps(this->vdiff + i),
_mm512_loadu_ps(this->vmin + i)));
}
};
/**********************************************************
* FP16 Quantizer
**********************************************************/
template <>
struct QuantizerFP16<SIMDLevel::AVX512> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m256i codei = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
return simd16float32(_mm512_cvtph_ps(codei));
}
};
/**********************************************************
* BF16 Quantizer
**********************************************************/
template <>
struct QuantizerBF16<SIMDLevel::AVX512> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m256i code_256i = _mm256_loadu_si256((const __m256i*)(code + 2 * i));
__m512i code_512i = _mm512_cvtepu16_epi32(code_256i);
code_512i = _mm512_slli_epi32(code_512i, 16);
return simd16float32(_mm512_castsi512_ps(code_512i));
}
};
/**********************************************************
* 8bit Direct Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirect<SIMDLevel::AVX512>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
__m512i y16 = _mm512_cvtepu8_epi32(x16); // 16 * int32
return simd16float32(_mm512_cvtepi32_ps(y16)); // 16 * float32
}
};
/**********************************************************
* 8bit Direct Signed Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::AVX512>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {
assert(d % 16 == 0);
}
FAISS_ALWAYS_INLINE simd16float32
reconstruct_16_components(const uint8_t* code, int i) const {
__m128i x16 = _mm_loadu_si128((__m128i*)(code + i)); // 16 * int8
__m512i y16 = _mm512_cvtepu8_epi32(x16); // 16 * int32
__m512i c16 = _mm512_set1_epi32(128);
__m512i z16 = _mm512_sub_epi32(y16, c16); // subtract 128 from all lanes
return simd16float32(_mm512_cvtepi32_ps(z16)); // 16 * float32
}
};
/**********************************************************
* Similarities (L2 and IP)
**********************************************************/
template <>
struct SimilarityL2<SIMDLevel::AVX512> {
static constexpr int simdwidth = 16;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
simd16float32 accu16;
FAISS_ALWAYS_INLINE void begin_16() {
accu16.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
simd16float32 yiv(yi);
yi += 16;
simd16float32 tmp = yiv - x;
accu16 = accu16 + tmp * tmp;
}
FAISS_ALWAYS_INLINE void add_16_components_2(
simd16float32 x,
simd16float32 y_2) {
simd16float32 tmp = y_2 - x;
accu16 = accu16 + tmp * tmp;
}
FAISS_ALWAYS_INLINE float result_16() {
return horizontal_add(accu16);
}
};
template <>
struct SimilarityIP<SIMDLevel::AVX512> {
static constexpr int simdwidth = 16;
static constexpr SIMDLevel simd_level = SIMDLevel::AVX512;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr) {}
simd16float32 accu16;
FAISS_ALWAYS_INLINE void begin_16() {
accu16.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_16_components(simd16float32 x) {
simd16float32 yiv(yi);
yi += 16;
accu16 = accu16 + yiv * x;
}
FAISS_ALWAYS_INLINE void add_16_components_2(
simd16float32 x1,
simd16float32 x2) {
accu16 = accu16 + x1 * x2;
}
FAISS_ALWAYS_INLINE float result_16() {
return horizontal_add(accu16);
}
};
/**********************************************************
* Distance Computers
**********************************************************/
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::AVX512>
: SQDistanceComputer {
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_16();
for (size_t i = 0; i < quant.d; i += 16) {
simd16float32 xi = quant.reconstruct_16_components(code, i);
sim.add_16_components(xi);
}
return sim.result_16();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_16();
for (size_t i = 0; i < quant.d; i += 16) {
simd16float32 x1 = quant.reconstruct_16_components(code1, i);
simd16float32 x2 = quant.reconstruct_16_components(code2, i);
sim.add_16_components_2(x1, x2);
}
return sim.result_16();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::AVX512>
: SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
// compute 16 lanes of 32-bit products (16-bytes) at once for
// the supported metrics
__m512i accu = _mm512_setzero_si512();
constexpr int kLanes = 16;
for (int i = 0; i < d; i += kLanes) {
__m128i c1 = _mm_loadu_si128((__m128i*)(code1 + i));
__m128i c2 = _mm_loadu_si128((__m128i*)(code2 + i));
__m512i c1i = _mm512_cvtepu8_epi32(c1);
__m512i c2i = _mm512_cvtepu8_epi32(c2);
__m512i v;
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
v = _mm512_mullo_epi32(c1i, c2i);
} else {
__m512i diff = _mm512_sub_epi32(c1i, c2i);
v = _mm512_mullo_epi32(diff, diff);
}
accu = _mm512_add_epi32(accu, v);
}
return _mm512_reduce_add_epi32(accu);
}
void set_query(const float* x) final {
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
} // namespace scalar_quantizer
} // namespace faiss
#define THE_LEVEL_TO_DISPATCH SIMDLevel::AVX512
#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
#endif // COMPILE_SIMD_AVX512
faiss/impl/scalar_quantizer/sq-dispatch.h
+329
View File
@@ -0,0 +1,329 @@
/*
* 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/impl/scalar_quantizer/codecs.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
#ifndef THE_LEVEL_TO_DISPATCH
#error "THE_LEVEL_TO_DISPATCH should be set on input to this header"
#endif
namespace faiss {
namespace scalar_quantizer {
// Define SL as alias for THE_LEVEL_TO_DISPATCH for use in this file
constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH;
// Returns true if dimension d is compatible with the given SIMD level
template <SIMDLevel SL2>
constexpr bool is_dimension_compatible(size_t d) {
if constexpr (SL2 == SIMDLevel::AVX512) {
return d % 16 == 0;
} else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
return d % 8 == 0;
} else {
return true; // SIMDLevel::NONE has no alignment requirements
}
}
/*******************************************************************
* sq_select_quantizer: the big switch returning SQuantizer*
*******************************************************************/
template <>
ScalarQuantizer::SQuantizer* sq_select_quantizer<THE_LEVEL_TO_DISPATCH>(
QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
// Return nullptr for incompatible dimensions in SIMD cases
if constexpr (SL != SIMDLevel::NONE) {
if (!is_dimension_compatible<SL>(d)) {
return nullptr;
}
}
switch (qtype) {
case ScalarQuantizer::QT_8bit:
return new QuantizerTemplate<
Codec8bit<SL>,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_6bit:
return new QuantizerTemplate<
Codec6bit<SL>,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_4bit:
return new QuantizerTemplate<
Codec4bit<SL>,
QuantizerTemplateScaling::NON_UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_8bit_uniform:
return new QuantizerTemplate<
Codec8bit<SL>,
QuantizerTemplateScaling::UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_4bit_uniform:
return new QuantizerTemplate<
Codec4bit<SL>,
QuantizerTemplateScaling::UNIFORM,
SL>(d, trained);
case ScalarQuantizer::QT_fp16:
return new QuantizerFP16<SL>(d, trained);
case ScalarQuantizer::QT_bf16:
return new QuantizerBF16<SL>(d, trained);
case ScalarQuantizer::QT_8bit_direct:
return new Quantizer8bitDirect<SL>(d, trained);
case ScalarQuantizer::QT_8bit_direct_signed:
return new Quantizer8bitDirectSigned<SL>(d, trained);
default:
FAISS_THROW_MSG("unknown qtype");
}
}
/*******************************************************************
* select_distance_computer_body: helper for sq_select_distance_computer
*******************************************************************/
template <class Sim, SIMDLevel SL2>
SQDistanceComputer* select_distance_computer_body(
ScalarQuantizer::QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
// Return nullptr for incompatible dimensions in SIMD cases
if constexpr (SL2 != SIMDLevel::NONE) {
if (!is_dimension_compatible<SL2>(d)) {
return nullptr;
}
}
switch (qtype) {
case ScalarQuantizer::QT_8bit_uniform:
return new DCTemplate<
QuantizerTemplate<
Codec8bit<SL2>,
QuantizerTemplateScaling::UNIFORM,
SL2>,
Sim,
SL2>(d, trained);
case ScalarQuantizer::QT_4bit_uniform:
return new DCTemplate<
QuantizerTemplate<
Codec4bit<SL2>,
QuantizerTemplateScaling::UNIFORM,
SL2>,
Sim,
SL2>(d, trained);
case ScalarQuantizer::QT_8bit:
return new DCTemplate<
QuantizerTemplate<
Codec8bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Sim,
SL2>(d, trained);
case ScalarQuantizer::QT_6bit:
return new DCTemplate<
QuantizerTemplate<
Codec6bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Sim,
SL2>(d, trained);
case ScalarQuantizer::QT_4bit:
return new DCTemplate<
QuantizerTemplate<
Codec4bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Sim,
SL2>(d, trained);
case ScalarQuantizer::QT_fp16:
return new DCTemplate<QuantizerFP16<SL2>, Sim, SL2>(d, trained);
case ScalarQuantizer::QT_bf16:
return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
case ScalarQuantizer::QT_8bit_direct:
if constexpr (SL2 == SIMDLevel::AVX512) {
if (d % 32 == 0) {
return new DistanceComputerByte<Sim, SL2>(
static_cast<int>(d), trained);
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
if (d % 16 == 0) {
return new DistanceComputerByte<Sim, SL2>(
static_cast<int>(d), trained);
}
}
return new DCTemplate<Quantizer8bitDirect<SL2>, Sim, SL2>(
d, trained);
case ScalarQuantizer::QT_8bit_direct_signed:
return new DCTemplate<Quantizer8bitDirectSigned<SL2>, Sim, SL2>(
d, trained);
default:
FAISS_THROW_MSG("unknown qtype");
}
}
/*******************************************************************
* sq_select_distance_computer: returns SQDistanceComputer*
*******************************************************************/
template <>
SQDistanceComputer* sq_select_distance_computer<THE_LEVEL_TO_DISPATCH>(
MetricType metric,
ScalarQuantizer::QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
if (metric == METRIC_L2) {
return select_distance_computer_body<SimilarityL2<SL>, SL>(
qtype, d, trained);
} else {
return select_distance_computer_body<SimilarityIP<SL>, SL>(
qtype, d, trained);
}
}
/*******************************************************************
* sq_select_InvertedListScanner: returns InvertedListScanner*
*******************************************************************/
template <>
InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
QuantizerType qtype,
MetricType mt,
size_t d,
size_t code_size,
const std::vector<float>& trained,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual) {
auto scan = [&]<class DCClass>() -> InvertedListScanner* {
if constexpr (DCClass::Sim::metric_type == METRIC_L2) {
return new IVFSQScannerL2<DCClass>(
int(d),
trained,
code_size,
quantizer,
store_pairs,
sel,
by_residual);
} else if constexpr (
DCClass::Sim::metric_type == METRIC_INNER_PRODUCT) {
return new IVFSQScannerIP<DCClass>(
int(d), trained, code_size, store_pairs, sel, by_residual);
} else {
FAISS_THROW_MSG("unsupported metric type");
}
};
auto select_by_simd_and_metric =
[&]<SIMDLevel SL2, class Similarity>() -> InvertedListScanner* {
// Return nullptr for incompatible dimensions in SIMD cases
if constexpr (SL2 != SIMDLevel::NONE) {
if (!is_dimension_compatible<SL2>(d)) {
return nullptr;
}
}
switch (qtype) {
case ScalarQuantizer::QT_8bit_uniform:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec8bit<SL2>,
QuantizerTemplateScaling::UNIFORM,
SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_4bit_uniform:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec4bit<SL2>,
QuantizerTemplateScaling::UNIFORM,
SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_8bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec8bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_4bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec4bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_6bit:
return scan.template operator()<DCTemplate<
QuantizerTemplate<
Codec6bit<SL2>,
QuantizerTemplateScaling::NON_UNIFORM,
SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_fp16:
return scan.template
operator()<DCTemplate<QuantizerFP16<SL2>, Similarity, SL2>>();
case ScalarQuantizer::QT_bf16:
return scan.template
operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
case ScalarQuantizer::QT_8bit_direct:
if constexpr (SL2 == SIMDLevel::AVX512) {
if (d % 32 == 0) {
return scan.template
operator()<DistanceComputerByte<Similarity, SL2>>();
}
} else if constexpr (SL2 == SIMDLevel::AVX2) {
if (d % 16 == 0) {
return scan.template
operator()<DistanceComputerByte<Similarity, SL2>>();
}
}
return scan.template operator()<DCTemplate<
Quantizer8bitDirect<SL2>,
Similarity,
SL2>>();
case ScalarQuantizer::QT_8bit_direct_signed:
return scan.template operator()<DCTemplate<
Quantizer8bitDirectSigned<SL2>,
Similarity,
SL2>>();
default:
FAISS_THROW_MSG("unknown qtype");
}
};
if (mt == METRIC_L2) {
return select_by_simd_and_metric
.template operator()<SL, SimilarityL2<SL>>();
} else if (mt == METRIC_INNER_PRODUCT) {
return select_by_simd_and_metric
.template operator()<SL, SimilarityIP<SL>>();
}
FAISS_THROW_MSG("unsupported metric type");
}
} // namespace scalar_quantizer
} // namespace faiss
faiss/impl/scalar_quantizer/sq-neon.cpp
+467
View File
@@ -0,0 +1,467 @@
/*
* 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.
*/
#ifdef COMPILE_SIMD_ARM_NEON
#include <faiss/impl/scalar_quantizer/codecs.h>
#include <faiss/impl/scalar_quantizer/distance_computers.h>
#include <faiss/impl/scalar_quantizer/quantizers.h>
#include <faiss/impl/scalar_quantizer/scanners.h>
#include <faiss/impl/scalar_quantizer/similarities.h>
namespace faiss {
namespace scalar_quantizer {
/**********************************************************
* Codecs
**********************************************************/
template <>
struct Codec8bit<SIMDLevel::ARM_NEON> : Codec8bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] =
Codec8bit<SIMDLevel::NONE>::decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
};
template <>
struct Codec4bit<SIMDLevel::ARM_NEON> : Codec4bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] =
Codec4bit<SIMDLevel::NONE>::decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
};
template <>
struct Codec6bit<SIMDLevel::ARM_NEON> : Codec6bit<SIMDLevel::NONE> {
static FAISS_ALWAYS_INLINE simd8float32
decode_8_components(const uint8_t* code, size_t i) {
float32_t result[8] = {};
for (size_t j = 0; j < 8; j++) {
result[j] =
Codec6bit<SIMDLevel::NONE>::decode_component(code, i + j);
}
float32x4_t res1 = vld1q_f32(result);
float32x4_t res2 = vld1q_f32(result + 4);
return simd8float32(float32x4x2_t{res1, res2});
}
};
/**********************************************************
* Quantizers (uniform and non-uniform)
**********************************************************/
template <class Codec>
struct QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::ARM_NEON>
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
simd8float32 xi = Codec::decode_8_components(code, i);
return simd8float32(
float32x4x2_t{
vfmaq_n_f32(
vdupq_n_f32(this->vmin),
xi.data.val[0],
this->vdiff),
vfmaq_n_f32(
vdupq_n_f32(this->vmin),
xi.data.val[1],
this->vdiff)});
}
};
template <class Codec>
struct QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::ARM_NEON>
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE> {
QuantizerTemplate(size_t d, const std::vector<float>& trained)
: QuantizerTemplate<
Codec,
scalar_quantizer::QuantizerTemplateScaling::NON_UNIFORM,
SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
simd8float32 xi = Codec::decode_8_components(code, i);
return simd8float32(
float32x4x2_t{
vfmaq_f32(
vld1q_f32(this->vmin + i),
xi.data.val[0],
vld1q_f32(this->vdiff + i)),
vfmaq_f32(
vld1q_f32(this->vmin + i + 4),
xi.data.val[1],
vld1q_f32(this->vdiff + i + 4))});
}
};
/**********************************************************
* FP16 Quantizer
**********************************************************/
template <>
struct QuantizerFP16<SIMDLevel::ARM_NEON> : QuantizerFP16<SIMDLevel::NONE> {
QuantizerFP16(size_t d, const std::vector<float>& trained)
: QuantizerFP16<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint16x4x2_t codei = vld1_u16_x2((const uint16_t*)(code + 2 * i));
return simd8float32(
float32x4x2_t{
vcvt_f32_f16(vreinterpret_f16_u16(codei.val[0])),
vcvt_f32_f16(vreinterpret_f16_u16(codei.val[1]))});
}
};
/**********************************************************
* BF16 Quantizer
**********************************************************/
template <>
struct QuantizerBF16<SIMDLevel::ARM_NEON> : QuantizerBF16<SIMDLevel::NONE> {
QuantizerBF16(size_t d, const std::vector<float>& trained)
: QuantizerBF16<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint16x4x2_t codei = vld1_u16_x2((const uint16_t*)(code + 2 * i));
return simd8float32(
float32x4x2_t{
vreinterpretq_f32_u32(
vshlq_n_u32(vmovl_u16(codei.val[0]), 16)),
vreinterpretq_f32_u32(
vshlq_n_u32(vmovl_u16(codei.val[1]), 16))});
}
};
/**********************************************************
* 8bit Direct Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirect<SIMDLevel::ARM_NEON>
: Quantizer8bitDirect<SIMDLevel::NONE> {
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirect<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint8x8_t x8 = vld1_u8((const uint8_t*)(code + i));
uint16x8_t y8 = vmovl_u8(x8);
uint16x4_t y8_0 = vget_low_u16(y8);
uint16x4_t y8_1 = vget_high_u16(y8);
return simd8float32(
float32x4x2_t{
vcvtq_f32_u32(vmovl_u16(y8_0)),
vcvtq_f32_u32(vmovl_u16(y8_1))});
}
};
/**********************************************************
* 8bit Direct Signed Quantizer
**********************************************************/
template <>
struct Quantizer8bitDirectSigned<SIMDLevel::ARM_NEON>
: Quantizer8bitDirectSigned<SIMDLevel::NONE> {
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
: Quantizer8bitDirectSigned<SIMDLevel::NONE>(d, trained) {
assert(d % 8 == 0);
}
FAISS_ALWAYS_INLINE simd8float32
reconstruct_8_components(const uint8_t* code, int i) const {
uint8x8_t x8 = vld1_u8((const uint8_t*)(code + i));
uint16x8_t y8 = vmovl_u8(x8);
int16x8_t z8 = vreinterpretq_s16_u16(
vsubq_u16(y8, vdupq_n_u16(128))); // subtract 128 from all lanes
int16x4_t z8_0 = vget_low_s16(z8);
int16x4_t z8_1 = vget_high_s16(z8);
return simd8float32(
float32x4x2_t{
vcvtq_f32_s32(vmovl_s16(z8_0)),
vcvtq_f32_s32(vmovl_s16(z8_1))});
}
};
/**********************************************************
* Similarities (L2 and IP)
**********************************************************/
template <>
struct SimilarityL2<SIMDLevel::ARM_NEON> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::ARM_NEON;
static constexpr MetricType metric_type = METRIC_L2;
const float *y, *yi;
explicit SimilarityL2(const float* y) : y(y), yi(nullptr) {}
simd8float32 accu8;
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
simd8float32 yiv(yi);
yi += 8;
simd8float32 tmp = yiv - x;
accu8 = accu8 + tmp * tmp;
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x,
simd8float32 y_2) {
simd8float32 tmp = y_2 - x;
accu8 = accu8 + tmp * tmp;
}
FAISS_ALWAYS_INLINE float result_8() {
return horizontal_add(accu8);
}
};
template <>
struct SimilarityIP<SIMDLevel::ARM_NEON> {
static constexpr int simdwidth = 8;
static constexpr SIMDLevel simd_level = SIMDLevel::ARM_NEON;
static constexpr MetricType metric_type = METRIC_INNER_PRODUCT;
const float *y, *yi;
explicit SimilarityIP(const float* y) : y(y), yi(nullptr) {}
simd8float32 accu8;
FAISS_ALWAYS_INLINE void begin_8() {
accu8.clear();
yi = y;
}
FAISS_ALWAYS_INLINE void add_8_components(simd8float32 x) {
simd8float32 yiv(yi);
yi += 8;
accu8 = accu8 + yiv * x;
}
FAISS_ALWAYS_INLINE void add_8_components_2(
simd8float32 x1,
simd8float32 x2) {
accu8 = accu8 + x1 * x2;
}
FAISS_ALWAYS_INLINE float result_8() {
return horizontal_add(accu8);
}
};
/**********************************************************
* Distance Computers
**********************************************************/
template <class Quantizer, class Similarity>
struct DCTemplate<Quantizer, Similarity, SIMDLevel::ARM_NEON>
: SQDistanceComputer {
using Sim = Similarity;
Quantizer quant;
DCTemplate(size_t d, const std::vector<float>& trained)
: quant(d, trained) {}
float compute_distance(const float* x, const uint8_t* code) const {
Similarity sim(x);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 xi = quant.reconstruct_8_components(code, i);
sim.add_8_components(xi);
}
return sim.result_8();
}
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
Similarity sim(nullptr);
sim.begin_8();
for (size_t i = 0; i < quant.d; i += 8) {
simd8float32 x1 = quant.reconstruct_8_components(code1, i);
simd8float32 x2 = quant.reconstruct_8_components(code2, i);
sim.add_8_components_2(x1, x2);
}
return sim.result_8();
}
void set_query(const float* x) final {
q = x;
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_distance(q, code);
}
};
template <class Similarity>
struct DistanceComputerByte<Similarity, SIMDLevel::ARM_NEON>
: SQDistanceComputer {
using Sim = Similarity;
int d;
std::vector<uint8_t> tmp;
DistanceComputerByte(int d, const std::vector<float>&) : d(d), tmp(d) {}
int compute_code_distance(const uint8_t* code1, const uint8_t* code2)
const {
int accu = 0;
for (int i = 0; i < d; i++) {
if (Sim::metric_type == METRIC_INNER_PRODUCT) {
accu += int(code1[i]) * code2[i];
} else {
int diff = int(code1[i]) - code2[i];
accu += diff * diff;
}
}
return accu;
}
void set_query(const float* x) final {
for (int i = 0; i < d; i++) {
tmp[i] = int(x[i]);
}
}
int compute_distance(const float* x, const uint8_t* code) {
set_query(x);
return compute_code_distance(tmp.data(), code);
}
float symmetric_dis(idx_t i, idx_t j) override {
return compute_code_distance(
codes + i * code_size, codes + j * code_size);
}
float query_to_code(const uint8_t* code) const final {
return compute_code_distance(tmp.data(), code);
}
};
} // namespace scalar_quantizer
} // namespace faiss
#define THE_LEVEL_TO_DISPATCH SIMDLevel::ARM_NEON
#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
#ifdef COMPILE_SIMD_ARM_SVE
// ARM_SVE: SVE is a superset of NEON. Forward to the NEON implementation
// until a dedicated SVE specialization is written.
namespace faiss {
namespace scalar_quantizer {
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
ScalarQuantizer::SQuantizer* sq_select_quantizer<SIMDLevel::ARM_SVE>(
QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
return sq_select_quantizer<SIMDLevel::ARM_NEON>(qtype, d, trained);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
SQDistanceComputer* sq_select_distance_computer<SIMDLevel::ARM_SVE>(
MetricType metric,
ScalarQuantizer::QuantizerType qtype,
size_t d,
const std::vector<float>& trained) {
return sq_select_distance_computer<SIMDLevel::ARM_NEON>(
metric, qtype, d, trained);
}
// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
template <>
InvertedListScanner* sq_select_InvertedListScanner<SIMDLevel::ARM_SVE>(
QuantizerType qtype,
MetricType mt,
size_t d,
size_t code_size,
const std::vector<float>& trained,
const Index* quantizer,
bool store_pairs,
const IDSelector* sel,
bool by_residual) {
return sq_select_InvertedListScanner<SIMDLevel::ARM_NEON>(
qtype,
mt,
d,
code_size,
trained,
quantizer,
store_pairs,
sel,
by_residual);
}
} // namespace scalar_quantizer
} // namespace faiss
#endif // COMPILE_SIMD_ARM_SVE
#endif // COMPILE_SIMD_ARM_NEON
faiss/utils/simdlib_avx2.h
+12 -1
View File
@@ -794,10 +794,21 @@ inline simd8float32 getlow128(simd8float32 a, simd8float32 b) {
return simd8float32(_mm256_permute2f128_ps(a.f, b.f, 0 | 2 << 4));
}
inline simd8float32 gethigh128(simd8float32 a, simd8float32 b) {
inline simd8float32 gethigh128(const simd8float32& a, const simd8float32& b) {
return simd8float32(_mm256_permute2f128_ps(a.f, b.f, 1 | 3 << 4));
}
// horizontal add: sum all 8 floats in the register
inline float horizontal_add(const simd8float32& a) {
__m128 sum = _mm_add_ps(
_mm256_castps256_ps128(a.f), _mm256_extractf128_ps(a.f, 1));
__m128 v0 = _mm_shuffle_ps(sum, sum, _MM_SHUFFLE(0, 0, 3, 2));
__m128 v1 = _mm_add_ps(sum, v0);
__m128 v2 = _mm_shuffle_ps(v1, v1, _MM_SHUFFLE(0, 0, 0, 1));
__m128 v3 = _mm_add_ps(v1, v2);
return _mm_cvtss_f32(v3);
}
} // namespace
} // namespace faiss
faiss/utils/simdlib_avx512.h
+5
View File
@@ -357,4 +357,9 @@ inline simd16float32 fmadd(
return simd16float32(_mm512_fmadd_ps(a.f, b.f, c.f));
}
// horizontal add: sum all 16 floats in the register
inline float horizontal_add(const simd16float32& a) {
return _mm512_reduce_add_ps(a.f);
}
} // namespace faiss
faiss/utils/simdlib_neon.h
+6
View File
@@ -1355,6 +1355,12 @@ simd8float32 gethigh128(const simd8float32& a, const simd8float32& b) {
return simd8float32{float32x4x2_t{a.data.val[1], b.data.val[1]}};
}
// horizontal add: sum all 8 floats in the register
inline float horizontal_add(const simd8float32& a) {
float32x4_t sum = vaddq_f32(a.data.val[0], a.data.val[1]);
return vaddvq_f32(sum);
}
} // namespace
} // namespace faiss