diff --git a/.gitignore b/.gitignore index caab1304c8..158c2cfc1e 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ /autom4te.cache/ /makefile.inc /bin/ +/build/ /c_api/bin/ /c_api/gpu/bin/ /tests/test diff --git a/CMakeLists.txt b/CMakeLists.txt index fa5ab4d1f5..39d562a8c7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -56,6 +56,7 @@ option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON) option(FAISS_ENABLE_RAFT "Enable RAFT for GPU indexes." OFF) option(FAISS_ENABLE_PYTHON "Build Python extension." ON) option(FAISS_ENABLE_C_API "Build C API." OFF) +option(FAISS_USE_JEMALLOC "Use jemalloc for memory allocation." OFF) # Force FAISS_OPT_LEVEL to "generic" if not building on x86_64. if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|AMD64") @@ -89,9 +90,9 @@ add_subdirectory(demos) add_subdirectory(benchs) add_subdirectory(tutorial/cpp) -# CTest must be included in the top level to enable `make test` target. -include(CTest) if(BUILD_TESTING) + # CTest must be included in the top level to enable `make test` target. + include(CTest) add_subdirectory(tests) if(FAISS_ENABLE_GPU) diff --git a/c_api/CMakeLists.txt b/c_api/CMakeLists.txt index 8af3ee8e4a..41b361a90f 100644 --- a/c_api/CMakeLists.txt +++ b/c_api/CMakeLists.txt @@ -28,6 +28,7 @@ set(FAISS_C_SRC MetaIndexes_c.cpp clone_index_c.cpp error_impl.cpp + jemalloc_override_c.cpp index_factory_c.cpp index_io_c.cpp index_io_c_ex.cpp @@ -42,6 +43,13 @@ elseif(FAISS_OPT_LEVEL STREQUAL "avx2") elseif(FAISS_OPT_LEVEL STREQUAL "avx512") target_link_libraries(faiss_c PRIVATE faiss_avx512) endif() + +if(FAISS_USE_JEMALLOC) + find_package(Jemalloc REQUIRED) + add_definitions(-DUSE_JEMALLOC) + target_link_libraries(faiss_c PRIVATE Jemalloc::jemalloc) +endif() + install(TARGETS faiss_c EXPORT faiss-targets RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} diff --git a/c_api/example_c.c b/c_api/example_c.c index 299305cf3f..59bd957c4f 100644 --- a/c_api/example_c.c +++ b/c_api/example_c.c @@ -20,6 +20,7 @@ #include "impl/AuxIndexStructures_c.h" #include "index_factory_c.h" #include "index_io_c.h" +#include "macros_impl.h" #define FAISS_TRY(C) \ { \ diff --git a/c_api/jemalloc_override_c.cpp b/c_api/jemalloc_override_c.cpp new file mode 100644 index 0000000000..b89c2fc7b0 --- /dev/null +++ b/c_api/jemalloc_override_c.cpp @@ -0,0 +1,15 @@ +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Copyright 2004-present Facebook. All Rights Reserved. +// -*- c++ -*- + +#ifdef USE_JEMALLOC + +#include + +#endif /* USE_JEMALLOC */ diff --git a/c_api/macros_impl.h b/c_api/macros_impl.h index 48c5efffbc..623047a648 100644 --- a/c_api/macros_impl.h +++ b/c_api/macros_impl.h @@ -18,6 +18,7 @@ #include #include "error_impl.h" #include "faiss_c.h" +#include #ifdef NDEBUG #define CATCH_AND_HANDLE \ diff --git a/cmake/FindJemalloc.cmake b/cmake/FindJemalloc.cmake new file mode 100644 index 0000000000..fe0a955455 --- /dev/null +++ b/cmake/FindJemalloc.cmake @@ -0,0 +1,63 @@ +# - Try to find jemalloc +# Once done this will define a target Jemalloc::jemalloc which will include all +# required definitions and libraries. + +if (NOT FindJemalloc_included) + + # First try finding Jemalloc using pkg-config. + find_package(PkgConfig QUIET) + if (PKG_CONFIG_FOUND) + pkg_check_modules(PC_JEMALLOC QUIET IMPORTED_TARGET GLOBAL jemalloc) + + if (PC_JEMALLOC_FOUND) + # Ensure it created the PkgConfig target + if(NOT TARGET PkgConfig::PC_JEMALLOC) + message(FATAL_ERROR + "Found Jemalloc via pkg-config, but it did not create the PkgConfig::PC_JEMALLOC target.") + endif() + + # Make the discovered target available as Jemalloc::jemalloc + add_library(Jemalloc::jemalloc ALIAS PkgConfig::PC_JEMALLOC) + set(_jemalloc_found TRUE) + endif() + endif() + + # If that didn't find it, try finding Jemalloc using CMake's config + # mode. + if(NOT _jemalloc_found) + find_package(Jemalloc CONFIG) + + if(Jemalloc_FOUND) + if(NOT TARGET Jemalloc::jemalloc) + # Ensure this found package created the standard target. + message(FATAL_ERROR + "Found Jemalloc, but it did not create the Jemalloc::jemalloc target.") + endif() + set(_jemalloc_found TRUE) + endif() + endif() + + # If still not found, report an error + if(NOT _jemalloc_found) + message(FATAL_ERROR "Could not find Jemalloc. Install Jemalloc or provide -DJemalloc_ROOT=.") + endif() + + # Determine the version of the found jemalloc. + get_target_property(_jemalloc_include_dirs Jemalloc::jemalloc INTERFACE_INCLUDE_DIRECTORIES) + list (GET _jemalloc_include_dirs 0 _jemalloc_include_dir) + set(_version_regex "^#define[ \t]+JEMALLOC_VERSION[ \t]+\"([^\"]+)\".*") + file(STRINGS "${_jemalloc_include_dir}/jemalloc/jemalloc.h" + JEMALLOC_VERSION REGEX "${_version_regex}") + string(REGEX REPLACE "${_version_regex}" "\\1" + JEMALLOC_VERSION "${JEMALLOC_VERSION}") + unset(_version_regex) + + # handle the QUIET and REQUIRED arguments, verify version (if + # necessary), report found status, and set JEMALLOC_FOUND to TRUE. + include(FindPackageHandleStandardArgs) + find_package_handle_standard_args( + Jemalloc VERSION_VAR JEMALLOC_VERSION + REQUIRED_VARS _jemalloc_found) + set(FindJemalloc_included TRUE) + +endif (NOT FindJemalloc_included) diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt index 1584023086..145659879e 100644 --- a/faiss/CMakeLists.txt +++ b/faiss/CMakeLists.txt @@ -46,6 +46,7 @@ set(FAISS_SRC IndexScalarQuantizer.cpp IndexShards.cpp IndexShardsIVF.cpp + jemalloc_override.cpp MatrixStats.cpp MetaIndexes.cpp OMPConfig.cpp @@ -332,6 +333,14 @@ else() target_link_libraries(faiss_avx512 PRIVATE ${LAPACK_LIBRARIES}) endif() +if(FAISS_USE_JEMALLOC) + find_package(Jemalloc REQUIRED) + add_definitions(-DUSE_JEMALLOC) + target_link_libraries(faiss PRIVATE Jemalloc::jemalloc) + target_link_libraries(faiss_avx2 PRIVATE Jemalloc::jemalloc) + target_link_libraries(faiss_avx512 PRIVATE Jemalloc::jemalloc) +endif() + if(FAISS_OPT_LEVEL STREQUAL "generic") install(TARGETS faiss diff --git a/faiss/impl/FaissException.cpp b/faiss/impl/FaissException.cpp index 3dcf47a5ea..fb0cb61671 100644 --- a/faiss/impl/FaissException.cpp +++ b/faiss/impl/FaissException.cpp @@ -14,6 +14,7 @@ #include #endif +#include namespace faiss { FaissException::FaissException(const std::string& m) : msg(m) {} diff --git a/faiss/impl/platform_macros.h b/faiss/impl/platform_macros.h index a6c9cf27ce..e3f2817275 100644 --- a/faiss/impl/platform_macros.h +++ b/faiss/impl/platform_macros.h @@ -30,9 +30,21 @@ #define __PRETTY_FUNCTION__ __FUNCSIG__ +#ifndef USE_JEMALLOC +// no jemalloc on windows, and windows does not have posix_memalign +// _aligned_malloc is the equivalent #define posix_memalign(p, a, s) \ (((*(p)) = _aligned_malloc((s), (a))), *(p) ? 0 : errno) +// _aligned_free is the equivalent of free for _aligned_malloc #define posix_memalign_free _aligned_free +#else +// je_posix_memalign is available on jemalloc +// posix_memalign_free is defined as je_free +// it MUST be ensured that the jemalloc_override.h file is imported +// AFTER whichever file imported the platform_macros.h file +// (which is where posix_memalign_free -> free -> je_free is defined) +#define posix_memalign_free free +#endif // aligned should be in front of the declaration #define ALIGNED(x) __declspec(align(x)) diff --git a/faiss/jemalloc_override.cpp b/faiss/jemalloc_override.cpp new file mode 100644 index 0000000000..96ad1a696d --- /dev/null +++ b/faiss/jemalloc_override.cpp @@ -0,0 +1,140 @@ +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +#ifdef USE_JEMALLOC +// ---------------------------------------------------------------------------- +// This header provides convenient overrides for the new and +// delete operations in C++. +// +// This header should be included in only one source file! +// +// See +// --------------------------------------------------------------------------- +#include "jemalloc_override.h" +#include + +#if __cplusplus > 202002L +#warning Only operator new/delete up to C++20 overridden. If later standards add additional overrides they should be added here. +#endif + + +// (new 1) +[[nodiscard]] void* operator new(std::size_t count) { + void* result = malloc(count); + if (result == nullptr) { + throw std::bad_alloc(); + } + return result; +} + +// (new 2) +[[nodiscard]] void* operator new[](std::size_t count) { + void* result = malloc(count); + if (result == nullptr) { + throw std::bad_alloc(); + } + return result; +} + +// (new 3) +[[nodiscard]] void* operator new(std::size_t count, std::align_val_t al) { + void* result = aligned_alloc(static_cast(al), count); + if (result == nullptr) { + throw std::bad_alloc(); + } + return result; +} + +// (new 4) +[[nodiscard]] void* operator new[](std::size_t count, std::align_val_t al) { + void* result = aligned_alloc(static_cast(al), count); + if (result == nullptr) { + throw std::bad_alloc(); + } + return result; +} + +// (new 5) +[[nodiscard]] void* operator new(std::size_t count, + const std::nothrow_t& tag) noexcept { + void* result = malloc(count); + return result; +} + +// (new 6) +[[nodiscard]] void* operator new[](std::size_t count, + const std::nothrow_t& tag) noexcept { + void* result = malloc(count); + return result; +} + +// (new 7) +[[nodiscard]] void* operator new(std::size_t count, + std::align_val_t al, + const std::nothrow_t& tag) noexcept { + void* result = aligned_alloc(static_cast(al), count); + return result; +} + +// (new 8) +[[nodiscard]] void* operator new[](std::size_t count, + std::align_val_t al, + const std::nothrow_t& tag) noexcept { + void* result = aligned_alloc(static_cast(al), count); + return result; +} + +// (del 1) +void operator delete(void* ptr) noexcept { + free(ptr); +} + +// (del 2) +void operator delete[](void* ptr) noexcept { + free(ptr); +} + +// (del 3) +void operator delete(void* ptr, std::align_val_t) noexcept { + free(ptr); +} + +// (del 4) +void operator delete[](void* ptr, std::align_val_t) noexcept { + free(ptr); +} + +// (del 6) +void operator delete[](void* ptr, std::size_t size) noexcept { + free(ptr); +} + +// (del 5) +void operator delete(void* ptr, std::size_t size) noexcept { + free(ptr); +} + +// (del 7) +void operator delete(void* ptr, std::size_t size, std::align_val_t) noexcept { + free(ptr); +} + +// (del 8) +void operator delete[](void* ptr, std::size_t size, std::align_val_t) noexcept { + free(ptr); +} + +// (del 9) +void operator delete(void* ptr, const std::nothrow_t& tag) noexcept { + free(ptr); +} + +// (del 10) +void operator delete[](void* ptr, const std::nothrow_t& tag) noexcept { + free(ptr); +} + +#endif // USE_JEMALLOC diff --git a/faiss/jemalloc_override.h b/faiss/jemalloc_override.h new file mode 100644 index 0000000000..8796f8cdea --- /dev/null +++ b/faiss/jemalloc_override.h @@ -0,0 +1,20 @@ +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +// -*- c++ -*- + +#ifndef JEMALLOC_OVERRIDE_H +#define JEMALLOC_OVERRIDE_H + +#ifdef USE_JEMALLOC + +#define JEMALLOC_MANGLE +#include + +#endif /* USE_JEMALLOC */ + +#endif /* JEMALLOC_OVERRIDE_H */ diff --git a/faiss/utils/AlignedTable.h b/faiss/utils/AlignedTable.h index 05adb1c0d0..5a0e903ffc 100644 --- a/faiss/utils/AlignedTable.h +++ b/faiss/utils/AlignedTable.h @@ -15,6 +15,7 @@ #include #include +#include namespace faiss {