faiss 0.1.0 → 0.1.1

data/vendor/faiss/c_api/gpu/GpuIndex_c.cpp

@@ -0,0 +1,17 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#include "gpu/GpuIndex.h"
+#include "GpuIndex_c.h"
+#include "macros_impl.h"
+
+using faiss::gpu::GpuIndexConfig;
+
+DEFINE_GETTER(GpuIndexConfig, int, device)

data/vendor/faiss/c_api/gpu/GpuIndex_c.h

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c -*-
+
+#ifndef FAISS_GPU_INDEX_C_H
+#define FAISS_GPU_INDEX_C_H
+
+#include "faiss_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+FAISS_DECLARE_CLASS(GpuIndexConfig)
+
+FAISS_DECLARE_GETTER(GpuIndexConfig, int, device)
+
+FAISS_DECLARE_CLASS_INHERITED(GpuIndex, Index)
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

data/vendor/faiss/c_api/gpu/GpuIndicesOptions_c.h

@@ -0,0 +1,38 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c -*-
+
+#ifndef FAISS_GPU_INDICES_OPTIONS_C_H
+#define FAISS_GPU_INDICES_OPTIONS_C_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// How user vector index data is stored on the GPU
+typedef enum FaissIndicesOptions {
+  /// The user indices are only stored on the CPU; the GPU returns
+  /// (inverted list, offset) to the CPU which is then translated to
+  /// the real user index.
+  INDICES_CPU = 0,
+  /// The indices are not stored at all, on either the CPU or
+  /// GPU. Only (inverted list, offset) is returned to the user as the
+  /// index.
+  INDICES_IVF = 1,
+  /// Indices are stored as 32 bit integers on the GPU, but returned
+  /// as 64 bit integers
+  INDICES_32_BIT = 2,
+  /// Indices are stored as 64 bit integers on the GPU
+  INDICES_64_BIT = 3,
+} FaissIndicesOptions;
+
+#ifdef __cplusplus
+}
+#endif
+#endif

data/vendor/faiss/c_api/gpu/GpuResources_c.cpp

