microsoft · Mark Hildebrand (hildebrandmw) · May 3, 2021 · May 4, 2021 · May 12, 2021 · May 14, 2021
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -43,53 +43,72 @@ function(checkEnvAndSetLocalVar env_var msg local_var)
 	endif()
 endfunction()
 
-	
+
 
 #MKL Config
 if (MSVC)
 	checkEnvAndSetLocalVar("INTEL_ROOT" "Please install Intel MKL libraries and set the env variable INTEL_ROOT to the intel software directory. Should be similar to: C:\\Program Files (x86)\\IntelSWTools\\compilers_and_libraries\\windows\\. " "INTEL_ROOT")
-	set(MKL_ROOT ${INTEL_ROOT}/mkl)
+	set(MKL_ROOT ${INTEL_ROOT}/mkl/latest/)
 	add_compile_options(/arch:AVX2 /Qpar)
-	link_libraries("${INTEL_ROOT}/mkl/lib/intel64/mkl_core_dll.lib" "${INTEL_ROOT}/mkl/lib/intel64/mkl_rt.lib"  "${INTEL_ROOT}/mkl/lib/intel64/mkl_intel_thread_dll.lib" "${INTEL_ROOT}/compiler/lib/intel64/libiomp5md.lib"  "${INTEL_ROOT}/mkl/lib/intel64/mkl_intel_ilp64_dll.lib" "${INTEL_ROOT}/mkl/lib/intel64/mkl_sequential_dll.lib")
+	link_libraries("${INTEL_ROOT}/mkl/latest/lib/intel64/mkl_core_dll.lib" "${INTEL_ROOT}/mkl/latest/lib/intel64/mkl_rt.lib"  "${INTEL_ROOT}/mkl/latest/lib/intel64/mkl_intel_thread_dll.lib" "${INTEL_ROOT}/compiler/latest/windows/compiler/lib/intel64_win/libiomp5md.lib"  "${INTEL_ROOT}/mkl/latest/lib/intel64/mkl_intel_ilp64_dll.lib" "${INTEL_ROOT}/mkl/latest/lib/intel64/mkl_sequential_dll.lib")
 	checkEnvAndSetLocalVar("BOOST_ROOT" "Please install Boost (1.71 or greater) from www.boost.org and set the env var BOOST_ROOT to the boost directory." "BOOST_ROOT")
 
 else()
-	set(INTEL_ROOT /opt/intel/compilers_and_libraries/linux)
+	set(INTEL_ROOT /home/stg/intel/compilers_and_libraries_2020.2.254/linux)
 	set(MKL_ROOT ${INTEL_ROOT}/mkl)
-        add_compile_options(-m64 -Wl,--no-as-needed) 
+        add_compile_options(-m64 -Wl,--no-as-needed)
 	link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core iomp5 pthread m dl)
 	link_directories(${INTEL_ROOT}/lib/intel64 ${MKL_ROOT}/lib/intel64)
 endif()
 
+#function(checkPMem env_var local_var)
+#	if (NOT EXISTS "$ENV{${env_var}}")
+#		message ("PMem variable not found defaulting to DRAM")
+#		set(${local_var} "DRAM" PARENT_SCOPE)
+#	else()
+#		message (STATUS "using ${env_var}")
+#		set(${local_var} "$ENV{${env_var}}" PARENT_SCOPE)
+#		if ($ENV{${env_var}} MATCHES "PMEM")
+#			link_libraries(memkind)
+#		endif()
+#	endif()
+#endfunction()
+
+# Include PM
+if (MSVC)
+else()
+link_libraries(memkind)
+endif()
+
 add_definitions(-DMKL_ILP64)
 include_directories(include ${INTEL_ROOT}/include ${MKL_ROOT}/include ${BOOST_ROOT})
 
 
-#Main compiler/linker settings 
+#Main compiler/linker settings
 if(MSVC)
 	#language options
-	add_compile_options(/permissive- /openmp:experimental /Zc:wchar_t /Zc:twoPhase- /Zc:forScope /Zc:inline /WX- /std:c++14 /Gd /W3 /MP /Zi /FC /nologo /diagnostics:classic) 
+	add_compile_options(/permissive- /openmp:experimental /Zc:wchar_t /Zc:twoPhase- /Zc:forScope /Zc:inline /WX- /std:c++14 /Gd /W3 /MP /Zi /FC /nologo /diagnostics:classic)
 	#code generation options
 	add_compile_options(/Qpar /fp:fast /Zp8 /fp:except- /EHsc /GS- /Gm- /Gy )
 	#optimization options
 	add_compile_options(/Ot /Oy /Oi)
 	#path options
 	#add_compile_options(/Fdx64/Release/vc141.pdb /Fox64/Release/)
 	add_definitions(-DUSE_AVX2 -DUSE_ACCELERATED_PQ -D_WINDOWS -DNOMINMAX -DUNICODE)
