update cmake for diskann_build, update cpp with malloc extension

microsoft · Aug 6, 2024 · cb3c441 · cb3c441
1 parent aa7faac
commit cb3c441
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 11 deletions.
diff --git a/src/dll/CMakeLists.txt b/src/dll/CMakeLists.txt
@@ -6,13 +6,21 @@ if (DISKANN_USE_STATIC_LIB)
         ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp 
         ../in_mem_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp
         ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp)
+
+    # tcmalloc verison
+    add_library(${PROJECT_NAME}_build STATIC ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp
+        ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../../gperftools/src/gperftools/malloc_extension.h 
+        ../in_mem_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp
+        ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp)
 else()
     add_library(${PROJECT_NAME} SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp
         ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp 
         ../in_mem_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp
         ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp)
+
+    # tcmalloc version
     add_library(${PROJECT_NAME}_build SHARED dllmain.cpp ../abstract_data_store.cpp ../partition.cpp ../pq.cpp ../pq_flash_index.cpp ../logger.cpp ../utils.cpp
-        ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp 
+        ../windows_aligned_file_reader.cpp ../distance.cpp ../memory_mapper.cpp ../index.cpp ../../gperftools/src/gperftools/malloc_extension.h 
         ../in_mem_data_store.cpp ../in_mem_graph_store.cpp ../math_utils.cpp ../disk_utils.cpp ../filter_utils.cpp
         ../ann_exception.cpp ../natural_number_set.cpp ../natural_number_map.cpp ../scratch.cpp ../index_factory.cpp ../abstract_index.cpp)
 endif()
@@ -21,20 +29,21 @@ set(TARGET_DIR "$<$<CONFIG:Debug>:${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}>$<$<CO
 
 if (DISKANN_USE_STATIC_LIB)
     target_compile_definitions(${PROJECT_NAME} PRIVATE _NODLL)
+
+    target_compile_definitions(${PROJECT_NAME}_build PRIVATE _NODLL DISKANN_BUILD RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS)
 else()
     set(DISKANN_DLL_IMPLIB "${TARGET_DIR}/${PROJECT_NAME}.lib")
     target_compile_definitions(${PROJECT_NAME} PRIVATE _USRDLL _WINDLL)
 
     set(DISKANN_DLL_IMPLIB_BUILD "${TARGET_DIR}/${PROJECT_NAME}_build.lib")
-    target_compile_definitions(${PROJECT_NAME}_build PRIVATE _USRDLL _WINDLL)
-    target_compile_definitions(${PROJECT_NAME}_build PRIVATE DISKANN_BUILD)
+    target_compile_definitions(${PROJECT_NAME}_build PRIVATE _USRDLL _WINDLL DISKANN_BUILD RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS)
 endif()
 
 target_compile_options(${PROJECT_NAME} PRIVATE /GL)
 target_include_directories(${PROJECT_NAME} PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES})
 
 target_compile_options(${PROJECT_NAME}_build PRIVATE /GL)
-target_include_directories(${PROJECT_NAME}_build PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES})
+target_include_directories(${PROJECT_NAME}_build PRIVATE ${DISKANN_MKL_INCLUDE_DIRECTORIES} ../../gperftools/src/)
 
 
 if (NOT DEFINED DISKANN_USE_STATIC_LIB)
@@ -44,19 +53,14 @@ endif()
 
 target_link_libraries(${PROJECT_NAME} PRIVATE ${DISKANN_MKL_LINK_LIBRARIES})
 target_link_libraries(${PROJECT_NAME} PRIVATE synchronization.lib)
+
 target_link_libraries(${PROJECT_NAME}_build PRIVATE ${DISKANN_MKL_LINK_LIBRARIES})
 target_link_libraries(${PROJECT_NAME}_build PRIVATE synchronization.lib)
-
 #This is the crux of the build dll
 target_link_libraries(${PROJECT_NAME}_build PUBLIC ${DISKANN_DLL_TCMALLOC_LINK_OPTIONS})
+target_link_libraries(${PROJECT_NAME}_build PRIVATE ${TCMALLOC_LINK_LIBRARY})
 set_target_properties(${PROJECT_NAME}_build PROPERTIES LINK_FLAGS /INCLUDE:_tcmalloc)
 
-
-
-if (DISKANN_DLL_TCMALLOC_LINK_OPTIONS)
-    target_link_libraries(${PROJECT_NAME} PUBLIC ${DISKANN_DLL_TCMALLOC_LINK_OPTIONS})
-endif()
-
 # Copy OpenMP DLL and PDB.
 set(RUNTIME_FILES_TO_COPY ${OPENMP_WINDOWS_RUNTIME_FILES} ${TCMALLOC_WINDOWS_RUNTIME_FILES})
 

diff --git a/src/index.cpp b/src/index.cpp
@@ -1601,6 +1601,10 @@ template <typename T, typename TagT, typename LabelT> void Index<T, TagT, LabelT
         }
     }
 
+#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD)
+    MallocExtension::instance()->ReleaseFreeMemory();
+#endif
+
     if (_nd > 0)
     {
         diskann::cout << "Starting final cleanup.." << std::flush;

diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
@@ -15,6 +15,10 @@
 #include "linux_aligned_file_reader.h"
 #endif
 
+#if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD)
+#include "gperftools/malloc_extension.h"
+#endif
+
 #define READ_U64(stream, val) stream.read((char *)&val, sizeof(uint64_t))
 #define READ_U32(stream, val) stream.read((char *)&val, sizeof(uint32_t))
 #define READ_UNSIGNED(stream, val) stream.read((char *)&val, sizeof(unsigned))
@@ -92,6 +96,10 @@ template <typename T, typename LabelT> PQFlashIndex<T, LabelT>::~PQFlashIndex()
     {
         delete[] _pts_to_labels;
     }
+
+    #if defined(RELEASE_UNUSED_TCMALLOC_MEMORY_AT_CHECKPOINTS) && defined(DISKANN_BUILD)
+    MallocExtension::instance()->ReleaseFreeMemory();
+    #endif
 }
 
 template <typename T, typename LabelT> inline uint64_t PQFlashIndex<T, LabelT>::get_node_sector(uint64_t node_id)