Skip to content

Commit 847cde8

Browse files
Michael Norrisfacebook-github-bot
Michael Norris
authored and
facebook-github-bot
committed
Add example of how to build, link, and test an external SWIG module (facebookresearch#3922)
Summary: Pull Request resolved: facebookresearch#3922 We need to be able to build external modules into FAISS, but don't have an example yet. This diff shows what CMakeLists.txt changes need to happen to incorporate an external module. Reference: facebookresearch#3699 Reviewed By: mdouze Differential Revision: D63991471 fbshipit-source-id: 0c1cd25eabbffb01d2a7170d6725a0c4a13c5bf0
1 parent a99dbcd commit 847cde8

File tree

9 files changed

+301
-29
lines changed

9 files changed

+301
-29
lines changed

‎conda/faiss-gpu-raft/meta.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ outputs:
112112
- pytorch
113113
- pytorch-cuda {{ cuda_constraints }}
114114
commands:
115-
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
115+
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
116116
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
117117
- cp tests/common_faiss_tests.py faiss/gpu/test
118118
- python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*"

‎conda/faiss-gpu/meta.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ outputs:
101101
- pytorch
102102
- pytorch-cuda {{ cuda_constraints }}
103103
commands:
104-
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
104+
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
105105
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
106106
- cp tests/common_faiss_tests.py faiss/gpu/test
107107
- python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*"

‎conda/faiss/meta.yaml

+1-1
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ outputs:
8686
- scipy
8787
- pytorch
8888
commands:
89-
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
89+
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
9090
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
9191
- sh test_cpu_dispatch.sh # [linux64]
9292
files:

‎faiss/python/CMakeLists.txt

+29
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ configure_swigfaiss(swigfaiss.swig)
6161
configure_swigfaiss(swigfaiss_avx2.swig)
6262
configure_swigfaiss(swigfaiss_avx512.swig)
6363
configure_swigfaiss(swigfaiss_sve.swig)
64+
configure_swigfaiss(faiss_example_external_module.swig)
6465

6566
if(TARGET faiss)
6667
# Manually add headers as extra dependencies of swigfaiss.
@@ -74,6 +75,8 @@ if(TARGET faiss)
7475
"${faiss_SOURCE_DIR}/faiss/${h}")
7576
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
7677
"${faiss_SOURCE_DIR}/faiss/${h}")
78+
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
79+
"${faiss_SOURCE_DIR}/faiss/${h}")
7780
endforeach()
7881
if(FAISS_ENABLE_ROCM)
7982
foreach(h ${FAISS_GPU_HEADERS})
@@ -83,6 +86,8 @@ if(TARGET faiss)
8386
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
8487
list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
8588
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
89+
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
90+
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
8691
endforeach()
8792
else()
8893
foreach(h ${FAISS_GPU_HEADERS})
@@ -94,6 +99,8 @@ if(TARGET faiss)
9499
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
95100
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
96101
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
102+
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
103+
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
97104
endforeach()
98105
endif()
99106
else()
@@ -152,25 +159,37 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
152159
set_target_properties(swigfaiss_sve PROPERTIES EXCLUDE_FROM_ALL TRUE)
153160
endif()
154161

162+
set_property(SOURCE faiss_example_external_module.swig
163+
PROPERTY SWIG_MODULE_NAME faiss_example_external_module)
164+
swig_add_library(faiss_example_external_module
165+
TYPE SHARED
166+
LANGUAGE python
167+
SOURCES faiss_example_external_module.swig
168+
)
169+
set_property(TARGET faiss_example_external_module PROPERTY SWIG_COMPILE_OPTIONS -doxygen)
170+
155171
if(NOT WIN32)
156172
# NOTE: Python does not recognize the dylib extension.
157173
set_target_properties(swigfaiss PROPERTIES SUFFIX .so)
158174
set_target_properties(swigfaiss_avx2 PROPERTIES SUFFIX .so)
159175
set_target_properties(swigfaiss_avx512 PROPERTIES SUFFIX .so)
160176
set_target_properties(swigfaiss_sve PROPERTIES SUFFIX .so)
177+
set_target_properties(faiss_example_external_module PROPERTIES SUFFIX .so)
161178
else()
162179
# we need bigobj for the swig wrapper
163180
target_compile_options(swigfaiss PRIVATE /bigobj)
164181
target_compile_options(swigfaiss_avx2 PRIVATE /bigobj)
165182
target_compile_options(swigfaiss_avx512 PRIVATE /bigobj)
166183
target_compile_options(swigfaiss_sve PRIVATE /bigobj)
184+
target_compile_options(faiss_example_external_module PRIVATE /bigobj)
167185
endif()
168186