-	
+
 	set(CMAKE_SHARED_LIBRARY_CXX_LINK_FLAGS "/MANIFEST /MACHINE:X64 /DEBUG:FULL /LTCG:incremental /NXCOMPAT /DYNAMICBASE /OPT:REF /SUBSYSTEM:CONSOLE /MANIFESTUAC:\"level='asInvoker' uiAccess='false'\"")
 	set(CMAKE_EXECUTABLE_CXX_LINK_FLAGS "/MANIFEST /MACHINE:X64 /DEBUG:FULL /LTCG:incremental /NXCOMPAT /DYNAMICBASE /OPT:REF /SUBSYSTEM:CONSOLE /MANIFESTUAC:\"level='asInvoker' uiAccess='false'\"")
 
 	set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} /D_DEBUG")
 	set(CMAKE_SHARED_LIBRARY_CXX_LINK_FLAGS_DEBUG "${CMAKE_SHARED_LIBRARY_CXX_LINK_FLAGS_DEBUG} /DEBUG")
 
-	set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug)
-	set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug)
-	set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/x64/Debug)
+	set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/out/build/Debug)
+	set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/out/build/Debug)
+	set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_DEBUG ${PROJECT_SOURCE_DIR}/out/build/Debug)
 
-	set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release)
-	set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release)
-	set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release)
+	set(CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/out/build/Release)
+	set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/out/build/Release)
+	set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/out/build/Release)
 else()
 	set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000)
     #	set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -O0 -fsanitize=address -fsanitize=leak -fsanitize=undefined")

diff --git a/CMakeSettings.json b/CMakeSettings.json
@@ -0,0 +1,17 @@
+{
+  "configurations": [
+    {
+      "name": "Release",
+      "generator": "Visual Studio 16 2019 Win64",
+      "configurationType": "Release",
+      "inheritEnvironments": [ "msvc_x64" ],
+      "buildRoot": "${projectDir}\\out\\build\\${name}",
+      "installRoot": "${projectDir}\\out\\install\\${name}",
+      "buildCommandArgs": "",
+      "ctestCommandArgs": "",
+      "cmakeExecutable": "C:\\Program Files\\CMake\\bin\\cmake.exe",
+      "intelliSenseMode": "windows-msvc-x64",
+      "cmakeCommandArgs": "-DCMAKE_BUILD_TYPE=Release"
+    }
+  ]
+}
diff --git a/include/index.h b/include/index.h
@@ -5,6 +5,7 @@
 
 #include <cassert>
 #include <map>
+#include <memory>
 #include <sstream>
 #include <stack>
 #include <string>
