c_api: expose ScalarQuantizer codec and complete FaissQuantizerType enum (#5477)

Summary:
Hi,

the `c_api` is missing exports for the FaissScalarQuantizer scalar quantizer. I've added them.

Additionally I synced the `c_api/IndexScalarQuantizer_c.h` `FaissQuantizerType` enum with the cpp side which is required because `faiss_IndexScalarQuantizer_sq` may return an enum value which is not contained in the `c_api` enum definition.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5477

Reviewed By: trang-nm-nguyen

Differential Revision: D114077316

Pulled By: mnorris11

fbshipit-source-id: 6c714a86daa20a490b3892c3a78356b7beebc92e

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