faiss 0.6.0 → 0.6.1

data/vendor/faiss/faiss/impl/simdlib/simdlib_avx512.h

@@ -0,0 +1,414 @@
+/*
+ * 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 <cstdint>
+#include <string>
+
+#include <immintrin.h>
+
+#include <faiss/impl/platform_macros.h>
+#include <faiss/impl/simdlib/simdlib.h>
+
+#include <faiss/impl/simdlib/simdlib_avx2.h>
+
+namespace faiss {
+
+/** Simple wrapper around the AVX 512-bit registers
+ *
+ * The objective is to separate the different interpretations of the same
+ * registers (as a vector of uint8, uint16 or uint32), to provide printing
+ * functions, and to give more readable names to the AVX intrinsics. It does not
+ * pretend to be exhaustive, functions are added as needed.
+ */
+
+/// 512-bit representation without interpretation as a vector
+template <>
+struct simd512bit_tpl<SIMDLevel::AVX512> {
+    union {
+        __m512i i;
+        __m512 f;
+    };
+
+    simd512bit_tpl() {}
+
+    explicit simd512bit_tpl(__m512i i) : i(i) {}
+
+    explicit simd512bit_tpl(__m512 f) : f(f) {}
+
+    explicit simd512bit_tpl(const void* x)
+            : i(_mm512_loadu_si512((__m512i const*)x)) {}
+
+    // sets up a lower half of the register while keeping upper one as zero
+    explicit simd512bit_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
+            : simd512bit_tpl(_mm512_inserti32x8(
+                      _mm512_castsi256_si512(lo.i),
+                      _mm256_setzero_si256(),
+                      1)) {}
+
+    // constructs from lower and upper halves
+    explicit simd512bit_tpl(
+            simd256bit_tpl<SIMDLevel::AVX2> lo,
+            simd256bit_tpl<SIMDLevel::AVX2> hi)
+            : simd512bit_tpl(_mm512_inserti32x8(
+                      _mm512_castsi256_si512(lo.i),
+                      hi.i,
+                      1)) {}
+
+    void clear() {
+        i = _mm512_setzero_si512();
+    }
+
+    void storeu(void* ptr) const {
+        _mm512_storeu_si512((__m512i*)ptr, i);
+    }
+
+    void loadu(const void* ptr) {
+        i = _mm512_loadu_si512((__m512i*)ptr);
+    }
+
+    void store(void* ptr) const {
+        _mm512_storeu_si512((__m512i*)ptr, i);
+    }
+
+    void bin(char bits[513]) const {
+        char bytes[64];
+        storeu((void*)bytes);
+        for (int i = 0; i < 512; i++) {
+            bits[i] = '0' + ((bytes[i / 8] >> (i % 8)) & 1);
+        }
+        bits[512] = 0;
+    }
+
+    std::string bin() const {
+        char bits[513];
+        bin(bits);
+        return std::string(bits);
+    }
+};
+
+/// vector of 32 elements in uint16
+template <>
+struct simd32uint16_tpl<SIMDLevel::AVX512> : simd512bit_tpl<SIMDLevel::AVX512> {
+    simd32uint16_tpl() {}
+
+    explicit simd32uint16_tpl(__m512i i)
+            : simd512bit_tpl<SIMDLevel::AVX512>(i) {}
+
+    explicit simd32uint16_tpl(int x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi16(x)) {}
+
+    explicit simd32uint16_tpl(uint16_t x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi16(x)) {}
+
+    explicit simd32uint16_tpl(simd512bit_tpl<SIMDLevel::AVX512> x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(x) {}
+
+    explicit simd32uint16_tpl(const uint16_t* x)
+            : simd512bit_tpl<SIMDLevel::AVX512>((const void*)x) {}
+
+    // sets up a lower half of the register
+    explicit simd32uint16_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
+            : simd512bit_tpl<SIMDLevel::AVX512>(lo) {}
+
+    // constructs from lower and upper halves
+    explicit simd32uint16_tpl(
+            simd256bit_tpl<SIMDLevel::AVX2> lo,
+            simd256bit_tpl<SIMDLevel::AVX2> hi)
+            : simd512bit_tpl<SIMDLevel::AVX512>(lo, hi) {}
+
+    std::string elements_to_string(const char* fmt) const {
+        uint16_t bytes[32];
+        storeu((void*)bytes);
+        char res[2000];
+        char* ptr = res;
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wformat-nonliteral"
+#endif
+        for (int i = 0; i < 32; i++) {
+            ptr += snprintf(
+                    ptr, (size_t)(res + sizeof(res) - ptr), fmt, bytes[i]);
+        }
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic pop
+#endif
+        // strip last ,
+        ptr[-1] = 0;
+        return std::string(res);
+    }
+
+    std::string hex() const {
+        return elements_to_string("%02x,");
+    }
+
+    std::string dec() const {
+        return elements_to_string("%3d,");
+    }
+
+    void set1(uint16_t x) {
+        i = _mm512_set1_epi16((short)x);
+    }
+
+    simd32uint16_tpl operator*(const simd32uint16_tpl& other) const {
+        return simd32uint16_tpl(_mm512_mullo_epi16(i, other.i));
+    }
+
+    // shift must be known at compile time
+    simd32uint16_tpl operator>>(const int shift) const {
+        return simd32uint16_tpl(_mm512_srli_epi16(i, shift));
+    }
+
+    // shift must be known at compile time
+    simd32uint16_tpl operator<<(const int shift) const {
+        return simd32uint16_tpl(_mm512_slli_epi16(i, shift));
+    }
+
+    simd32uint16_tpl& operator+=(simd32uint16_tpl other) {
+        i = _mm512_add_epi16(i, other.i);
+        return *this;
+    }
+
+    simd32uint16_tpl& operator-=(simd32uint16_tpl other) {
+        i = _mm512_sub_epi16(i, other.i);
+        return *this;
+    }
+
+    simd32uint16_tpl operator+(simd32uint16_tpl other) const {
+        return simd32uint16_tpl(_mm512_add_epi16(i, other.i));
+    }
+
+    simd32uint16_tpl operator-(simd32uint16_tpl other) const {
+        return simd32uint16_tpl(_mm512_sub_epi16(i, other.i));
+    }
+
+    simd32uint16_tpl operator&(simd512bit_tpl<SIMDLevel::AVX512> other) const {
+        return simd32uint16_tpl(_mm512_and_si512(i, other.i));
+    }
+
+    simd32uint16_tpl operator|(simd512bit_tpl<SIMDLevel::AVX512> other) const {
+        return simd32uint16_tpl(_mm512_or_si512(i, other.i));
+    }
+
+    simd32uint16_tpl operator^(simd512bit_tpl<SIMDLevel::AVX512> other) const {
+        return simd32uint16_tpl(_mm512_xor_si512(i, other.i));
+    }
+
+    simd32uint16_tpl operator~() const {
+        return simd32uint16_tpl(_mm512_xor_si512(i, _mm512_set1_epi32(-1)));
+    }
+
+    simd16uint16_tpl<SIMDLevel::AVX2> low() const {
+        return simd16uint16_tpl<SIMDLevel::AVX2>(_mm512_castsi512_si256(i));
+    }
+
+    simd16uint16_tpl<SIMDLevel::AVX2> high() const {
+        return simd16uint16_tpl<SIMDLevel::AVX2>(
+                _mm512_extracti32x8_epi32(i, 1));
+    }
+
+    // for debugging only
+    uint16_t operator[](int i) const {
+        ALIGNED(64) uint16_t tab[32];
+        store(tab);
+        return tab[i];
+    }
+
+    void accu_min(simd32uint16_tpl incoming) {
+        i = _mm512_min_epu16(i, incoming.i);
+    }
+
+    void accu_max(simd32uint16_tpl incoming) {
+        i = _mm512_max_epu16(i, incoming.i);
+    }
+};
+
+// decompose in 128-lanes: a = (a0, a1, a2, a3), b = (b0, b1, b2, b3)
+// return (a0 + a1 + a2 + a3, b0 + b1 + b2 + b3)
+inline simd16uint16_tpl<SIMDLevel::AVX2> combine4x2(
+        simd32uint16_tpl<SIMDLevel::AVX512> a,
+        simd32uint16_tpl<SIMDLevel::AVX512> b) {
+    return combine2x2(a.low(), b.low()) + combine2x2(a.high(), b.high());
+}
+
+// vector of 32 unsigned 8-bit integers
+template <>
+struct simd64uint8_tpl<SIMDLevel::AVX512> : simd512bit_tpl<SIMDLevel::AVX512> {
+    simd64uint8_tpl() {}
+
+    explicit simd64uint8_tpl(__m512i i)
+            : simd512bit_tpl<SIMDLevel::AVX512>(i) {}
+
+    explicit simd64uint8_tpl(int x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi8(x)) {}
+
+    explicit simd64uint8_tpl(uint8_t x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_epi8(x)) {}
+
+    // sets up a lower half of the register
+    explicit simd64uint8_tpl(simd256bit_tpl<SIMDLevel::AVX2> lo)
+            : simd512bit_tpl<SIMDLevel::AVX512>(lo) {}
+
+    // constructs from lower and upper halves
+    explicit simd64uint8_tpl(
+            simd256bit_tpl<SIMDLevel::AVX2> lo,
+            simd256bit_tpl<SIMDLevel::AVX2> hi)
+            : simd512bit_tpl<SIMDLevel::AVX512>(lo, hi) {}
+
+    explicit simd64uint8_tpl(simd512bit_tpl<SIMDLevel::AVX512> x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(x) {}
+
+    explicit simd64uint8_tpl(const uint8_t* x)
+            : simd512bit_tpl<SIMDLevel::AVX512>((const void*)x) {}
+
+    std::string elements_to_string(const char* fmt) const {
+        uint8_t bytes[64];
+        storeu((void*)bytes);
+        char res[2000];
+        char* ptr = res;
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wformat-nonliteral"
+#endif
+        for (int i = 0; i < 64; i++) {
+            ptr += snprintf(
+                    ptr, (size_t)(res + sizeof(res) - ptr), fmt, bytes[i]);
+        }
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic pop
+#endif
+        // strip last ,
+        ptr[-1] = 0;
+        return std::string(res);
+    }
+
+    std::string hex() const {
+        return elements_to_string("%02x,");
+    }
+
+    std::string dec() const {
+        return elements_to_string("%3d,");
+    }
+
+    void set1(uint8_t x) {
+        i = _mm512_set1_epi8((char)x);
+    }
+
+    simd64uint8_tpl operator&(simd512bit_tpl<SIMDLevel::AVX512> other) const {
+        return simd64uint8_tpl(_mm512_and_si512(i, other.i));
+    }
+
+    simd64uint8_tpl operator+(simd64uint8_tpl other) const {
+        return simd64uint8_tpl(_mm512_add_epi8(i, other.i));
+    }
+
+    simd64uint8_tpl lookup_4_lanes(simd64uint8_tpl idx) const {
+        return simd64uint8_tpl(_mm512_shuffle_epi8(i, idx.i));
+    }
+
+    // extract + 0-extend lane
+    // this operation is slow (3 cycles)
+    simd32uint16_tpl<SIMDLevel::AVX512> lane0_as_uint16() const {
+        __m256i x = _mm512_extracti32x8_epi32(i, 0);
+        return simd32uint16_tpl<SIMDLevel::AVX512>(_mm512_cvtepu8_epi16(x));
+    }
+
+    simd32uint16_tpl<SIMDLevel::AVX512> lane1_as_uint16() const {
+        __m256i x = _mm512_extracti32x8_epi32(i, 1);
+        return simd32uint16_tpl<SIMDLevel::AVX512>(_mm512_cvtepu8_epi16(x));
+    }
+
+    simd64uint8_tpl& operator+=(simd64uint8_tpl other) {
+        i = _mm512_add_epi8(i, other.i);
+        return *this;
+    }
+
+    // for debugging only
+    uint8_t operator[](int i) const {
+        ALIGNED(64) uint8_t tab[64];
+        store(tab);
+        return tab[i];
+    }
+};
+
+/// vector of 16 32-bit floats
+template <>
+struct simd16float32_tpl<SIMDLevel::AVX512>
+        : simd512bit_tpl<SIMDLevel::AVX512> {
+    simd16float32_tpl() {}
+
+    explicit simd16float32_tpl(__m512 f)
+            : simd512bit_tpl<SIMDLevel::AVX512>(f) {}
+
+    explicit simd16float32_tpl(float x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_set1_ps(x)) {}
+
+    explicit simd16float32_tpl(const float* x)
+            : simd512bit_tpl<SIMDLevel::AVX512>(_mm512_loadu_ps(x)) {}
+
+    void clear() {
+        f = _mm512_setzero_ps();
+    }
+
+    void storeu(float* ptr) const {
+        _mm512_storeu_ps(ptr, f);
+    }
+
+    void loadu(const float* ptr) {
+        f = _mm512_loadu_ps(ptr);
+    }
+
+    void store(float* ptr) const {
+        _mm512_storeu_ps(ptr, f);
+    }
+
+    simd16float32_tpl operator*(const simd16float32_tpl& other) const {
+        return simd16float32_tpl(_mm512_mul_ps(f, other.f));
+    }
+
+    simd16float32_tpl operator+(const simd16float32_tpl& other) const {
+        return simd16float32_tpl(_mm512_add_ps(f, other.f));
+    }
+
+    simd16float32_tpl operator-(const simd16float32_tpl& other) const {
+        return simd16float32_tpl(_mm512_sub_ps(f, other.f));
+    }
+
+    simd16float32_tpl& operator+=(const simd16float32_tpl& other) {
+        f = _mm512_add_ps(f, other.f);
+        return *this;
+    }
+
+    std::string tostring() const {
+        float tab[16];
+        storeu(tab);
+        char res[1000];
+        char* ptr = res;
+        for (int i = 0; i < 16; i++) {
+            ptr += snprintf(
+                    ptr, (size_t)(res + sizeof(res) - ptr), "%g,", tab[i]);
+        }
+        ptr[-1] = 0;
+        return std::string(res);
+    }
+};
+
+// compute a * b + c
+inline simd16float32_tpl<SIMDLevel::AVX512> fmadd(
+        const simd16float32_tpl<SIMDLevel::AVX512>& a,
+        const simd16float32_tpl<SIMDLevel::AVX512>& b,
+        const simd16float32_tpl<SIMDLevel::AVX512>& c) {
+    return simd16float32_tpl<SIMDLevel::AVX512>(_mm512_fmadd_ps(a.f, b.f, c.f));
+}
+
+// horizontal add: sum all 16 floats in the register
+inline float horizontal_add(const simd16float32_tpl<SIMDLevel::AVX512>& a) {
+    return _mm512_reduce_add_ps(a.f);
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/simdlib/simdlib_dispatch.h

@@ -0,0 +1,44 @@
+/*
+ * 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
+
+/** Includes simdlib.h (primary templates + NONE/emulated specialization)
+ * plus the platform specialization for the current compilation context.
+ *
+ * Generic code should include this header.
+ * Per-SIMD TUs should include the concrete header directly
+ * (e.g., simdlib_avx2.h).
+ */
+
+#include <faiss/impl/simdlib/simdlib.h>
+
+// Platform specializations — guarded by COMPILE_SIMD_* AND compiler macros.
+// In DD mode: COMPILE_SIMD_* are target-wide, compiler macros are per-file.
+// Only per-SIMD TUs (compiled with -mavx2 etc.) see the platform
+// specializations. In static mode: only the compiled-in level is available.
+
+#if (defined(COMPILE_SIMD_AVX512) || defined(COMPILE_SIMD_AVX512_SPR)) && \
+        defined(__AVX512F__)
+
+// AVX512 includes AVX2 (simdlib_avx512.h includes simdlib_avx2.h)
+#include <faiss/impl/simdlib/simdlib_avx512.h>
+
+#elif defined(COMPILE_SIMD_AVX2) && defined(__AVX2__)
+
+#include <faiss/impl/simdlib/simdlib_avx2.h>
+
+#elif defined(COMPILE_SIMD_ARM_NEON) && defined(__aarch64__)
+
+#include <faiss/impl/simdlib/simdlib_neon.h>
+
+#endif
+
+// No global bare-name aliases (simd16uint16, simd32uint8, etc.) — each file
+// that needs them must declare its own `using` with an explicit SIMD level.
+// This prevents per-ISA TUs from accidentally picking up SINGLE_SIMD_LEVEL
+// (= NONE in DD mode) when they should use THE_SIMD_LEVEL.