@@ -24,10 +25,11 @@
            ((double) size * degree) * sizeof(unsigned) * SLACK_FACTOR))
 
 namespace diskann {
-  template<typename T, typename TagT = int>
+  template<typename T, typename TagT = int, typename Allocator = std::allocator<unsigned>>
   class Index {
    public:
-    DISKANN_DLLEXPORT Index(Metric m, const char *filename,
+    DISKANN_DLLEXPORT Index(Metric m, const char *filename, bool data_in_pm = false,
+                            const Allocator& allocator = std::allocator<unsigned>(),
                             const size_t max_points = 0, const size_t nd = 0,
                             const size_t num_frozen_pts = 0,
                             const bool   enable_tags = false,
@@ -98,15 +100,16 @@ namespace diskann {
     DISKANN_DLLEXPORT int eager_delete(const TagT        tag,
                                        const Parameters &parameters);
 
-    DISKANN_DLLEXPORT void optimize_graph();
+    DISKANN_DLLEXPORT void optimize_graph(bool use_pm = false);
 
     DISKANN_DLLEXPORT void search_with_opt_graph(const T *query, size_t K,
                                                  size_t L, unsigned *indices);
 
     /*  Internals of the library */
    protected:
+    const Allocator _allocator;
     typedef std::vector<SimpleNeighbor>        vecNgh;
-    typedef std::vector<std::vector<unsigned>> CompactGraph;
+    typedef std::vector<std::vector<unsigned, Allocator>> CompactGraph;
     CompactGraph                               _final_graph;
     CompactGraph                               _in_graph;
 

diff --git a/include/pm.h b/include/pm.h
@@ -0,0 +1,27 @@
+#pragma once
+#ifndef _WINDOWS
+#include <cassert>
+#include <string>
+#include <memkind.h>
+#include <pmem_allocator.h>
+
+namespace diskann {
+  bool is_pm_init();
+  void init_pm(const std::string& dir);
+  void* __alloc(size_t size, size_t align, bool pm = false);
+  void __free(void* ptr);
+
+  template<typename T>
+  using pmem_allocator = libmemkind::pmem::allocator<T>;
+
+  // Bootstrap copying allocators with different type parameters.
+  const pmem_allocator<uint8_t>& _pm_allocator();
+
+  // Get a `pmem::allocator` for any type.
+  template<typename T>
+  pmem_allocator<T> pm_allocator()
+  {
+    return pmem_allocator<T>(_pm_allocator());
+  }
+} // namespace diskann
+#endif
diff --git a/include/utils.h b/include/utils.h
@@ -30,6 +30,7 @@ typedef int FileHandle;
 #include "cached_io.h"
 #include "common_includes.h"
 #include "windows_customizations.h"
+#include "pm.h" // PM modifications
 
 #ifdef EXEC_ENV_OLS
 #include "content_buf.h"
@@ -66,11 +67,11 @@ namespace diskann {
 
   enum Metric { L2 = 0, INNER_PRODUCT = 1, FAST_L2 = 2, PQ = 3 };
 
-  inline void alloc_aligned(void** ptr, size_t size, size_t align) {
+  inline void alloc_aligned(void** ptr, size_t size, size_t align, bool persistent = false) {
     *ptr = nullptr;
     assert(IS_ALIGNED(size, align));
 #ifndef _WINDOWS
-    *ptr = ::aligned_alloc(align, size);
+    *ptr = __alloc(size, align, persistent);
 #else
     *ptr = ::_aligned_malloc(size, align);  // note the swapped arguments!
 #endif
@@ -84,7 +85,7 @@ namespace diskann {
       return;
     }
 #ifndef _WINDOWS
-    free(ptr);
+    __free(ptr);
 #else
     ::_aligned_free(ptr);
 #endif
@@ -331,7 +332,8 @@ namespace diskann {
   inline void load_aligned_bin_impl(std::basic_istream<char>& reader,
                                     size_t actual_file_size, T*& data,
                                     size_t& npts, size_t& dim,
-                                    size_t& rounded_dim) {
+                                    size_t& rounded_dim,
+                                    bool use_pm = false) {
     int npts_i32, dim_i32;
     reader.read((char*) &npts_i32, sizeof(int));
     reader.read((char*) &dim_i32, sizeof(int));
@@ -356,7 +358,7 @@ namespace diskann {
     size_t allocSize = npts * rounded_dim * sizeof(T);
     diskann::cout << "allocating aligned memory, " << allocSize << " bytes..."
                   << std::flush;
-    alloc_aligned(((void**) &data), allocSize, 8 * sizeof(T));
+    alloc_aligned(((void**) &data), allocSize, 8 * sizeof(T), use_pm);
     diskann::cout << "done. Copying data..." << std::flush;
 
     for (size_t i = 0; i < npts; i++) {
@@ -384,7 +386,7 @@ namespace diskann {
 
   template<typename T>
   inline void load_aligned_bin(const std::string& bin_file, T*& data,
-                               size_t& npts, size_t& dim, size_t& rounded_dim) {
+                               size_t& npts, size_t& dim, size_t& rounded_dim, bool use_pm = false) {
     diskann::cout << "Reading bin file " << bin_file << " ..." << std::flush;
     // START OLS
     //_u64            read_blk_size = 64 * 1024 * 1024;
@@ -396,7 +398,7 @@ namespace diskann {
     uint64_t      fsize = reader.tellg();
     reader.seekg(0);
 
-    load_aligned_bin_impl(reader, fsize, data, npts, dim, rounded_dim);
+    load_aligned_bin_impl(reader, fsize, data, npts, dim, rounded_dim, use_pm);
   }
 
   template<typename InType, typename OutType>

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -9,7 +9,7 @@ else()
 	#file(GLOB CPP_SOURCES *.cpp)
 	set(CPP_SOURCES ann_exception.cpp aux_utils.cpp index.cpp
         linux_aligned_file_reader.cpp math_utils.cpp memory_mapper.cpp
-        partition_and_pq.cpp  pq_flash_index.cpp logger.cpp utils.cpp)
+        partition_and_pq.cpp pm.cpp pq_flash_index.cpp logger.cpp utils.cpp)
 	add_library(${PROJECT_NAME} ${CPP_SOURCES})
 	add_library(${PROJECT_NAME}_s STATIC ${CPP_SOURCES})
 endif()