MB-71397: Add C bindings for GPU Index class

c_api/IndexBinary_c.cpp

@@ -23,6 +23,8 @@ DEFINE_GETTER(IndexBinary, idx_t, ntotal)
 
 DEFINE_GETTER(IndexBinary, FaissMetricType, metric_type)
 
+DEFINE_GETTER(IndexBinary, int, code_size)
+
 DEFINE_GETTER(IndexBinary, int, verbose);
 DEFINE_SETTER(IndexBinary, int, verbose);
 

c_api/IndexBinary_c.h

@@ -40,6 +40,9 @@ FAISS_DECLARE_GETTER(IndexBinary, idx_t, ntotal)
 /// Getter for metric_type
 FAISS_DECLARE_GETTER(IndexBinary, FaissMetricType, metric_type)
 
+/// Getter for code_size (number of bytes per vector = d / 8)
+FAISS_DECLARE_GETTER(IndexBinary, int, code_size)
+
 FAISS_DECLARE_GETTER_SETTER(IndexBinary, int, verbose)
 
 /** Perform training on a representative set of vectors

c_api/gpu/CMakeLists.txt

@@ -8,6 +8,7 @@ target_sources(faiss_c PRIVATE
   GpuAutoTune_c.cpp
   GpuClonerOptions_c.cpp
   GpuIndex_c.cpp
+  GpuIndexIVF_c.cpp
   GpuResources_c.cpp
   StandardGpuResources_c.cpp
 )

c_api/gpu/GpuIndexIVF_c.cpp

@@ -0,0 +1,45 @@
+/*
+ * 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.
+ */
+
+// -*- c++ -*-
+
+#include "GpuIndexIVF_c.h"
+#include <faiss/gpu/GpuIndexIVF.h>
+#include <faiss/gpu/GpuIndexIVFFlat.h>
+#include <faiss/gpu/GpuIndexIVFPQ.h>
+#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
+#include "macros_impl.h"
+
+using faiss::gpu::GpuIndex;
+using faiss::gpu::GpuIndexIVF;
+using faiss::gpu::GpuIndexIVFFlat;
+using faiss::gpu::GpuIndexIVFPQ;
+using faiss::gpu::GpuIndexIVFScalarQuantizer;
+
+int faiss_GpuIndexIVF_reserve_memory(
+        FaissGpuIndexIVF* index,
+        size_t num_vectors) {
+    try {
+        auto ivf = reinterpret_cast<GpuIndexIVF*>(index);
+        if (auto flat = dynamic_cast<GpuIndexIVFFlat*>(ivf)) {
+            flat->reserveMemory(num_vectors);
+        } else if (auto pq = dynamic_cast<GpuIndexIVFPQ*>(ivf)) {
+            pq->reserveMemory(num_vectors);
+        } else if (auto sq = dynamic_cast<GpuIndexIVFScalarQuantizer*>(ivf)) {
+            sq->reserveMemory(num_vectors);
+        } else {
+            throw std::runtime_error(
+                    "reserve_memory not supported for this GpuIndexIVF subtype");
+        }
+    }
+    CATCH_AND_HANDLE
+}
+
+FaissGpuIndexIVF* faiss_GpuIndexIVF_cast(FaissGpuIndex* index) {
+    return reinterpret_cast<FaissGpuIndexIVF*>(
+            dynamic_cast<GpuIndexIVF*>(reinterpret_cast<GpuIndex*>(index)));
+}

c_api/gpu/GpuIndexIVF_c.h

@@ -0,0 +1,35 @@
+/*
+ * 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.
+ */
+
+// -*- c -*-
+
+#ifndef FAISS_GPU_INDEX_IVF_C_H
+#define FAISS_GPU_INDEX_IVF_C_H
+
+#include <stddef.h>
+
+#include "GpuIndex_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+FAISS_DECLARE_CLASS_INHERITED(GpuIndexIVF, GpuIndex)
+
+/// Reserve memory for a given number of vectors, before adding them to the index.
+int faiss_GpuIndexIVF_reserve_memory(
+        FaissGpuIndexIVF* index,
+        size_t num_vectors);
+
+/// Downcast a GpuIndex to GpuIndexIVF (returns NULL if not a GpuIndexIVF)
+FaissGpuIndexIVF* faiss_GpuIndexIVF_cast(FaissGpuIndex* index);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

c_api/gpu/GpuIndex_c.cpp

@@ -8,9 +8,45 @@
 // -*- c++ -*-
 
 #include "GpuIndex_c.h"
+#include "GpuIndexIVF_c.h"
 #include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/GpuIndexIVF.h>
 #include "macros_impl.h"
 
+using faiss::gpu::GpuIndex;
 using faiss::gpu::GpuIndexConfig;
+using faiss::gpu::GpuIndexIVF;
 
 DEFINE_GETTER(GpuIndexConfig, int, device)
+
+DEFINE_GETTER(GpuIndex, int, d)
+
+int faiss_GpuIndex_train(FaissGpuIndex* index, idx_t n, const float* x) {
+    try {
+        reinterpret_cast<GpuIndex*>(index)->train(n, x);
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_GpuIndex_add(FaissGpuIndex* index, idx_t n, const float* x) {
+    try {
+        reinterpret_cast<GpuIndex*>(index)->add(n, x);
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_GpuIndex_search(
+        const FaissGpuIndex* index,
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) {
+    try {
+        reinterpret_cast<const GpuIndex*>(index)->search(
+                n, x, k, distances, labels);
+    }
+    CATCH_AND_HANDLE
+}
+
+DEFINE_DESTRUCTOR(GpuIndex)

c_api/gpu/GpuIndex_c.h

@@ -22,6 +22,45 @@ FAISS_DECLARE_GETTER(GpuIndexConfig, int, device)
 
 FAISS_DECLARE_CLASS_INHERITED(GpuIndex, Index)
 
+/// Getter for d (dimension)
+FAISS_DECLARE_GETTER(GpuIndex, int, d)
+
+/** Perform training on a representative set of vectors
+ *
+ * @param index  opaque pointer to GPU index object
+ * @param n      nb of training vectors
+ * @param x      training vectors, size n * d
+ */
+int faiss_GpuIndex_train(FaissGpuIndex* index, idx_t n, const float* x);
+
+/** Add n vectors of dimension d to the index.
+ *
+ * @param index  opaque pointer to GPU index object
+ * @param n      nb of vectors to add
+ * @param x      input vectors, size n * d
+ */
+int faiss_GpuIndex_add(FaissGpuIndex* index, idx_t n, const float* x);
+
+/** Query n vectors of dimension d to the index.
+ *
+ * @param index       opaque pointer to GPU index object
+ * @param n           number of vectors to search
+ * @param x           input vectors to search, size n * d
+ * @param k           number of nearest neighbors to return
+ * @param distances   output pairwise distances, size n * k
+ * @param labels      output labels of the NNs, size n * k
+ */
+int faiss_GpuIndex_search(
+        const FaissGpuIndex* index,
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels);
+
+/// Free the GPU index
+FAISS_DECLARE_DESTRUCTOR(GpuIndex)
+
 #ifdef __cplusplus
 }
 #endif