@@ -0,0 +1,86 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#include "gpu/GpuResources_c.h"
+#include "gpu/GpuResources.h"
+#include "macros_impl.h"
+
+using faiss::gpu::GpuResources;
+
+DEFINE_DESTRUCTOR(GpuResources)
+
+int faiss_GpuResources_initializeForDevice(FaissGpuResources* res, int device) {
+    try {
+        reinterpret_cast<GpuResources*>(res)->initializeForDevice(device);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getBlasHandle(FaissGpuResources* res, int device, cublasHandle_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getBlasHandle(device);
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getDefaultStream(FaissGpuResources* res, int device, cudaStream_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getDefaultStream(device);
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getPinnedMemory(FaissGpuResources* res, void** p_buffer, size_t* p_size) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getPinnedMemory();
+        *p_buffer = o.first;
+        *p_size = o.second;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getAsyncCopyStream(FaissGpuResources* res, int device, cudaStream_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getAsyncCopyStream(device);
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getBlasHandleCurrentDevice(FaissGpuResources* res, cublasHandle_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getBlasHandleCurrentDevice();
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getDefaultStreamCurrentDevice(FaissGpuResources* res, cudaStream_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getDefaultStreamCurrentDevice();
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_syncDefaultStream(FaissGpuResources* res, int device) {
+    try {
+        reinterpret_cast<GpuResources*>(res)->syncDefaultStream(device);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_syncDefaultStreamCurrentDevice(FaissGpuResources* res) {
+    try {
+        reinterpret_cast<GpuResources*>(res)->syncDefaultStreamCurrentDevice();
+    } CATCH_AND_HANDLE
+}
+
+int faiss_GpuResources_getAsyncCopyStreamCurrentDevice(FaissGpuResources* res, cudaStream_t* out) {
+    try {
+        auto o = reinterpret_cast<GpuResources*>(res)->getAsyncCopyStreamCurrentDevice();
+        *out = o;
+    } CATCH_AND_HANDLE
+}
+

data/vendor/faiss/c_api/gpu/GpuResources_c.h

@@ -0,0 +1,66 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c -*-
+
+#ifndef FAISS_GPU_RESOURCES_C_H
+#define FAISS_GPU_RESOURCES_C_H
+
+#include <cuda_runtime_api.h>
+#include <cublas.h>
+#include "faiss_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Base class of GPU-side resource provider; hides provision of
+/// cuBLAS handles, CUDA streams and a temporary memory manager
+FAISS_DECLARE_CLASS(GpuResources)
+
+FAISS_DECLARE_DESTRUCTOR(GpuResources)
+
+/// Call to pre-allocate resources for a particular device. If this is
+/// not called, then resources will be allocated at the first time
+/// of demand
+int faiss_GpuResources_initializeForDevice(FaissGpuResources*, int);
+
+/// Returns the cuBLAS handle that we use for the given device
+int faiss_GpuResources_getBlasHandle(FaissGpuResources*, int, cublasHandle_t*);
+
+/// Returns the stream that we order all computation on for the
+/// given device
+int faiss_GpuResources_getDefaultStream(FaissGpuResources*, int, cudaStream_t*);
+
+/// Returns the available CPU pinned memory buffer
+int faiss_GpuResources_getPinnedMemory(FaissGpuResources*, void**, size_t*);
+
+/// Returns the stream on which we perform async CPU <-> GPU copies
+int faiss_GpuResources_getAsyncCopyStream(FaissGpuResources*, int, cudaStream_t*);
+
+/// Calls getBlasHandle with the current device
+int faiss_GpuResources_getBlasHandleCurrentDevice(FaissGpuResources*, cublasHandle_t*);
+
+/// Calls getDefaultStream with the current device
+int faiss_GpuResources_getDefaultStreamCurrentDevice(FaissGpuResources*, cudaStream_t*);
+
+/// Synchronizes the CPU with respect to the default stream for the
+/// given device
+// equivalent to cudaDeviceSynchronize(getDefaultStream(device))
+int faiss_GpuResources_syncDefaultStream(FaissGpuResources*, int);
+
+/// Calls syncDefaultStream for the current device
+int faiss_GpuResources_syncDefaultStreamCurrentDevice(FaissGpuResources*);
+
+/// Calls getAsyncCopyStream for the current device
+int faiss_GpuResources_getAsyncCopyStreamCurrentDevice(FaissGpuResources*, cudaStream_t*);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

data/vendor/faiss/c_api/gpu/StandardGpuResources_c.cpp

@@ -0,0 +1,54 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#include "gpu/StandardGpuResources_c.h"
+#include "gpu/StandardGpuResources.h"
+#include "macros_impl.h"
+
+using faiss::gpu::StandardGpuResources;
+
+DEFINE_DESTRUCTOR(StandardGpuResources)
+
+int faiss_StandardGpuResources_new(FaissStandardGpuResources** p_res) {
+    try {
+        auto p = new StandardGpuResources();
+        *p_res = reinterpret_cast<FaissStandardGpuResources*>(p);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_StandardGpuResources_noTempMemory(FaissStandardGpuResources* res) {
+    try {
+        reinterpret_cast<StandardGpuResources*>(res)->noTempMemory();
+    } CATCH_AND_HANDLE
+}
+
+int faiss_StandardGpuResources_setTempMemory(FaissStandardGpuResources* res, size_t size) {
+    try {
+        reinterpret_cast<StandardGpuResources*>(res)->setTempMemory(size);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_StandardGpuResources_setPinnedMemory(FaissStandardGpuResources* res, size_t size) {
+    try {
+        reinterpret_cast<StandardGpuResources*>(res)->setPinnedMemory(size);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_StandardGpuResources_setDefaultStream(FaissStandardGpuResources* res, int device, cudaStream_t stream) {
+    try {
+        reinterpret_cast<StandardGpuResources*>(res)->setDefaultStream(device, stream);
+    } CATCH_AND_HANDLE
+}
+
+int faiss_StandardGpuResources_setDefaultNullStreamAllDevices(FaissStandardGpuResources* res) {
+    try {
+        reinterpret_cast<StandardGpuResources*>(res)->setDefaultNullStreamAllDevices();
+    } CATCH_AND_HANDLE
+}

data/vendor/faiss/c_api/gpu/StandardGpuResources_c.h

@@ -0,0 +1,53 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c -*-
+
+#ifndef FAISS_STANDARD_GPURESOURCES_C_H
+#define FAISS_STANDARD_GPURESOURCES_C_H
+
+#include <cuda_runtime_api.h>
+#include "faiss_c.h"
+#include "gpu/GpuResources_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Default implementation of GpuResources that allocates a cuBLAS
+/// stream and 2 streams for use, as well as temporary memory
+FAISS_DECLARE_CLASS_INHERITED(StandardGpuResources, GpuResources)
+
+FAISS_DECLARE_DESTRUCTOR(StandardGpuResources)
+
+/// Default constructor for StandardGpuResources
+int faiss_StandardGpuResources_new(FaissStandardGpuResources**);
+
+/// Disable allocation of temporary memory; all temporary memory
+/// requests will call cudaMalloc / cudaFree at the point of use
+int faiss_StandardGpuResources_noTempMemory(FaissStandardGpuResources*);
+
+/// Specify that we wish to use a certain fixed size of memory on
+/// all devices as temporary memory
+int faiss_StandardGpuResources_setTempMemory(FaissStandardGpuResources*, size_t size);
+
+/// Set amount of pinned memory to allocate, for async GPU <-> CPU
+/// transfers
+int faiss_StandardGpuResources_setPinnedMemory(FaissStandardGpuResources*, size_t size);
+
+/// Called to change the stream for work ordering
+int faiss_StandardGpuResources_setDefaultStream(FaissStandardGpuResources*, int device, cudaStream_t stream);
+
+/// Called to change the work ordering streams to the null stream
+/// for all devices
+int faiss_StandardGpuResources_setDefaultNullStreamAllDevices(FaissStandardGpuResources*);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

data/vendor/faiss/c_api/gpu/macros_impl.h

@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef GPU_MACROS_IMPL_H
+#define GPU_MACROS_IMPL_H
+#include "../macros_impl.h"
+
+#undef DEFINE_GETTER
+#define DEFINE_GETTER(clazz, ty, name)                               \
+    ty faiss_ ## clazz ## _ ## name (const Faiss ## clazz *obj) {    \
+        return static_cast< ty >(                                    \
+            reinterpret_cast< const faiss::gpu::clazz *>(obj)-> name \
+        );                                                           \
+    }
+
+#undef DEFINE_SETTER
+#define DEFINE_SETTER(clazz, ty, name)                                    \
+    void faiss_ ## clazz ## _set_ ## name (Faiss ## clazz *obj, ty val) { \
+        reinterpret_cast< faiss::gpu::clazz *>(obj)-> name = val;              \
+    }
+
+#undef DEFINE_SETTER_STATIC
+#define DEFINE_SETTER_STATIC(clazz, ty_to, ty_from, name)                      \
+    void faiss_ ## clazz ## _set_ ## name (Faiss ## clazz *obj, ty_from val) { \
+        reinterpret_cast< faiss::gpu::clazz *>(obj)-> name =                   \
+            static_cast< ty_to >(val);                                         \
+    }
+
+#undef DEFINE_DESTRUCTOR
+#define DEFINE_DESTRUCTOR(clazz)                           \
+    void faiss_ ## clazz ## _free (Faiss ## clazz *obj) {  \
+        delete reinterpret_cast<faiss::gpu::clazz *>(obj); \
+    }
+
+#endif

data/vendor/faiss/c_api/impl/AuxIndexStructures_c.cpp

@@ -0,0 +1,220 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#include "AuxIndexStructures_c.h"
+#include "../../impl/AuxIndexStructures.h"
+#include "../macros_impl.h"
+#include <iostream>
+
+using faiss::BufferList;
+using faiss::IDSelector;
+using faiss::IDSelectorBatch;
+using faiss::IDSelectorRange;
+using faiss::RangeSearchResult;
+using faiss::RangeSearchPartialResult;
+using faiss::RangeQueryResult;
+using faiss::DistanceComputer;
+
+DEFINE_GETTER(RangeSearchResult, size_t, nq)
+
+int faiss_RangeSearchResult_new(FaissRangeSearchResult** p_rsr, idx_t nq) {
+    try {
+        *p_rsr = reinterpret_cast<FaissRangeSearchResult*>(
+            new RangeSearchResult(nq));
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_RangeSearchResult_new_with(FaissRangeSearchResult** p_rsr, idx_t nq, int alloc_lims) {
+    try {
+        *p_rsr = reinterpret_cast<FaissRangeSearchResult*>(
+            new RangeSearchResult(nq, static_cast<bool>(alloc_lims)));
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+/// called when lims contains the nb of elements result entries
+/// for each query
+int faiss_RangeSearchResult_do_allocation(FaissRangeSearchResult* rsr) {
+    try {
+        reinterpret_cast<RangeSearchResult*>(rsr)->do_allocation();
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+DEFINE_DESTRUCTOR(RangeSearchResult)
+
+/// getter for buffer_size
+DEFINE_GETTER(RangeSearchResult, size_t, buffer_size)
+
+/// getter for lims: size (nq + 1)
+void faiss_RangeSearchResult_lims(FaissRangeSearchResult* rsr, size_t** lims) {
+    *lims = reinterpret_cast<RangeSearchResult*>(rsr)->lims;
+}
+
+/// getter for labels and respective distances (not sorted):
+/// result for query i is labels[lims[i]:lims[i+1]]
+void faiss_RangeSearchResult_labels(FaissRangeSearchResult* rsr, idx_t** labels, float** distances) {
+    auto sr = reinterpret_cast<RangeSearchResult*>(rsr);
+    *labels = sr->labels;
+    *distances = sr->distances;
+}
+
+DEFINE_DESTRUCTOR(IDSelector)
+
+int faiss_IDSelector_is_member(const FaissIDSelector* sel, idx_t id)  {
+    return reinterpret_cast<const IDSelector*>(sel)->is_member(id);
+}
+
+DEFINE_DESTRUCTOR(IDSelectorRange)
+
+DEFINE_GETTER(IDSelectorRange, idx_t, imin)
+DEFINE_GETTER(IDSelectorRange, idx_t, imax)
+
+int faiss_IDSelectorRange_new(FaissIDSelectorRange** p_sel, idx_t imin, idx_t imax) {
+    try {
+        *p_sel = reinterpret_cast<FaissIDSelectorRange*>(
+            new IDSelectorRange(imin, imax)
+        );
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+DEFINE_GETTER(IDSelectorBatch, int, nbits)
+DEFINE_GETTER(IDSelectorBatch, idx_t, mask)
+
+int faiss_IDSelectorBatch_new(FaissIDSelectorBatch** p_sel, size_t n, const idx_t* indices) {
+    try {
+        *p_sel = reinterpret_cast<FaissIDSelectorBatch*>(
+            new IDSelectorBatch(n, indices)
+        );
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+// Below are structures used only by Index implementations
+
+DEFINE_DESTRUCTOR(BufferList)
+
+DEFINE_GETTER(BufferList, size_t, buffer_size)
+DEFINE_GETTER(BufferList, size_t, wp)
+
+int faiss_BufferList_append_buffer(FaissBufferList* bl) {
+    try {
+        reinterpret_cast<BufferList*>(bl)->append_buffer();
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_BufferList_new(FaissBufferList** p_bl, size_t buffer_size) {
+    try {
+        *p_bl = reinterpret_cast<FaissBufferList*>(
+            new BufferList(buffer_size)
+        );
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_BufferList_add(FaissBufferList* bl, idx_t id, float dis) {
+    try {
+        reinterpret_cast<BufferList*>(bl)->add(id, dis);
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+/// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
+/// tables dest_ids, dest_dis
+int faiss_BufferList_copy_range(
+    FaissBufferList* bl, size_t ofs, size_t n, idx_t *dest_ids, float *dest_dis) {
+    try {
+        reinterpret_cast<BufferList*>(bl)->copy_range(ofs, n, dest_ids, dest_dis);
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+DEFINE_GETTER(RangeQueryResult, idx_t, qno)
+DEFINE_GETTER(RangeQueryResult, size_t, nres)
+DEFINE_GETTER_PERMISSIVE(RangeQueryResult, FaissRangeSearchPartialResult*, pres)
+
+int faiss_RangeQueryResult_add(FaissRangeQueryResult* qr, float dis, idx_t id) {
+    try {
+        reinterpret_cast<RangeQueryResult*>(qr)->add(dis, id);
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+DEFINE_GETTER_PERMISSIVE(RangeSearchPartialResult, FaissRangeSearchResult*, res)
+
+int faiss_RangeSearchPartialResult_new(
+    FaissRangeSearchPartialResult** p_res, FaissRangeSearchResult* res_in) {
+    try {
+        *p_res = reinterpret_cast<FaissRangeSearchPartialResult*>(
+            new RangeSearchPartialResult(
+                reinterpret_cast<RangeSearchResult*>(res_in))
+        );
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_RangeSearchPartialResult_finalize(
+    FaissRangeSearchPartialResult* res) {
+    try {
+        reinterpret_cast<RangeSearchPartialResult*>(res)->finalize();
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+/// called by range_search before do_allocation
+int faiss_RangeSearchPartialResult_set_lims(
+    FaissRangeSearchPartialResult* res) {
+    try {
+        reinterpret_cast<RangeSearchPartialResult*>(res)->set_lims();
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+int faiss_RangeSearchPartialResult_new_result(
+    FaissRangeSearchPartialResult* res, idx_t qno, FaissRangeQueryResult** qr) {
+
+    try {
+        auto q = 
+            &reinterpret_cast<RangeSearchPartialResult*>(res)->new_result(qno);
+        if (qr) {
+            *qr = reinterpret_cast<FaissRangeQueryResult*>(&q);
+        }
+        return 0;
+    } CATCH_AND_HANDLE
+}
+
+DEFINE_DESTRUCTOR(DistanceComputer)
+
+int faiss_DistanceComputer_set_query(FaissDistanceComputer *dc, const float *x) {
+    try {
+        reinterpret_cast<DistanceComputer*>(dc)->set_query(x);
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_DistanceComputer_vector_to_query_dis(FaissDistanceComputer *dc, idx_t i, float *qd) {
+    try {
+        *qd = reinterpret_cast<DistanceComputer*>(dc)->operator()(i);
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}
+
+int faiss_DistanceComputer_symmetric_dis(FaissDistanceComputer *dc, idx_t i, idx_t j, float *vd) {
+    try {
+        *vd = reinterpret_cast<DistanceComputer*>(dc)->symmetric_dis(i, j);
+        return 0;
+    }
+    CATCH_AND_HANDLE
+}