From 692070f501e16629223d62f46c3b895b8775707a Mon Sep 17 00:00:00 2001 From: "Hildebrand, Mark" Date: Mon, 3 May 2021 11:42:20 -0700 Subject: [PATCH 1/4] [PoC] Initial support for Optane PM. Minimal changes to the DiskANN code to facilitate placing the dataset and graph data structures in Optane PM for benchmarking purposes. Additional flags are added to the `build_memory_index` and `search_memory_index` executables to provide control over data structures placement. Notes: * This change is only implemented for Linux based systems. * This change adds an additional dependency of `libmemkind` to support heap allocations in either DRAM or PM. On Ubuntu systems, this library can be installed system wide with ``` apt install libmemkind-dev libmemkind0 ``` Other distributions require building `libmemkind` from scratch. * The directory provided as a commandline argument must point to a properly configured "DAX" (direct-access) filesystem backed by NVDIMMs. --- CMakeLists.txt | 15 +-- include/index.h | 11 ++- include/pm.h | 27 ++++++ include/utils.h | 16 ++-- src/CMakeLists.txt | 2 +- src/index.cpp | 175 ++++++++++++++++++++-------------- src/pm.cpp | 62 ++++++++++++ tests/build_memory_index.cpp | 62 ++++++++---- tests/search_memory_index.cpp | 75 +++++++++++---- 9 files changed, 322 insertions(+), 123 deletions(-) create mode 100644 include/pm.h create mode 100644 src/pm.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd3886f5b..02d5a268d2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,7 +14,7 @@ if(MSVC) set(CMAKE_CXX_COMPILER $ENV{VCToolsInstallDir}/bin/Hostx64/x64/cl.exe) set(CMAKE_CXX_LINK_EXECUTABLE $ENV{VCToolsInstallDir}/bin/Hostx64/x64/link.exe) else() - set(CMAKE_CXX_COMPILER g++) + set(CMAKE_CXX_COMPILER clang++) endif() project(diskann) @@ -43,7 +43,7 @@ function(checkEnvAndSetLocalVar env_var msg local_var) endif() endfunction() - + #MKL Config if (MSVC) @@ -56,19 +56,22 @@ if (MSVC) else() set(INTEL_ROOT /opt/intel/compilers_and_libraries/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() +# Include PM +link_libraries(memkind) + 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 @@ -76,7 +79,7 @@ if(MSVC) #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'\"") diff --git a/include/index.h b/include/index.h index eedbb1491f..ee557eb38b 100644 --- a/include/index.h +++ b/include/index.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -24,10 +25,11 @@ ((double) size * degree) * sizeof(unsigned) * SLACK_FACTOR)) namespace diskann { - template + template> 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(), 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 ¶meters); - 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 vecNgh; - typedef std::vector> CompactGraph; + typedef std::vector> CompactGraph; CompactGraph _final_graph; CompactGraph _in_graph; diff --git a/include/pm.h b/include/pm.h new file mode 100644 index 0000000000..b1800fc251 --- /dev/null +++ b/include/pm.h @@ -0,0 +1,27 @@ +#pragma once + +#include // For some reason, "pmem_allocator.h" doesn't include this. +#include + +#include +#include + +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 + using pmem_allocator = libmemkind::pmem::allocator; + + // Bootstrap copying allocators with different type parameters. + const pmem_allocator& _pm_allocator(); + + // Get a `pmem::allocator` for any type. + template + pmem_allocator pm_allocator() + { + return pmem_allocator(_pm_allocator()); + } +} diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..3c02e50f07 100644 --- 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& 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 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 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9cd8a86bc4..cacfce0fc4 100644 --- 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() diff --git a/src/index.cpp b/src/index.cpp index 2405878669..3cf149a110 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -32,6 +32,7 @@ #include "memory_mapper.h" #include "parameters.h" #include "partition_and_pq.h" +#include "pm.h" #include "timer.h" #include "utils.h" #include "windows_customizations.h" @@ -47,7 +48,6 @@ namespace { template<> diskann::Distance *get_distance_function(diskann::Metric m) { if (m == diskann::Metric::FAST_L2) { - std::cout << "Here" << std::endl; return new diskann::DistanceFastL2(); } else if (m == diskann::Metric::L2) { if (Avx2SupportedCPU) { @@ -124,12 +124,15 @@ namespace diskann { // Initialize an index with metric m, load the data of type T with filename // (bin), and initialize max_points - template - Index::Index(Metric m, const char *filename, const size_t max_points, + template + Index::Index(Metric m, const char *filename, + bool data_in_pm, + const Allocator& allocator, + const size_t max_points, const size_t nd, const size_t num_frozen_pts, const bool enable_tags, const bool store_data, const bool support_eager_delete) - : _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), + : _allocator(allocator), _num_frozen_pts(num_frozen_pts), _has_built(false), _width(0), _can_delete(false), _eager_done(true), _lazy_done(true), _compacted_order(true), _enable_tags(enable_tags), _consolidated_order(true), _support_eager_delete(support_eager_delete), @@ -138,7 +141,7 @@ namespace diskann { // zero-padding diskann::cout << "Number of frozen points = " << _num_frozen_pts << std::endl; - load_aligned_bin(std::string(filename), _data, _nd, _dim, _aligned_dim); + load_aligned_bin(std::string(filename), _data, _nd, _dim, _aligned_dim, data_in_pm); if (nd > 0) { if (_nd >= nd) @@ -187,25 +190,49 @@ namespace diskann { Index::~Index() { delete this->_distance; aligned_free(_data); + aligned_free(_opt_graph); + } + + template<> + Index>::~Index() { + delete this->_distance; + aligned_free(_data); + aligned_free(_opt_graph); } template<> Index<_s8>::~Index() { delete this->_distance; aligned_free(_data); + aligned_free(_opt_graph); + } + + template<> + Index<_s8,int,diskann::pmem_allocator>::~Index() { + delete this->_distance; + aligned_free(_data); + aligned_free(_opt_graph); } template<> Index<_u8>::~Index() { delete this->_distance; aligned_free(_data); + aligned_free(_opt_graph); + } + + template<> + Index<_u8,int,diskann::pmem_allocator>::~Index() { + delete this->_distance; + aligned_free(_data); + aligned_free(_opt_graph); } // save the graph index on a file as an adjacency list. For each point, // first store the number of neighbors, and then the neighbor list (each as // 4 byte unsigned) - template - void Index::save(const char *filename) { + template + void Index::save(const char *filename) { long long total_gr_edges = 0; size_t index_size = 0; std::ofstream out(std::string(filename), std::ios::binary | std::ios::out); @@ -287,8 +314,8 @@ namespace diskann { // load the index from file and update the width (max_degree), ep // (navigating node id), and _final_graph (adjacency list) - template - void Index::load(const char *filename, const bool load_tags, + template + void Index::load(const char *filename, const bool load_tags, const char *tag_filename) { if (!validate_file_size(filename)) { return; @@ -309,7 +336,7 @@ namespace diskann { break; cc += k; ++nodes; - std::vector tmp(k); + std::vector tmp(k, _allocator); in.read((char *) tmp.data(), k * sizeof(unsigned)); _final_graph.emplace_back(tmp); @@ -357,8 +384,8 @@ namespace diskann { /* This function finds out the navigating node, which is the medoid node * in the graph. */ - template - unsigned Index::calculate_entry_point() { + template + unsigned Index::calculate_entry_point() { // allocate and init centroid float *center = new float[_aligned_dim](); for (size_t j = 0; j < _aligned_dim; j++) @@ -411,8 +438,8 @@ namespace diskann { * search. * best_L_nodes: ids of closest L nodes in list */ - template - std::pair Index::iterate_to_fixed_point( + template + std::pair Index::iterate_to_fixed_point( const T *node_coords, const unsigned Lsize, const std::vector &init_ids, std::vector & expanded_nodes_info, @@ -495,8 +522,8 @@ namespace diskann { return std::make_pair(hops, cmps); } - template - void Index::get_expanded_nodes( + template + void Index::get_expanded_nodes( const size_t node_id, const unsigned Lindex, std::vector init_ids, std::vector & expanded_nodes_info, @@ -511,8 +538,8 @@ namespace diskann { expanded_nodes_ids, best_L_nodes); } - template - void Index::occlude_list(std::vector &pool, + template + void Index::occlude_list(std::vector &pool, const float alpha, const unsigned degree, const unsigned maxc, std::vector &result) { @@ -521,8 +548,8 @@ namespace diskann { occlude_list(pool, alpha, degree, maxc, result, occlude_factor); } - template - void Index::occlude_list(std::vector &pool, + template + void Index::occlude_list(std::vector &pool, const float alpha, const unsigned degree, const unsigned maxc, std::vector &result, @@ -559,8 +586,8 @@ namespace diskann { } } - template - void Index::prune_neighbors(const unsigned location, + template + void Index::prune_neighbors(const unsigned location, std::vector &pool, const Parameters & parameter, std::vector &pruned_list) { @@ -606,8 +633,8 @@ namespace diskann { * This function tries to add reverse links from all the visited nodes to * the current node n. */ - template - void Index::batch_inter_insert( + template + void Index::batch_inter_insert( unsigned n, const std::vector &pruned_list, const Parameters ¶meter, std::vector &need_to_sync) { const auto range = parameter.Get("R"); @@ -639,8 +666,8 @@ namespace diskann { * This function tries to add reverse links from all the visited nodes to * the current node n. */ - template - void Index::inter_insert(unsigned n, + template + void Index::inter_insert(unsigned n, std::vector &pruned_list, const Parameters & parameter, bool update_in_graph) { @@ -671,7 +698,7 @@ namespace diskann { } prune_needed = false; } else { - copy_of_neighbors = des_pool; + copy_of_neighbors.assign(des_pool.begin(), des_pool.end()); prune_needed = true; } } @@ -717,8 +744,8 @@ namespace diskann { * The graph creation function. * The graph will be updated periodically in NUM_SYNCS batches */ - template - void Index::link(Parameters ¶meters) { + template + void Index::link(Parameters ¶meters) { unsigned NUM_THREADS = parameters.Get("num_threads"); if (NUM_THREADS != 0) omp_set_num_threads(NUM_THREADS); @@ -765,11 +792,14 @@ namespace diskann { else _ep = calculate_entry_point(); + // Since `diskann::pmem_allocator` is not default constructable, we need to create + // an empty vector and use the copying version of `resize`. + auto copy_val = std::vector(_allocator); _final_graph.reserve(_max_points + _num_frozen_pts); - _final_graph.resize(_max_points + _num_frozen_pts); + _final_graph.resize(_max_points + _num_frozen_pts, copy_val); if (_support_eager_delete) { _in_graph.reserve(_max_points + _num_frozen_pts); - _in_graph.resize(_max_points + _num_frozen_pts); + _in_graph.resize(_max_points + _num_frozen_pts, copy_val); } for (uint64_t p = 0; p < _max_points + _num_frozen_pts; p++) { @@ -970,8 +1000,8 @@ namespace diskann { << std::endl; } - template - void Index::build(Parameters ¶meters, + template + void Index::build(Parameters ¶meters, const std::vector &tags) { if (_enable_tags) { if (tags.size() != _nd) { @@ -1008,8 +1038,8 @@ namespace diskann { _has_built = true; } - template - std::pair Index::search(const T *query, + template + std::pair Index::search(const T *query, const size_t K, const unsigned L, unsigned * indices) { @@ -1035,8 +1065,8 @@ namespace diskann { return retval; } - template - std::pair Index::search( + template + std::pair Index::search( const T *query, const uint64_t K, const unsigned L, std::vector init_ids, uint64_t *indices, float *distances) { tsl::robin_set visited(10 * L); @@ -1061,8 +1091,8 @@ namespace diskann { return retval; } - template - std::pair Index::search_with_tags( + template + std::pair Index::search_with_tags( const T *query, const size_t K, const unsigned L, TagT *tags, unsigned frozen_pts, unsigned *indices_buffer) { const bool alloc = indices_buffer == NULL; @@ -1075,12 +1105,15 @@ namespace diskann { return ret; } - template - void Index::optimize_graph() { // use after build or load + template + void Index::optimize_graph(bool use_pm) { // use after build or load _data_len = (_aligned_dim + 1) * sizeof(float); _neighbor_len = (_width + 1) * sizeof(unsigned); _node_size = _data_len + _neighbor_len; - _opt_graph = (char *) malloc(_node_size * _nd); + void* temp_ptr; + diskann::alloc_aligned(&temp_ptr, _node_size * _nd, sizeof(uint8_t), use_pm); + _opt_graph = (char *) temp_ptr; + //_opt_graph = (char *) malloc(_node_size * _nd); DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; for (unsigned i = 0; i < _nd; i++) { char *cur_node_offset = _opt_graph + i * _node_size; @@ -1094,14 +1127,14 @@ namespace diskann { std::memcpy(cur_node_offset, &k, sizeof(unsigned)); std::memcpy(cur_node_offset + sizeof(unsigned), _final_graph[i].data(), k * sizeof(unsigned)); - std::vector().swap(_final_graph[i]); + std::vector(_allocator).swap(_final_graph[i]); } _final_graph.clear(); _final_graph.shrink_to_fit(); } - template - void Index::search_with_opt_graph(const T *query, size_t K, size_t L, + template + void Index::search_with_opt_graph(const T *query, size_t K, size_t L, unsigned *indices) { DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; @@ -1205,8 +1238,8 @@ namespace diskann { // in case we add ''frozen'' auxiliary points to the dataset, these are not // visible to external world, we generate them here and update our dataset - template - int Index::generate_random_frozen_points(const char *filename) { + template + int Index::generate_random_frozen_points(const char *filename) { if (_has_built) { diskann::cout << "Index already built. Cannot add more points" << std::endl; @@ -1242,8 +1275,8 @@ namespace diskann { return 0; } - template - int Index::enable_delete() { + template + int Index::enable_delete() { LockGuard guard(_change_lock); assert(!_can_delete); assert(_enable_tags); @@ -1272,8 +1305,8 @@ namespace diskann { return 0; } - template - int Index::eager_delete(const TagT tag, + template + int Index::eager_delete(const TagT tag, const Parameters ¶meters) { if (_lazy_done && (!_consolidated_order)) { diskann::cout << "Lazy delete reuests issued but data not consolidated, " @@ -1384,8 +1417,8 @@ namespace diskann { return 0; } - template - void Index::update_in_graph() { + template + void Index::update_in_graph() { diskann::cout << "Updating in_graph....." << std::flush; for (unsigned i = 0; i < _in_graph.size(); i++) _in_graph[i].clear(); @@ -1422,8 +1455,8 @@ namespace diskann { // Do not call consolidate_deletes() if you have not locked _change_lock. // Returns number of live points left after consolidation - template - size_t Index::consolidate_deletes(const Parameters ¶meters) { + template + size_t Index::consolidate_deletes(const Parameters ¶meters) { if (_eager_done) { diskann::cout << "No consolidation required, eager deletes done" << std::endl; @@ -1503,8 +1536,8 @@ namespace diskann { return _nd; } - template - std::vector Index::get_new_location(unsigned &active) { + template + std::vector Index::get_new_location(unsigned &active) { std::vector new_location; new_location.resize(_max_points + _num_frozen_pts, (unsigned) (_max_points + _num_frozen_pts)); @@ -1518,8 +1551,8 @@ namespace diskann { return new_location; } - template - void Index::compact_data(std::vector new_location, + template + void Index::compact_data(std::vector new_location, unsigned active, bool &mode) { // If start node is removed, replace it. assert(!mode); @@ -1611,8 +1644,8 @@ namespace diskann { // Do not call reserve_location() if you have not locked _change_lock. // It is not thread safe. - template - unsigned Index::reserve_location() { + template + unsigned Index::reserve_location() { assert(_nd < _max_points); unsigned location; @@ -1632,8 +1665,8 @@ namespace diskann { return location; } - template - void Index::readjust_data(unsigned _num_frozen_pts) { + template + void Index::readjust_data(unsigned _num_frozen_pts) { if (_num_frozen_pts > 0) { if (_final_graph[_max_points].empty()) { diskann::cout << "Readjusting data to correctly position frozen point" @@ -1669,8 +1702,8 @@ namespace diskann { << std::endl; } - template - int Index::insert_point(const T *point, const Parameters ¶meters, + template + int Index::insert_point(const T *point, const Parameters ¶meters, std::vector & pool, std::vector & tmp, tsl::robin_set &visited, @@ -1747,8 +1780,8 @@ namespace diskann { return 0; } - template - int Index::disable_delete(const Parameters ¶meters, + template + int Index::disable_delete(const Parameters ¶meters, const bool consolidate) { LockGuard guard(_change_lock); if (!_can_delete) { @@ -1792,8 +1825,8 @@ namespace diskann { return 0; } - template - int Index::delete_point(const TagT tag) { + template + int Index::delete_point(const TagT tag) { if ((_eager_done) && (!_compacted_order)) { diskann::cout << "Eager delete requests were issued but data was not " "compacted, cannot proceed with lazy_deletes" @@ -1816,4 +1849,8 @@ namespace diskann { template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; + + template DISKANN_DLLEXPORT class Index>; + template DISKANN_DLLEXPORT class Index>; + template DISKANN_DLLEXPORT class Index>; } // namespace diskann diff --git a/src/pm.cpp b/src/pm.cpp new file mode 100644 index 0000000000..524f9f48e3 --- /dev/null +++ b/src/pm.cpp @@ -0,0 +1,62 @@ +#include "pm.h" + +#include +#include +#include +#include + +#include + +// Must be initialized before use. +// Keep the allocation type of the base allocator to `uint8_t`. +// +// Templated function in the header file can copy-construct allocators with other +// type parameters. +static std::unique_ptr> __allocator; +static std::unordered_map __pm_pointers; + +const libmemkind::pmem::allocator& diskann::_pm_allocator() +{ + assert(diskann::is_pm_init()); + return *__allocator.get(); +} + +bool diskann::is_pm_init() +{ + return bool(__allocator); +} + +void diskann::init_pm(const std::string& dir) +{ + if (!is_pm_init()) { + std::cout << "[PM} Initializing PM: " << dir << std::endl; + //auto policy = libmemkind::allocation_policy::CONSERVATIVE; + auto policy = libmemkind::allocation_policy::DEFAULT; + __allocator = std::make_unique>(dir, 0, policy); + } +} + +void* diskann::__alloc(size_t size, size_t align, bool pm) +{ + void* ptr; + if (pm) { + assert(is_pm_init()); + ptr = (void*) __allocator->allocate(size); + __pm_pointers.insert({ptr, size}); + } else { + ptr = ::aligned_alloc(align, size); + } + return ptr; +} + +void diskann::__free(void* ptr) +{ + auto search = __pm_pointers.find(ptr); + if (search != __pm_pointers.end()) { + __allocator->deallocate((uint8_t*) ptr, search->second); + __pm_pointers.erase(search); + } else { + free(ptr); + } +} + diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..3c0fe87b47 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -5,6 +5,7 @@ #include #include #include "utils.h" +#include "pm.h" #ifndef _WINDOWS #include @@ -15,11 +16,13 @@ #include "memory_mapper.h" -template +template> int build_in_memory_index(const std::string& data_path, const unsigned R, const unsigned L, const float alpha, const std::string& save_path, - const unsigned num_threads) { + const unsigned num_threads, + const bool data_in_pm, + const Allocator& allocator = std::allocator()) { diskann::Parameters paras; paras.Set("R", R); paras.Set("L", L); @@ -29,7 +32,7 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, paras.Set("saturate_graph", 0); paras.Set("num_threads", num_threads); - diskann::Index index(diskann::L2, data_path.c_str()); + diskann::Index index(diskann::L2, data_path.c_str(), data_in_pm, allocator); auto s = std::chrono::high_resolution_clock::now(); index.build(paras); std::chrono::duration diff = @@ -42,13 +45,13 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, } int main(int argc, char** argv) { - if (argc != 8) { + if (argc != 11) { std::cout << "Usage: " << argv[0] - << " [data_type] [data_file.bin] " - "[output_index_file] " - << "[R] [L] [alpha]" - << " [num_threads_to_use]. See README for more information on " - "parameters." + << " [data_type] [data_file.bin] [output_index_file]" + << " [R] [L] [alpha]" + << " [num_threads_to_use]" + << " [PM directory (use \"null\" for none)] [Data in PM] [Graph in PM]." + << " See README for more information on parameters." << std::endl; exit(-1); } @@ -59,16 +62,35 @@ int main(int argc, char** argv) { const unsigned L = (unsigned) atoi(argv[5]); const float alpha = (float) atof(argv[6]); const unsigned num_threads = (unsigned) atoi(argv[7]); + const std::string pm_directory(argv[8]); + const bool data_in_pm = std::atoi(argv[9]); + const bool graph_in_pm = std::atoi(argv[10]); - if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); - else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); - else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, - num_threads); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + if (pm_directory != "null") { + diskann::init_pm(pm_directory); + } else if (graph_in_pm || data_in_pm) { + std::cout << "Please set the PM Directory in order to put the graph in PM" << std::endl; + exit(-1); + } + + if (!graph_in_pm) { + if (std::string(argv[1]) == std::string("int8")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + else if (std::string(argv[1]) == std::string("uint8")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + else if (std::string(argv[1]) == std::string("float")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } else { + auto allocator = diskann::pm_allocator(); + if (std::string(argv[1]) == std::string("int8")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + else if (std::string(argv[1]) == std::string("uint8")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + else if (std::string(argv[1]) == std::string("float")) + build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } } diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 32592710e5..513c872800 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -17,10 +18,14 @@ #include "aux_utils.h" #include "index.h" #include "memory_mapper.h" +#include "pm.h" #include "utils.h" -template -int search_memory_index(int argc, char** argv) { +template> +int search_memory_index(int argc, + char** argv, + const Allocator& allocator = std::allocator() +) { T* query = nullptr; unsigned* gt_ids = nullptr; float* gt_dists = nullptr; @@ -35,6 +40,10 @@ int search_memory_index(int argc, char** argv) { _u64 recall_at = std::atoi(argv[7]); std::string result_output_prefix(argv[8]); bool use_optimized_search = std::atoi(argv[9]); + std::string pm_directory(argv[10]); + bool data_in_pm = std::atoi(argv[11]); + // `graph_in_pm` checked by caller. + //bool graph_in_pm = std::atoi(argv[12]); if ((std::string(argv[1]) != std::string("float")) && (use_optimized_search == true)) { @@ -46,7 +55,7 @@ int search_memory_index(int argc, char** argv) { bool calc_recall_flag = false; - for (int ctr = 10; ctr < argc; ctr++) { + for (int ctr = 13; ctr < argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); @@ -58,8 +67,11 @@ int search_memory_index(int argc, char** argv) { return -1; } - diskann::load_aligned_bin(query_bin, query, query_num, query_dim, - query_aligned_dim); + if (data_in_pm && pm_directory == "null") { + std::cout << "Please set a PM directory to use PM!" << std::endl; + return -1; + } + diskann::load_aligned_bin(query_bin, query, query_num, query_dim, query_aligned_dim); if (file_exists(truthset_bin)) { diskann::load_truthset(truthset_bin, gt_ids, gt_dists, gt_num, gt_dim); @@ -76,12 +88,17 @@ int search_memory_index(int argc, char** argv) { auto metric = diskann::L2; if (use_optimized_search) metric = diskann::FAST_L2; - diskann::Index index(metric, data_file.c_str()); + diskann::Index index( + metric, + data_file.c_str(), + data_in_pm, + allocator + ); index.load(memory_index_file.c_str()); // to load NSG std::cout << "Index loaded" << std::endl; if (use_optimized_search) - index.optimize_graph(); + index.optimize_graph(data_in_pm); diskann::Parameters paras; std::string recall_string = "Recall@" + std::to_string(recall_at); @@ -157,7 +174,7 @@ int search_memory_index(int argc, char** argv) { } int main(int argc, char** argv) { - if (argc < 11) { + if (argc < 14) { std::cout << "Usage: " << argv[0] << " [index_type] [data_file.bin] " @@ -165,16 +182,42 @@ int main(int argc, char** argv) { "[query_file.bin] [truthset.bin (use \"null\" for none)] " " [K] [result_output_prefix] [use_optimized_search (for small ~1M " "data)] " + " [PM directory (use \"null\" for none)] [Data in PM] [Graph in PM]" " [L1] [L2] etc. See README for more information on parameters. " << std::endl; exit(-1); } - if (std::string(argv[1]) == std::string("int8")) - search_memory_index(argc, argv); - else if (std::string(argv[1]) == std::string("uint8")) - search_memory_index(argc, argv); - else if (std::string(argv[1]) == std::string("float")) - search_memory_index(argc, argv); - else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + + // Parse some PM options here so the `diskann::pmem_allocator` can be constructed + // if needed. + std::string pm_directory(argv[10]); + bool graph_in_pm = std::atoi(argv[12]); + + if (pm_directory != "null") { + diskann::init_pm(pm_directory); + } else if (graph_in_pm) { + std::cout << "Please set the PM Directory in order to put the graph in PM" << std::endl; + exit(-1); + } + + if (!graph_in_pm) { + if (std::string(argv[1]) == std::string("int8")) + search_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("uint8")) + search_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("float")) + search_memory_index(argc, argv); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } else { + auto allocator = diskann::pm_allocator(); + if (std::string(argv[1]) == std::string("int8")) + search_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("uint8")) + search_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("float")) + search_memory_index(argc, argv); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + } } From fd3da70b12d9885b441cdb91857c9703acfc9070 Mon Sep 17 00:00:00 2001 From: "Hildebrand, Mark" Date: Tue, 4 May 2021 01:01:40 -0700 Subject: [PATCH 2/4] Set compiler back to `g++` --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 02d5a268d2..1df0e25f74 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,7 +14,7 @@ if(MSVC) set(CMAKE_CXX_COMPILER $ENV{VCToolsInstallDir}/bin/Hostx64/x64/cl.exe) set(CMAKE_CXX_LINK_EXECUTABLE $ENV{VCToolsInstallDir}/bin/Hostx64/x64/link.exe) else() - set(CMAKE_CXX_COMPILER clang++) + set(CMAKE_CXX_COMPILER g++) endif() project(diskann) From b8cc19b7f78d2250361b084599c9d0aa3a9d9fea Mon Sep 17 00:00:00 2001 From: sdongaon Date: Wed, 12 May 2021 15:30:31 -0700 Subject: [PATCH 3/4] added ifdefs to bypass PMem specific changes on windows --- CMakeLists.txt | 32 ++++++++++++---- CMakeSettings.json | 17 +++++++++ include/pm.h | 10 ++--- src/index.cpp | 30 ++++++++------- src/pm.cpp | 10 ++++- tests/build_memory_index.cpp | 72 ++++++++++++++++++++++++++++------- tests/search_memory_index.cpp | 72 +++++++++++++++++++++++++---------- 7 files changed, 182 insertions(+), 61 deletions(-) create mode 100644 CMakeSettings.json diff --git a/CMakeLists.txt b/CMakeLists.txt index 1df0e25f74..9bac2c277f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -48,9 +48,9 @@ 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() @@ -61,8 +61,24 @@ else() 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}) @@ -86,13 +102,13 @@ if(MSVC) 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 new file mode 100644 index 0000000000..9d3f1c813e --- /dev/null +++ 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" + } + ] +} \ No newline at end of file diff --git a/include/pm.h b/include/pm.h index b1800fc251..22babb1e13 100644 --- a/include/pm.h +++ b/include/pm.h @@ -1,10 +1,9 @@ #pragma once - -#include // For some reason, "pmem_allocator.h" doesn't include this. -#include - +#ifndef _WINDOWS #include #include +#include +#include namespace diskann { bool is_pm_init(); @@ -24,4 +23,5 @@ namespace diskann { { return pmem_allocator(_pm_allocator()); } -} +} // namespace diskann +#endif diff --git a/src/index.cpp b/src/index.cpp index 3cf149a110..c9877f4118 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -32,7 +32,9 @@ #include "memory_mapper.h" #include "parameters.h" #include "partition_and_pq.h" +#ifndef _WINDOWS #include "pm.h" +#endif #include "timer.h" #include "utils.h" #include "windows_customizations.h" @@ -193,29 +195,29 @@ namespace diskann { aligned_free(_opt_graph); } - template<> - Index>::~Index() { - delete this->_distance; - aligned_free(_data); - aligned_free(_opt_graph); + + template <> Index<_s8>::~Index() { + delete this->_distance; + aligned_free(_data); + aligned_free(_opt_graph); } - template<> - Index<_s8>::~Index() { - delete this->_distance; - aligned_free(_data); - aligned_free(_opt_graph); + template <> Index<_u8>::~Index() { + delete this->_distance; + aligned_free(_data); + aligned_free(_opt_graph); } + #ifndef _WINDOWS template<> - Index<_s8,int,diskann::pmem_allocator>::~Index() { + Index>::~Index() { delete this->_distance; aligned_free(_data); aligned_free(_opt_graph); } template<> - Index<_u8>::~Index() { + Index<_s8,int,diskann::pmem_allocator>::~Index() { delete this->_distance; aligned_free(_data); aligned_free(_opt_graph); @@ -227,6 +229,7 @@ namespace diskann { aligned_free(_data); aligned_free(_opt_graph); } + #endif // save the graph index on a file as an adjacency list. For each point, // first store the number of neighbors, and then the neighbor list (each as @@ -1849,8 +1852,9 @@ namespace diskann { template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; template DISKANN_DLLEXPORT class Index; - + #ifndef _WINDOWS template DISKANN_DLLEXPORT class Index>; template DISKANN_DLLEXPORT class Index>; template DISKANN_DLLEXPORT class Index>; + #endif } // namespace diskann diff --git a/src/pm.cpp b/src/pm.cpp index 524f9f48e3..107336390f 100644 --- a/src/pm.cpp +++ b/src/pm.cpp @@ -1,11 +1,16 @@ -#include "pm.h" +#pragma once + +#ifndef _WINDOWS #include #include #include #include -#include + + #include "pm.h" + #include + // Must be initialized before use. // Keep the allocation type of the base allocator to `uint8_t`. @@ -60,3 +65,4 @@ void diskann::__free(void* ptr) } } +#endif \ No newline at end of file diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index 3c0fe87b47..4b07eec5a2 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -45,6 +45,7 @@ int build_in_memory_index(const std::string& data_path, const unsigned R, } int main(int argc, char** argv) { + #ifndef _WINDOWS if (argc != 11) { std::cout << "Usage: " << argv[0] << " [data_type] [data_file.bin] [output_index_file]" @@ -58,10 +59,10 @@ int main(int argc, char** argv) { const std::string data_path(argv[2]); const std::string save_path(argv[3]); - const unsigned R = (unsigned) atoi(argv[4]); - const unsigned L = (unsigned) atoi(argv[5]); - const float alpha = (float) atof(argv[6]); - const unsigned num_threads = (unsigned) atoi(argv[7]); + const unsigned R = (unsigned)atoi(argv[4]); + const unsigned L = (unsigned)atoi(argv[5]); + const float alpha = (float)atof(argv[6]); + const unsigned num_threads = (unsigned)atoi(argv[7]); const std::string pm_directory(argv[8]); const bool data_in_pm = std::atoi(argv[9]); const bool graph_in_pm = std::atoi(argv[10]); @@ -69,28 +70,73 @@ int main(int argc, char** argv) { if (pm_directory != "null") { diskann::init_pm(pm_directory); } else if (graph_in_pm || data_in_pm) { - std::cout << "Please set the PM Directory in order to put the graph in PM" << std::endl; + std::cout << "Please set the PM Directory in order to put the graph in PM" + << std::endl; exit(-1); } if (!graph_in_pm) { if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } else { auto allocator = diskann::pm_allocator(); if (std::string(argv[1]) == std::string("int8")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm, allocator); else if (std::string(argv[1]) == std::string("uint8")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm, allocator); else if (std::string(argv[1]) == std::string("float")) - build_in_memory_index(data_path, R, L, alpha, save_path, num_threads, data_in_pm, allocator); + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm, allocator); else - std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } + +#else + if (argc != 8) { + std::cout + << "Usage: " << argv[0] + << " [data_type] [data_file.bin] " + "[output_index_file] " + << "[R] [L] [alpha]" + << " [num_threads_to_use]. See README for more information on " + "parameters." + << std::endl; + exit(-1); + } + + const std::string data_path(argv[2]); + const std::string save_path(argv[3]); + const unsigned R = (unsigned)atoi(argv[4]); + const unsigned L = (unsigned)atoi(argv[5]); + const float alpha = (float)atof(argv[6]); + const unsigned num_threads = (unsigned)atoi(argv[7]); + const std::string pm_directory = "null"; + const bool data_in_pm = 0; + const bool graph_in_pm = 0; + + if (std::string(argv[1]) == std::string("int8")) + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); + else if (std::string(argv[1]) == std::string("uint8")) + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); + else if (std::string(argv[1]) == std::string("float")) + build_in_memory_index(data_path, R, L, alpha, save_path, + num_threads, data_in_pm); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + #endif + + } diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 513c872800..87d367cee0 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -13,17 +13,16 @@ #include #include #include +#include "pm.h" #endif #include "aux_utils.h" #include "index.h" #include "memory_mapper.h" -#include "pm.h" #include "utils.h" template> -int search_memory_index(int argc, - char** argv, +int search_memory_index(int argc, char** argv, const Allocator& allocator = std::allocator() ) { T* query = nullptr; @@ -34,16 +33,13 @@ int search_memory_index(int argc, std::string data_file(argv[2]); std::string memory_index_file(argv[3]); - _u64 num_threads = std::atoi(argv[4]); + _u64 num_threads = std::atoi(argv[4]); std::string query_bin(argv[5]); std::string truthset_bin(argv[6]); - _u64 recall_at = std::atoi(argv[7]); + _u64 recall_at = std::atoi(argv[7]); std::string result_output_prefix(argv[8]); - bool use_optimized_search = std::atoi(argv[9]); - std::string pm_directory(argv[10]); - bool data_in_pm = std::atoi(argv[11]); - // `graph_in_pm` checked by caller. - //bool graph_in_pm = std::atoi(argv[12]); + bool use_optimized_search = std::atoi(argv[9]); + if ((std::string(argv[1]) != std::string("float")) && (use_optimized_search == true)) { @@ -52,14 +48,27 @@ int search_memory_index(int argc, << std::endl; use_optimized_search = false; } - - bool calc_recall_flag = false; + #ifndef _WINDOWS + std::string pm_directory(argv[10]); + bool data_in_pm = std::atoi(argv[11]); for (int ctr = 13; ctr < argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) Lvec.push_back(curL); } + #else + std::string pm_directory = "null"; + bool data_in_pm = 0; + + for (int ctr = 10; ctr < argc; ctr++) { + _u64 curL = std::atoi(argv[ctr]); + if (curL >= recall_at) + Lvec.push_back(curL); + } + #endif + + bool calc_recall_flag = false; if (Lvec.size() == 0) { std::cout << "No valid Lsearch found. Lsearch must be at least recall_at." @@ -174,6 +183,7 @@ int search_memory_index(int argc, } int main(int argc, char** argv) { + #ifndef _WINDOWS if (argc < 14) { std::cout << "Usage: " << argv[0] @@ -190,8 +200,6 @@ int main(int argc, char** argv) { // Parse some PM options here so the `diskann::pmem_allocator` can be constructed // if needed. - std::string pm_directory(argv[10]); - bool graph_in_pm = std::atoi(argv[12]); if (pm_directory != "null") { diskann::init_pm(pm_directory); @@ -202,22 +210,46 @@ int main(int argc, char** argv) { if (!graph_in_pm) { if (std::string(argv[1]) == std::string("int8")) - search_memory_index(argc, argv); + search_memory_index(argc,argv); else if (std::string(argv[1]) == std::string("uint8")) - search_memory_index(argc, argv); + search_memory_index(argc, argv); else if (std::string(argv[1]) == std::string("float")) - search_memory_index(argc, argv); + search_memory_index(argc, argv); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } else { auto allocator = diskann::pm_allocator(); if (std::string(argv[1]) == std::string("int8")) - search_memory_index(argc, argv); + search_memory_index(argc, argv, allocator); else if (std::string(argv[1]) == std::string("uint8")) - search_memory_index(argc, argv); + search_memory_index(argc, argv, allocator); else if (std::string(argv[1]) == std::string("float")) - search_memory_index(argc, argv); + search_memory_index(argc, argv, allocator); else std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; } + #else + if (argc < 11) { + std::cout + << "Usage: " << argv[0] + << " [index_type] [data_file.bin] " + "[memory_index_path] [num_threads] " + "[query_file.bin] [truthset.bin (use \"null\" for none)] " + " [K] [result_output_prefix] [use_optimized_search (for small ~1M " + "data)] " + " [L1] [L2] etc. See README for more information on parameters. " + << std::endl; + exit(-1); + } + + + if (std::string(argv[1]) == std::string("int8")) + search_memory_index(argc,argv); + else if (std::string(argv[1]) == std::string("uint8")) + search_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("float")) + search_memory_index(argc, argv); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; + #endif } From 643debc58ff0293c5b5cfc969e47308588d6d0cc Mon Sep 17 00:00:00 2001 From: sdongaon Date: Fri, 14 May 2021 06:00:54 -0700 Subject: [PATCH 4/4] fixed argumnet passing in search on linux --- CMakeLists.txt | 2 +- src/pm.cpp | 4 +--- tests/search_memory_index.cpp | 13 +++++++------ 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9bac2c277f..7893b7269e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -54,7 +54,7 @@ if (MSVC) 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) link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core iomp5 pthread m dl) diff --git a/src/pm.cpp b/src/pm.cpp index 107336390f..dac1f5f1dc 100644 --- a/src/pm.cpp +++ b/src/pm.cpp @@ -1,5 +1,3 @@ -#pragma once - #ifndef _WINDOWS #include @@ -65,4 +63,4 @@ void diskann::__free(void* ptr) } } -#endif \ No newline at end of file +#endif diff --git a/tests/search_memory_index.cpp b/tests/search_memory_index.cpp index 87d367cee0..51ce72b4ba 100644 --- a/tests/search_memory_index.cpp +++ b/tests/search_memory_index.cpp @@ -49,9 +49,13 @@ int search_memory_index(int argc, char** argv, use_optimized_search = false; } #ifndef _WINDOWS - std::string pm_directory(argv[10]); bool data_in_pm = std::atoi(argv[11]); + std::string pm_directory(argv[10]); + if (data_in_pm && pm_directory == "null") { + std::cout << "Please set a PM directory to use PM!" << std::endl; + return -1; + } for (int ctr = 13; ctr < argc; ctr++) { _u64 curL = std::atoi(argv[ctr]); if (curL >= recall_at) @@ -76,10 +80,6 @@ int search_memory_index(int argc, char** argv, return -1; } - if (data_in_pm && pm_directory == "null") { - std::cout << "Please set a PM directory to use PM!" << std::endl; - return -1; - } diskann::load_aligned_bin(query_bin, query, query_num, query_dim, query_aligned_dim); if (file_exists(truthset_bin)) { @@ -200,7 +200,8 @@ int main(int argc, char** argv) { // Parse some PM options here so the `diskann::pmem_allocator` can be constructed // if needed. - + std::string pm_directory(argv[10]); + bool graph_in_pm = std::atoi(argv[12]); if (pm_directory != "null") { diskann::init_pm(pm_directory); } else if (graph_in_pm) {