169187
if(FAISS_ENABLE_GPU)
170188
if(FAISS_ENABLE_ROCM)
171189
target_link_libraries(swigfaiss PRIVATE hip::host)
172190
target_link_libraries(swigfaiss_avx2 PRIVATE hip::host)
173191
target_link_libraries(swigfaiss_avx512 PRIVATE hip::host)
192+
target_link_libraries(faiss_example_external_module PRIVATE hip::host)
174193
else()
175194
find_package(CUDAToolkit REQUIRED)
176195
if(FAISS_ENABLE_RAFT)
@@ -221,12 +240,21 @@ target_link_libraries(swigfaiss_sve PRIVATE
221240
OpenMP::OpenMP_CXX
222241
)
223242

243+
target_link_libraries(faiss_example_external_module PRIVATE
244+
Python::Module
245+
Python::NumPy
246+
OpenMP::OpenMP_CXX
247+
swigfaiss
248+
faiss
249+
)
250+
224251
# Hack so that python_callbacks.h can be included as
225252
# `#include <faiss/python/python_callbacks.h>`.
226253
target_include_directories(swigfaiss PRIVATE ${PROJECT_SOURCE_DIR}/../..)
227254
target_include_directories(swigfaiss_avx2 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
228255
target_include_directories(swigfaiss_avx512 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
229256
target_include_directories(swigfaiss_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..)
257+
target_include_directories(faiss_example_external_module PRIVATE ${PROJECT_SOURCE_DIR}/../..)
230258

231259
find_package(Python REQUIRED
232260
COMPONENTS Development NumPy
@@ -252,6 +280,7 @@ target_link_libraries(swigfaiss PRIVATE faiss_python_callbacks)
252280
target_link_libraries(swigfaiss_avx2 PRIVATE faiss_python_callbacks)
253281
target_link_libraries(swigfaiss_avx512 PRIVATE faiss_python_callbacks)
254282
target_link_libraries(swigfaiss_sve PRIVATE faiss_python_callbacks)
283+
target_link_libraries(faiss_example_external_module PRIVATE faiss_python_callbacks)
255284

256285
configure_file(setup.py setup.py COPYONLY)
257286
configure_file(__init__.py __init__.py COPYONLY)

‎faiss/python/faiss_example_external_module.swig

+133
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
2+
%module faiss_example_external_module;
3+
4+
5+
// Put C++ includes here
6+
%{
7+
8+
#include <faiss/impl/FaissException.h>
9+
#include <faiss/impl/IDSelector.h>
10+
11+
%}
12+
13+
#pragma SWIG nowarn=322
14+
15+
typedef unsigned char uint8_t;
16+
typedef unsigned short uint16_t;
17+
typedef unsigned int uint32_t;
18+
19+
typedef signed char int8_t;
20+
typedef short int16_t;
21+
typedef int int32_t;
22+
23+
#ifdef SWIGWORDSIZE64
24+
typedef unsigned long uint64_t;
25+
typedef long int64_t;
26+
#else
27+
typedef unsigned long long uint64_t;
28+
typedef long long int64_t;
29+
#endif
30+
31+
typedef uint64_t size_t;
32+
33+
// This means: assume what's declared in these .h files is provided
34+
// by the Faiss module.
35+
%import(module="faiss") "faiss/MetricType.h"
36+
%import(module="faiss") "faiss/impl/IDSelector.h"
37+
38+
// functions to be parsed here
39+
40+
// This is important to release GIL and do Faiss exception handing
41+
%exception {
42+
Py_BEGIN_ALLOW_THREADS
43+
try {
44+
$action
45+
} catch(faiss::FaissException & e) {
46+
PyEval_RestoreThread(_save);
47+
48+
if (PyErr_Occurred()) {
49+
// some previous code already set the error type.
50+
} else {
51+
PyErr_SetString(PyExc_RuntimeError, e.what());
52+
}
53+
SWIG_fail;
54+
} catch(std::bad_alloc & ba) {
55+
PyEval_RestoreThread(_save);
56+
PyErr_SetString(PyExc_MemoryError, "std::bad_alloc");
57+
SWIG_fail;
58+
}
59+
Py_END_ALLOW_THREADS
60+
}
61+
62+
63+
// any class or function declared below will be made available
64+
// in the module.
65+
%inline %{
66+
67+
struct IDSelectorModulo : faiss::IDSelector {
68+
int mod;
69+
70+
IDSelectorModulo(int mod): mod(mod) {}
71+
72+
bool is_member(faiss::idx_t id) const {
73+
return id % mod == 0;
74+
}
75+
76+
~IDSelectorModulo() override {}
77+
};
78+
79+
faiss::idx_t sum_of_idx(size_t n, const faiss::idx_t *tab) {
80+
faiss::idx_t sum = 0;
81+
for(size_t i = 0; i < n; i++) {
82+
sum += tab[i];
83+
}
84+
return sum;
85+
}
86+
87+
float sum_of_float32(size_t n, const float *tab) {
88+
float sum = 0;
89+
for(size_t i = 0; i < n; i++) {
90+
sum += tab[i];
91+
}
92+
return sum;
93+
}
94+
95+
double sum_of_float64(size_t n, const double *tab) {
96+
double sum = 0;
97+
for(size_t i = 0; i < n; i++) {
98+
sum += tab[i];
99+
}
100+
return sum;
101+
}
102+
103+
%}
104+
105+
/**********************************************
106+
* To test if passing a swig_ptr on all array types works
107+
**********************************************/
108+
109+
%define SUM_OF_TYPE(ty)
110+
111+
%inline %{
112+
113+
ty##_t sum_of_##ty (size_t n, const ty##_t * tab) {
114+
ty##_t sum = 0;
115+
for(size_t i = 0; i < n; i++) {
116+
sum += tab[i];
117+
}
118+
return sum;
119+
}
120+
121+
%}
122+
123+
%enddef
124+
125+
SUM_OF_TYPE(uint8);
126+
SUM_OF_TYPE(uint16);
127+
SUM_OF_TYPE(uint32);
128+
SUM_OF_TYPE(uint64);
129+
130+
SUM_OF_TYPE(int8);
131+
SUM_OF_TYPE(int16);
132+
SUM_OF_TYPE(int32);
133+
SUM_OF_TYPE(int64);

‎faiss/python/setup.py

+38-20
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44
# LICENSE file in the root directory of this source tree.
55

66
from __future__ import print_function
7-
from setuptools import setup, find_packages
7+
88
import os
9-
import shutil
109
import platform
10+
import shutil
11+
12+
from setuptools import find_packages, setup
1113

1214
# make the faiss python package dir
1315
shutil.rmtree("faiss", ignore_errors=True)
@@ -20,25 +22,32 @@
2022
shutil.copyfile("extra_wrappers.py", "faiss/extra_wrappers.py")
2123
shutil.copyfile("array_conversions.py", "faiss/array_conversions.py")
2224

23-
ext = ".pyd" if platform.system() == 'Windows' else ".so"
24-
prefix = "Release/" * (platform.system() == 'Windows')
25+
ext = ".pyd" if platform.system() == "Windows" else ".so"
26+
prefix = "Release/" * (platform.system() == "Windows")
2527

2628
swigfaiss_generic_lib = f"{prefix}_swigfaiss{ext}"
2729
swigfaiss_avx2_lib = f"{prefix}_swigfaiss_avx2{ext}"
2830
swigfaiss_avx512_lib = f"{prefix}_swigfaiss_avx512{ext}"
2931
callbacks_lib = f"{prefix}libfaiss_python_callbacks{ext}"
3032
swigfaiss_sve_lib = f"{prefix}_swigfaiss_sve{ext}"
33+
faiss_example_external_module_lib = f"_faiss_example_external_module{ext}"
3134

3235
found_swigfaiss_generic = os.path.exists(swigfaiss_generic_lib)
3336
found_swigfaiss_avx2 = os.path.exists(swigfaiss_avx2_lib)
3437
found_swigfaiss_avx512 = os.path.exists(swigfaiss_avx512_lib)
3538
found_callbacks = os.path.exists(callbacks_lib)
3639
found_swigfaiss_sve = os.path.exists(swigfaiss_sve_lib)
40+
found_faiss_example_external_module_lib = os.path.exists(
41+
faiss_example_external_module_lib
42+
)
3743

38-
assert (found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512 or found_swigfaiss_sve), \
39-
f"Could not find {swigfaiss_generic_lib} or " \
40-
f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib} or {swigfaiss_sve_lib}. " \
44+
assert (
45+
found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512 or found_swigfaiss_sve or found_faiss_example_external_module_lib
46+
), (
47+
f"Could not find {swigfaiss_generic_lib} or "
48+
f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib} or {swigfaiss_sve_lib} or {faiss_example_external_module_lib}. "
4149
f"Faiss may not be compiled yet."
50+
)
4251

