faiss 0.3.1 → 0.3.3

data/vendor/faiss/faiss/impl/code_distance/code_distance-sve.h

@@ -0,0 +1,440 @@
+/*
+ * 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
+
+#ifdef __ARM_FEATURE_SVE
+
+#include <arm_sve.h>
+
+#include <tuple>
+#include <type_traits>
+
+#include <faiss/impl/code_distance/code_distance-generic.h>
+
+namespace faiss {
+
+template <typename PQDecoderT>
+std::enable_if_t<!std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_single_code_sve(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        const uint8_t* code) {
+    // default implementation
+    return distance_single_code_generic<PQDecoderT>(M, nbits, sim_table, code);
+}
+
+static inline void distance_codes_kernel(
+        svbool_t pg,
+        svuint32_t idx1,
+        svuint32_t offsets_0,
+        const float* tab,
+        svfloat32_t& partialSum) {
+    // add offset
+    const auto indices_to_read_from = svadd_u32_x(pg, idx1, offsets_0);
+
+    // gather values, similar to some operations of tab[index]
+    const auto collected =
+            svld1_gather_u32index_f32(pg, tab, indices_to_read_from);
+
+    // collect partial sum
+    partialSum = svadd_f32_m(pg, partialSum, collected);
+}
+
+static float distance_single_code_sve_for_small_m(
+        // the product quantizer
+        const size_t M,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code) {
+    constexpr size_t nbits = 8u;
+
+    const size_t ksub = 1 << nbits;
+
+    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
+
+    // loop
+    const auto pg = svwhilelt_b32_u64(0, M);
+
+    auto mm1 = svld1ub_u32(pg, code);
+    mm1 = svadd_u32_x(pg, mm1, offsets_0);
+    const auto collected0 = svld1_gather_u32index_f32(pg, sim_table, mm1);
+    return svaddv_f32(pg, collected0);
+}
+
+template <typename PQDecoderT>
+std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, float> inline distance_single_code_sve(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        const uint8_t* code) {
+    if (M <= svcntw())
+        return distance_single_code_sve_for_small_m(M, sim_table, code);
+
+    const float* tab = sim_table;
+
+    const size_t ksub = 1 << nbits;
+
+    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
+
+    // accumulators of partial sums
+    auto partialSum = svdup_n_f32(0.f);
+
+    const auto lanes = svcntb();
+    const auto quad_lanes = lanes / 4;
+
+    // loop
+    for (std::size_t m = 0; m < M;) {
+        const auto pg = svwhilelt_b8_u64(m, M);
+
+        const auto mm1 = svld1_u8(pg, code + m);
+        {
+            const auto mm1lo = svunpklo_u16(mm1);
+            const auto pglo = svunpklo_b(pg);
+
+            {
+                // convert uint8 values to uint32 values
+                const auto idx1 = svunpklo_u32(mm1lo);
+                const auto pglolo = svunpklo_b(pglo);
+
+                distance_codes_kernel(pglolo, idx1, offsets_0, tab, partialSum);
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                // convert uint8 values to uint32 values
+                const auto idx1 = svunpkhi_u32(mm1lo);
+                const auto pglohi = svunpkhi_b(pglo);
+
+                distance_codes_kernel(pglohi, idx1, offsets_0, tab, partialSum);
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+        }
+
+        {
+            const auto mm1hi = svunpkhi_u16(mm1);
+            const auto pghi = svunpkhi_b(pg);
+
+            {
+                // convert uint8 values to uint32 values
+                const auto idx1 = svunpklo_u32(mm1hi);
+                const auto pghilo = svunpklo_b(pghi);
+
+                distance_codes_kernel(pghilo, idx1, offsets_0, tab, partialSum);
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                // convert uint8 values to uint32 values
+                const auto idx1 = svunpkhi_u32(mm1hi);
+                const auto pghihi = svunpkhi_b(pghi);
+
+                distance_codes_kernel(pghihi, idx1, offsets_0, tab, partialSum);
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+        }
+    }
+
+    return svaddv_f32(svptrue_b32(), partialSum);
+}
+
+template <typename PQDecoderT>
+std::enable_if_t<!std::is_same_v<PQDecoderT, PQDecoder8>, void>
+distance_four_codes_sve(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    distance_four_codes_generic<PQDecoderT>(
+            M,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+static void distance_four_codes_sve_for_small_m(
+        // the product quantizer
+        const size_t M,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    constexpr size_t nbits = 8u;
+
+    const size_t ksub = 1 << nbits;
+
+    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
+
+    const auto quad_lanes = svcntw();
+
+    // loop
+    const auto pg = svwhilelt_b32_u64(0, M);
+
+    auto mm10 = svld1ub_u32(pg, code0);
+    auto mm11 = svld1ub_u32(pg, code1);
+    auto mm12 = svld1ub_u32(pg, code2);
+    auto mm13 = svld1ub_u32(pg, code3);
+    mm10 = svadd_u32_x(pg, mm10, offsets_0);
+    mm11 = svadd_u32_x(pg, mm11, offsets_0);
+    mm12 = svadd_u32_x(pg, mm12, offsets_0);
+    mm13 = svadd_u32_x(pg, mm13, offsets_0);
+    const auto collected0 = svld1_gather_u32index_f32(pg, sim_table, mm10);
+    const auto collected1 = svld1_gather_u32index_f32(pg, sim_table, mm11);
+    const auto collected2 = svld1_gather_u32index_f32(pg, sim_table, mm12);
+    const auto collected3 = svld1_gather_u32index_f32(pg, sim_table, mm13);
+    result0 = svaddv_f32(pg, collected0);
+    result1 = svaddv_f32(pg, collected1);
+    result2 = svaddv_f32(pg, collected2);
+    result3 = svaddv_f32(pg, collected3);
+}
+
+// Combines 4 operations of distance_single_code()
+template <typename PQDecoderT>
+std::enable_if_t<std::is_same_v<PQDecoderT, PQDecoder8>, void>
+distance_four_codes_sve(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    if (M <= svcntw()) {
+        distance_four_codes_sve_for_small_m(
+                M,
+                sim_table,
+                code0,
+                code1,
+                code2,
+                code3,
+                result0,
+                result1,
+                result2,
+                result3);
+        return;
+    }
+
+    const float* tab = sim_table;
+
+    const size_t ksub = 1 << nbits;
+
+    const auto offsets_0 = svindex_u32(0, static_cast<uint32_t>(ksub));
+
+    // accumulators of partial sums
+    auto partialSum0 = svdup_n_f32(0.f);
+    auto partialSum1 = svdup_n_f32(0.f);
+    auto partialSum2 = svdup_n_f32(0.f);
+    auto partialSum3 = svdup_n_f32(0.f);
+
+    const auto lanes = svcntb();
+    const auto quad_lanes = lanes / 4;
+
+    // loop
+    for (std::size_t m = 0; m < M;) {
+        const auto pg = svwhilelt_b8_u64(m, M);
+
+        const auto mm10 = svld1_u8(pg, code0 + m);
+        const auto mm11 = svld1_u8(pg, code1 + m);
+        const auto mm12 = svld1_u8(pg, code2 + m);
+        const auto mm13 = svld1_u8(pg, code3 + m);
+        {
+            const auto mm10lo = svunpklo_u16(mm10);
+            const auto mm11lo = svunpklo_u16(mm11);
+            const auto mm12lo = svunpklo_u16(mm12);
+            const auto mm13lo = svunpklo_u16(mm13);
+            const auto pglo = svunpklo_b(pg);
+
+            {
+                const auto pglolo = svunpklo_b(pglo);
+                {
+                    const auto idx1 = svunpklo_u32(mm10lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm11lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm12lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm13lo);
+                    distance_codes_kernel(
+                            pglolo, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                const auto pglohi = svunpkhi_b(pglo);
+                {
+                    const auto idx1 = svunpkhi_u32(mm10lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm11lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm12lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm13lo);
+                    distance_codes_kernel(
+                            pglohi, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+        }
+
+        {
+            const auto mm10hi = svunpkhi_u16(mm10);
+            const auto mm11hi = svunpkhi_u16(mm11);
+            const auto mm12hi = svunpkhi_u16(mm12);
+            const auto mm13hi = svunpkhi_u16(mm13);
+            const auto pghi = svunpkhi_b(pg);
+
+            {
+                const auto pghilo = svunpklo_b(pghi);
+                {
+                    const auto idx1 = svunpklo_u32(mm10hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm11hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm12hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpklo_u32(mm13hi);
+                    distance_codes_kernel(
+                            pghilo, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+            if (m >= M)
+                break;
+
+            {
+                const auto pghihi = svunpkhi_b(pghi);
+                {
+                    const auto idx1 = svunpkhi_u32(mm10hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum0);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm11hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum1);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm12hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum2);
+                }
+                {
+                    const auto idx1 = svunpkhi_u32(mm13hi);
+                    distance_codes_kernel(
+                            pghihi, idx1, offsets_0, tab, partialSum3);
+                }
+                tab += ksub * quad_lanes;
+            }
+
+            m += quad_lanes;
+        }
+    }
+
+    result0 = svaddv_f32(svptrue_b32(), partialSum0);
+    result1 = svaddv_f32(svptrue_b32(), partialSum1);
+    result2 = svaddv_f32(svptrue_b32(), partialSum2);
+    result3 = svaddv_f32(svptrue_b32(), partialSum3);
+}
+
+} // namespace faiss
+
+#endif

data/vendor/faiss/faiss/impl/code_distance/code_distance.h

@@ -1,5 +1,5 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
@@ -77,6 +77,59 @@ inline void distance_four_codes(
 
 } // namespace faiss
 
+#elif defined(__ARM_FEATURE_SVE)
+
+#include <faiss/impl/code_distance/code_distance-sve.h>
+
+namespace faiss {
+
+template <typename PQDecoderT>
+inline float distance_single_code(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // the code
+        const uint8_t* code) {
+    return distance_single_code_sve<PQDecoderT>(M, nbits, sim_table, code);
+}
+
+template <typename PQDecoderT>
+inline void distance_four_codes(
+        // the product quantizer
+        const size_t M,
+        // number of bits per quantization index
+        const size_t nbits,
+        // precomputed distances, layout (M, ksub)
+        const float* sim_table,
+        // codes
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        // computed distances
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    distance_four_codes_sve<PQDecoderT>(
+            M,
+            nbits,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+} // namespace faiss
+
 #else
 
 #include <faiss/impl/code_distance/code_distance-generic.h>

data/vendor/faiss/faiss/impl/index_read.cpp

@@ -1,12 +1,11 @@
-/**
- * Copyright (c) Facebook, Inc. and its affiliates.
+/*
+ * 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.
  */
 
