faiss 0.6.4 → 0.6.5

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

@@ -45,6 +45,7 @@
 #include <faiss/IndexIVFPQR.h>
 #include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/IndexIVFRaBitQFastScan.h>
+#include <faiss/IndexIVFSQFastScan.h>
 #include <faiss/IndexIVFSpectralHash.h>
 #include <faiss/IndexLSH.h>
 #include <faiss/IndexLattice.h>
@@ -52,11 +53,13 @@
 #include <faiss/IndexNSG.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexPiPNN.h>
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
 #include <faiss/IndexRefine.h>
 #include <faiss/IndexRowwiseMinMax.h>
+#include <faiss/IndexSQFastScan.h>
 #ifdef FAISS_ENABLE_SVS
 #include <faiss/impl/svs_io.h>
 #include <faiss/svs/IndexSVSFlat.h>
@@ -765,6 +768,10 @@ void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
             pq->d * pq->ksub);
 }
 
+static void validate_codebooks_size(
+        const AdditiveQuantizer& aq,
+        size_t effective_d);
+
 static void read_ResidualQuantizer_old(ResidualQuantizer& rq, IOReader* f) {
     READ1(rq.d);
     FAISS_THROW_IF_NOT_FMT(
@@ -786,6 +793,7 @@ static void read_ResidualQuantizer_old(ResidualQuantizer& rq, IOReader* f) {
     READ1(rq.norm_min);
     READ1(rq.norm_max);
     rq.set_derived_values();
+    validate_codebooks_size(rq, rq.d);
 }
 
 static void read_AdditiveQuantizer(AdditiveQuantizer& aq, IOReader* f) {
@@ -847,6 +855,9 @@ static void read_AdditiveQuantizer(AdditiveQuantizer& aq, IOReader* f) {
 static void validate_codebooks_size(
         const AdditiveQuantizer& aq,
         size_t effective_d) {
+    if (!aq.is_trained && aq.codebooks.empty()) {
+        return;
+    }
     size_t required = mul_no_overflow(
             effective_d, aq.total_codebook_size, "codebooks validation");
     FAISS_THROW_IF_NOT_FMT(
@@ -1359,6 +1370,59 @@ static void read_NSG(NSG& nsg, IOReader* f) {
     FAISS_CHECK_RANGE(nsg.enterpoint, 0, N);
 }
 
+static void read_PiPNN(PiPNN& pipnn, idx_t index_ntotal, IOReader* f) {
+    READ1(pipnn.R);
+    READ1(pipnn.leaf_k);
+    READ1(pipnn.hash_bits);
+    READ1(pipnn.reservoir_size);
+    READ1(pipnn.c_max);
+    READ1(pipnn.c_min);
+    READ1(pipnn.leader_cap);
+    // Former replicas field: build metadata that search does not need.
+    int replicas;
+    READ1(replicas);
+    READ1(pipnn.leader_fraction);
+    READ1(pipnn.alpha);
+    READVECTOR(pipnn.fanout);
+    READ1(pipnn.seed);
+    READ1_BOOL(pipnn.ip_partition_by_angle);
+    READ1(pipnn.search_L);
+    pipnn.check_params();
+    READ1(pipnn.ntotal);
+    READ1(pipnn.entry_point);
+    // Checked before the graph is allocated; ntotal == 0 means unbuilt.
+    if (pipnn.ntotal != 0) {
+        FAISS_THROW_IF_NOT_FMT(
+                pipnn.ntotal == index_ntotal,
+                "PiPNN ntotal %" PRId64 " != index ntotal %" PRId64,
+                pipnn.ntotal,
+                index_ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                pipnn.ntotal <= INT32_MAX,
+                "built PiPNN ntotal %" PRId64 " exceeds INT32_MAX",
+                pipnn.ntotal);
+        FAISS_CHECK_RANGE(pipnn.entry_point, 0, pipnn.ntotal);
+    }
+    const size_t graph_size = mul_no_overflow(
+            (size_t)pipnn.ntotal, (size_t)pipnn.R, "PiPNN graph");
+    // Views (IO_FLAG_MMAP_IFC, zero-copy) skip the scan: it would fault in
+    // the whole graph. Search bounds-checks every neighbour id instead.
+    if (!read_vector_base(pipnn.graph, f, graph_size, std::nullopt)) {
+        FAISS_THROW_IF_NOT(
+                graph_size <
+                get_deserialization_vector_byte_limit() / sizeof(int32_t));
+        pipnn.graph.resize(graph_size);
+        READANDCHECK(pipnn.graph.data(), graph_size);
+        for (int32_t id : pipnn.graph) {
+            FAISS_CHECK_RANGE(id, -1, pipnn.ntotal);
+        }
+    }
+    read_vector(pipnn.next_copy, f);
+    FAISS_THROW_IF_NOT(
+            pipnn.next_copy.size() == 0 ||
+            pipnn.next_copy.size() == size_t(pipnn.ntotal));
+}
+
 static void read_NNDescent(NNDescent& nnd, IOReader* f) {
     READ1(nnd.ntotal);
     READ1(nnd.d);
@@ -1841,6 +1905,14 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         } else {
             read_ResidualQuantizer(idxr->rq, f, io_flags);
         }
+        FAISS_THROW_IF_NOT_MSG(
+                idxr->rq.is_trained == idxr->is_trained,
+                "IndexResidualQuantizer is_trained does not match "
+                "ResidualQuantizer is_trained");
+        FAISS_THROW_IF_NOT_MSG(
+                idxr->ntotal == 0 || idxr->rq.is_trained,
+                "IndexResidualQuantizer with ntotal > 0 must have trained "
+                "ResidualQuantizer");
         validate_aq_dimension_match(
                 idxr->rq, idxr->d, "IndexResidualQuantizer");
         READ1(idxr->code_size);
@@ -2760,6 +2832,66 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
                     idxnnd->d);
         }
         idx = std::move(idxnnd);
+    } else if (h == fourcc("IPNf") || h == fourcc("IPNs")) {
+        std::unique_ptr<IndexPiPNN> idxpipnn;
+        if (h == fourcc("IPNf")) {
+            idxpipnn = std::make_unique<IndexPiPNNFlat>();
+        } else {
+            idxpipnn = std::make_unique<IndexPiPNNSQ>();
+        }
+        read_index_header(*idxpipnn, f);
+        READ1_BOOL(idxpipnn->defer_build);
+        read_PiPNN(idxpipnn->pipnn, idxpipnn->ntotal, f);
+        idxpipnn->storage = read_index(f, io_flags);
+        idxpipnn->own_fields = true;
+        FAISS_THROW_IF_NOT_MSG(
+                idxpipnn->storage, "IndexPiPNN requires non-null storage");
+        IndexPiPNN::check_storage_supported(*idxpipnn->storage);
+        FAISS_THROW_IF_NOT_FMT(
+                idxpipnn->storage->ntotal == idxpipnn->ntotal,
+                "PiPNN storage ntotal %" PRId64 " != index ntotal %" PRId64,
+                idxpipnn->storage->ntotal,
+                idxpipnn->ntotal);
+        FAISS_THROW_IF_NOT_FMT(
+                idxpipnn->storage->d == idxpipnn->d,
+                "PiPNN storage d %d != index d %d",
+                idxpipnn->storage->d,
+                idxpipnn->d);
+        FAISS_THROW_IF_NOT_FMT(
+                idxpipnn->storage->metric_type == idxpipnn->metric_type,
+                "PiPNN storage metric %d != index metric %d",
+                (int)idxpipnn->storage->metric_type,
+                (int)idxpipnn->metric_type);
+        idx = std::move(idxpipnn);
+    } else if (h == fourcc("ISfs")) {
+        auto idxsqfs = std::make_unique<IndexSQFastScan>();
+        read_index_header(*idxsqfs, f);
+        read_ScalarQuantizer(&idxsqfs->sq, f, *idxsqfs);
+        READ1(idxsqfs->implem);
+        READ1(idxsqfs->bbs);
+        READ1(idxsqfs->qbs);
+        FAISS_THROW_IF_NOT_MSG(idxsqfs->qbs >= 0, "qbs must be non-negative");
+        READ1(idxsqfs->ntotal2);
+        READ1(idxsqfs->M2);
+        READVECTOR(idxsqfs->codes);
+        READ1(idxsqfs->rerank_factor);
+        READVECTOR(idxsqfs->lo_codes);
+
+        // Restore FastScan base-class fields from the SQ
+        idxsqfs->M = idxsqfs->sq.d;
+        idxsqfs->nbits = 4;
+        idxsqfs->ksub = 16;
+        idxsqfs->code_size = idxsqfs->M2 / 2;
+
+        validate_fastscan_fields(
+                idxsqfs->M,
+                idxsqfs->M2,
+                idxsqfs->ksub,
+                idxsqfs->bbs,
+                "IndexSQFastScan");
+
+        idx = std::move(idxsqfs);
+
     } else if (h == fourcc("IPfs")) {
         auto idxpqfs = std::make_unique<IndexPQFastScan>();
         read_index_header(*idxpqfs, f);
@@ -2787,6 +2919,46 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
 
         idx = std::move(idxpqfs);
 
+    } else if (h == fourcc("IwSf")) {
+        auto ivfsqfs = std::make_unique<IndexIVFSQFastScan>();
+        read_ivf_header(ivfsqfs.get(), f);
+        READ1_BOOL(ivfsqfs->by_residual);
+        READ1(ivfsqfs->code_size);
+        READ1(ivfsqfs->bbs);
+        READ1(ivfsqfs->M2);
+        READ1(ivfsqfs->implem);
+        READ1(ivfsqfs->rerank_factor);
+        read_ScalarQuantizer(&ivfsqfs->sq, f, *ivfsqfs);
+        read_InvertedLists(*ivfsqfs, f, io_flags);
+
+        ivfsqfs->M = ivfsqfs->d;
+        ivfsqfs->nbits = 4;
+        ivfsqfs->ksub = 16;
+        ivfsqfs->init_code_packer();
+
+        bool has_orig;
+        READ1(has_orig);
+        if (has_orig) {
+            ivfsqfs->lo_codes_invlists = read_InvertedLists(f, io_flags);
+            // Rebuild direct_map from lo_codes_invlists for reranking
+            ivfsqfs->direct_map.set_type(
+                    DirectMap::Hashtable, ivfsqfs->invlists, 0);
+            for (size_t list_no = 0; list_no < ivfsqfs->nlist; list_no++) {
+                size_t list_size =
+                        ivfsqfs->lo_codes_invlists->list_size(list_no);
+                if (list_size == 0) {
+                    continue;
+                }
+                InvertedLists::ScopedIds ids(
+                        ivfsqfs->lo_codes_invlists, list_no);
+                for (size_t j = 0; j < list_size; j++) {
+                    ivfsqfs->direct_map.add_single_id(ids[j], list_no, j);
+                }
+            }
+        }
+
+        idx = std::move(ivfsqfs);
+
     } else if (h == fourcc("IwPf")) {
         auto ivpq = std::make_unique<IndexIVFPQFastScan>();
         read_ivf_header(ivpq.get(), f);

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

@@ -40,6 +40,7 @@
 #include <faiss/IndexIVFPQR.h>
 #include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/IndexIVFRaBitQFastScan.h>
+#include <faiss/IndexIVFSQFastScan.h>
 #include <faiss/IndexIVFSpectralHash.h>
 #include <faiss/IndexLSH.h>
 #include <faiss/IndexLattice.h>
@@ -47,11 +48,13 @@
 #include <faiss/IndexNSG.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexPQFastScan.h>
+#include <faiss/IndexPiPNN.h>
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
 #include <faiss/IndexRefine.h>
 #include <faiss/IndexRowwiseMinMax.h>
+#include <faiss/IndexSQFastScan.h>
 #ifdef FAISS_ENABLE_SVS
 #include <faiss/impl/svs_io.h>
 #include <faiss/svs/IndexSVSFlat.h>
@@ -421,6 +424,29 @@ static void write_NNDescent(const NNDescent* nnd, IOWriter* f) {
     WRITEVECTOR(nnd->final_graph);
 }
 
+static void write_PiPNN(const PiPNN& pipnn, IOWriter* f) {
+    WRITE1(pipnn.R);
+    WRITE1(pipnn.leaf_k);
+    WRITE1(pipnn.hash_bits);
+    WRITE1(pipnn.reservoir_size);
+    WRITE1(pipnn.c_max);
+    WRITE1(pipnn.c_min);
+    WRITE1(pipnn.leader_cap);
+    // Former replicas field, kept so that the format does not change.
+    const int replicas = 1;
+    WRITE1(replicas);
+    WRITE1(pipnn.leader_fraction);
+    WRITE1(pipnn.alpha);
+    WRITEVECTOR(pipnn.fanout);
+    WRITE1(pipnn.seed);
+    WRITE1(pipnn.ip_partition_by_angle);
+    WRITE1(pipnn.search_L);
+    WRITE1(pipnn.ntotal);
+    WRITE1(pipnn.entry_point);
+    WRITEANDCHECK(pipnn.graph.data(), pipnn.graph.size());
+    WRITEVECTOR(pipnn.next_copy);
+}
+
 // Write RaBitQuantizer for 1-bit format (backward compatible)
 static void write_RaBitQuantizer(
         const RaBitQuantizer* rabitq,
@@ -938,6 +964,40 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(idxnsg->nndescent_iter);
         write_NSG(&idxnsg->nsg, f);
         write_index(idxnsg->storage, f);
+    } else if (
+            const IndexPiPNN* idxpipnn = dynamic_cast<const IndexPiPNN*>(idx)) {
+        FAISS_THROW_IF_NOT_FMT(
+                typeid(*idx) == typeid(IndexPiPNN) ||
+                        typeid(*idx) == typeid(IndexPiPNNFlat) ||
+                        typeid(*idx) == typeid(IndexPiPNNSQ),
+                "don't know how to serialize this IndexPiPNN subtype: %s",
+                typeid(*idx).name());
+        const Index* storage = idxpipnn->storage;
+        const PiPNN& state = idxpipnn->pipnn;
+        // The reader rejects these states, so they are not written either.
+        FAISS_THROW_IF_NOT_MSG(
+                storage, "cannot serialize an IndexPiPNN without storage");
+        IndexPiPNN::check_storage_supported(*storage);
+        state.check_params();
+        FAISS_THROW_IF_NOT_MSG(
+                storage->ntotal == idxpipnn->ntotal &&
+                        storage->d == idxpipnn->d &&
+                        storage->metric_type == idxpipnn->metric_type,
+                "cannot serialize IndexPiPNN: the storage is inconsistent "
+                "with the index (was the storage modified directly?)");
+        FAISS_THROW_IF_NOT_MSG(
+                state.graph.size() == size_t(state.ntotal) * size_t(state.R),
+                "cannot serialize IndexPiPNN: graph size != ntotal * R");
+        // The subclasses differ only by their constructors: the fourcc names
+        // the storage.
+        const uint32_t h = dynamic_cast<const IndexScalarQuantizer*>(storage)
+                ? fourcc("IPNs")
+                : fourcc("IPNf");
+        WRITE1(h);
+        write_index_header(idxpipnn, f);
+        WRITE1(idxpipnn->defer_build);
+        write_PiPNN(state, f);
+        write_index(storage, f);
     } else if (
             const IndexNNDescent* idxnnd =
                     dynamic_cast<const IndexNNDescent*>(idx)) {
@@ -951,6 +1011,21 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         write_index_header(idxnnd, f);
         write_NNDescent(&idxnnd->nndescent, f);
         write_index(idxnnd->storage, f);
+    } else if (
+            const IndexSQFastScan* idxsqfs =
+                    dynamic_cast<const IndexSQFastScan*>(idx)) {
+        uint32_t h = fourcc("ISfs");
+        WRITE1(h);
+        write_index_header(idxsqfs, f);
+        write_ScalarQuantizer(&idxsqfs->sq, f);
+        WRITE1(idxsqfs->implem);
+        WRITE1(idxsqfs->bbs);
+        WRITE1(idxsqfs->qbs);
+        WRITE1(idxsqfs->ntotal2);
+        WRITE1(idxsqfs->M2);
+        WRITEVECTOR(idxsqfs->codes);
+        WRITE1(idxsqfs->rerank_factor);
+        WRITEVECTOR(idxsqfs->lo_codes);
     } else if (
             const IndexPQFastScan* idxpqfs =
                     dynamic_cast<const IndexPQFastScan*>(idx)) {
@@ -964,6 +1039,26 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
         WRITE1(idxpqfs->ntotal2);
         WRITE1(idxpqfs->M2);
         WRITEVECTOR(idxpqfs->codes);
+    } else if (
+            const IndexIVFSQFastScan* ivfsqfs =
+                    dynamic_cast<const IndexIVFSQFastScan*>(idx)) {
+        uint32_t h = fourcc("IwSf");
+        WRITE1(h);
+        write_ivf_header(ivfsqfs, f);
+        WRITE1(ivfsqfs->by_residual);
+        WRITE1(ivfsqfs->code_size);
+        WRITE1(ivfsqfs->bbs);
+        WRITE1(ivfsqfs->M2);
+        WRITE1(ivfsqfs->implem);
+        WRITE1(ivfsqfs->rerank_factor);
+        write_ScalarQuantizer(&ivfsqfs->sq, f);
+        write_InvertedLists(ivfsqfs->invlists, f);
+        // Write lo_codes_invlists if present
+        bool has_orig = (ivfsqfs->lo_codes_invlists != nullptr);
+        WRITE1(has_orig);
+        if (has_orig) {
+            write_InvertedLists(ivfsqfs->lo_codes_invlists, f);
+        }
     } else if (
             const IndexIVFPQFastScan* ivpq_2 =
                     dynamic_cast<const IndexIVFPQFastScan*>(idx)) {

data/vendor/faiss/faiss/impl/pipnn/HashPrune.cpp

@@ -0,0 +1,159 @@
+/*
+ * 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.
+ */
+
+#include <faiss/impl/pipnn/HashPrune.h>
+
+#include <cstring>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/bf16.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+namespace pipnn {
+
+uint16_t distance_key(float v) {
+    // Normalized on the bit pattern, so -ffast-math cannot remove the checks.
+    uint32_t u;
+    std::memcpy(&u, &v, sizeof(u));
+    if ((u & 0x7fffffffu) == 0) {
+        u = 0; // -0 -> +0
+    } else if ((u & 0x7f800000u) == 0x7f800000u && (u & 0x007fffffu) != 0) {
+        u = 0x7f800000u; // NaN -> +inf
+    }
+    std::memcpy(&v, &u, sizeof(v));
+    // Scalar `encode_bf16` only: the SIMD path rounds differently, so keys
+    // would depend on the CPU.
+    const uint16_t b = encode_bf16(v);
+    return (b & 0x8000) ? static_cast<uint16_t>(~b)
+                        : static_cast<uint16_t>(b | 0x8000);
+}
+
+bool ReservoirRef::offer(int32_t id, uint16_t hash, uint16_t key) {
+    const int n = *count;
+    // A full reservoir rejects anything not below its maximum (rules 1 and 3).
+    if (n == capacity && *farthest != kNoFarthest) {
+        const Slot& f = slots[*farthest];
+        if (!slot_less(key, id, f.key, f.id)) {
+            return false;
+        }
+    }
+
+    int lo = 0;
+    int hi = n;
+    while (lo < hi) {
+        const int mid = (lo + hi) / 2;
+        if (slots[mid].hash < hash) {
+            lo = mid + 1;
+        } else {
+            hi = mid;
+        }
+    }
+
+    // Rule 1: bucket hit. The cached farthest stays valid only when a
+    // non-farthest slot is lowered.
+    if (lo < n && slots[lo].hash == hash) {
+        Slot& s = slots[lo];
+        if (!slot_less(key, id, s.key, s.id)) {
+            return false;
+        }
+        s.id = id;
+        s.key = key;
+        if (*farthest == lo) {
+            *farthest = kNoFarthest;
+        }
+        return true;
+    }
+
+    // Rule 2: miss with room.
+    if (n < capacity) {
+        std::memmove(
+                slots + lo + 1,
+                slots + lo,
+                sizeof(Slot) * static_cast<size_t>(n - lo));
+        slots[lo] = Slot{id, hash, key};
+        *count = static_cast<uint8_t>(n + 1);
+        *farthest = kNoFarthest;
+        return true;
+    }
+
+    // Rule 3: miss, full. Compare against the slot_less-maximum.
+    int z = *farthest;
+    if (z == kNoFarthest) {
+        z = 0;
+        for (int i = 1; i < n; i++) {
+            if (slot_less(
+                        slots[z].key, slots[z].id, slots[i].key, slots[i].id)) {
+                z = i;
+            }
+        }
+        *farthest = static_cast<uint8_t>(z);
+    }
+    const Slot& zs = slots[z];
+    if (!slot_less(key, id, zs.key, zs.id)) {
+        return false;
+    }
+    if (z < lo) {
+        std::memmove(
+                slots + z,
+                slots + z + 1,
+                sizeof(Slot) * static_cast<size_t>(lo - 1 - z));
+        slots[lo - 1] = Slot{id, hash, key};
+    } else {
+        std::memmove(
+                slots + lo + 1,
+                slots + lo,
+                sizeof(Slot) * static_cast<size_t>(z - lo));
+        slots[lo] = Slot{id, hash, key};
+    }
+    *farthest = kNoFarthest;
+    return true;
+}
+
+Hyperplanes::Hyperplanes(int d, int m, int64_t seed) : d(d), m(m) {
+    FAISS_THROW_IF_NOT_FMT(d >= 1, "Hyperplanes: d must be >= 1, got %d", d);
+    FAISS_THROW_IF_NOT_FMT(
+            m >= 1 && m <= 16, "Hyperplanes: m must be in [1, 16], got %d", m);
+    std::vector<float> h(static_cast<size_t>(m) * d);
+    float_randn(h.data(), h.size(), seed);
+    ht.assign(static_cast<size_t>(d) * 16, 0.0f);
+    for (int i = 0; i < m; i++) {
+        for (int j = 0; j < d; j++) {
+            ht[static_cast<size_t>(j) * 16 + i] =
+                    h[static_cast<size_t>(i) * d + j];
+        }
+    }
+}
+
+FAISS_NOINLINE void Hyperplanes::sketch(const float* x, float* out16) const {
+    // Local accumulators, j ascending: the result depends only on the values
+    // of x, never on its buffer, alignment or caller.
+    float acc[16] = {};
+    const float* h = ht.data();
+    for (int j = 0; j < d; j++) {
+        const float xj = x[j];
+        const float* hj = h + static_cast<size_t>(j) * 16;
+        for (int i = 0; i < 16; i++) {
+            acc[i] += xj * hj[i];
+        }
+    }
+    std::memcpy(out16, acc, sizeof(acc));
+}
+
+uint16_t residual_hash(const float* sketch_p, const float* sketch_c, int m) {
+    uint16_t h = 0;
+    for (int i = 0; i < m; i++) {
+        if (sketch_c[i] - sketch_p[i] >= 0) {
+            h = static_cast<uint16_t>(h | (1u << i));
+        }
+    }
+    return h;
+}
+
+} // namespace pipnn
+} // namespace faiss

data/vendor/faiss/faiss/impl/pipnn/HashPrune.h

@@ -0,0 +1,69 @@
+/*
+ * 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 <vector>
+
+namespace faiss {
+namespace pipnn {
+
+/// One reservoir slot: candidate id, residual-LSH bucket, order-preserving key.
+struct Slot {
+    int32_t id;
+    uint16_t hash;
+    uint16_t key;
+};
+static_assert(sizeof(Slot) == 8, "PiPNN reservoir slots must be 8 bytes");
+
+/// "No cached farthest slot"; never a slot index since capacity <= 255.
+constexpr uint8_t kNoFarthest = 0xFF;
+
+/// Monotone 16-bit key of a distance, which may be negative, from its scalar
+/// bf16 encoding. NaN maps to +inf and -0 to +0.
+uint16_t distance_key(float v);
+
+/// Strict (key, id) order. Deciding key ties by id makes the reservoir
+/// independent of the offer order.
+inline bool slot_less(uint16_t ka, int32_t ia, uint16_t kb, int32_t ib) {
+    return ka < kb || (ka == kb && ia < ib);
+}
+
+/// One point's reservoir in caller-owned storage; not thread-safe. Slots are
+/// sorted by distinct hash, and *farthest caches the slot_less-maximum.
+struct ReservoirRef {
+    Slot* slots;
+    uint8_t* count;    // 0 when fresh
+    uint8_t* farthest; // kNoFarthest when fresh or unknown
+    int capacity;      // 1..255
+
+    /// Keeps the smaller offer on a bucket hit, else inserts or evicts the
+    /// maximum. An id must always come with the same hash. True iff changed.
+    bool offer(int32_t id, uint16_t hash, uint16_t key);
+};
+
+/// m <= 16 Gaussian hyperplanes from float_randn(seed), stored transposed
+/// and zero-padded: ht[j * 16 + i].
+struct Hyperplanes {
+    int d = 0;
+    int m = 0;
+    std::vector<float> ht; // d * 16
+
+    Hyperplanes() = default;
+    Hyperplanes(int d, int m, int64_t seed);
+
+    /// out16[i] = sum over ascending j of x[j] * ht[j * 16 + i]. Not inlined,
+    /// so all call sites produce bitwise-identical sketches.
+    void sketch(const float* x, float* out16) const;
+};
+
+/// Bit i (i < m <= 16) is set iff sketch_c[i] - sketch_p[i] >= 0.
+uint16_t residual_hash(const float* sketch_p, const float* sketch_c, int m);
+
+} // namespace pipnn
+} // namespace faiss