Skip to content

Commit

Permalink
Add example of how to build, link, and test an external SWIG module (#…
Browse files Browse the repository at this point in the history
…3922)

Summary:
Pull Request resolved: #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: #3699

Reviewed By: mdouze

Differential Revision: D63991471

fbshipit-source-id: 0c1cd25eabbffb01d2a7170d6725a0c4a13c5bf0
  • Loading branch information
Michael Norris authored and facebook-github-bot committed Oct 11, 2024
1 parent a99dbcd commit 847cde8
Show file tree
Hide file tree
Showing 9 changed files with 301 additions and 29 deletions.
2 changes: 1 addition & 1 deletion conda/faiss-gpu-raft/meta.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ outputs:
- pytorch
- pytorch-cuda {{ cuda_constraints }}
commands:
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
- cp tests/common_faiss_tests.py faiss/gpu/test
- python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*"
Expand Down
2 changes: 1 addition & 1 deletion conda/faiss-gpu/meta.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ outputs:
- pytorch
- pytorch-cuda {{ cuda_constraints }}
commands:
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
- cp tests/common_faiss_tests.py faiss/gpu/test
- python -X faulthandler -m unittest discover -v -s faiss/gpu/test/ -p "test_*"
Expand Down
2 changes: 1 addition & 1 deletion conda/faiss/meta.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ outputs:
- scipy
- pytorch
commands:
- python -X faulthandler -m unittest discover -v -s tests/ -p "test_*"
- python -X faulthandler -m unittest discover -v -s tests/ -p "(?!.*test_external_module\.py)test_.*py"
- python -X faulthandler -m unittest discover -v -s tests/ -p "torch_*"
- sh test_cpu_dispatch.sh # [linux64]
files:
Expand Down
29 changes: 29 additions & 0 deletions faiss/python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ configure_swigfaiss(swigfaiss.swig)
configure_swigfaiss(swigfaiss_avx2.swig)
configure_swigfaiss(swigfaiss_avx512.swig)
configure_swigfaiss(swigfaiss_sve.swig)
configure_swigfaiss(faiss_example_external_module.swig)