4352
if found_swigfaiss_generic:
4453
print(f"Copying {swigfaiss_generic_lib}")
@@ -64,7 +73,17 @@
6473
shutil.copyfile("swigfaiss_sve.py", "faiss/swigfaiss_sve.py")
6574
shutil.copyfile(swigfaiss_sve_lib, f"faiss/_swigfaiss_sve{ext}")
6675

67-
long_description="""
76+
if found_faiss_example_external_module_lib:
77+
print(f"Copying {faiss_example_external_module_lib}")
78+
shutil.copyfile(
79+
"faiss_example_external_module.py", "faiss/faiss_example_external_module.py"
80+
)
81+
shutil.copyfile(
82+
faiss_example_external_module_lib,
83+
f"faiss/_faiss_example_external_module{ext}",
84+
)
85+
86+
long_description = """
6887
Faiss is a library for efficient similarity search and clustering of dense
6988
vectors. It contains algorithms that search in sets of vectors of any size,
7089
up to ones that possibly do not fit in RAM. It also contains supporting
@@ -73,20 +92,19 @@
7392
are implemented on the GPU. It is developed by Facebook AI Research.
7493
"""
7594
setup(
76-
name='faiss',
77-
version='1.9.0',
78-
description='A library for efficient similarity search and clustering of dense vectors',
95+
name="faiss",
96+
version="1.9.0",
97+
description="A library for efficient similarity search and clustering of dense vectors",
7998
long_description=long_description,
80-
url='https://github.com/facebookresearch/faiss',
81-
author='Matthijs Douze, Jeff Johnson, Herve Jegou, Lucas Hosseini',
82-
author_email='matthijs@meta.com',
83-
license='MIT',
84-
keywords='search nearest neighbors',
85-
86-
install_requires=['numpy', 'packaging'],
87-
packages=['faiss', 'faiss.contrib', 'faiss.contrib.torch'],
99+
url="https://github.com/facebookresearch/faiss",
100+
author="Matthijs Douze, Jeff Johnson, Herve Jegou, Lucas Hosseini",
101+
author_email="matthijs@meta.com",
102+
license="MIT",
103+
keywords="search nearest neighbors",
104+
install_requires=["numpy", "packaging"],
105+
packages=["faiss", "faiss.contrib", "faiss.contrib.torch"],
88106
package_data={
89-
'faiss': ['*.so', '*.pyd'],
107+
"faiss": ["*.so", "*.pyd"],
90108
},
91109
zip_safe=False,
92110
)

0 commit comments

Comments
 (0)