Merge branch 'main' into panorama-fp16-storage

Author: mnorris11

benchs/bench_cppcontrib_sa_decode.cpp

@@ -1472,7 +1472,7 @@ void testMinMaxIndexPQDecoder(
 }
 
 //
-int main(int argc, char** argv) {
+int main() {
     // 1 MB points
     const uint64_t INDEX_SIZE = 65536 * 16;
     const uint64_t N_ITERATIONS = 18;

benchs/bench_rabitq_simd.cpp

@@ -82,6 +82,18 @@ void bench_rabitq_and_dot_product_with_sum(benchmark::State& state) {
             });
 }
 
+void bench_rabitq_and_dot_product_with_sum_fused(benchmark::State& state) {
+    bench_rabitq_generic(
+            state,
+            [](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
+                    -> int64_t {
+                auto result = rabitq::bitwise_and_dot_product_with_popcount(
+                        q, x, size, qb);
+                // Synthetic operation using both inputs for benchmarking.
+                return result.popcount + result.dot_product;
+            });
+}
+
 template <SIMDLevel SL>
 void bench_rabitq_rearrange_impl(benchmark::State& state) {
     size_t qb = state.range(0);
@@ -127,6 +139,9 @@ BENCHMARK(bench_rabitq_xor_dot_product)
 BENCHMARK(bench_rabitq_and_dot_product_with_sum)
         ->ArgsProduct({qbs, dims})
         ->ArgNames({"qb", "d"});
+BENCHMARK(bench_rabitq_and_dot_product_with_sum_fused)
+        ->ArgsProduct({qbs, dims})
+        ->ArgNames({"qb", "d"});
 BENCHMARK(bench_rabitq_rearrange_scalar)
         ->ArgsProduct({qbs, dims})
         ->ArgNames({"qb", "d"});

benchs/bench_result_handler_overhead.cpp

@@ -46,8 +46,7 @@ std::pair<double, double> run_search(
         IndexData& data,
         int d,
         int k,
-        int nprobe,
-        const char* factory_string) {
+        int nprobe) {
     ParameterSpace().set_index_parameter(data.index.get(), "nprobe", nprobe);
 
     omp_set_num_threads(1);
@@ -184,8 +183,7 @@ int main() {
         std::vector<faiss::BenchmarkResult> results;
         for (int k : ks) {
             for (int nprobe : nprobes) {
-                auto [mean, std] = faiss::run_search(
-                        data, d, k, nprobe, index_factory.c_str());
+                auto [mean, std] = faiss::run_search(data, d, k, nprobe);
                 results.push_back({index_factory, d, k, nprobe, mean, std});
             }
         }

c_api/IndexScalarQuantizer_c.cpp

@@ -109,3 +109,42 @@ int faiss_IndexIVFScalarQuantizer_add_core(
     }
     CATCH_AND_HANDLE
 }
+
+/** The returned pointer is a borrowed reference to a member of the
+ * owning index; must not be freed.*/
+FaissScalarQuantizer* faiss_IndexScalarQuantizer_sq(
+        FaissIndexScalarQuantizer* index) {
+    return reinterpret_cast<FaissScalarQuantizer*>(
+            &reinterpret_cast<IndexScalarQuantizer*>(index)->sq);
+}
+
+/** The returned pointer is a borrowed reference to a member of the
+ * owning index; must not be freed.*/
+FaissScalarQuantizer* faiss_IndexIVFScalarQuantizer_sq(
+        FaissIndexIVFScalarQuantizer* index) {
+    return reinterpret_cast<FaissScalarQuantizer*>(
+            &reinterpret_cast<IndexIVFScalarQuantizer*>(index)->sq);
+}
+
+DEFINE_GETTER(ScalarQuantizer, FaissQuantizerType, qtype)
+DEFINE_GETTER(ScalarQuantizer, size_t, bits)
+DEFINE_GETTER(ScalarQuantizer, size_t, d)
+DEFINE_GETTER(ScalarQuantizer, size_t, code_size)
+DEFINE_GETTER(ScalarQuantizer, FaissRangeStat, rangestat)
+DEFINE_SETTER_STATIC(
+        ScalarQuantizer,
+        faiss::ScalarQuantizer::RangeStat,
+        FaissRangeStat,
+        rangestat)
+DEFINE_GETTER(ScalarQuantizer, float, rangestat_arg)
+DEFINE_SETTER(ScalarQuantizer, float, rangestat_arg)
+
+size_t faiss_ScalarQuantizer_trained_size(const FaissScalarQuantizer* sq) {
+    return reinterpret_cast<const faiss::ScalarQuantizer*>(sq)->trained.size();
+}
+
+void faiss_ScalarQuantizer_trained(const FaissScalarQuantizer* sq, float* out) {
+    const auto& trained =
+            reinterpret_cast<const faiss::ScalarQuantizer*>(sq)->trained;
+    memcpy(out, trained.data(), trained.size() * sizeof(float));
+}

c_api/IndexScalarQuantizer_c.h

@@ -28,11 +28,43 @@ typedef enum FaissQuantizerType {
     QT_bf16,
     QT_8bit_direct_signed, ///< fast indexing of signed int8s ranging from [-128
                            ///< to 127]
+    QT_0bit, ///< 0 bits per component, centroid-only distance (IVF only; a
+             ///< flat index with QT_0bit stores nothing)
+    QT_1bit_tqmse, ///< TurboQuant MSE-optimized, 1 bit per component
+    QT_2bit_tqmse, ///< TurboQuant MSE-optimized, 2 bits per component
+    QT_3bit_tqmse, ///< TurboQuant MSE-optimized, 3 bits per component
+    QT_4bit_tqmse, ///< TurboQuant MSE-optimized, 4 bits per component
+    QT_8bit_tqmse, ///< TurboQuant MSE-optimized, 8 bits per component
+    QT_2bit_tq,    ///< Full TurboQuant (1-bit MSE + 1-bit QJL + factors)
+    QT_3bit_tq,    ///< Full TurboQuant (2-bit MSE + 1-bit QJL + factors)
+    QT_4bit_tq,    ///< Full TurboQuant (3-bit MSE + 1-bit QJL + factors)
+    QT_5bit_tq,    ///< Full TurboQuant (4-bit MSE + 1-bit QJL + factors)
+    QT_1bit_eden,  ///< EDEN Lloyd-Max scalar code, 1 bit per component
+    QT_2bit_eden,  ///< EDEN Lloyd-Max scalar code, 2 bits per component
+    QT_3bit_eden,  ///< EDEN Lloyd-Max scalar code, 3 bits per component
+    QT_4bit_eden,  ///< EDEN Lloyd-Max scalar code, 4 bits per component
+    QT_5bit_eden,  ///< EDEN Lloyd-Max scalar code, 5 bits per component
+    QT_6bit_eden,  ///< EDEN Lloyd-Max scalar code, 6 bits per component
+    QT_7bit_eden,  ///< EDEN Lloyd-Max scalar code, 7 bits per component
+    QT_8bit_eden,  ///< EDEN Lloyd-Max scalar code, 8 bits per component
+    QT_count
 } FaissQuantizerType;
 
+typedef enum FaissRangeStat {
+    RS_minmax,    ///< [min - rs*(max-min), max + rs*(max-min)]
+    RS_meanstd,   ///< [mean - std * rs, mean + std * rs]
+    RS_quantiles, ///< [Q(rs), Q(1-rs)]
+    RS_optim,     ///< alternate optimization of reconstruction error
+} FaissRangeStat;
+
 // forward declaration
 typedef enum FaissMetricType FaissMetricType;
 
+/** Opaque type for the ScalarQuantizer codec owned by an index.
+ * Valid only while the owning index is alive and must not be freed.
+ */
+FAISS_DECLARE_CLASS(ScalarQuantizer)
+
 /** Opaque type for IndexScalarQuantizer */
 FAISS_DECLARE_CLASS_INHERITED(IndexScalarQuantizer, Index)
 
@@ -90,6 +122,39 @@ int faiss_IndexIVFScalarQuantizer_add_core(
         const idx_t* xids,
         const idx_t* precomputed_idx);
 
+/** Access the scalar quantizer codec of an IndexScalarQuantizer.
+ * Borrowed reference, owned by the index; do not free.
+ */
+FaissScalarQuantizer* faiss_IndexScalarQuantizer_sq(
+        FaissIndexScalarQuantizer* index);
+
+/** Access the scalar quantizer codec of an IndexIVFScalarQuantizer.
+ * Borrowed reference, owned by the index; do not free.
+ */
+FaissScalarQuantizer* faiss_IndexIVFScalarQuantizer_sq(
+        FaissIndexIVFScalarQuantizer* index);
+
+/// quantizer type of this codec
+FAISS_DECLARE_GETTER(ScalarQuantizer, FaissQuantizerType, qtype)
+/// bits per scalar code
+FAISS_DECLARE_GETTER(ScalarQuantizer, size_t, bits)
+/// size of the input vectors
+FAISS_DECLARE_GETTER(ScalarQuantizer, size_t, d)
+/// bytes per encoded vector
+FAISS_DECLARE_GETTER(ScalarQuantizer, size_t, code_size)
+/// range estimation strategy (uniform encoder)
+FAISS_DECLARE_GETTER_SETTER(ScalarQuantizer, FaissRangeStat, rangestat)
+/// argument to the range estimation strategy (rs)
+FAISS_DECLARE_GETTER_SETTER(ScalarQuantizer, float, rangestat_arg)
+
+/// Number of trained values
+size_t faiss_ScalarQuantizer_trained_size(const FaissScalarQuantizer* sq);
+
+/** Copy the trained values into out, which must hold at least
+ * faiss_ScalarQuantizer_trained_size() floats.
+ */
+void faiss_ScalarQuantizer_trained(const FaissScalarQuantizer* sq, float* out);
+
 #ifdef __cplusplus
 }
 #endif

faiss/IndexAdditiveQuantizer.cpp

@@ -235,8 +235,7 @@ void IndexAdditiveQuantizer::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
 
     if (aq->search_type == AdditiveQuantizer::ST_decompress) {
         with_VectorDistance(d, metric_type, metric_arg, [&](auto vd) {
@@ -459,8 +458,7 @@ void AdditiveCoarseQuantizer::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
 
     if (metric_type == METRIC_INNER_PRODUCT) {
         aq->knn_centroids_inner_product(n, x, k, distances, labels);

faiss/IndexAdditiveQuantizerFastScan.cpp

@@ -36,7 +36,7 @@ void IndexAdditiveQuantizerFastScan::init(
         MetricType metric,
         int bbs_) {
     FAISS_THROW_IF_NOT(aq_init != nullptr);
-    FAISS_THROW_IF_NOT(!aq_init->nbits.empty());
+    FAISS_THROW_IF_MSG(aq_init->nbits.empty(), "nbits must not be empty");
     FAISS_THROW_IF_NOT(aq_init->nbits[0] == 4);
     if (metric == METRIC_INNER_PRODUCT) {
         FAISS_THROW_IF_NOT_MSG(
@@ -192,8 +192,7 @@ void IndexAdditiveQuantizerFastScan::search(
         float* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
     bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2);
     if (!rescale) {

faiss/IndexBinaryFromFloat.cpp

@@ -55,8 +55,7 @@ void IndexBinaryFromFloat::search(
         int32_t* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     constexpr idx_t bs = 32768;

faiss/IndexBinaryHNSW.cpp

@@ -292,8 +292,8 @@ void IndexBinaryHNSW::reconstruct(idx_t key, uint8_t* recons) const {
 
 DistanceComputer* IndexBinaryHNSW::get_distance_computer() const {
     IndexBinaryFlat* flat_storage = dynamic_cast<IndexBinaryFlat*>(storage);
-    FAISS_THROW_IF_NOT_MSG(
-            flat_storage != nullptr,
+    FAISS_THROW_IF_MSG(
+            flat_storage == nullptr,
             "IndexBinaryHNSW requires IndexBinaryFlat storage");
     return with_simd_level([&]<SIMDLevel SL>() {
         return make_binary_hnsw_distance_computer_fixSL<SL>(
@@ -316,8 +316,8 @@ IndexBinaryHNSWCagra::IndexBinaryHNSWCagra(int d_, int M)
 }
 
 void IndexBinaryHNSWCagra::add(idx_t n, const uint8_t* x) {
-    FAISS_THROW_IF_NOT_MSG(
-            !base_level_only,
+    FAISS_THROW_IF_MSG(
+            base_level_only,
             "Cannot add vectors when base_level_only is set to True");
 
     IndexBinaryHNSW::add(n, x);

faiss/IndexBinaryHash.cpp

@@ -83,8 +83,7 @@ void IndexBinaryHash::range_search(
         int radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
 #pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
@@ -116,8 +115,7 @@ void IndexBinaryHash::search(
         int32_t* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     using HeapForL2 = CMax<int32_t, idx_t>;
@@ -216,8 +214,7 @@ void IndexBinaryMultiHash::range_search(
         int radius,
         RangeSearchResult* result,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     size_t nlist = 0, ndis = 0, n0 = 0;
 
 #pragma omp parallel if (n > 100) reduction(+ : ndis, n0, nlist)
@@ -249,8 +246,7 @@ void IndexBinaryMultiHash::search(
         int32_t* distances,
         idx_t* labels,
         const SearchParameters* params) const {
-    FAISS_THROW_IF_NOT_MSG(
-            !params, "search params not supported for this index");
+    FAISS_THROW_IF_MSG(params, "search params not supported for this index");
     FAISS_THROW_IF_NOT(k > 0);
 
     using HeapForL2 = CMax<int32_t, idx_t>;