Skip to content

Commit 6f83aea

Browse files
jmazanec15jmazanec15
committed
Add faiss support in jni (#28)
Signed-off-by: Jack Mazanec <[email protected]>
1 parent d7c55e9 commit 6f83aea

32 files changed

+3969
-1091
lines changed

‎.gitmodules

+3
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
11
[submodule "jni/external/nmslib"]
22
path = jni/external/nmslib
33
url = https://github.com/nmslib/nmslib.git
4+
[submodule "jni/external/faiss"]
5+
path = jni/external/faiss
6+
url = https://github.com/facebookresearch/faiss.git

‎jni/CMakeLists.txt

+144-40
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,11 @@
2525

2626
cmake_minimum_required(VERSION 2.8)
2727

28-
project(KNNIndexV2_0_11)
28+
project(KNNPlugin_JNI)
29+
30+
# ---------------------------------- SETUP ----------------------------------
31+
# Target library to be compiled
32+
set(TARGET_LIB OpensearchKNN)
2933

3034
# Corner case. For CMake 2.8, there is no option to specify set(CMAKE_CXX_STANDARD 11). Instead, the flag manually needs
3135
# to be set.
@@ -36,15 +40,32 @@ else()
3640
set(CMAKE_CXX_STANDARD_REQUIRED True)
3741
endif()
3842

39-
# Target Library to be built
40-
set(KNN_INDEX KNNIndexV2_0_11)
41-
set(KNN_PACKAGE_NAME opensearch-knnlib)
43+
# Set OS specific variables
44+
if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
45+
set(CMAKE_MACOSX_RPATH 1)
46+
set(JVM_OS_TYPE darwin)
47+
set(LIB_EXT .jnilib)
48+
elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux)
49+
set(JVM_OS_TYPE linux)
50+
set(LIB_EXT .so)
51+
else()
52+
message(FATAL_ERROR "Unable to run on system: ${CMAKE_SYSTEM_NAME}")
53+
endif()
4254

4355
if(NOT KNN_PLUGIN_VERSION)
4456
set(KNN_PLUGIN_VERSION "1.2.0.0")
4557
endif()
4658

47-
# Check if similarity search exists
59+
# Set architecture specific variables
60+
if (${CMAKE_SYSTEM_PROCESSOR} STREQUAL aarch64)
61+
set(MACH_ARCH arm64)
62+
elseif(${CMAKE_SYSTEM_PROCESSOR} STREQUAL x86_64)
63+
set(MACH_ARCH x64)
64+
endif()
65+
# ----------------------------------------------------------------------------
66+
67+
# ---------------------------------- NMSLIB ----------------------------------
68+
# Check if nmslib exists
4869
find_path(NMS_REPO_DIR NAMES similarity_search PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib)
4970

5071
# If not, pull the updated submodule
@@ -53,67 +74,150 @@ if (NOT EXISTS ${NMS_REPO_DIR})
5374
execute_process(COMMAND git submodule update --init -- external/nmslib WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
5475
endif ()
5576

56-
# Add the subdirectory so it is possible to use its targets
57-
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search EXCLUDE_FROM_ALL)
77+
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search)
78+
# ---------------------------------------------------------------------------
79+
80+
# ---------------------------------- FAISS ----------------------------------
81+
# Avoid building faiss tests
82+
set(BUILD_TESTING OFF)
5883

59-
# Set OS specific variables
6084
if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
61-
set(CMAKE_MACOSX_RPATH 1)
62-
set(JVM_OS_TYPE darwin)
63-
set(LIB_EXT .jnilib)
64-
elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux)
65-
set(JVM_OS_TYPE linux)
66-
set(LIB_EXT .so)
67-
else()
68-
message( FATAL_ERROR "Unable to run on system: ${CMAKE_SYSTEM_NAME}")
85+
if(CMAKE_C_COMPILER_ID MATCHES "Clang\$")
86+
set(OpenMP_C_FLAGS "-Xpreprocessor -fopenmp")
87+
set(OpenMP_C_LIB_NAMES "omp")
88+
set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
89+
endif()
90+
91+
if(CMAKE_CXX_COMPILER_ID MATCHES "Clang\$")
92+
set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp -I/usr/local/opt/libomp/include")
93+
set(OpenMP_CXX_LIB_NAMES "omp")
94+
set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
95+
endif()
6996
endif()
7097

