MB-63430: Add jemalloc option

.gitignore

@@ -12,6 +12,7 @@
 /autom4te.cache/
 /makefile.inc
 /bin/
+/build/
 /c_api/bin/
 /c_api/gpu/bin/
 /tests/test

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)

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}

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)                                       \
     {                                                      \

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 <faiss/jemalloc_override.cpp>
+
+#endif /* USE_JEMALLOC */

c_api/macros_impl.h

@@ -18,6 +18,7 @@
 #include <stdexcept>
 #include "error_impl.h"
 #include "faiss_c.h"
+#include <faiss/jemalloc_override.h>
 
 #ifdef NDEBUG
 #define CATCH_AND_HANDLE                                                      \

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=<path>.")
+  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)

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

faiss/impl/FaissException.cpp

@@ -14,6 +14,7 @@
 #include <cxxabi.h>
 #endif
 
+#include <faiss/jemalloc_override.h>
 namespace faiss {
 
 FaissException::FaissException(const std::string& m) : msg(m) {}

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))

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 <https://en.cppreference.com/w/cpp/memory/new/operator_new>
+// ---------------------------------------------------------------------------
+#include "jemalloc_override.h"
+#include <new>
+
+#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<size_t>(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<size_t>(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<size_t>(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<size_t>(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

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 <jemalloc/jemalloc.h>
+
+#endif /* USE_JEMALLOC */
+
+#endif /* JEMALLOC_OVERRIDE_H */

faiss/utils/AlignedTable.h

@@ -15,6 +15,7 @@
 #include <algorithm>
 
 #include <faiss/impl/platform_macros.h>
+#include <faiss/jemalloc_override.h>
 
 namespace faiss {