if(TARGET faiss)
# Manually add headers as extra dependencies of swigfaiss.
Expand All @@ -74,6 +75,8 @@ if(TARGET faiss)
"${faiss_SOURCE_DIR}/faiss/${h}")
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/${h}")
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/${h}")
endforeach()
if(FAISS_ENABLE_ROCM)
foreach(h ${FAISS_GPU_HEADERS})
Expand All @@ -83,6 +86,8 @@ if(TARGET faiss)
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
list(APPEND SWIG_MODULE_swigfaiss_avx512_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/gpu-rocm/${h}")
endforeach()
else()
foreach(h ${FAISS_GPU_HEADERS})
Expand All @@ -94,6 +99,8 @@ if(TARGET faiss)
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
list(APPEND SWIG_MODULE_swigfaiss_sve_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
list(APPEND SWIG_MODULE_faiss_example_external_module_EXTRA_DEPS
"${faiss_SOURCE_DIR}/faiss/gpu/${h}")
endforeach()
endif()
else()
Expand Down Expand Up @@ -152,25 +159,37 @@ if(NOT FAISS_OPT_LEVEL STREQUAL "sve")
set_target_properties(swigfaiss_sve PROPERTIES EXCLUDE_FROM_ALL TRUE)
endif()

set_property(SOURCE faiss_example_external_module.swig
PROPERTY SWIG_MODULE_NAME faiss_example_external_module)
swig_add_library(faiss_example_external_module
TYPE SHARED
LANGUAGE python
SOURCES faiss_example_external_module.swig
)
set_property(TARGET faiss_example_external_module PROPERTY SWIG_COMPILE_OPTIONS -doxygen)

if(NOT WIN32)
# NOTE: Python does not recognize the dylib extension.
set_target_properties(swigfaiss PROPERTIES SUFFIX .so)
set_target_properties(swigfaiss_avx2 PROPERTIES SUFFIX .so)
set_target_properties(swigfaiss_avx512 PROPERTIES SUFFIX .so)
set_target_properties(swigfaiss_sve PROPERTIES SUFFIX .so)
set_target_properties(faiss_example_external_module PROPERTIES SUFFIX .so)
else()
# we need bigobj for the swig wrapper
target_compile_options(swigfaiss PRIVATE /bigobj)
target_compile_options(swigfaiss_avx2 PRIVATE /bigobj)
target_compile_options(swigfaiss_avx512 PRIVATE /bigobj)
target_compile_options(swigfaiss_sve PRIVATE /bigobj)
target_compile_options(faiss_example_external_module PRIVATE /bigobj)
endif()

if(FAISS_ENABLE_GPU)
if(FAISS_ENABLE_ROCM)
target_link_libraries(swigfaiss PRIVATE hip::host)
target_link_libraries(swigfaiss_avx2 PRIVATE hip::host)
target_link_libraries(swigfaiss_avx512 PRIVATE hip::host)
target_link_libraries(faiss_example_external_module PRIVATE hip::host)
else()
find_package(CUDAToolkit REQUIRED)
if(FAISS_ENABLE_RAFT)
Expand Down Expand Up @@ -221,12 +240,21 @@ target_link_libraries(swigfaiss_sve PRIVATE
OpenMP::OpenMP_CXX
)

target_link_libraries(faiss_example_external_module PRIVATE
Python::Module
Python::NumPy
OpenMP::OpenMP_CXX
swigfaiss
faiss
)

# Hack so that python_callbacks.h can be included as
# `#include <faiss/python/python_callbacks.h>`.
target_include_directories(swigfaiss PRIVATE ${PROJECT_SOURCE_DIR}/../..)
target_include_directories(swigfaiss_avx2 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
target_include_directories(swigfaiss_avx512 PRIVATE ${PROJECT_SOURCE_DIR}/../..)
target_include_directories(swigfaiss_sve PRIVATE ${PROJECT_SOURCE_DIR}/../..)
target_include_directories(faiss_example_external_module PRIVATE ${PROJECT_SOURCE_DIR}/../..)

find_package(Python REQUIRED
COMPONENTS Development NumPy
Expand All @@ -252,6 +280,7 @@ target_link_libraries(swigfaiss PRIVATE faiss_python_callbacks)
target_link_libraries(swigfaiss_avx2 PRIVATE faiss_python_callbacks)
target_link_libraries(swigfaiss_avx512 PRIVATE faiss_python_callbacks)
target_link_libraries(swigfaiss_sve PRIVATE faiss_python_callbacks)
target_link_libraries(faiss_example_external_module PRIVATE faiss_python_callbacks)

configure_file(setup.py setup.py COPYONLY)
configure_file(__init__.py __init__.py COPYONLY)
Expand Down
133 changes: 133 additions & 0 deletions faiss/python/faiss_example_external_module.swig
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@

%module faiss_example_external_module;


// Put C++ includes here
%{

#include <faiss/impl/FaissException.h>
#include <faiss/impl/IDSelector.h>

%}

#pragma SWIG nowarn=322

typedef unsigned char uint8_t;
typedef unsigned short uint16_t;
typedef unsigned int uint32_t;

typedef signed char int8_t;
typedef short int16_t;
typedef int int32_t;

#ifdef SWIGWORDSIZE64
typedef unsigned long uint64_t;
typedef long int64_t;
#else
typedef unsigned long long uint64_t;
typedef long long int64_t;
#endif

typedef uint64_t size_t;

// This means: assume what's declared in these .h files is provided
// by the Faiss module.
%import(module="faiss") "faiss/MetricType.h"
%import(module="faiss") "faiss/impl/IDSelector.h"

// functions to be parsed here

// This is important to release GIL and do Faiss exception handing
%exception {
Py_BEGIN_ALLOW_THREADS
try {
$action
} catch(faiss::FaissException & e) {
PyEval_RestoreThread(_save);

if (PyErr_Occurred()) {
// some previous code already set the error type.
} else {
PyErr_SetString(PyExc_RuntimeError, e.what());
}
SWIG_fail;
} catch(std::bad_alloc & ba) {
PyEval_RestoreThread(_save);
PyErr_SetString(PyExc_MemoryError, "std::bad_alloc");
SWIG_fail;
}
Py_END_ALLOW_THREADS
}


// any class or function declared below will be made available
// in the module.
%inline %{

struct IDSelectorModulo : faiss::IDSelector {
int mod;

IDSelectorModulo(int mod): mod(mod) {}

bool is_member(faiss::idx_t id) const {
return id % mod == 0;
}

~IDSelectorModulo() override {}
};

faiss::idx_t sum_of_idx(size_t n, const faiss::idx_t *tab) {
faiss::idx_t sum = 0;
for(size_t i = 0; i < n; i++) {
sum += tab[i];
}
return sum;
}

float sum_of_float32(size_t n, const float *tab) {
float sum = 0;
for(size_t i = 0; i < n; i++) {
sum += tab[i];
}
return sum;
}

double sum_of_float64(size_t n, const double *tab) {
double sum = 0;
for(size_t i = 0; i < n; i++) {
sum += tab[i];
}
return sum;
}

%}

/**********************************************
* To test if passing a swig_ptr on all array types works
**********************************************/

%define SUM_OF_TYPE(ty)

%inline %{

ty##_t sum_of_##ty (size_t n, const ty##_t * tab) {
ty##_t sum = 0;
for(size_t i = 0; i < n; i++) {
sum += tab[i];
}
return sum;
}

%}

%enddef

SUM_OF_TYPE(uint8);
SUM_OF_TYPE(uint16);
SUM_OF_TYPE(uint32);
SUM_OF_TYPE(uint64);

SUM_OF_TYPE(int8);
SUM_OF_TYPE(int16);
SUM_OF_TYPE(int32);
SUM_OF_TYPE(int64);
58 changes: 38 additions & 20 deletions faiss/python/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@
# LICENSE file in the root directory of this source tree.

from __future__ import print_function
from setuptools import setup, find_packages

import os
import shutil
import platform
import shutil

from setuptools import find_packages, setup

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

ext = ".pyd" if platform.system() == 'Windows' else ".so"
prefix = "Release/" * (platform.system() == 'Windows')
ext = ".pyd" if platform.system() == "Windows" else ".so"
prefix = "Release/" * (platform.system() == "Windows")

swigfaiss_generic_lib = f"{prefix}_swigfaiss{ext}"
swigfaiss_avx2_lib = f"{prefix}_swigfaiss_avx2{ext}"
swigfaiss_avx512_lib = f"{prefix}_swigfaiss_avx512{ext}"
callbacks_lib = f"{prefix}libfaiss_python_callbacks{ext}"
swigfaiss_sve_lib = f"{prefix}_swigfaiss_sve{ext}"
faiss_example_external_module_lib = f"_faiss_example_external_module{ext}"

found_swigfaiss_generic = os.path.exists(swigfaiss_generic_lib)
found_swigfaiss_avx2 = os.path.exists(swigfaiss_avx2_lib)
found_swigfaiss_avx512 = os.path.exists(swigfaiss_avx512_lib)
found_callbacks = os.path.exists(callbacks_lib)
found_swigfaiss_sve = os.path.exists(swigfaiss_sve_lib)
found_faiss_example_external_module_lib = os.path.exists(
faiss_example_external_module_lib
)

assert (found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512 or found_swigfaiss_sve), \
f"Could not find {swigfaiss_generic_lib} or " \
f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib} or {swigfaiss_sve_lib}. " \
assert (
found_swigfaiss_generic or found_swigfaiss_avx2 or found_swigfaiss_avx512 or found_swigfaiss_sve or found_faiss_example_external_module_lib
), (
f"Could not find {swigfaiss_generic_lib} or "
f"{swigfaiss_avx2_lib} or {swigfaiss_avx512_lib} or {swigfaiss_sve_lib} or {faiss_example_external_module_lib}. "
f"Faiss may not be compiled yet."
)

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

long_description="""
if found_faiss_example_external_module_lib:
print(f"Copying {faiss_example_external_module_lib}")
shutil.copyfile(
"faiss_example_external_module.py", "faiss/faiss_example_external_module.py"
)
shutil.copyfile(
faiss_example_external_module_lib,
f"faiss/_faiss_example_external_module{ext}",
)

long_description = """
Faiss is a library for efficient similarity search and clustering of dense
vectors. It contains algorithms that search in sets of vectors of any size,
up to ones that possibly do not fit in RAM. It also contains supporting
Expand All @@ -73,20 +92,19 @@
are implemented on the GPU. It is developed by Facebook AI Research.
"""
setup(
name='faiss',
version='1.9.0',
description='A library for efficient similarity search and clustering of dense vectors',
name="faiss",
version="1.9.0",
description="A library for efficient similarity search and clustering of dense vectors",
long_description=long_description,
url='https://github.com/facebookresearch/faiss',
author='Matthijs Douze, Jeff Johnson, Herve Jegou, Lucas Hosseini',
author_email='[email protected]',
license='MIT',
keywords='search nearest neighbors',

install_requires=['numpy', 'packaging'],
packages=['faiss', 'faiss.contrib', 'faiss.contrib.torch'],
url="https://github.com/facebookresearch/faiss",
author="Matthijs Douze, Jeff Johnson, Herve Jegou, Lucas Hosseini",
author_email="[email protected]",
license="MIT",
keywords="search nearest neighbors",
install_requires=["numpy", "packaging"],
packages=["faiss", "faiss.contrib", "faiss.contrib.torch"],
package_data={
'faiss': ['*.so', '*.pyd'],
"faiss": ["*.so", "*.pyd"],
},
zip_safe=False,
)
Loading

0 comments on commit 847cde8

Please sign in to comment.