71-
# Compile the library
72-
add_library(${KNN_INDEX} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_index_v2011_KNNIndex.cpp)
73-
target_link_libraries(${KNN_INDEX} NonMetricSpaceLib)
74-
target_include_directories(${KNN_INDEX} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include)
98+
find_package(OpenMP REQUIRED)
99+
find_package(ZLIB REQUIRED)
100+
find_package(BLAS REQUIRED)
101+
find_package(LAPACK REQUIRED)
75102

76-
set_target_properties(${KNN_INDEX} PROPERTIES SUFFIX ${LIB_EXT})
77-
set_target_properties(${KNN_INDEX} PROPERTIES POSITION_INDEPENDENT_CODE ON)
78-
set_target_properties(${KNN_INDEX} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
103+
# Check if faiss exists
104+
find_path(FAISS_REPO_DIR NAMES faiss PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss)
79105

106+
# If not, pull the updated submodule
107+
if (NOT EXISTS ${FAISS_REPO_DIR})
108+
message(STATUS "Could not find faiss. Pulling updated submodule.")
109+
execute_process(COMMAND git submodule update --init -- external/faiss WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
110+
endif ()
111+
112+
set(FAISS_ENABLE_GPU OFF)
113+
set(FAISS_ENABLE_PYTHON OFF)
114+
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/faiss EXCLUDE_FROM_ALL)
115+
# ---------------------------------------------------------------------------
116+
117+
# ------------------------------ Lib Compiling ------------------------------
118+
add_library(${TARGET_LIB} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_index_JNIService.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/jni_util.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/nmslib_wrapper.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/faiss_wrapper.cpp)
119+
target_link_libraries(${TARGET_LIB} faiss NonMetricSpaceLib ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} OpenMP::OpenMP_CXX)
120+
target_include_directories(${TARGET_LIB} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include)
121+
122+
set_target_properties(${TARGET_LIB} PROPERTIES SUFFIX ${LIB_EXT})
123+
set_target_properties(${TARGET_LIB} PROPERTIES POSITION_INDEPENDENT_CODE ON)
124+
set_target_properties(${TARGET_LIB} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
125+
# ---------------------------------------------------------------------------
126+
127+
# --------------------------------- TESTS -----------------------------------
128+
# Reference - https://crascit.com/2015/07/25/cmake-gtest/
129+
configure_file(CMakeLists.txt.in googletest-download/CMakeLists.txt)
130+
execute_process(COMMAND "${CMAKE_COMMAND}" -G "${CMAKE_GENERATOR}" .
131+
WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
132+
)
133+
execute_process(COMMAND "${CMAKE_COMMAND}" --build .
134+
WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/googletest-download"
135+
)
136+
set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
137+
138+
add_subdirectory("${CMAKE_BINARY_DIR}/googletest-src"
139+
"${CMAKE_BINARY_DIR}/googletest-build"
140+
)
141+
add_executable(
142+
jni_test
143+
tests/faiss_wrapper_test.cpp
144+
tests/nmslib_wrapper_test.cpp
145+
tests/test_util.cpp)
146+
147+
target_link_libraries(
148+
jni_test
149+
gtest_main
150+
gmock_main
151+
faiss
152+
NonMetricSpaceLib
153+
${BLAS_LIBRARIES}
154+
${LAPACK_LIBRARIES}
155+
OpenMP::OpenMP_CXX
156+
${TARGET_LIB}
157+
)
158+
159+
target_include_directories(jni_test PRIVATE
160+
${CMAKE_CURRENT_SOURCE_DIR}/tests
161+
${CMAKE_CURRENT_SOURCE_DIR}/include
162+
$ENV{JAVA_HOME}/include
163+
$ENV{JAVA_HOME}/include/${JVM_OS_TYPE}
164+
${CMAKE_CURRENT_SOURCE_DIR}/external/faiss
165+
${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include
166+
${gtest_SOURCE_DIR}/include
167+
${gmock_SOURCE_DIR}/include)
168+
169+
170+
set_target_properties(jni_test PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/bin)
171+
172+
# ---------------------------------------------------------------------------
173+
174+
# -------------------------------- INSTALL ----------------------------------
80175
# Installation rules for shared library
81-
install(TARGETS ${KNN_INDEX}
176+
install(TARGETS ${TARGET_LIB}
82177
LIBRARY DESTINATION lib
83178
COMPONENT library)
84179

85-
# CPack section to build artifacts
86-
if (${CMAKE_SYSTEM_PROCESSOR} STREQUAL aarch64)
87-
set(MACH_ARCH arm64)
88-
elseif(${CMAKE_SYSTEM_PROCESSOR} STREQUAL x86_64)
89-
set(MACH_ARCH x64)
90-
endif()
91-
92180
set(KNN_MAINTAINER "OpenSearch Team <[email protected]>")
93181
set(OPENSEARCH_DOWNLOAD_URL "https://opensearch.org/downloads.html")
94182
set(CPACK_PACKAGE_NAME ${KNN_PACKAGE_NAME})
95183
set(CPACK_PACKAGE_VERSION ${KNN_PLUGIN_VERSION})
96184
set(CMAKE_INSTALL_PREFIX /usr)
97185
set(CPACK_GENERATOR "RPM;DEB")
98-
SET(CPACK_OUTPUT_FILE_PREFIX packages)
186+
set(CPACK_OUTPUT_FILE_PREFIX packages)
99187
set(CPACK_PACKAGE_RELEASE 1)
100-
set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "KNN JNI library built off of NMSLIB for OpenSearch.")
101188
set(CPACK_PACKAGE_VENDOR "Amazon")
102189
set(CPACK_PACKAGE_CONTACT "Maintainer: ${KNN_MAINTAINER}")
103190
set(CPACK_PACKAGING_INSTALL_PREFIX ${CMAKE_INSTALL_PREFIX})
104-
set(CPACK_PACKAGE_FILE_NAME "${CPACK_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${JVM_OS_TYPE}-${MACH_ARCH}")
191+
set(CPACK_COMPONENTS_GROUPING IGNORE)
192+
get_cmake_property(CPACK_COMPONENTS_ALL COMPONENTS)
193+
list(REMOVE_ITEM CPACK_COMPONENTS_ALL "Unspecified")
105194

106-
# RPM Specific variables
107-
set(CPACK_RPM_PACKAGE_RELEASE ${CPACK_PACKAGE_RELEASE})
108-
set(CPACK_RPM_PACKAGE_URL ${OPENSEARCH_DOWNLOAD_URL})
109-
set(CPACK_RPM_PACKAGE_DESCRIPTION "OpenSearch KNN JNI Library")
195+
# Component variable
196+
set(KNN_PACKAGE_NAME opensearch-knnlib)
197+
set(KNN_PACKAGE_DESCRIPTION "KNN JNI library built off of nmslib and faiss for OpenSearch")
198+
199+
# RPM
110200
set(CPACK_RPM_PACKAGE_LICENSE "ASL-2.0")
201+
set(CPACK_RPM_COMPONENT_INSTALL ON)
202+
set(CPACK_RPM_PACKAGE_URL ${OPENSEARCH_DOWNLOAD_URL})
203+
set(CPACK_RPM_PACKAGE_RELEASE ${CPACK_PACKAGE_RELEASE})
111204

112-
# DEB Specific variables
205+
set(CPACK_RPM_PACKAGE_NAME ${KNN_PACKAGE_NAME})
206+
set(CPACK_RPM_FILE_NAME "${CPACK_RPM_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${JVM_OS_TYPE}-${MACH_ARCH}.rpm")
207+
set(CPACK_RPM_PACKAGE_DESCRIPTION ${KNN_PACKAGE_DESCRIPTION})
208+
set(CPACK_RPM_PACKAGE_SUMMARY "OpenSearch k-NN JNI Library with nmslib and faiss")
209+
210+
# DEB
113211
set(CPACK_DEBIAN_PACKAGE_HOMEPAGE ${OPENSEARCH_DOWNLOAD_URL})
114212
set(CPACK_DEBIAN_PACKAGE_MAINTAINER ${KNN_MAINTAINER})
115-
set(CPACK_DEBIAN_PACKAGE_SOURCE ${CPACK_PACKAGE_NAME})
116213
set(CPACK_DEBIAN_PACKAGE_VERSION ${CPACK_PACKAGE_VERSION})
117214
set(CPACK_DEBIAN_PACKAGE_SECTION "libs")
215+
set(CPACK_DEB_COMPONENT_INSTALL ON)
216+
217+
set(CPACK_DEBIAN_PACKAGE_NAME ${KNN_PACKAGE_NAME})
218+
set(CPACK_DEBIAN_FILE_NAME "${CPACK_DEBIAN_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${JVM_OS_TYPE}-${MACH_ARCH}.deb")
219+
set(CPACK_DEBIAN_DESCRIPTION ${KNN_PACKAGE_DESCRIPTION})
220+
set(CPACK_DEBIAN_PACKAGE_SOURCE ${CPACK_DEBIAN_PACKAGE_NAME})
118221

119222
include(CPack)
223+
# ---------------------------------------------------------------------------

‎jni/external/faiss

Submodule faiss added at 88eabe9

‎jni/include/faiss_wrapper.h

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*
8+
* Modifications Copyright OpenSearch Contributors. See
9+
* GitHub history for details.
10+
*/
11+
12+
#ifndef OPENSEARCH_KNN_FAISS_WRAPPER_H
13+
#define OPENSEARCH_KNN_FAISS_WRAPPER_H
14+
15+
#include <jni.h>
16+
17+
namespace knn_jni {
18+
namespace faiss_wrapper {
19+
// Create an index with ids and vectors. The configuration is defined by values in the Java map, parametersJ.
20+
// The index is serialized to indexPathJ.
21+
void CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ, jobjectArray vectorsJ,
22+
jstring indexPathJ, jobject parametersJ);
23+
24+
// Create an index with ids and vectors. Instead of creating a new index, this function creates the index
25+
// based off of the template index passed in. The index is serialized to indexPathJ.
26+
void CreateIndexFromTemplate(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,
27+
jobjectArray vectorsJ, jstring indexPathJ, jbyteArray templateIndexJ);
28+
29+
// Load an index from indexPathJ into memory.
30+
//
31+
// Return a pointer to the loaded index
32+
jlong LoadIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jstring indexPathJ);
33+
34+
// Execute a query against the index located in memory at indexPointerJ.
35+
//
36+
// Return an array of KNNQueryResults
37+
jobjectArray QueryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
38+
jfloatArray queryVectorJ, jint kJ);
39+
40+
// Free the index located in memory at indexPointerJ
41+
void Free(jlong indexPointer);
42+
43+
// Perform initilization operations for the library
44+
void InitLibrary();
45+
46+
// Create an empty index defined by the values in the Java map, parametersJ. Train the index with
47+
// the vector of floats located at trainVectorsPointerJ.
48+
//
49+
// Return the serialized representation
50+
jbyteArray TrainIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jobject parametersJ, jint dimension,
51+
jlong trainVectorsPointerJ);
52+
}
53+
}
54+
55+
#endif //OPENSEARCH_KNN_FAISS_WRAPPER_H

0 commit comments

Comments
 (0)