-// -*- c++ -*-
-
+#include <faiss/impl/index_read_utils.h>
 #include <faiss/index_io.h>
 
 #include <faiss/impl/io_macros.h>
@@ -14,9 +13,6 @@
 #include <cstdio>
 #include <cstdlib>
 
-#include <sys/stat.h>
-#include <sys/types.h>
-
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/io.h>
 #include <faiss/impl/io_macros.h>
@@ -63,7 +59,7 @@ namespace faiss {
  * Read
  **************************************************************/
 
-static void read_index_header(Index* idx, IOReader* f) {
+void read_index_header(Index* idx, IOReader* f) {
     READ1(idx->d);
     READ1(idx->ntotal);
     idx_t dummy;
@@ -232,7 +228,7 @@ InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
     }
 }
 
-static void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
+void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
     InvertedLists* ils = read_InvertedLists(f, io_flags);
     if (ils) {
         FAISS_THROW_IF_NOT(ils->nlist == ivf->nlist);
@@ -244,7 +240,7 @@ static void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags) {
     ivf->own_invlists = true;
 }
 
-static void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
+void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
     READ1(pq->d);
     READ1(pq->M);
     READ1(pq->nbits);
@@ -354,7 +350,7 @@ static void read_ProductLocalSearchQuantizer(
     }
 }
 
