make sure swig_ptr and rev_swig_ptr work on all primitive types (#1382)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1382

The array types supported for swig_ptr were not complete. This diff fixes that.

Reviewed By: wickedfoo

Differential Revision: D23411297

fbshipit-source-id: d94249b140aeb8c8179d9da3fcbc97eb034eac91
This commit is contained in:
Matthijs Douze
2020-08-28 23:43:52 -07:00
committed by Facebook GitHub Bot
parent 8dfd51d3a7
commit c97f890651
2 changed files with 43 additions and 5 deletions
+23 -5
View File
@@ -210,6 +210,7 @@ namespace std {
%include <std_pair.i>
%include <std_map.i>
// primitive array types
%template(FloatVector) std::vector<float>;
%template(DoubleVector) std::vector<double>;
%template(ByteVector) std::vector<uint8_t>;
@@ -220,6 +221,10 @@ namespace std {
%template(LongVector) std::vector<long>;
%template(LongLongVector) std::vector<long long>;
%template(IntVector) std::vector<int>;
%template(UInt32Vector) std::vector<uint32_t>;
%template(Int16Vector) std::vector<int16_t>;
%template(UInt16Vector) std::vector<uint16_t>;
%template(FloatVectorVector) std::vector<std::vector<float> >;
%template(ByteVectorVector) std::vector<std::vector<unsigned char> >;
%template(LongVectorVector) std::vector<std::vector<long> >;
@@ -799,15 +804,24 @@ PyObject *swig_ptr (PyObject *a)
if(PyArray_TYPE(ao) == NPY_FLOAT64) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_double, 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_unsigned_char, 0);
}
if(PyArray_TYPE(ao) == NPY_INT8) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_char, 0);
}
if(PyArray_TYPE(ao) == NPY_UINT16) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_short, 0);
}
if(PyArray_TYPE(ao) == NPY_INT16) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_short, 0);
}
if(PyArray_TYPE(ao) == NPY_UINT32) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_int, 0);
}
if(PyArray_TYPE(ao) == NPY_INT32) {
return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
}
if(PyArray_TYPE(ao) == NPY_UINT64) {
#ifdef SWIGWORDSIZE64
return SWIG_NewPointerObj(data, SWIGTYPE_p_unsigned_long, 0);
@@ -870,8 +884,13 @@ PyObject * rev_swig_ptr(ctype *src, size_t size);
%enddef
REV_SWIG_PTR(float, NPY_FLOAT32);
REV_SWIG_PTR(int, NPY_INT32);
REV_SWIG_PTR(double, NPY_FLOAT64);
REV_SWIG_PTR(unsigned char, NPY_UINT8);
REV_SWIG_PTR(char, NPY_INT8);
REV_SWIG_PTR(unsigned short, NPY_UINT16);
REV_SWIG_PTR(short, NPY_INT16);
REV_SWIG_PTR(int, NPY_INT32);
REV_SWIG_PTR(unsigned int, NPY_UINT32);
REV_SWIG_PTR(int64_t, NPY_INT64);
REV_SWIG_PTR(uint64_t, NPY_UINT64);
@@ -1038,4 +1057,3 @@ struct MapLong2Long {
%}
// End of file...
+20
View File
@@ -628,5 +628,25 @@ class TestSWIGWrap(unittest.TestCase):
assert isinstance(index2, faiss.IndexRefineFlat)
def do_test_array_type(self, dtype):
""" tests swig_ptr and rev_swig_ptr for this type of array """
a = np.arange(12).astype(dtype)
ptr = faiss.swig_ptr(a)
print(ptr)
a2 = faiss.rev_swig_ptr(ptr, 12)
np.testing.assert_array_equal(a, a2)
def test_all_array_types(self):
self.do_test_array_type('float32')
self.do_test_array_type('float64')
self.do_test_array_type('int8')
self.do_test_array_type('uint8')
self.do_test_array_type('int16')
self.do_test_array_type('uint16')
self.do_test_array_type('int32')
self.do_test_array_type('uint32')
self.do_test_array_type('int64')
self.do_test_array_type('uint64')
if __name__ == '__main__':
unittest.main()