639 lines
13 KiB
Plaintext
639 lines
13 KiB
Plaintext
|
|
/**
|
|
* Copyright (c) 2015-present, Facebook, Inc.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the CC-by-NC license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
// -*- C++ -*-
|
|
|
|
// This file describes the C++-scripting language bridge for both Lua
|
|
// and Python It contains mainly includes and a few macros. There are
|
|
// 3 preprocessor macros of interest:
|
|
// SWIGLUA: Lua-specific code
|
|
// SWIGPYTHON: Python-specific code
|
|
// GPU_WRAPPER: also compile interfaces for GPU.
|
|
|
|
|
|
#ifdef GPU_WRAPPER
|
|
%module swigfaiss_gpu;
|
|
#else
|
|
%module swigfaiss;
|
|
#endif
|
|
|
|
typedef unsigned long uint64_t;
|
|
typedef uint64_t size_t;
|
|
typedef int int32_t ;
|
|
|
|
#define __restrict
|
|
|
|
|
|
/*******************************************************************
|
|
* Copied verbatim to wrapper. Contains the C++-visible includes, and
|
|
* the language includes for their respective matrix libraries.
|
|
*******************************************************************/
|
|
|
|
%{
|
|
|
|
|
|
#include <stdint.h>
|
|
#include <omp.h>
|
|
|
|
|
|
#ifdef SWIGLUA
|
|
|
|
#include <pthread.h>
|
|
|
|
extern "C" {
|
|
|
|
#include <TH/TH.h>
|
|
#include <luaT.h>
|
|
#undef THTensor
|
|
|
|
}
|
|
|
|
#endif
|
|
|
|
|
|
#ifdef SWIGPYTHON
|
|
|
|
#undef popcount64
|
|
|
|
#define SWIG_FILE_WITH_INIT
|
|
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
|
|
#include <numpy/arrayobject.h>
|
|
|
|
#endif
|
|
|
|
|
|
#include "IndexFlat.h"
|
|
#include "VectorTransform.h"
|
|
|
|
#include "IndexLSH.h"
|
|
#include "IndexPQ.h"
|
|
#include "IndexIVF.h"
|
|
#include "IndexIVFPQ.h"
|
|
|
|
#include "MetaIndexes.h"
|
|
|
|
#include "index_io.h"
|
|
|
|
#include "utils.h"
|
|
#include "Heap.h"
|
|
#include "AuxIndexStructures.h"
|
|
|
|
|
|
#include "Clustering.h"
|
|
|
|
#include "hamming.h"
|
|
|
|
#include "AutoTune.h"
|
|
|
|
|
|
|
|
%}
|
|
|
|
/*******************************************************************
|
|
* Types of vectors we want to manipulate at the scripting language
|
|
* level.
|
|
*******************************************************************/
|
|
|
|
// simplified interface for vector
|
|
namespace std {
|
|
|
|
template<class T>
|
|
class vector {
|
|
public:
|
|
vector();
|
|
void push_back(T);
|
|
T * data();
|
|
size_t size();
|
|
T at (size_t n) const;
|
|
void resize (size_t n);
|
|
};
|
|
};
|
|
|
|
|
|
%template(FloatVector) std::vector<float>;
|
|
%template(ByteVector) std::vector<uint8_t>;
|
|
%template(Uint64Vector) std::vector<uint64_t>;
|
|
%template(LongVector) std::vector<long>;
|
|
%template(IntVector) std::vector<int>;
|
|
%template(VectorTransformVector) std::vector<faiss::VectorTransform*>;
|
|
%template(OperatingPointVector) std::vector<faiss::OperatingPoint>;
|
|
|
|
#ifdef GPU_WRAPPER
|
|
%template(GpuResourcesVector) std::vector<faiss::gpu::GpuResources*>;
|
|
#endif
|
|
|
|
%include <std_string.i>
|
|
|
|
// produces an error on the Mac
|
|
%ignore faiss::hamming;
|
|
|
|
|
|
|
|
/*******************************************************************
|
|
* Parse headers
|
|
*******************************************************************/
|
|
|
|
#ifdef SWIGPYTHON
|
|
// Python-specific: release GIL by default for all functions
|
|
%exception {
|
|
Py_BEGIN_ALLOW_THREADS
|
|
$action
|
|
Py_END_ALLOW_THREADS
|
|
}
|
|
#endif
|
|
|
|
|
|
|
|
%ignore *::cmp;
|
|
|
|
%include "Heap.h"
|
|
%include "hamming.h"
|
|
|
|
int get_num_gpus();
|
|
|
|
#ifdef GPU_WRAPPER
|
|
|
|
%{
|
|
|
|
#include "gpu/StandardGpuResources.h"
|
|
#include "gpu/GpuIndicesOptions.h"
|
|
#include "gpu/GpuIndex.h"
|
|
#include "gpu/GpuIndexFlat.h"
|
|
#include "gpu/GpuIndexIVF.h"
|
|
#include "gpu/GpuIndexIVFPQ.h"
|
|
#include "gpu/GpuIndexIVFFlat.h"
|
|
#include "gpu/IndexProxy.h"
|
|
#include "gpu/GpuAutoTune.h"
|
|
|
|
int get_num_gpus()
|
|
{
|
|
return faiss::gpu::getNumDevices();
|
|
}
|
|
|
|
%}
|
|
|
|
// causes weird wrapper bug
|
|
%ignore *::getMemoryManager;
|
|
%ignore *::getMemoryManagerCurrentDevice;
|
|
|
|
%include "gpu/GpuResources.h"
|
|
%include "gpu/StandardGpuResources.h"
|
|
|
|
#else
|
|
|
|
%{
|
|
int get_num_gpus()
|
|
{
|
|
return 0;
|
|
}
|
|
%}
|
|
|
|
|
|
#endif
|
|
|
|
|
|
%include "utils.h"
|
|
|
|
%include "Index.h"
|
|
%include "Clustering.h"
|
|
|
|
%ignore faiss::ProductQuantizer::get_centroids(size_t,size_t) const;
|
|
|
|
%include "ProductQuantizer.h"
|
|
|
|
%include "VectorTransform.h"
|
|
|
|
%include "IndexFlat.h"
|
|
%include "IndexLSH.h"
|
|
%include "PolysemousTraining.h"
|
|
%include "IndexPQ.h"
|
|
%include "IndexIVF.h"
|
|
|
|
%ignore faiss::IndexIVFPQ::alloc_type;
|
|
%include "IndexIVFPQ.h"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
%include "MetaIndexes.h"
|
|
|
|
|
|
#ifdef GPU_WRAPPER
|
|
|
|
// quiet SWIG warnings
|
|
%ignore faiss::gpu::GpuIndexIVF::GpuIndexIVF;
|
|
%ignore faiss::gpu::IndexProxy::at(int) const;
|
|
|
|
%include "gpu/GpuIndicesOptions.h"
|
|
%include "gpu/GpuIndex.h"
|
|
%include "gpu/GpuIndexFlat.h"
|
|
%include "gpu/GpuIndexIVF.h"
|
|
%include "gpu/GpuIndexIVFPQ.h"
|
|
%include "gpu/GpuIndexIVFFlat.h"
|
|
%include "gpu/IndexProxy.h"
|
|
|
|
#ifdef SWIGLUA
|
|
|
|
/// in Lua, swigfaiss_gpu is known as swigfaiss
|
|
%luacode {
|
|
local swigfaiss = swigfaiss_gpu
|
|
}
|
|
|
|
#endif
|
|
|
|
|
|
#endif
|
|
|
|
|
|
// Python-specific: do not release GIL any more, as functions below
|
|
// use the Python/C API
|
|
#ifdef SWIGPYTHON
|
|
%exception;
|
|
#endif
|
|
|
|
|
|
/*******************************************************************
|
|
* Lua-specific: support async execution of searches in an index
|
|
* Python equivalent is just to use Python threads.
|
|
*******************************************************************/
|
|
|
|
|
|
#ifdef SWIGLUA
|
|
|
|
%{
|
|
|
|
|
|
namespace faiss {
|
|
|
|
struct AsyncIndexSearchC {
|
|
typedef Index::idx_t idx_t;
|
|
const Index * index;
|
|
|
|
idx_t n;
|
|
const float *x;
|
|
idx_t k;
|
|
float *distances;
|
|
idx_t *labels;
|
|
|
|
bool is_finished;
|
|
|
|
pthread_t thread;
|
|
|
|
|
|
AsyncIndexSearchC (const Index *index,
|
|
idx_t n, const float *x, idx_t k,
|
|
float *distances, idx_t *labels):
|
|
index(index), n(n), x(x), k(k), distances(distances),
|
|
labels(labels)
|
|
{
|
|
is_finished = false;
|
|
pthread_create (&thread, NULL, &AsyncIndexSearchC::callback,
|
|
this);
|
|
}
|
|
|
|
static void *callback (void *arg)
|
|
{
|
|
AsyncIndexSearchC *aidx = (AsyncIndexSearchC *)arg;
|
|
aidx->do_search();
|
|
return NULL;
|
|
}
|
|
|
|
void do_search ()
|
|
{
|
|
index->search (n, x, k, distances, labels);
|
|
}
|
|
void join ()
|
|
{
|
|
pthread_join (thread, NULL);
|
|
}
|
|
|
|
};
|
|
|
|
}
|
|
|
|
%}
|
|
|
|
// re-decrlare only what we need
|
|
namespace faiss {
|
|
|
|
struct AsyncIndexSearchC {
|
|
typedef Index::idx_t idx_t;
|
|
bool is_finished;
|
|
AsyncIndexSearchC (const Index *index,
|
|
idx_t n, const float *x, idx_t k,
|
|
float *distances, idx_t *labels);
|
|
|
|
|
|
void join ();
|
|
};
|
|
|
|
}
|
|
|
|
|
|
#endif
|
|
|
|
|
|
|
|
|
|
/*******************************************************************
|
|
* downcast return of some functions so that the sub-class is used
|
|
* instead of the generic upper-class.
|
|
*******************************************************************/
|
|
|
|
#ifdef SWIGLUA
|
|
|
|
%define DOWNCAST(subclass)
|
|
if (dynamic_cast<faiss::subclass *> ($1)) {
|
|
SWIG_NewPointerObj(L,$1,SWIGTYPE_p_faiss__ ## subclass,$owner);
|
|
} else
|
|
%enddef
|
|
|
|
%define DOWNCAST_GPU(subclass)
|
|
if (dynamic_cast<faiss::gpu::subclass *> ($1)) {
|
|
SWIG_NewPointerObj(L,$1,SWIGTYPE_p_faiss__gpu__ ## subclass,$owner);
|
|
} else
|
|
%enddef
|
|
|
|
#endif
|
|
|
|
|
|
#ifdef SWIGPYTHON
|
|
|
|
%define DOWNCAST(subclass)
|
|
if (dynamic_cast<faiss::subclass *> ($1)) {
|
|
$result = SWIG_NewPointerObj($1,SWIGTYPE_p_faiss__ ## subclass,$owner);
|
|
} else
|
|
%enddef
|
|
|
|
%define DOWNCAST_GPU(subclass)
|
|
if (dynamic_cast<faiss::gpu::subclass *> ($1)) {
|
|
$result = SWIG_NewPointerObj($1,SWIGTYPE_p_faiss__gpu__ ## subclass,$owner);
|
|
} else
|
|
%enddef
|
|
|
|
#endif
|
|
|
|
%newobject read_index;
|
|
%newobject read_VectorTransform;
|
|
%newobject read_ProductQuantizer;
|
|
%newobject clone_index;
|
|
%newobject clone_VectorTransform;
|
|
|
|
// Subclasses should appear before their parent
|
|
%typemap(out) faiss::Index * {
|
|
DOWNCAST ( IndexIDMap )
|
|
DOWNCAST ( IndexShards )
|
|
DOWNCAST ( IndexIVFPQCompact )
|
|
DOWNCAST ( IndexIVFPQR )
|
|
DOWNCAST ( IndexIVFPQ )
|
|
DOWNCAST ( IndexIVFFlat )
|
|
DOWNCAST ( IndexIVF )
|
|
DOWNCAST ( IndexFlat )
|
|
DOWNCAST ( IndexPQ )
|
|
DOWNCAST ( IndexLSH )
|
|
DOWNCAST ( IndexPreTransform )
|
|
DOWNCAST ( MultiIndexQuantizer )
|
|
#ifdef GPU_WRAPPER
|
|
DOWNCAST_GPU ( IndexProxy )
|
|
DOWNCAST_GPU ( GpuIndexIVFPQ )
|
|
DOWNCAST_GPU ( GpuIndexIVFFlat )
|
|
DOWNCAST_GPU ( GpuIndexFlat )
|
|
#endif
|
|
// default for non-recognized classes
|
|
DOWNCAST( Index )
|
|
{
|
|
assert(false);
|
|
}
|
|
#ifdef SWIGLUA
|
|
SWIG_arg++;
|
|
#endif
|
|
}
|
|
|
|
%typemap(out) faiss::VectorTransform * {
|
|
DOWNCAST (RemapDimensionsTransform)
|
|
DOWNCAST (OPQMatrix)
|
|
DOWNCAST (PCAMatrix)
|
|
DOWNCAST (RandomRotationMatrix)
|
|
DOWNCAST (LinearTransform)
|
|
DOWNCAST (VectorTransform)
|
|
{
|
|
assert(false);
|
|
}
|
|
#ifdef SWIGLUA
|
|
SWIG_arg++;
|
|
#endif
|
|
}
|
|
|
|
|
|
// just to downcast pointers that come from elsewhere (eg. direct
|
|
// access to object fields)
|
|
%{
|
|
faiss::Index * downcast_index (faiss::Index *index)
|
|
{
|
|
return index;
|
|
}
|
|
faiss::VectorTransform * downcast_VectorTransform (faiss::VectorTransform *vt)
|
|
{
|
|
return vt;
|
|
}
|
|
%}
|
|
|
|
faiss::Index * downcast_index (faiss::Index *);
|
|
faiss::VectorTransform * downcast_VectorTransform (faiss::VectorTransform *vt);
|
|
|
|
|
|
%include "index_io.h"
|
|
|
|
%newobject index_factory;
|
|
|
|
%include "AutoTune.h"
|
|
|
|
|
|
#ifdef GPU_WRAPPER
|
|
|
|
%newobject index_gpu_to_cpu;
|
|
%newobject index_cpu_to_gpu;
|
|
%newobject index_cpu_to_gpu_multiple;
|
|
|
|
%include "gpu/GpuAutoTune.h"
|
|
|
|
#endif
|
|
|
|
|
|
|
|
|
|
|
|
|
|
/*******************************************************************
|
|
* Python specific: numpy array <-> C++ pointer interface
|
|
*******************************************************************/
|
|
|
|
#ifdef SWIGPYTHON
|
|
|
|
%{
|
|
PyObject *swig_ptr (PyObject *a)
|
|
{
|
|
if(!PyArray_Check(a)) {
|
|
PyErr_SetString(PyExc_ValueError, "input not a numpy array");
|
|
return NULL;
|
|
}
|
|
PyArrayObject *ao = (PyArrayObject *)a;
|
|
|
|
if(!PyArray_ISCONTIGUOUS(ao)) {
|
|
PyErr_SetString(PyExc_ValueError, "array is not C-contiguous");
|
|
return NULL;
|
|
}
|
|
void * data = PyArray_DATA(ao);
|
|
if(PyArray_TYPE(ao) == NPY_FLOAT32) {
|
|
return SWIG_NewPointerObj(data, SWIGTYPE_p_float, 0);
|
|
}
|
|
if(PyArray_TYPE(ao) == NPY_INT32) {
|
|
return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
|
|
}
|
|
if(PyArray_TYPE(ao) == NPY_UINT8) {
|
|
return SWIG_NewPointerObj(data, SWIGTYPE_p_uint8_t, 0);
|
|
}
|
|
if(PyArray_TYPE(ao) == NPY_UINT64) {
|
|
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_long, 0);
|
|
}
|
|
if(PyArray_TYPE(ao) == NPY_INT64) {
|
|
return SWIG_NewPointerObj(data, SWIGTYPE_p_long, 0);
|
|
}
|
|
PyErr_SetString(PyExc_ValueError, "did not recognize array type");
|
|
return NULL;
|
|
}
|
|
|
|
|
|
|
|
%}
|
|
|
|
|
|
%init %{
|
|
/* needed, else crash at runtime */
|
|
import_array();
|
|
%}
|
|
|
|
// return a pointer usable as input for functions that expect pointers
|
|
PyObject *swig_ptr (PyObject *a);
|
|
|
|
#endif
|
|
|
|
|
|
|
|
/*******************************************************************
|
|
* Lua specific: Torch tensor <-> C++ pointer interface
|
|
*******************************************************************/
|
|
|
|
#ifdef SWIGLUA
|
|
|
|
|
|
// provide a XXX_ptr function to convert Lua XXXTensor -> C++ XXX*
|
|
|
|
%define TYPE_CONVERSION(ctype, tensortype)
|
|
|
|
// typemap for the *_ptr_from_cdata function
|
|
%typemap(in) ctype** {
|
|
if(lua_type(L, $input) != 10) {
|
|
fprintf(stderr, "not cdata input\n");
|
|
SWIG_fail;
|
|
}
|
|
$1 = (ctype**)lua_topointer(L, $input);
|
|
}
|
|
|
|
|
|
// SWIG and C declaration for the *_ptr_from_cdata function
|
|
%{
|
|
ctype * ctype ## _ptr_from_cdata(ctype **x, long ofs) {
|
|
return *x + ofs;
|
|
}
|
|
%}
|
|
ctype * ctype ## _ptr_from_cdata(ctype **x, long ofs);
|
|
|
|
// the *_ptr function
|
|
%luacode {
|
|
|
|
function swigfaiss. ctype ## _ptr(tensor)
|
|
assert(tensor:type() == "torch." .. # tensortype, "need a " .. # tensortype)
|
|
assert(tensor:isContiguous(), "requires contiguous tensor")
|
|
return swigfaiss. ctype ## _ptr_from_cdata(
|
|
tensor:storage():data(),
|
|
tensor:storageOffset() - 1)
|
|
end
|
|
|
|
}
|
|
|
|
%enddef
|
|
|
|
TYPE_CONVERSION (int, IntTensor)
|
|
TYPE_CONVERSION (float, FloatTensor)
|
|
TYPE_CONVERSION (long, LongTensor)
|
|
TYPE_CONVERSION (uint64_t, LongTensor)
|
|
TYPE_CONVERSION (uint8_t, ByteTensor)
|
|
|
|
#endif
|
|
|
|
/*******************************************************************
|
|
* How should the template objects apprear in the scripting language?
|
|
*******************************************************************/
|
|
|
|
// answer: the same as the C++ typedefs, but we still have to redefine them
|
|
|
|
%template() faiss::CMin<float, long>;
|
|
%template() faiss::CMin<int, long>;
|
|
%template() faiss::CMax<float, long>;
|
|
%template() faiss::CMax<int, long>;
|
|
|
|
%template(float_minheap_array_t) faiss::HeapArray<faiss::CMin<float, long> >;
|
|
%template(int_minheap_array_t) faiss::HeapArray<faiss::CMin<int, long> >;
|
|
|
|
%template(float_maxheap_array_t) faiss::HeapArray<faiss::CMax<float, long> >;
|
|
%template(int_maxheap_array_t) faiss::HeapArray<faiss::CMax<int, long> >;
|
|
|
|
|
|
/*******************************************************************
|
|
* Expose a few basic functions
|
|
*******************************************************************/
|
|
|
|
|
|
void omp_set_num_threads (int num_threads);
|
|
int omp_get_max_threads ();
|
|
void *memcpy(void *dest, const void *src, size_t n);
|
|
|
|
|
|
|
|
/*******************************************************************
|
|
* Range search interface
|
|
*******************************************************************/
|
|
|
|
%ignore faiss::BufferList::Buffer;
|
|
%ignore faiss::RangeSearchPartialResult::QueryResult;
|
|
%ignore faiss::IDSelectorBatch::set;
|
|
%ignore faiss::IDSelectorBatch::bloom;
|
|
|
|
%include "AuxIndexStructures.h"
|
|
|
|
|
|
|
|
|
|
%{
|
|
// may be useful for lua code launched in background from shell
|
|
|
|
#include <signal.h>
|
|
void ignore_SIGTTIN() {
|
|
signal(SIGTTIN, SIG_IGN);
|
|
}
|
|
%}
|
|
|
|
void ignore_SIGTTIN();
|