Add support for IndexIDMap with Cagra fp16

faiss/Index.h

@@ -61,6 +61,7 @@ struct DistanceComputer;
 enum NumericType {
     Float32,
     Float16,
+    UInt8,
 };
 
 inline size_t get_numeric_type_size(NumericType numeric_type) {

faiss/IndexBinary.h

@@ -54,13 +54,31 @@ struct IndexBinary {
      * @param x      training vecors, size n * d / 8
      */
     virtual void train(idx_t n, const uint8_t* x);
+    // This typed train function is a dummy to enable overriding of typed train
+    // function in the templated IndexIDMap struct.
+    virtual void train(idx_t n, const void* x, NumericType numeric_type) {
+        if (numeric_type == NumericType::UInt8) {
+            train(n, static_cast<const uint8_t*>(x));
+        } else {
+            FAISS_THROW_MSG("IndexBinary::train: unsupported numeric type");
+        }
+    }
 
     /** Add n vectors of dimension d to the index.
      *
      * Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
      * @param x      input matrix, size n * d / 8
      */
     virtual void add(idx_t n, const uint8_t* x) = 0;
+    // This typed add function is a dummy to enable overriding of typed add
+    // function in the templated IndexIDMap struct.
+    virtual void add(idx_t n, const void* x, NumericType numeric_type) {
+        if (numeric_type == NumericType::UInt8) {
+            add(n, static_cast<const uint8_t*>(x));
+        } else {
+            FAISS_THROW_MSG("IndexBinary::add: unsupported numeric type");
+        }
+    }
 
     /** Same as add, but stores xids instead of sequential ids.
      *
@@ -70,6 +88,20 @@ struct IndexBinary {
      * @param xids if non-null, ids to store for the vectors (size n)
      */
     virtual void add_with_ids(idx_t n, const uint8_t* x, const idx_t* xids);
+    // This typed add_with_ids function is a dummy to enable overriding of typed
+    // add_with_ids function in the templated IndexIDMap struct.
+    virtual void add_with_ids(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* xids) {
+        if (numeric_type == NumericType::UInt8) {
+            add_with_ids(n, static_cast<const uint8_t*>(x), xids);
+        } else {
+            FAISS_THROW_MSG(
+                    "IndexBinary::add_with_ids: unsupported numeric type");
+        }
+    }
 
     /** Query n vectors of dimension d to the index.
      *
@@ -87,6 +119,27 @@ struct IndexBinary {
             int32_t* distances,
             idx_t* labels,
             const SearchParameters* params = nullptr) const = 0;
+    // This typed search function is a dummy to enable overriding of typed
+    // search function in the templated IndexIDMap struct.
+    virtual void search(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            idx_t k,
+            int32_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const {
+        if (numeric_type == NumericType::UInt8) {
+            search(n,
+                   static_cast<const uint8_t*>(x),
+                   k,
+                   distances,
+                   labels,
+                   params);
+        } else {
+            FAISS_THROW_MSG("IndexBinary::search: unsupported numeric type");
+        }
+    }
 
     /** Query n vectors of dimension d to the index.
      *

faiss/IndexIDMap.cpp

@@ -12,6 +12,7 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstdio>
+#include "faiss/Index.h"
 
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
@@ -33,6 +34,17 @@ void sync_d(IndexBinary* index) {
 
 } // anonymous namespace
 
+template <typename componentT>
+NumericType component_t_to_numeric() {
+    if constexpr (std::is_same<componentT, float>::value) {
+        return NumericType::Float32;
+    } else if constexpr (std::is_same<componentT, uint8_t>::value) {
+        return NumericType::UInt8;
+    } else {
+        FAISS_THROW_MSG("Unsupported component_t");
+    }
+}
+
 /*****************************************************
  * IndexIDMap implementation
  *******************************************************/
@@ -59,11 +71,21 @@ void IndexIDMapTemplate<IndexT>::add(
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::train(
         idx_t n,
-        const typename IndexT::component_t* x) {
-    index->train(n, x);
+        const void* x,
+        NumericType numeric_type) {
+    index->train(n, x, numeric_type);
     this->is_trained = index->is_trained;
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::train(
+        idx_t n,
+        const typename IndexT::component_t* x) {
+    train(n,
+          static_cast<const void*>(x),
+          component_t_to_numeric<typename IndexT::component_t>());
+}
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::reset() {
     index->reset();
@@ -74,14 +96,27 @@ void IndexIDMapTemplate<IndexT>::reset() {
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::add_with_ids(
         idx_t n,
-        const typename IndexT::component_t* x,
+        const void* x,
+        NumericType numeric_type,
         const idx_t* xids) {
-    index->add(n, x);
+    index->add(n, x, numeric_type);
     for (idx_t i = 0; i < n; i++)
         id_map.push_back(xids[i]);
     this->ntotal = index->ntotal;
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::add_with_ids(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        const idx_t* xids) {
+    add_with_ids(
+            n,
+            static_cast<const void*>(x),
+            component_t_to_numeric<typename IndexT::component_t>(),
+            xids);
+}
+
 template <typename IndexT>
 size_t IndexIDMapTemplate<IndexT>::sa_code_size() const {
     return index->sa_code_size();
@@ -123,7 +158,8 @@ struct ScopedSelChange {
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::search(
         idx_t n,
-        const typename IndexT::component_t* x,
+        const void* x,
+        NumericType numeric_type,
         idx_t k,
         typename IndexT::distance_t* distances,
         idx_t* labels,
@@ -147,14 +183,31 @@ void IndexIDMapTemplate<IndexT>::search(
             sel_change.set(params_non_const, &this_idtrans);
         }
     }
-    index->search(n, x, k, distances, labels, params);
+    index->search(n, x, numeric_type, k, distances, labels, params);
     idx_t* li = labels;
 #pragma omp parallel for
     for (idx_t i = 0; i < n * k; i++) {
         li[i] = li[i] < 0 ? li[i] : id_map[li[i]];
     }
 }
 
+template <typename IndexT>
+void IndexIDMapTemplate<IndexT>::search(
+        idx_t n,
+        const typename IndexT::component_t* x,
+        idx_t k,
+        typename IndexT::distance_t* distances,
+        idx_t* labels,
+        const SearchParameters* params) const {
+    search(n,
+           static_cast<const void*>(x),
+           component_t_to_numeric<typename IndexT::component_t>(),
+           k,
+           distances,
+           labels,
+           params);
+}
+
 template <typename IndexT>
 void IndexIDMapTemplate<IndexT>::range_search(
         idx_t n,

faiss/IndexIDMap.h

@@ -31,6 +31,11 @@ struct IndexIDMapTemplate : IndexT {
     /// @param xids if non-null, ids to store for the vectors (size n)
     void add_with_ids(idx_t n, const component_t* x, const idx_t* xids)
             override;
+    void add_with_ids(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            const idx_t* xids) override;
 
     /// this will fail. Use add_with_ids
     void add(idx_t n, const component_t* x) override;
@@ -42,8 +47,17 @@ struct IndexIDMapTemplate : IndexT {
             distance_t* distances,
             idx_t* labels,
             const SearchParameters* params = nullptr) const override;
+    void search(
+            idx_t n,
+            const void* x,
+            NumericType numeric_type,
+            idx_t k,
+            distance_t* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
 
     void train(idx_t n, const component_t* x) override;
+    void train(idx_t n, const void* x, NumericType numeric_type) override;
 
     void reset() override;