faiss 0.5.2 → 0.6.0

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

@@ -19,7 +19,6 @@ 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,

data/vendor/faiss/faiss/impl/index_write.cpp

@@ -12,10 +12,14 @@
 
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 
 #include <faiss/invlists/InvertedListsIOHook.h>
 
+#include <faiss/invlists/BlockInvertedLists.h>
+
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/RaBitQUtils.h>
 #include <faiss/utils/hamming.h>
 
 #include <faiss/Index2Layer.h>
@@ -101,7 +105,14 @@ static void write_index_header(const Index* idx, IOWriter* f) {
 }
 
 void write_VectorTransform(const VectorTransform* vt, IOWriter* f) {
-    if (const LinearTransform* lt = dynamic_cast<const LinearTransform*>(vt)) {
+    if (const HadamardRotation* hr =
+                dynamic_cast<const HadamardRotation*>(vt)) {
+        uint32_t h = fourcc("HRot");
+        WRITE1(h);
+        WRITE1(hr->seed);
+    } else if (
+            const LinearTransform* lt =
+                    dynamic_cast<const LinearTransform*>(vt)) {
         if (dynamic_cast<const RandomRotationMatrix*>(lt)) {
             uint32_t h = fourcc("rrot");
             WRITE1(h);
@@ -446,9 +457,9 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         uint32_t h = fourcc("null");
         WRITE1(h);
     } else if (
-            const IndexFlatL2Panorama* idxpan =
-                    dynamic_cast<const IndexFlatL2Panorama*>(idx)) {
-        uint32_t h = fourcc("IxFP");
+            const IndexFlatPanorama* idxpan =
+                    dynamic_cast<const IndexFlatPanorama*>(idx)) {
+        uint32_t h = fourcc(idxpan->metric_type == METRIC_L2 ? "IxFP" : "IxFp");
         WRITE1(h);
         WRITE1(idxpan->d);
         WRITE1(idxpan->n_levels);
@@ -937,13 +948,12 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
     } else if (
             const IndexRaBitQFastScan* idxqfs =
                     dynamic_cast<const IndexRaBitQFastScan*>(idx)) {
-        uint32_t h = fourcc("Irfs");
+        uint32_t h = fourcc("Irfn");
         WRITE1(h);
         write_index_header(idx, f);
         write_RaBitQuantizer(&idxqfs->rabitq, f);
         WRITEVECTOR(idxqfs->center);
         WRITE1(idxqfs->qb);
-        WRITEVECTOR(idxqfs->flat_storage);
         WRITE1(idxqfs->bbs);
         WRITE1(idxqfs->ntotal2);
         WRITE1(idxqfs->M2);
@@ -1060,7 +1070,7 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
     else if (
             const IndexIVFRaBitQFastScan* ivrqfs =
                     dynamic_cast<const IndexIVFRaBitQFastScan*>(idx)) {
-        uint32_t h = fourcc("Iwrf");
+        uint32_t h = fourcc("Iwrn");
         WRITE1(h);
         write_ivf_header(ivrqfs, f);
         write_RaBitQuantizer(&ivrqfs->rabitq, f);
@@ -1072,7 +1082,6 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(ivrqfs->implem);
         WRITE1(ivrqfs->qb);
         WRITE1(ivrqfs->centered);
-        WRITEVECTOR(ivrqfs->flat_storage);
         write_InvertedLists(ivrqfs->invlists, f);
     } else {
         FAISS_THROW_MSG("don't know how to serialize this type of index");

data/vendor/faiss/faiss/impl/lattice_Zn.cpp

@@ -18,6 +18,9 @@
 #include <queue>
 #include <unordered_set>
 
+#include <faiss/impl/FaissAssert.h>
+
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
 namespace faiss {
@@ -302,18 +305,20 @@ void EnumeratedVectors::find_nn(
     }
 
     std::vector<float> c(dim);
-    for (size_t i = 0; i < nc; i++) {
-        uint64_t code = codes[nc];
-        decode(code, c.data());
-        for (size_t j = 0; j < nq; j++) {
-            const float* x = xq + j * dim;
-            float dis = fvec_inner_product(x, c.data(), dim);
-            if (dis > distances[j]) {
-                distances[j] = dis;
-                labels[j] = i;
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t i = 0; i < nc; i++) {
+            uint64_t code = codes[nc];
+            decode(code, c.data());
+            for (size_t j = 0; j < nq; j++) {
+                const float* x = xq + j * dim;
+                float dis = fvec_inner_product<SL>(x, c.data(), dim);
+                if (dis > distances[j]) {
+                    distances[j] = dis;
+                    labels[j] = i;
+                }
             }
         }
-    }
+    });
 }
 
 /**********************************************************
@@ -321,6 +326,12 @@ void EnumeratedVectors::find_nn(
  **********************************************************/
 
 ZnSphereSearch::ZnSphereSearch(int dim, int r2) : dimS(dim), r2(r2) {
+    FAISS_THROW_IF_NOT_MSG(
+            dim > 0 && dim <= 64, "ZnSphereSearch: dim must be in [1, 64]");
+    FAISS_THROW_IF_NOT_MSG(
+            r2 >= 0 && r2 <= 512,
+            "ZnSphereSearch: r2 must be in [0, 512] to avoid"
+            " excessive computation in sum_of_sq");
     voc = sum_of_sq(r2, int(ceil(sqrt(r2)) + 1), dim);
     natom = voc.size() / dim;
 }
@@ -355,13 +366,15 @@ float ZnSphereSearch::search(
     // find best
     int ibest = -1;
     float dpbest = -100;
-    for (int i = 0; i < natom; i++) {
-        float dp = fvec_inner_product(voc.data() + i * dim, xperm, dim);
-        if (dp > dpbest) {
-            dpbest = dp;
-            ibest = i;
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (int i = 0; i < natom; i++) {
+            float dp = fvec_inner_product<SL>(voc.data() + i * dim, xperm, dim);
+            if (dp > dpbest) {
+                dpbest = dp;
+                ibest = i;
+            }
         }
-    }
+    });
     // revert sort
     const float* cin = voc.data() + ibest * dim;
     for (int i = 0; i < dim; i++) {
@@ -486,14 +499,28 @@ void ZnSphereCodecRec::set_nv_cum(int ld, int r2t, int r2a, uint64_t cum) {
 
 ZnSphereCodecRec::ZnSphereCodecRec(int dim, int r2)
         : EnumeratedVectors(dim), r2(r2) {
+    FAISS_THROW_IF_NOT_MSG(
+            dim > 0 && r2 >= 0, "invalid ZnSphereCodecRec parameters");
     log2_dim = 0;
     while (dim > (1 << log2_dim)) {
         log2_dim++;
     }
-    assert(dim == (1 << log2_dim) || !"dimension must be a power of 2");
-
-    all_nv.resize((log2_dim + 1) * (r2 + 1));
-    all_nv_cum.resize((log2_dim + 1) * (r2 + 1) * (r2 + 1));
+    assert(dim == (1 << log2_dim) && "dimension must be a power of 2");
+
+    // Validate allocation sizes to avoid null pointer dereference on
+    // allocation failure. The cumulative table has O(r2^2) entries.
+    size_t nv_size = (size_t)(log2_dim + 1) * (r2 + 1);
+    size_t nv_cum_size = nv_size * (r2 + 1);
+    FAISS_THROW_IF_NOT_MSG(
+            nv_cum_size / (r2 + 1) == nv_size,
+            "ZnSphereCodecRec: allocation size overflow");
+    // Cap at ~1GB worth of uint64_t entries
+    FAISS_THROW_IF_NOT_MSG(
+            nv_cum_size <= (size_t(1) << 27),
+            "ZnSphereCodecRec: r2 too large, would require excessive memory");
+
+    all_nv.resize(nv_size);
+    all_nv_cum.resize(nv_cum_size);
 
     for (int r2a = 0; r2a <= r2; r2a++) {
         int r = int(sqrt(r2a));

data/vendor/faiss/faiss/impl/mapped_io.cpp

@@ -34,8 +34,15 @@ struct MmappedFileMappingOwner::PImpl {
     size_t ptr_size = 0;
 
     explicit PImpl(const std::string& filename) {
-        auto f = std::unique_ptr<FILE, decltype(&fclose)>(
-                fopen(filename.c_str(), "r"), &fclose);
+        struct FileDeleter {
+            void operator()(FILE* f) const {
+                if (f)
+                    fclose(f);
+            }
+        };
+
+        auto f = std::unique_ptr<FILE, FileDeleter>(
+                fopen(filename.c_str(), "r"), FileDeleter{});
         FAISS_THROW_IF_NOT_FMT(
                 f.get(),
                 "could not open %s for reading: %s",

data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp

@@ -111,7 +111,8 @@ void pq4_pack_codes_range(
         size_t bbs,
         size_t nsq,
         uint8_t* blocks,
-        size_t code_stride) {
+        size_t code_stride,
+        size_t block_stride) {
     // Determine stride: use custom if provided, otherwise use legacy
     // calculation
     size_t actual_stride = (code_stride == 0) ? (M + 1) / 2 : code_stride;
@@ -136,7 +137,7 @@ void pq4_pack_codes_range(
     size_t block1 = ((i1 - 1) / bbs) + 1;
 
     for (size_t b = block0; b < block1; b++) {
-        uint8_t* codes2 = blocks + b * bbs * nsq / 2;
+        uint8_t* codes2 = blocks + b * block_stride;
         int64_t i_base = b * bbs - i0;
         for (int sq = 0; sq < nsq; sq += 2) {
             for (size_t i = 0; i < bbs; i += 32) {
@@ -272,6 +273,10 @@ void CodePackerPQ4::unpack_1(
     }
 }
 
+CodePacker* CodePackerPQ4::clone() const {
+    return new CodePackerPQ4(*this);
+}
+
 /***************************************************************
  * Packing functions for Look-Up Tables (LUT)
  ***************************************************************/

data/vendor/faiss/faiss/impl/pq4_fast_scan.h

@@ -59,6 +59,7 @@ void pq4_pack_codes(
  * @param blocks  output array, size at least ceil(i1 / bbs) * bbs * nsq / 2
  * @param code_stride  optional stride between consecutive codes (0 = use
  * default (M + 1) / 2)
+ * @param block_stride  stride in bytes between consecutive blocks.
  */
 void pq4_pack_codes_range(
         const uint8_t* codes,
@@ -68,7 +69,8 @@ void pq4_pack_codes_range(
         size_t bbs,
         size_t nsq,
         uint8_t* blocks,
-        size_t code_stride = 0);
+        size_t code_stride,
+        size_t block_stride);
 
 /** get a single element from a packed codes table
  *
@@ -101,6 +103,8 @@ struct CodePackerPQ4 : CodePacker {
 
     CodePackerPQ4(size_t nsq, size_t bbs);
 
+    CodePacker* clone() const final;
+
     void pack_1(const uint8_t* flat_code, size_t offset, uint8_t* block)
             const final;
     void unpack_1(const uint8_t* block, size_t offset, uint8_t* flat_code)
@@ -125,6 +129,7 @@ void pq4_pack_LUT(int nq, int nsq, const uint8_t* src, uint8_t* dest);
  * @param codes   packed codes array
  * @param LUT     packed look-up table
  * @param scaler  scaler to scale the encoded norm
+ * @param block_stride  stride in bytes between consecutive blocks.
  */
 void pq4_accumulate_loop(
         int nq,
@@ -134,7 +139,8 @@ void pq4_accumulate_loop(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler);
+        const NormTableScaler* scaler,
+        size_t block_stride);
 
 /* qbs versions, supported only for bbs=32.
  *
@@ -185,6 +191,7 @@ int pq4_pack_LUT_qbs_q_map(
  * @param LUT     look-up table (packed)
  * @param res     call-back for the results
  * @param scaler  scaler to scale the encoded norm
+ * @param block_stride  stride in bytes between consecutive blocks.
  */
 void pq4_accumulate_loop_qbs(
         int qbs,
@@ -193,7 +200,8 @@ void pq4_accumulate_loop_qbs(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler = nullptr);
+        const NormTableScaler* scaler,
+        size_t block_stride);
 
 /** Wrapper of pq4_accumulate_loop_qbs using simple StoreResultHandler
  *  and DummyScaler

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp

@@ -123,14 +123,15 @@ void accumulate_fixed_blocks(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride) {
     constexpr int bbs = 32 * BB;
     for (size_t j0 = 0; j0 < nb; j0 += bbs) {
         FixedStorageHandler<NQ, 2 * BB> res2;
         kernel_accumulate_block<NQ, BB>(nsq, codes, LUT, res2, scaler);
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
-        codes += bbs * nsq / 2;
+        codes += block_stride;
     }
 }
 
@@ -143,15 +144,17 @@ void pq4_accumulate_loop_fixed_scaler(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride) {
     FAISS_THROW_IF_NOT(is_aligned_pointer(codes));
     FAISS_THROW_IF_NOT(is_aligned_pointer(LUT));
     FAISS_THROW_IF_NOT(bbs % 32 == 0);
     FAISS_THROW_IF_NOT(nb % bbs == 0);
 
-#define DISPATCH(NQ, BB)                                                   \
-    case NQ * 1000 + BB:                                                   \
-        accumulate_fixed_blocks<NQ, BB>(nb, nsq, codes, LUT, res, scaler); \
+#define DISPATCH(NQ, BB)                                         \
+    case NQ * 1000 + BB:                                         \
+        accumulate_fixed_blocks<NQ, BB>(                         \
+                nb, nsq, codes, LUT, res, scaler, block_stride); \
         break
 
     switch (nq * 1000 + bbs / 32) {
@@ -179,14 +182,15 @@ void pq4_accumulate_loop_fixed_handler(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const NormTableScaler* scaler) {
+        const NormTableScaler* scaler,
+        size_t block_stride) {
     if (scaler) {
         pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, *scaler);
+                nq, nb, bbs, nsq, codes, LUT, res, *scaler, block_stride);
     } else {
         DummyScaler dscaler;
         pq4_accumulate_loop_fixed_scaler(
-                nq, nb, bbs, nsq, codes, LUT, res, dscaler);
+                nq, nb, bbs, nsq, codes, LUT, res, dscaler, block_stride);
     }
 }
 
@@ -199,9 +203,10 @@ struct Run_pq4_accumulate_loop {
            int nsq,
            const uint8_t* codes,
            const uint8_t* LUT,
-           const NormTableScaler* scaler) {
+           const NormTableScaler* scaler,
+           size_t block_stride) {
         pq4_accumulate_loop_fixed_handler(
-                nq, nb, bbs, nsq, codes, LUT, res, scaler);
+                nq, nb, bbs, nsq, codes, LUT, res, scaler, block_stride);
     }
 };
 
@@ -215,10 +220,11 @@ void pq4_accumulate_loop(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler) {
+        const NormTableScaler* scaler,
+        size_t block_stride) {
     Run_pq4_accumulate_loop consumer;
     dispatch_SIMDResultHandler(
-            res, consumer, nq, nb, bbs, nsq, codes, LUT, scaler);
+            res, consumer, nq, nb, bbs, nsq, codes, LUT, scaler, block_stride);
 }
 
 } // namespace faiss

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp

@@ -565,7 +565,8 @@ void accumulate_q_4step(
         const uint8_t* codes,
         const uint8_t* LUT0,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride) {
     constexpr int Q1 = QBS & 15;
     constexpr int Q2 = (QBS >> 4) & 15;
     constexpr int Q3 = (QBS >> 8) & 15;
@@ -593,7 +594,7 @@ void accumulate_q_4step(
         }
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
-        codes += 32 * nsq / 2;
+        codes += block_stride;
     }
 }
 
@@ -604,11 +605,13 @@ void kernel_accumulate_block_loop(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride) {
     for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
         res.set_block_origin(0, j0);
         kernel_accumulate_block<NQ, ResultHandler>(
-                nsq, codes + j0 * nsq / 2, LUT, res, scaler);
+                nsq, codes, LUT, res, scaler);
+        codes += block_stride;
     }
 }
 
@@ -621,14 +624,15 @@ void accumulate(
         const uint8_t* codes,
         const uint8_t* LUT,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(LUT));
 
-#define DISPATCH(NQ)                                     \
-    case NQ:                                             \
-        kernel_accumulate_block_loop<NQ, ResultHandler>( \
-                ntotal2, nsq, codes, LUT, res, scaler);  \
+#define DISPATCH(NQ)                                                  \
+    case NQ:                                                          \
+        kernel_accumulate_block_loop<NQ, ResultHandler>(              \
+                ntotal2, nsq, codes, LUT, res, scaler, block_stride); \
         return
 
     switch (nq) {
@@ -650,16 +654,18 @@ void pq4_accumulate_loop_qbs_fixed_scaler(
         const uint8_t* codes,
         const uint8_t* LUT0,
         ResultHandler& res,
-        const Scaler& scaler) {
+        const Scaler& scaler,
+        size_t block_stride = 0) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(codes));
     assert(is_aligned_pointer(LUT0));
 
     // try out optimized versions
     switch (qbs) {
-#define DISPATCH(QBS)                                                    \
-    case QBS:                                                            \
-        accumulate_q_4step<QBS>(ntotal2, nsq, codes, LUT0, res, scaler); \
+#define DISPATCH(QBS)                                                  \
+    case QBS:                                                          \
+        accumulate_q_4step<QBS>(                                       \
+                ntotal2, nsq, codes, LUT0, res, scaler, block_stride); \
         return;
         DISPATCH(0x3333); // 12
         DISPATCH(0x2333); // 11
@@ -688,7 +694,6 @@ void pq4_accumulate_loop_qbs_fixed_scaler(
     }
 
     // default implementation where qbs is not known at compile time
-
     for (size_t j0 = 0; j0 < ntotal2; j0 += 32) {
         const uint8_t* LUT = LUT0;
         int qi = qbs;
@@ -714,7 +719,7 @@ void pq4_accumulate_loop_qbs_fixed_scaler(
             i0 += nq;
             LUT += nq * nsq * 16;
         }
-        codes += 32 * nsq / 2;
+        codes += block_stride;
     }
 }
 
@@ -726,14 +731,15 @@ struct Run_pq4_accumulate_loop_qbs {
            int nsq,
            const uint8_t* codes,
            const uint8_t* LUT,
-           const NormTableScaler* scaler) {
+           const NormTableScaler* scaler,
+           size_t block_stride) {
         if (scaler) {
             pq4_accumulate_loop_qbs_fixed_scaler(
-                    qbs, nb, nsq, codes, LUT, res, *scaler);
+                    qbs, nb, nsq, codes, LUT, res, *scaler, block_stride);
         } else {
             DummyScaler dummy;
             pq4_accumulate_loop_qbs_fixed_scaler(
-                    qbs, nb, nsq, codes, LUT, res, dummy);
+                    qbs, nb, nsq, codes, LUT, res, dummy, block_stride);
         }
     }
 };
@@ -747,9 +753,11 @@ void pq4_accumulate_loop_qbs(
         const uint8_t* codes,
         const uint8_t* LUT,
         SIMDResultHandler& res,
-        const NormTableScaler* scaler) {
+        const NormTableScaler* scaler,
+        size_t block_stride) {
     Run_pq4_accumulate_loop_qbs consumer;
-    dispatch_SIMDResultHandler(res, consumer, qbs, nb, nsq, codes, LUT, scaler);
+    dispatch_SIMDResultHandler(
+            res, consumer, qbs, nb, nsq, codes, LUT, scaler, block_stride);
 }
 
 /***************************************************************
@@ -777,7 +785,7 @@ void accumulate_to_mem(
     FAISS_THROW_IF_NOT(ntotal2 % 32 == 0);
     StoreResultHandler handler(accu, ntotal2);
     DummyScaler scaler;
-    accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler);
+    accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler, 32 * nsq / 2);
 }
 
 int pq4_preferred_qbs(int n) {