swigfaiss.swig

/**
 * 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();