-static void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f) {
+void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f) {
     READ1(ivsc->qtype);
     READ1(ivsc->rangestat);
     READ1(ivsc->rangestat_arg);
@@ -375,7 +371,10 @@ static void read_HNSW(HNSW* hnsw, IOReader* f) {
     READ1(hnsw->max_level);
     READ1(hnsw->efConstruction);
     READ1(hnsw->efSearch);
-    READ1(hnsw->upper_beam);
+
+    // // deprecated field
+    // READ1(hnsw->upper_beam);
+    READ1_DUMMY(int)
 }
 
 static void read_NSG(NSG* nsg, IOReader* f) {
@@ -440,7 +439,7 @@ ProductQuantizer* read_ProductQuantizer(IOReader* reader) {
     return pq;
 }
 
-static void read_direct_map(DirectMap* dm, IOReader* f) {
+void read_direct_map(DirectMap* dm, IOReader* f) {
     char maintain_direct_map;
     READ1(maintain_direct_map);
     dm->type = (DirectMap::Type)maintain_direct_map;
@@ -456,10 +455,10 @@ static void read_direct_map(DirectMap* dm, IOReader* f) {
     }
 }
 
-static void read_ivf_header(
+void read_ivf_header(
         IndexIVF* ivf,
         IOReader* f,
-        std::vector<std::vector<idx_t>>* ids = nullptr) {
+        std::vector<std::vector<idx_t>>* ids) {
     read_index_header(ivf, f);
     READ1(ivf->nlist);
     READ1(ivf->nprobe);
@@ -474,7 +473,7 @@ static void read_ivf_header(
 }
 
 // used for legacy formats
-static ArrayInvertedLists* set_array_invlist(
+ArrayInvertedLists* set_array_invlist(
         IndexIVF* ivf,
         std::vector<std::vector<idx_t>>& ids) {
     ArrayInvertedLists* ail =
@@ -531,7 +530,11 @@ Index* read_index(IOReader* f, int io_flags) {
     Index* idx = nullptr;
     uint32_t h;
     READ1(h);
-    if (h == fourcc("IxFI") || h == fourcc("IxF2") || h == fourcc("IxFl")) {
+    if (h == fourcc("null")) {
+        // denotes a missing index, useful for some cases
+        return nullptr;
+    } else if (
+            h == fourcc("IxFI") || h == fourcc("IxF2") || h == fourcc("IxFl")) {
         IndexFlat* idxf;
         if (h == fourcc("IxFI")) {
             idxf = new IndexFlatIP();
@@ -948,7 +951,7 @@ Index* read_index(IOReader* f, int io_flags) {
         idx = idxp;
     } else if (
             h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") ||
-            h == fourcc("IHN2")) {
+            h == fourcc("IHN2") || h == fourcc("IHNc")) {
         IndexHNSW* idxhnsw = nullptr;
         if (h == fourcc("IHNf"))
             idxhnsw = new IndexHNSWFlat();
@@ -958,11 +961,19 @@ Index* read_index(IOReader* f, int io_flags) {
             idxhnsw = new IndexHNSWSQ();
         if (h == fourcc("IHN2"))
             idxhnsw = new IndexHNSW2Level();
+        if (h == fourcc("IHNc"))
+            idxhnsw = new IndexHNSWCagra();
         read_index_header(idxhnsw, f);
+        if (h == fourcc("IHNc")) {
+            READ1(idxhnsw->keep_max_size_level0);
+            auto idx_hnsw_cagra = dynamic_cast<IndexHNSWCagra*>(idxhnsw);
+            READ1(idx_hnsw_cagra->base_level_only);
+            READ1(idx_hnsw_cagra->num_base_level_search_entrypoints);
+        }
         read_HNSW(&idxhnsw->hnsw, f);
         idxhnsw->storage = read_index(f, io_flags);
-        idxhnsw->own_fields = true;
-        if (h == fourcc("IHNp")) {
+        idxhnsw->own_fields = idxhnsw->storage != nullptr;
+        if (h == fourcc("IHNp") && !(io_flags & IO_FLAG_PQ_SKIP_SDC_TABLE)) {
             dynamic_cast<IndexPQ*>(idxhnsw->storage)->pq.compute_sdc_table();
         }
         idx = idxhnsw;

data/vendor/faiss/faiss/impl/index_read_utils.h

@@ -0,0 +1,37 @@
+/*
+ * 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.
+ */
+
+// Utils for index_read
+
+#ifndef FAISS_INDEX_READ_UTILS_H
+#define FAISS_INDEX_READ_UTILS_H
+
+#include <faiss/IndexIVF.h>
+#include <faiss/impl/io.h>
+
+#pragma once
+
+namespace faiss {
+struct ProductQuantizer;
+struct ScalarQuantizer;
+
+void read_index_header(Index* idx, IOReader* f);
+void read_direct_map(DirectMap* dm, IOReader* f);
+void read_ivf_header(
+        IndexIVF* ivf,
+        IOReader* f,
+        std::vector<std::vector<idx_t>>* ids = nullptr);
+void read_InvertedLists(IndexIVF* ivf, IOReader* f, int io_flags);
+ArrayInvertedLists* set_array_invlist(
+        IndexIVF* ivf,
+        std::vector<std::vector<idx_t>>& ids);
+void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f);
+void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f);
+
+} // namespace faiss
+
+#endif