diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd3886f5b..3240bc4309 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -95,7 +95,7 @@ else() # set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -O0 -fsanitize=address -fsanitize=leak -fsanitize=undefined") set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG -Wall -Wextra") set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -Ofast -DNDEBUG -march=native -mtune=native -ftree-vectorize") - add_compile_options(-march=native -Wall -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_ACCELERATED_PQ -DUSE_AVX2) + add_compile_options(-march=native -Wall -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_ACCELERATED_PQ -fPIC) endif() add_subdirectory(src) diff --git a/include/distance.h b/include/distance.h index a4f311e510..b037d1042a 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,7 @@ #pragma once #include +#include #ifdef _WINDOWS #include #include @@ -174,7 +175,54 @@ namespace diskann { } }; - class DistanceL2 : public Distance { + class AVX512DistanceL2Float : public Distance { + public: +#ifndef _WINDOWS + float compare(const float *a, const float *b, unsigned size) const + __attribute__((hot)) { + a = (const float *) __builtin_assume_aligned(a, 32); + b = (const float *) __builtin_assume_aligned(b, 32); +#else + float compare(const float *a, const float *b, unsigned size) const { +#endif + + float result = 0; +#ifdef USE_AVX512 + // assume size is divisible by 16 + _u16 niters = size / 16; + __m512 sum = _mm512_setzero_ps(); + for (_u16 j = 0; j < niters; j++) { + // scope is a[16j:16j+15], b[16j:16j+15] + // load a_vec + if (j < (niters - 1)) { + _mm_prefetch((char *) (a + 16 * (j + 1)), _MM_HINT_T0); + _mm_prefetch((char *) (b + 16 * (j + 1)), _MM_HINT_T0); + } + __m512 a_vec = _mm512_load_ps(a + 16 * j); + // load b_vec + __m512 b_vec = _mm512_load_ps(b + 16 * j); + // a_vec - b_vec + __m512 tmp_vec = _mm512_sub_ps(a_vec, b_vec); + + // sum = (tmp_vec**2) + sum + sum = _mm512_fmadd_ps(tmp_vec, tmp_vec, sum); + } + + // horizontal add sum + result = _mm512_reduce_add_ps(sum); +#else +#ifndef _WINDOWS +#pragma omp simd reduction(+ : result) aligned(a, b : 32) +#endif + for (_s32 i = 0; i < (_s32) size; i++) { + result += (a[i] - b[i]) * (a[i] - b[i]); + } +#endif + return result; + } + }; + + class AVX2DistanceL2Float : public Distance { public: #ifndef _WINDOWS float compare(const float *a, const float *b, unsigned size) const @@ -330,202 +378,299 @@ namespace diskann { template class DistanceInnerProduct : public Distance { public: - float compare(const T *a, const T *b, unsigned size) const { + virtual float norm(const T *a, unsigned size) const { float result = 0; #ifdef __GNUC__ -#ifdef __AVX__ -#define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm256_loadu_ps(addr1); \ - tmp2 = _mm256_loadu_ps(addr2); \ - tmp1 = _mm256_mul_ps(tmp1, tmp2); \ - dest = _mm256_add_ps(dest, tmp1); +#ifdef __AVX512F__ +#define AVX512_L2NORM(addr, dest, tmp) \ + tmp = _mm512_loadu_ps(addr); \ + dest = _mm512_fmadd_ps(tmp, tmp, dest); + + __m512 sum; + __m512 l0, l1; + unsigned D = (size + 15) & ~15U; + unsigned DR = D % 32; + unsigned DD = D - DR; + const float *l = (float *) a; + const float *e_l = l + DD; + + sum = _mm512_setzero_ps(); + if (DR) { + AVX512_L2NORM(e_l, sum, l0); + } + + for (unsigned i = 0; i < DD; i += 32, l += 32) { + AVX512_L2NORM(l, sum, l0); + AVX512_L2NORM(l + 16, sum, l1); + } + + result = _mm512_reduce_add_ps(sum); +#elif defined(__AVX__) +#define AVX_L2NORM(addr, dest, tmp) \ + tmp = _mm256_loadu_ps(addr); \ + dest = _mm256_fmadd_ps(tmp, tmp, dest); __m256 sum; __m256 l0, l1; - __m256 r0, r1; unsigned D = (size + 7) & ~7U; unsigned DR = D % 16; unsigned DD = D - DR; const float *l = (float *) a; - const float *r = (float *) b; const float *e_l = l + DD; - const float *e_r = r + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; - sum = _mm256_loadu_ps(unpack); + sum = _mm256_setzero_ps(); if (DR) { - AVX_DOT(e_l, e_r, sum, l0, r0); + AVX_L2NORM(e_l, sum, l0); } - for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { - AVX_DOT(l, r, sum, l0, r0); - AVX_DOT(l + 8, r + 8, sum, l1, r1); + for (unsigned i = 0; i < DD; i += 16, l += 16) { + AVX_L2NORM(l, sum, l0); + AVX_L2NORM(l + 8, sum, l1); } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + - unpack[5] + unpack[6] + unpack[7]; + result = _mm256_reduce_add_ps(sum); #else #ifdef __SSE2__ -#define SSE_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm128_loadu_ps(addr1); \ - tmp2 = _mm128_loadu_ps(addr2); \ - tmp1 = _mm128_mul_ps(tmp1, tmp2); \ - dest = _mm128_add_ps(dest, tmp1); +#define SSE_L2NORM(addr, dest, tmp) \ + tmp = _mm128_loadu_ps(addr); \ + tmp = _mm128_mul_ps(tmp, tmp); \ + dest = _mm128_add_ps(dest, tmp); + __m128 sum; __m128 l0, l1, l2, l3; - __m128 r0, r1, r2, r3; unsigned D = (size + 3) & ~3U; unsigned DR = D % 16; unsigned DD = D - DR; const float *l = a; - const float *r = b; const float *e_l = l + DD; - const float *e_r = r + DD; float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; sum = _mm_load_ps(unpack); switch (DR) { case 12: - SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); + SSE_L2NORM(e_l + 8, sum, l2); case 8: - SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); + SSE_L2NORM(e_l + 4, sum, l1); case 4: - SSE_DOT(e_l, e_r, sum, l0, r0); + SSE_L2NORM(e_l, sum, l0); default: break; } - for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { - SSE_DOT(l, r, sum, l0, r0); - SSE_DOT(l + 4, r + 4, sum, l1, r1); - SSE_DOT(l + 8, r + 8, sum, l2, r2); - SSE_DOT(l + 12, r + 12, sum, l3, r3); + for (unsigned i = 0; i < DD; i += 16, l += 16) { + SSE_L2NORM(l, sum, l0); + SSE_L2NORM(l + 4, sum, l1); + SSE_L2NORM(l + 8, sum, l2); + SSE_L2NORM(l + 12, sum, l3); } _mm_storeu_ps(unpack, sum); result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; #else - float dot0, dot1, dot2, dot3; const float *last = a + size; const float *unroll_group = last - 3; /* Process 4 items with each loop for efficiency. */ while (a < unroll_group) { - dot0 = a[0] * b[0]; - dot1 = a[1] * b[1]; - dot2 = a[2] * b[2]; - dot3 = a[3] * b[3]; + dot0 = a[0] * a[0]; + dot1 = a[1] * a[1]; + dot2 = a[2] * a[2]; + dot3 = a[3] * a[3]; result += dot0 + dot1 + dot2 + dot3; a += 4; - b += 4; } /* Process last 0-3 pixels. Not needed for standard vector lengths. */ while (a < last) { - result += *a++ * *b++; + result += (*a) * (*a); + a++; } #endif #endif #endif return result; } - }; - template - class DistanceFastL2 : public DistanceInnerProduct { - public: - float norm(const T *a, unsigned size) const { + virtual float compare(const T *a, const T *b, float norm, + unsigned size) const = 0; + +#ifndef _WINDOWS + virtual float compare(const T *a, const T *b, unsigned size) const + __attribute__((hot)) { + const float *l = (const float *) __builtin_assume_aligned(a, 32); + const float *r = (const float *) __builtin_assume_aligned(b, 32); +#else + virtual float compare(const T *a, const T *b, unsigned size) const { + const float *l = (float *) a; + const float *r = (float *) b; +#endif float result = 0; #ifdef __GNUC__ -#ifdef __AVX__ -#define AVX_L2NORM(addr, dest, tmp) \ - tmp = _mm256_loadu_ps(addr); \ - tmp = _mm256_mul_ps(tmp, tmp); \ - dest = _mm256_add_ps(dest, tmp); +#ifdef __AVX512F__ +#define AVX512_DOT(addr1, addr2, dest, tmp1, tmp2) \ + tmp1 = _mm512_loadu_ps(addr1); \ + tmp2 = _mm512_loadu_ps(addr2); \ + dest = _mm512_fmadd_ps(tmp1, tmp2, dest); + + __m512 sum; + __m512 l0, l1; + __m512 r0, r1; + unsigned D = (size + 15) & ~15U; + unsigned DR = D % 32; + unsigned DD = D - DR; + const float *e_l = l + DD; + const float *e_r = r + DD; + + sum = _mm512_setzero_ps(); + if (DR) { + AVX512_DOT(e_l, e_r, sum, l0, r0); + } + + for (unsigned i = 0; i < DD; i += 32, l += 32, r += 32) { + AVX512_DOT(l, r, sum, l0, r0); + AVX512_DOT(l + 16, r + 16, sum, l1, r1); + } + + result = _mm512_reduce_add_ps(sum); +#elif defined(__AVX__) +#define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \ + tmp1 = _mm256_loadu_ps(addr1); \ + tmp2 = _mm256_loadu_ps(addr2); \ + dest = _mm256_fmadd_ps(tmp1, tmp2, dest); __m256 sum; __m256 l0, l1; + __m256 r0, r1; unsigned D = (size + 7) & ~7U; unsigned DR = D % 16; unsigned DD = D - DR; - const float *l = (float *) a; const float *e_l = l + DD; - float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0}; + const float *e_r = r + DD; - sum = _mm256_loadu_ps(unpack); + sum = _mm256_setzero_ps(); if (DR) { - AVX_L2NORM(e_l, sum, l0); + AVX_DOT(e_l, e_r, sum, l0, r0); } - for (unsigned i = 0; i < DD; i += 16, l += 16) { - AVX_L2NORM(l, sum, l0); - AVX_L2NORM(l + 8, sum, l1); + + for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { + AVX_DOT(l, r, sum, l0, r0); + AVX_DOT(l + 8, r + 8, sum, l1, r1); } - _mm256_storeu_ps(unpack, sum); - result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] + - unpack[5] + unpack[6] + unpack[7]; + + result = _mm256_reduce_add_ps(sum); #else #ifdef __SSE2__ -#define SSE_L2NORM(addr, dest, tmp) \ - tmp = _mm128_loadu_ps(addr); \ - tmp = _mm128_mul_ps(tmp, tmp); \ - dest = _mm128_add_ps(dest, tmp); - +#define SSE_DOT(addr1, addr2, dest, tmp1, tmp2) \ + tmp1 = _mm128_loadu_ps(addr1); \ + tmp2 = _mm128_loadu_ps(addr2); \ + tmp1 = _mm128_mul_ps(tmp1, tmp2); \ + dest = _mm128_add_ps(dest, tmp1); __m128 sum; __m128 l0, l1, l2, l3; + __m128 r0, r1, r2, r3; unsigned D = (size + 3) & ~3U; unsigned DR = D % 16; unsigned DD = D - DR; - const float *l = a; const float *e_l = l + DD; + const float *e_r = r + DD; float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0}; sum = _mm_load_ps(unpack); switch (DR) { case 12: - SSE_L2NORM(e_l + 8, sum, l2); + SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2); case 8: - SSE_L2NORM(e_l + 4, sum, l1); + SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1); case 4: - SSE_L2NORM(e_l, sum, l0); + SSE_DOT(e_l, e_r, sum, l0, r0); default: break; } - for (unsigned i = 0; i < DD; i += 16, l += 16) { - SSE_L2NORM(l, sum, l0); - SSE_L2NORM(l + 4, sum, l1); - SSE_L2NORM(l + 8, sum, l2); - SSE_L2NORM(l + 12, sum, l3); + for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) { + SSE_DOT(l, r, sum, l0, r0); + SSE_DOT(l + 4, r + 4, sum, l1, r1); + SSE_DOT(l + 8, r + 8, sum, l2, r2); + SSE_DOT(l + 12, r + 12, sum, l3, r3); } _mm_storeu_ps(unpack, sum); result += unpack[0] + unpack[1] + unpack[2] + unpack[3]; #else + float dot0, dot1, dot2, dot3; const float *last = a + size; const float *unroll_group = last - 3; /* Process 4 items with each loop for efficiency. */ while (a < unroll_group) { - dot0 = a[0] * a[0]; - dot1 = a[1] * a[1]; - dot2 = a[2] * a[2]; - dot3 = a[3] * a[3]; + dot0 = a[0] * b[0]; + dot1 = a[1] * b[1]; + dot2 = a[2] * b[2]; + dot3 = a[3] * b[3]; result += dot0 + dot1 + dot2 + dot3; a += 4; + b += 4; } /* Process last 0-3 pixels. Not needed for standard vector lengths. */ while (a < last) { - result += (*a) * (*a); - a++; + result += *a++ * *b++; } #endif #endif #endif return result; } - using DistanceInnerProduct::compare; - float compare(const T *a, const T *b, float norm, - unsigned size) const { // not implement - float result = -2 * DistanceInnerProduct::compare(a, b, size); - result += norm; + }; + + template + class DistanceFastL2 : public DistanceInnerProduct { + public: + float norm(const T *a, unsigned size) const { + float norm = DistanceInnerProduct::norm(a, size); + if (norm == 0.0) { + return std::numeric_limits::max(); + } + return norm; + } + + float compare(const T *a, const T *b, unsigned size) const { + float norm_a = DistanceInnerProduct::norm(a, size); + float norm_b = DistanceInnerProduct::norm(b, size); + if (norm_a == 0.0 || norm_b == 0.0) { + return std::numeric_limits::max(); + } + float result = + norm_a + norm_b - (2 * DistanceInnerProduct::compare(a, b, size)); + return result; + } + + float compare(const T *a, const T *b, float norm, unsigned size) const { + float result = norm - (2 * DistanceInnerProduct::compare(a, b, size)); + return result; + } + }; + + template + class DistanceFastInnerProduct : public DistanceInnerProduct { + public: + float norm(const T *a, unsigned size) const { + float norm = std::sqrt(DistanceInnerProduct::norm(a, size)); + if (norm == 0.0) { + return std::numeric_limits::max(); + } + return 1 / norm; + } + + float compare(const T *a, const T *b, unsigned size) const { + float norm_a = std::sqrt(DistanceInnerProduct::norm(a, size)); + float norm_b = std::sqrt(DistanceInnerProduct::norm(b, size)); + if (norm_a == 0.0 || norm_b == 0.0) { + return std::numeric_limits::max(); + } + float result = + DistanceInnerProduct::compare(a, b, size) / (norm_a * norm_b); + return 1 - result; + } + + float compare(const T *a, const T *b, float norm, unsigned size) const { + float result = -DistanceInnerProduct::compare(a, b, size) * norm; return result; } }; diff --git a/include/index.h b/include/index.h index eedbb1491f..4236ce0582 100644 --- a/include/index.h +++ b/include/index.h @@ -17,6 +17,8 @@ #include "utils.h" #include "windows_customizations.h" +#include "pq_table.h" + #define SLACK_FACTOR 1.3 #define ESTIMATE_RAM_USAGE(size, dim, datasize, degree) \ @@ -41,6 +43,7 @@ namespace diskann { DISKANN_DLLEXPORT void load(const char *filename, const bool load_tags = false, const char *tag_filename = NULL); + DISKANN_DLLEXPORT void pq_load(const char *pq_prefix); // generates one or more frozen points that will never get deleted from the // graph DISKANN_DLLEXPORT int generate_random_frozen_points( @@ -49,6 +52,9 @@ namespace diskann { DISKANN_DLLEXPORT void build( Parameters & parameters, const std::vector &tags = std::vector()); + DISKANN_DLLEXPORT void pq_build(const char *dataFilePath, + const char *indexFilePath, + Parameters ¶meters); // Gopal. Added search overload that takes L as parameter, so that we // can customize L on a per-query basis without tampering with "Parameters" @@ -102,6 +108,8 @@ namespace diskann { DISKANN_DLLEXPORT void search_with_opt_graph(const T *query, size_t K, size_t L, unsigned *indices); + DISKANN_DLLEXPORT void pq_search(T *query, size_t K, size_t L, + unsigned *indices); /* Internals of the library */ protected: @@ -178,6 +186,7 @@ namespace diskann { unsigned _width; unsigned _ep; bool _saturate_graph = false; + bool _normalize = false; std::vector _locks; // Per node lock, cardinality=max_points_ char * _opt_graph; @@ -198,6 +207,17 @@ namespace diskann { // deletion bool _store_data; + // _pq_data: Stores the data points in compressed format (_u8 * n_chunks) + // _chunk_size = chunk size of each dimension chunk + // _n_chunks = # of bytes the data is compressed to + // _pq_table = [[2^8 * [chunk_size]] * n_chunks] + _u8 * _pq_data = nullptr; + _u64 _chunk_size; + _u64 _n_chunks; + FixedChunkPQTable _pq_table; + + float *_pq_table_dists = nullptr; // Must be atleast [256 * _n_chunks] + std::unordered_map _tag_to_location; std::unordered_map _location_to_tag; diff --git a/include/neighbor.h b/include/neighbor.h index 5c37df38e5..aba8266dc1 100644 --- a/include/neighbor.h +++ b/include/neighbor.h @@ -113,15 +113,15 @@ namespace diskann { Neighbor nn) { // find the location to insert unsigned left = 0, right = K - 1; + if (addr[right].distance < nn.distance) { + addr[K] = nn; + return K; + } if (addr[left].distance > nn.distance) { memmove((char *) &addr[left + 1], &addr[left], K * sizeof(Neighbor)); addr[left] = nn; return left; } - if (addr[right].distance < nn.distance) { - addr[K] = nn; - return K; - } while (right > 1 && left < right - 1) { unsigned mid = (left + right) / 2; if (addr[mid].distance > nn.distance) @@ -129,8 +129,8 @@ namespace diskann { else left = mid; } - // check equal ID + // check equal ID while (left > 0) { if (addr[left].distance < nn.distance) break; diff --git a/include/parameters.h b/include/parameters.h index 42e7d3a463..1d9f78d035 100644 --- a/include/parameters.h +++ b/include/parameters.h @@ -45,7 +45,7 @@ namespace diskann { const ParamType & default_value) { try { return Get(name); - } catch (std::invalid_argument e) { + } catch (std::invalid_argument &e) { return default_value; } } diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index c7601851c6..17f6478c06 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -24,6 +24,53 @@ #define MAX_PQ_CHUNKS 100 namespace diskann { + static inline void aggregate_coords(const unsigned *ids, const _u64 n_ids, + const _u8 *all_coords, const _u64 ndims, + _u8 *out) { + for (_u64 i = 0; i < n_ids; i++) { + memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(_u8)); + } + } + + static inline void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, + const _u64 pq_nchunks, + const float *pq_dists, float *dists_out) { + _mm_prefetch((char *) dists_out, _MM_HINT_T0); + memset(dists_out, 0, n_pts * sizeof(float)); + for (_u64 chunk = 0; chunk < pq_nchunks; chunk++) { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) { + _mm_prefetch((char *) (chunk_dists + 256), _MM_HINT_T0); + } + for (_u64 idx = 0; idx < n_pts; idx++) { + _u8 pq_centerid = pq_ids[pq_nchunks * idx + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; + } + } + } + + static inline void pq_dist_fast(const unsigned *ids, const _u8 *pq_coords, + const _u64 n_pts, const _u64 pq_nchunks, + const float *pq_dists, float *dists_out) { + _mm_prefetch((char *) dists_out, _MM_HINT_T0); + memset(dists_out, 0, n_pts * sizeof(float)); + + for (_u64 idx = 0; idx < n_pts; idx++) { + _mm_prefetch((char *) pq_coords + ids[idx] * pq_nchunks, _MM_HINT_T0); + } + + for (_u64 chunk = 0; chunk < pq_nchunks; chunk++) { + const float *chunk_dists = pq_dists + 256 * chunk; + if (chunk < pq_nchunks - 1) { + _mm_prefetch((char *) (chunk_dists + 256), _MM_HINT_T0); + } + for (_u64 idx = 0; idx < n_pts; idx++) { + _u8 pq_centerid = pq_coords[pq_nchunks * ids[idx] + chunk]; + dists_out[idx] += chunk_dists[pq_centerid]; + } + } + } + template struct QueryScratch { T * coord_scratch = nullptr; // MUST BE AT LEAST [MAX_N_CMPS * data_dim] @@ -100,16 +147,6 @@ namespace diskann { DISKANN_DLLEXPORT void cache_bfs_levels(_u64 num_nodes_to_cache, std::vector &node_list); - // DISKANN_DLLEXPORT void cache_from_samples(const std::string - // sample_file, _u64 num_nodes_to_cache, std::vector - // &node_list); - - // DISKANN_DLLEXPORT void save_cached_nodes(_u64 num_nodes, - // std::string cache_file_path); - - // setting up thread-specific data - - // implemented DISKANN_DLLEXPORT void cached_beam_search( const T *query, const _u64 k_search, const _u64 l_search, _u64 *res_ids, float *res_dists, const _u64 beam_width, QueryStats *stats = nullptr, diff --git a/include/pq_table.h b/include/pq_table.h index 3cac23c15a..d5475ba048 100644 --- a/include/pq_table.h +++ b/include/pq_table.h @@ -49,14 +49,15 @@ namespace diskann { std::string chunk_offset_file = std::string(pq_table_file) + "_chunk_offsets.bin"; std::string centroid_file = std::string(pq_table_file) + "_centroid.bin"; + std::string table_file = std::string(pq_table_file) + ".bin"; // bin structure: [256][ndims][ndims(float)] uint64_t numr, numc; size_t npts_u64, ndims_u64; #ifdef EXEC_ENV_OLS - diskann::load_bin(files, pq_table_file, tables, npts_u64, ndims_u64); + diskann::load_bin(files, table_file, tables, npts_u64, ndims_u64); #else - diskann::load_bin(pq_table_file, tables, npts_u64, ndims_u64); + diskann::load_bin(table_file, tables, npts_u64, ndims_u64); #endif this->ndims = ndims_u64; diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..3080d18edc 100644 --- a/include/utils.h +++ b/include/utils.h @@ -350,7 +350,7 @@ namespace diskann { throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, __LINE__); } - rounded_dim = ROUND_UP(dim, 8); + rounded_dim = ROUND_UP(dim, 16); diskann::cout << "Metadata: #pts = " << npts << ", #dims = " << dim << ", aligned_dim = " << rounded_dim << "..." << std::flush; size_t allocSize = npts * rounded_dim * sizeof(T); @@ -525,11 +525,6 @@ inline void printProcessMemory(const char* message) { } #else -// need to check and change this -inline bool avx2Supported() { - return true; -} - inline void printProcessMemory(const char* message) { diskann::cout << message << std::endl; } @@ -537,3 +532,4 @@ inline void printProcessMemory(const char* message) { extern bool AvxSupportedCPU; extern bool Avx2SupportedCPU; +extern bool Avx512SupportedCPU; diff --git a/python/setup.py b/python/setup.py new file mode 100644 index 0000000000..0a59ae9961 --- /dev/null +++ b/python/setup.py @@ -0,0 +1,86 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import os +import sys +import numpy +import pybind11 +from setuptools import setup, Extension +from pybind11.setup_helpers import Pybind11Extension, build_ext + + +__version__ = "0.1.0" + + +class BuildExt(build_ext): + """A custom build extension for adding compiler-specific options.""" + c_opts = {'unix': ['-ggdb', '-Ofast', '-DMKL_ILP64', '-m64', '-Wl,--no-as-needed']} + arch_list = '-march -msse -msse2 -msse3 -mssse3 -msse4 -msse4a -msse4.1 -msse4.2 -mavx -mavx2 -mavx512f'.split() + no_arch_flag = True + + if 'CFLAGS' in os.environ: + for flag in arch_list: + if flag in os.environ["CFLAGS"]: + no_arch_flag = False + break + + if no_arch_flag: + c_opts['unix'].append('-march=native') + + link_opts = {'unix': ['-L/opt/intel/compilers_and_libraries/linux/mkl/lib/intel64/', '-L/opt/intel/compilers_and_libraries/linux/lib/intel64/', '-lmkl_rt', '-lmkl_core', '-lmkl_intel_ilp64', '-lmkl_sequential', '-lmkl_intel_thread', '-liomp5', '-lpthread', '-lm', '-ldl']} + c_opts['unix'].append('-fopenmp') + link_opts['unix'].extend(['-fopenmp', '-pthread']) + + def build_extensions(self): + ct = 'unix' + opts = self.c_opts.get(ct, []) + opts.append('-DVERSION_INFO="%s"' % + self.distribution.get_version()) + opts.append('-std=c++14') + opts.append('-fvisibility=hidden') + print('Extra compilation arguments:', opts) + + for ext in self.extensions: + ext.extra_compile_args.extend(opts) + ext.extra_link_args.extend(self.link_opts.get(ct, [])) + ext.include_dirs.extend([ + # Path to pybind11 headers + pybind11.get_include(False), + pybind11.get_include(True), + # Path to numpy headers + numpy.get_include() + ]) + + build_ext.build_extensions(self) + + +ext_modules = [ + Extension( + 'diskannpy', + ['src/diskann_bindings.cpp'], + include_dirs=['../include/', + '/opt/intel/compilers_and_libraries/linux/mkl/include/', + '/usr/include', + pybind11.get_include(False), + pybind11.get_include(True)], + libraries=['aio'], + language='c++', + extra_objects=['../build/src/libdiskann_s.a'], + ) +] + + +setup( + name="diskannpy", + version=__version__, + author="Shikhar Jaiswal, Harsha Vardhan Simhadri", + author_email="t-sjaiswal@microsoft.com, harshasi@microsoft.com", + url="https://github.com/microsoft/diskann", + description="DiskANN Bindings using PyBind11", + long_description="", + ext_modules=ext_modules, + install_requires=['numpy', 'pybind11'], + cmdclass={"build_ext": BuildExt}, + test_suite="tests", + zip_safe=False, +) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp new file mode 100644 index 0000000000..3147b79130 --- /dev/null +++ b/python/src/diskann_bindings.cpp @@ -0,0 +1,440 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "linux_aligned_file_reader.h" +#include "aux_utils.h" +#include "pq_flash_index.h" + +PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); + + +namespace py = pybind11; +using namespace diskann; + +#ifdef __linux__ +template +struct DiskANNIndex { + PQFlashIndex * pq_flash_index; + std::shared_ptr reader; + + DiskANNIndex() { + reader = std::make_shared(); + pq_flash_index = new PQFlashIndex(reader); + } + + ~DiskANNIndex() { + delete pq_flash_index; + } + + int load_index(const std::string &index_path_prefix, const int num_threads) { + const std::string pq_path = index_path_prefix; + const std::string index_path = + index_path_prefix + std::string("_disk.index"); + int load_success = + pq_flash_index->load(num_threads, pq_path.c_str(), index_path.c_str()); + if (load_success != 0) { + std::cout << "Index load failed" << std::endl; + return load_success; + } + std::vector node_list; + _u64 num_nodes_to_cache = 1000; + pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); + std::cout << "loaded index, cached " << node_list.size() + << " nodes based on BFS" << std::endl; + return 0; + } + + void search(std::vector &query, const _u64 query_idx, const _u64 dim, + const _u64 num_queries, const _u64 knn, const _u64 l_search, + const _u64 beam_width, std::vector &ids, + std::vector &dists) { + QueryStats stats; + if (ids.size() < knn * num_queries) { + ids.resize(knn * num_queries); + dists.resize(knn * num_queries); + } + std::vector<_u64> _u64_ids(knn); + pq_flash_index->cached_beam_search( + query.data() + (query_idx * dim), knn, l_search, _u64_ids.data(), + dists.data() + (query_idx * knn), beam_width, &stats); + for (_u64 i = 0; i < knn; i++) + ids[(query_idx * knn) + i] = _u64_ids[i]; + } + + void batch_search(std::vector &queries, const _u64 dim, + const _u64 num_queries, const _u64 knn, const _u64 l_search, + const _u64 beam_width, std::vector &ids, + std::vector &dists, const int num_threads) { + if (ids.size() < knn * num_queries) { + ids.resize(knn * num_queries); + dists.resize(knn * num_queries); + } + omp_set_num_threads(num_threads); +#pragma omp parallel for schedule(dynamic, 1) + for (_u64 q = 0; q < num_queries; ++q) { + std::vector<_u64> u64_ids(knn); + + pq_flash_index->cached_beam_search(queries.data() + q * dim, knn, + l_search, u64_ids.data(), + dists.data() + q * knn, beam_width); + for (_u64 i = 0; i < knn; i++) + ids[(q * knn) + i] = u64_ids[i]; + } + } + + auto search_numpy_input( + py::array_t &query, + const _u64 dim, const _u64 knn, const _u64 l_search, + const _u64 beam_width) { + py::array_t ids(knn); + py::array_t dists(knn); + + std::vector u32_ids(knn); + std::vector<_u64> u64_ids(knn); + QueryStats stats; + + pq_flash_index->cached_beam_search(query.data(), knn, l_search, + u64_ids.data(), dists.mutable_data(), + beam_width, &stats); + + auto r = ids.mutable_unchecked<1>(); + for (_u64 i = 0; i < knn; ++i) + r(i) = (unsigned) u64_ids[i]; + + return std::make_pair(ids, dists); + } + + auto batch_search_numpy_input( + py::array_t &queries, + const _u64 dim, const _u64 num_queries, const _u64 knn, + const _u64 l_search, const _u64 beam_width, const int num_threads) { + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); + + std::vector<_u64> u64_ids(knn * num_queries); + +#pragma omp parallel for schedule(dynamic, 1) + for (_u64 i = 0; i < num_queries; i++) { + pq_flash_index->cached_beam_search(queries.data(i), knn, l_search, + u64_ids.data() + i * knn, + dists.mutable_data(i), beam_width); + } + + auto r = ids.mutable_unchecked(); + for (_u64 i = 0; i < num_queries; ++i) + for (_u64 j = 0; j < knn; ++j) + r(i, j) = (unsigned) u64_ids[i * knn + j]; + + return std::make_pair(ids, dists); + } +}; + +#endif + +PYBIND11_MODULE(diskannpy, m) { + m.doc() = "DiskANN Python Bindings"; + m.attr("__version__") = "0.1.1"; + + py::bind_vector>(m, "VectorUnsigned"); + py::bind_vector>(m, "VectorFloat"); + py::bind_vector>(m, "VectorInt8"); + py::bind_vector>(m, "VectorUInt8"); + + + py::enum_(m, "Metric").value("L2", Metric::L2).export_values(); + + py::class_(m, "Parameters") + .def(py::init<>()) + .def( + "set", + [](Parameters &self, const std::string &name, py::object value) { + if (py::isinstance(value)) { + return self.Set(name, py::cast(value)); + } else if (py::isinstance(value)) { + return self.Set(name, py::cast(value)); + } else if (py::isinstance(value)) { + return self.Set(name, py::cast(value)); + } + }, + py::arg("name"), py::arg("value")); + + py::class_(m, "Neighbor") + .def(py::init<>()) + .def(py::init()) + .def(py::self < py::self) + .def(py::self == py::self); + + py::class_(m, "SimpleNeighbor") + .def(py::init<>()) + .def(py::init()) + .def(py::self < py::self) + .def(py::self == py::self); + + py::class_(m, "AlignedFileReader"); + + py::class_(m, "LinuxAlignedFileReader") + .def(py::init<>()); + + m.def( + "omp_set_num_threads", + [](const size_t num_threads) { omp_set_num_threads(num_threads); }, + py::arg("num_threads") = 1); + + m.def("omp_get_max_threads", []() { return omp_get_max_threads(); }); + + m.def( + "load_aligned_bin_float", + [](const std::string &path, std::vector &data) { + float *data_ptr = nullptr; + size_t num, dims, aligned_dims; + load_aligned_bin(path, data_ptr, num, dims, aligned_dims); + // TODO: Remove redundant copy. + data.assign(data_ptr, data_ptr + num * dims); + auto l = py::list(3); + l[0] = py::int_(num); + l[1] = py::int_(dims); + l[2] = py::int_(aligned_dims); + aligned_free(data_ptr); + return l; + }, + py::arg("path"), py::arg("data")); + + m.def( + "load_truthset", + [](const std::string &path, std::vector &ids, + std::vector &distances) { + unsigned *id_ptr = nullptr; + float * dist_ptr = nullptr; + size_t num, dims; + load_truthset(path, id_ptr, dist_ptr, num, dims); + // TODO: Remove redundant copies. + ids.assign(id_ptr, id_ptr + num * dims); + distances.assign(dist_ptr, dist_ptr + num * dims); + auto l = py::list(2); + l[0] = py::int_(num); + l[1] = py::int_(dims); + delete[] id_ptr; + delete[] dist_ptr; + return l; + }, + py::arg("path"), py::arg("ids"), py::arg("distances")); + + m.def( + "calculate_recall", + [](const unsigned num_queries, std::vector &ground_truth_ids, + std::vector &ground_truth_dists, + const unsigned ground_truth_dims, std::vector &results, + const unsigned result_dims, const unsigned recall_at) { + unsigned *gti_ptr = ground_truth_ids.data(); + float * gtd_ptr = ground_truth_dists.data(); + unsigned *r_ptr = results.data(); + + double total_recall = 0; + std::set gt, res; + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + size_t tie_breaker = recall_at; + if (gtd_ptr != nullptr) { + tie_breaker = recall_at - 1; + float *gt_dist_vec = gtd_ptr + ground_truth_dims * i; + while (tie_breaker < ground_truth_dims && + gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) + tie_breaker++; + } + + gt.insert(gti_ptr + ground_truth_dims * i, + gti_ptr + ground_truth_dims * i + tie_breaker); + res.insert(r_ptr + result_dims * i, + r_ptr + result_dims * i + recall_at); + unsigned cur_recall = 0; + for (auto &v : gt) { + if (res.find(v) != res.end()) { + cur_recall++; + } + } + total_recall += cur_recall; + } + return py::float_(total_recall / (num_queries) * (100.0 / recall_at)); + }, + py::arg("num_queries"), py::arg("ground_truth_ids"), + py::arg("ground_truth_dists"), py::arg("ground_truth_dims"), + py::arg("results"), py::arg("result_dims"), py::arg("recall_at")); + + m.def( + "calculate_recall_numpy_input", + [](const unsigned num_queries, std::vector &ground_truth_ids, + std::vector &ground_truth_dists, + const unsigned ground_truth_dims, + py::array_t + & results, + const unsigned result_dims, const unsigned recall_at) { + unsigned *gti_ptr = ground_truth_ids.data(); + float * gtd_ptr = ground_truth_dists.data(); + unsigned *r_ptr = results.mutable_data(); + + double total_recall = 0; + std::set gt, res; + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + size_t tie_breaker = recall_at; + if (gtd_ptr != nullptr) { + tie_breaker = recall_at - 1; + float *gt_dist_vec = gtd_ptr + ground_truth_dims * i; + while (tie_breaker < ground_truth_dims && + gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) + tie_breaker++; + } + + gt.insert(gti_ptr + ground_truth_dims * i, + gti_ptr + ground_truth_dims * i + tie_breaker); + res.insert(r_ptr + result_dims * i, + r_ptr + result_dims * i + recall_at); + unsigned cur_recall = 0; + for (auto &v : gt) { + if (res.find(v) != res.end()) { + cur_recall++; + } + } + total_recall += cur_recall; + } + return py::float_(total_recall / (num_queries) * (100.0 / recall_at)); + }, + py::arg("num_queries"), py::arg("ground_truth_ids"), + py::arg("ground_truth_dists"), py::arg("ground_truth_dims"), + py::arg("results"), py::arg("result_dims"), py::arg("recall_at")); + + m.def( + "save_bin_u32", + [](const std::string &file_name, std::vector &data, size_t npts, + size_t dims) { save_bin<_u32>(file_name, data.data(), npts, dims); }, + py::arg("file_name"), py::arg("data"), py::arg("npts"), py::arg("dims")); + + py::class_>(m, "DiskANNFloatIndex") + .def(py::init([]() { return new DiskANNIndex(); })) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads")) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("query_idx"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists"), py::arg("num_threads")) + .def("search_numpy_input", &DiskANNIndex::search_numpy_input, + py::arg("query"), py::arg("dim"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width")) + .def("batch_search_numpy_input", + &DiskANNIndex::batch_search_numpy_input, py::arg("queries"), + py::arg("dim"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + diskann::build_disk_index(data_file_path, index_prefix_path, + params.c_str(), + diskann::Metric::L2); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads")); + + py::class_>(m, "DiskANNInt8Index") + .def(py::init([]() { return new DiskANNIndex(); })) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads")) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("query_idx"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists"), py::arg("num_threads")) + .def("search_numpy_input", &DiskANNIndex::search_numpy_input, + py::arg("query"), py::arg("dim"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width")) + .def("batch_search_numpy_input", + &DiskANNIndex::batch_search_numpy_input, py::arg("queries"), + py::arg("dim"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + diskann::build_disk_index(data_file_path, index_prefix_path, + params.c_str(), + diskann::Metric::L2); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads")); + + + + py::class_>(m, "DiskANNUInt8Index") + .def(py::init([]() { return new DiskANNIndex(); })) + .def("load_index", &DiskANNIndex::load_index, + py::arg("index_path_prefix"), py::arg("num_threads")) + .def("search", &DiskANNIndex::search, py::arg("query"), + py::arg("query_idx"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists")) + .def("batch_search", &DiskANNIndex::batch_search, + py::arg("queries"), py::arg("dim"), py::arg("num_queries"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists"), py::arg("num_threads")) + .def("search_numpy_input", &DiskANNIndex::search_numpy_input, + py::arg("query"), py::arg("dim"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width")) + .def("batch_search_numpy_input", + &DiskANNIndex::batch_search_numpy_input, py::arg("queries"), + py::arg("dim"), py::arg("num_queries"), py::arg("knn"), + py::arg("l_search"), py::arg("beam_width"), py::arg("num_threads")) + .def( + "build", + [](DiskANNIndex &self, const char *data_file_path, + const char *index_prefix_path, unsigned R, unsigned L, + double final_index_ram_limit, double indexing_ram_budget, + unsigned num_threads) { + std::string params = std::to_string(R) + " " + std::to_string(L) + + " " + std::to_string(final_index_ram_limit) + + " " + std::to_string(indexing_ram_budget) + + " " + std::to_string(num_threads); + diskann::build_disk_index( + data_file_path, index_prefix_path, + params.c_str(), + diskann::Metric::L2); + }, + py::arg("data_file_path"), py::arg("index_prefix_path"), py::arg("R"), + py::arg("L"), py::arg("final_index_ram_limit"), + py::arg("indexing_ram_limit"), py::arg("num_threads")); +} diff --git a/python/src/vamana_bindings.cpp b/python/src/vamana_bindings.cpp new file mode 100644 index 0000000000..1961144c63 --- /dev/null +++ b/python/src/vamana_bindings.cpp @@ -0,0 +1,326 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include +#include +#include +#include +#include + +#include "index.h" +#include "utils.h" +#include "memory_mapper.h" + +PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); + +namespace py = pybind11; +using namespace diskann; + +PYBIND11_MODULE(vamanapy, m) { + m.doc() = "Vamana Python Bindings"; + m.attr("__version__") = "0.1.0"; + + py::bind_vector>(m, "VectorUnsigned"); + py::bind_vector>(m, "VectorFloat"); + + py::enum_(m, "Metric") + .value("L2", Metric::L2) + .value("INNER_PRODUCT", Metric::INNER_PRODUCT) + .value("FAST_L2", Metric::FAST_L2) + .value("PQ", Metric::PQ) + .export_values(); + + py::class_(m, "Parameters") + .def(py::init<>()) + .def("set", [](Parameters &self, const std::string &name, py::object value) { + if (py::isinstance(value)) { + return self.Set(name, py::cast(value)); + } else if (py::isinstance(value)) { + return self.Set(name, py::cast(value)); + } else if (py::isinstance(value)){ + return self.Set(name, py::cast(value)); + } + }, py::arg("name"), py::arg("value")); + + py::class_(m, "Neighbor") + .def(py::init<>()) + .def(py::init()) + .def(py::self < py::self) + .def(py::self == py::self); + + py::class_(m, "SimpleNeighbor") + .def(py::init<>()) + .def(py::init()) + .def(py::self < py::self) + .def(py::self == py::self); + + m.def("set_num_threads", [](const size_t num_threads) { + omp_set_num_threads(num_threads); + }, py::arg("num_threads") = 1); + + m.def("load_aligned_bin_float", [](const std::string &path, + std::vector &data) { + float *data_ptr = nullptr; + size_t num, dims, aligned_dims; + load_aligned_bin(path, data_ptr, num, dims, aligned_dims); + // TODO: Remove redundant copy. + data.assign(data_ptr, data_ptr + num * dims); + auto l = py::list(3); + l[0] = py::int_(num); + l[1] = py::int_(dims); + l[2] = py::int_(aligned_dims); + aligned_free(data_ptr); + return l; + }, py::arg("path"), py::arg("data")); + + m.def("load_truthset", [](const std::string &path, + std::vector &ids, + std::vector &distances) { + unsigned *id_ptr = nullptr; + float *dist_ptr = nullptr; + size_t num, dims; + load_truthset(path, id_ptr, dist_ptr, num, dims); + // TODO: Remove redundant copies. + ids.assign(id_ptr, id_ptr + num * dims); + distances.assign(dist_ptr, dist_ptr + num * dims); + auto l = py::list(2); + l[0] = py::int_(num); + l[1] = py::int_(dims); + delete[] id_ptr; + delete[] dist_ptr; + return l; + }, py::arg("path"), py::arg("ids"), py::arg("distances")); + + m.def("calculate_recall", [](const unsigned num_queries, + std::vector &ground_truth_ids, + std::vector &ground_truth_dists, + const unsigned ground_truth_dims, + std::vector &results, + const unsigned result_dims, + const unsigned recall_at) { + unsigned *gti_ptr = ground_truth_ids.data(); + float *gtd_ptr = ground_truth_dists.data(); + unsigned *r_ptr = results.data(); + + double total_recall = 0; + std::set gt, res; + for (size_t i = 0; i < num_queries; i++) { + gt.clear(); + res.clear(); + size_t tie_breaker = recall_at; + if (gtd_ptr != nullptr) { + tie_breaker = recall_at - 1; + float *gt_dist_vec = gtd_ptr + ground_truth_dims * i; + while (tie_breaker < ground_truth_dims && + gt_dist_vec[tie_breaker] == gt_dist_vec[recall_at - 1]) + tie_breaker++; + } + + gt.insert(gti_ptr + ground_truth_dims * i, gti_ptr + ground_truth_dims * i + tie_breaker); + res.insert(r_ptr + result_dims * i, r_ptr + result_dims * i + recall_at); + unsigned cur_recall = 0; + for (auto &v : gt) { + if (res.find(v) != res.end()) { + cur_recall++; + } + } + total_recall += cur_recall; + } + return py::float_(total_recall / (num_queries) * (100.0 / recall_at)); + }, py::arg("num_queries"), py::arg("ground_truth_ids"), + py::arg("ground_truth_dists"), py::arg("ground_truth_dims"), + py::arg("results"), py::arg("result_dims"), py::arg("recall_at")); + + m.def("save_bin_u32", [](const std::string& file_name, + std::vector &data, size_t npts, + size_t dims) { + save_bin<_u32>(file_name, data.data(), npts, dims); + }, py::arg("file_name"), py::arg("data"), py::arg("npts"), + py::arg("dims")); + + py::class_>(m, "SinglePrecisionIndex") + .def(py::init(), + py::arg("m"), py::arg("filename"), + py::arg("max_points") = 0, py::arg("nd") = 0, + py::arg("num_frozen_pts") = 0, py::arg("enable_tags") = false, + py::arg("store_data") = true, py::arg("support_eager_delete") = false) + .def("save", [](Index &self, + const std::string file_name) { + return self.save(file_name.c_str()); + }, py::arg("file_name")) + .def("load", [](Index &self, + const std::string file_name, bool load_tags, + const std::string tag_file_name) { + if (tag_file_name == "") { + return self.load(file_name.c_str(), load_tags, NULL); + } else { + return self.load(file_name.c_str(), load_tags, tag_file_name.c_str()); + } + }, py::arg("file_name"), py::arg("load_tags") = false, + py::arg("tag_file_name") = "") + .def("pq_load", [](Index &self, + const std::string pq_prefix_path) { + return self.pq_load(pq_prefix_path.c_str()); + }, py::arg("pq_prefix_path")) + .def("generate_random_frozen_points", [](Index &self, + const std::string file_name) { + if (file_name == "") { + return self.generate_random_frozen_points(NULL); + } else { + return self.generate_random_frozen_points(file_name.c_str()); + } + }, py::arg("file_name") = "") + .def("build", [](Index &self, Parameters ¶meters, + const std::vector &tags) { + if (tags.size() == 0) { + return self.build(parameters); + } else { + return self.build(parameters, tags); + } + }, py::arg("parameters"), py::arg("tags")) + .def("pq_build", [](Index &self, const std::string file_name, + const std::string index_path, Parameters ¶meters) { + return self.pq_build(file_name.c_str(), index_path.c_str(), + parameters); + }, py::arg("file_name"), py::arg("index_path"), py::arg("parameters")) + .def("search", [](Index &self, std::vector &query, + const size_t query_index, const size_t knn, + const size_t num_queries, const size_t l_search, + std::vector &ids, const size_t id_index) { + if (ids.size() == 0) { + ids.resize(knn * num_queries); + } + + self.search(query.data() + query_index, knn, l_search, + ids.data() + id_index); + }, py::arg("query"), py::arg("query_index"), py::arg("knn") = 10, + py::arg("num_queries"), py::arg("l_search"), py::arg("ids"), + py::arg("id_index")) + .def("search_with_tags", [](Index &self, + std::vector &query, size_t knn, + size_t l_search, std::vector &tags, + unsigned num_frozen_pts, + std::vector &indices_buffer) { + if (indices_buffer.size() == 0) { + return self.search_with_tags(query.data(), knn, l_search, + tags.data(), num_frozen_pts, NULL); + } else { + return self.search_with_tags(query.data(), knn, l_search, + tags.data(), num_frozen_pts, + indices_buffer.data()); + } + }, py::arg("query"), py::arg("knn") = 10, py::arg("l_search"), + py::arg("tags"), py::arg("num_frozen_pts"), py::arg("indices_buffer")) + .def("read_just_data", &Index::readjust_data) + .def("insert_point", [](Index &self, + const std::vector &point, + const Parameters ¶meter, + std::vector &pool, + std::vector &tmp, + tsl::robin_set &visited, + std::vector &cut_graph, + const int tag) { + return self.insert_point(point.data(), parameter, pool, tmp, visited, + cut_graph, tag); + }, py::arg("point"), py::arg("parameter"), py::arg("pool"), + py::arg("tmp"), py::arg("visited"), py::arg("cut_graph"), py::arg("tag")) + .def("enable_delete", &Index::enable_delete) + .def("disable_delete", [](Index &self, + const Parameters ¶meters, + const bool consolidate) { + return self.disable_delete(parameters, consolidate); + }, py::arg("parameters"), py::arg("consolidate") = false) + .def("delete_point", &Index::delete_point) + .def("eager_delete", [](Index &self, const int tag, + const Parameters ¶meters) { + return self.eager_delete(tag, parameters); + }, py::arg("tag"), py::arg("parameters")) + .def("optimize_graph", &Index::optimize_graph) + .def("search_with_optimized_graph", [](Index &self, + std::vector &query, + const size_t query_index, + const size_t knn, + const size_t num_queries, + const size_t l_search, + std::vector &ids, + const size_t id_index) { + if (ids.size() == 0) { + ids.resize(knn * num_queries); + } + + self.search_with_opt_graph(query.data() + query_index, knn, + l_search, ids.data() + id_index); + }, py::arg("query"), py::arg("query_index"), py::arg("knn") = 10, + py::arg("num_queries"), py::arg("l_search"), py::arg("ids"), + py::arg("id_index")) + .def("single_numpy_query", [](Index &self, + py::array_t &query, + const size_t knn, + const size_t l_search) { + py::array_t ids(knn); + self.search_with_opt_graph(query.data(), knn, + l_search, ids.mutable_data()); + return ids; + }, py::arg("query"), py::arg("knn") = 10, py::arg("l_search")) + .def("batch_numpy_query", [](Index &self, + py::array_t &queries, + const size_t knn, + const size_t num_queries, + const size_t l_search) { + py::array_t ids(knn * num_queries); + #pragma omp parallel for schedule(dynamic, 1) + for (unsigned i = 0; i < num_queries; i++) { + self.search_with_opt_graph(queries.data(i), knn, + l_search, ids.mutable_data(i * knn)); + } + return ids; + }, py::arg("queries"), py::arg("knn") = 10, py::arg("num_queries"), + py::arg("l_search")) + .def("pq_search", [](Index &self, std::vector &query, + const size_t query_index, const size_t knn, + const size_t num_queries, const size_t l_search, + std::vector &ids, const size_t id_index) { + if (ids.size() == 0) { + ids.resize(knn * num_queries); + } + self.pq_search(query.data() + query_index, knn, l_search, + ids.data() + id_index); + }, py::arg("query"), py::arg("query_index"), py::arg("knn") = 10, + py::arg("num_queries"), py::arg("l_search"), py::arg("ids"), + py::arg("id_index")) + .def("pq_single_numpy_query", [](Index &self, + py::array_t &query, + const size_t knn, + const size_t l_search) { + py::array_t ids(knn); + self.pq_search(query.mutable_data(), knn, l_search, + ids.mutable_data()); + return ids; + }, py::arg("query"), py::arg("knn") = 10, py::arg("l_search")) + .def("pq_batch_numpy_query", [](Index &self, + py::array_t &queries, + const size_t knn, + const size_t num_queries, + const size_t l_search) { + py::array_t ids(knn * num_queries); + #pragma omp parallel for schedule(dynamic, 1) + for (unsigned i = 0; i < num_queries; i++) { + self.pq_search(queries.mutable_data(i), knn, l_search, + ids.mutable_data(i * knn)); + } + return ids; + }, py::arg("queries"), py::arg("knn") = 10, py::arg("num_queries"), + py::arg("l_search")); +} diff --git a/python/tests/test_build_disk_index.py b/python/tests/test_build_disk_index.py new file mode 100644 index 0000000000..df5602dc3b --- /dev/null +++ b/python/tests/test_build_disk_index.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import argparse +from diskannpy import Metric, Parameters, DiskANNFloatIndex + + +parser = argparse.ArgumentParser() +parser.add_argument('data_path', type=str, help='Path to the input base set of vectors.') +parser.add_argument('save_path', type=str, help='Path to the built index.') +parser.add_argument('R', type=int, help='Graph degree.') +parser.add_argument('L', type=int, help='Index build complexity.') +parser.add_argument('B', type=float, help='Memory budget in GB for the final index.') +parser.add_argument('M', type=float, help='Memory budget in GB for the index construction.') +parser.add_argument('T', type=int, help='Number of threads for index construction.') + +args = parser.parse_args() + +start = time.time() +index = DiskANNFloatIndex() +index.build(args.data_path, args.save_path, args.R, args.L, args.B, args.M, args.T) +end = time.time() + +print("Indexing Time: " + str(end - start) + " seconds") \ No newline at end of file diff --git a/python/tests/test_build_memory_index.py b/python/tests/test_build_memory_index.py new file mode 100644 index 0000000000..5bc58bc3fa --- /dev/null +++ b/python/tests/test_build_memory_index.py @@ -0,0 +1,29 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import argparse +from vamanapy import Metric, Parameters, SinglePrecisionIndex + + +parser = argparse.ArgumentParser() +parser.add_argument('data_path', type=str, help='Path to the input base set of vectors.') +parser.add_argument('save_path', type=str, help='Path to the built index.') +args = parser.parse_args() + +params = Parameters() +params.set("L", 125) +params.set("R", 32) +params.set("C", 750) +params.set("alpha", 1.2) +params.set("saturate_graph", False) +params.set("num_threads", 32) + +start = time.time() +index = SinglePrecisionIndex(Metric.FAST_L2, args.data_path) +index.build(params, []) +end = time.time() + +print("Indexing Time: " + str(end - start) + " seconds") + +index.save(args.save_path) diff --git a/python/tests/test_build_pq_memory_index.py b/python/tests/test_build_pq_memory_index.py new file mode 100644 index 0000000000..f076fd43f3 --- /dev/null +++ b/python/tests/test_build_pq_memory_index.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +from diskannpy import Metric, Parameters, SinglePrecisionIndex + + +data_path = "/mnt/SIFT1M/sift_base.bin" +pq_path = "/home/t-sjaiswal/diskann/build/tests/PQ_SIFT1M/test_pq_memory_index" +save_path = "/home/t-sjaiswal/diskann/build/tests/PQ_SIFT1M/test_build_pq_memory_index.bin" + +params = Parameters() +params.set("L", 125) +params.set("R", 32) +params.set("C", 750) +params.set("alpha", 1.2) +params.set("saturate_graph", False) +params.set("num_chunks", 32) +params.set("num_threads", 32) + +start = time.time() +index = SinglePrecisionIndex(Metric.FAST_L2, data_path) +index.pq_build(data_path, pq_path, params) +end = time.time() + +print("Indexing Time: " + str(end - start) + " seconds") + +index.save(save_path) diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py new file mode 100644 index 0000000000..5d44a83476 --- /dev/null +++ b/python/tests/test_search_disk_index.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import argparse +import numpy as np +import diskannpy + + +parser = argparse.ArgumentParser() +parser.add_argument('query_path', type=str, help='Path to the input query set of vectors.') +parser.add_argument('ground_truth_path', type=str, help='Path to the input groundtruth set.') +parser.add_argument('index_path_prefix', type=str, help='Path prefix for index files.') +parser.add_argument('output_path_prefix', type=str, help='Prefix for the generated output files.') +parser.add_argument('K', type=int, help='k value for recall@K.') +parser.add_argument('W', type=int, help='Beamwidth for search.') +parser.add_argument('T', type=int, help='Number of threads to use for search.') + +args = parser.parse_args() + +recall_at = args.K +W = args.W +# Use multi-threaded search only for batch mode. +num_threads = args.T +single_query_mode = False +l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] + + +query_data = diskannpy.VectorFloat() +ground_truth_ids = diskannpy.VectorUnsigned() +ground_truth_dists = diskannpy.VectorFloat() + +num_queries, query_dims, query_aligned_dims = diskannpy.load_aligned_bin_float(args.query_path, query_data) +num_ground_truth, ground_truth_dims = diskannpy.load_truthset(args.ground_truth_path, ground_truth_ids, ground_truth_dists) + +index = diskannpy.DiskANNFloatIndex() +index.load_index(args.index_path_prefix, num_threads) +print("Index Loaded") + +#index.optimize_graph() +#print("Graph Optimization Completed") + +if single_query_mode: + print("Ls QPS Mean Latency (mus) 99.9 Latency Recall@10") + print("================================================================") + for i, L in enumerate(l_search): + latency_stats = [] + query_result_ids = diskannpy.VectorUnsigned() + query_result_dists = diskannpy.VectorFloat() + + s = time.time() + + for j in range(num_queries): + qs = time.time() + index.search(query_data, j, query_aligned_dims, num_queries, + recall_at, L, W, query_result_ids, query_result_dists) + qe = time.time() + latency_stats.append(float((qe - qs) * 1000000)) + + e = time.time() + qps = (num_queries / (e - s)) + recall = diskannpy.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + latency_stats.sort() + mean_latency = sum(latency_stats) / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>20}".format("{:.2f}".format(latency_stats[int((0.999 * num_queries))])) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = args.output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + diskannpy.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) +else: + print("Ls QPS Mean Latency (mus) Recall@10") + print("=============================================") + for i, L in enumerate(l_search): + diskannpy.omp_set_num_threads(num_threads) + + query_result_ids = diskannpy.VectorUnsigned() + query_result_dists = diskannpy.VectorFloat() + + qs = time.time() + index.batch_search(query_data, query_aligned_dims, num_queries, + recall_at, L, W, + query_result_ids, query_result_dists, + num_threads) + qe = time.time() + latency_stats = float((qe - qs) * 1000000) + + qps = (num_queries / (qe - qs)) + recall = diskannpy.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + mean_latency = latency_stats / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = args.output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + diskannpy.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) diff --git a/python/tests/test_search_disk_index_numpy.py b/python/tests/test_search_disk_index_numpy.py new file mode 100644 index 0000000000..196541e713 --- /dev/null +++ b/python/tests/test_search_disk_index_numpy.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. +import time +import argparse +import numpy as np +import diskannpy + + +parser = argparse.ArgumentParser() +parser.add_argument('query_path', type=str, help='Path to the input query set of vectors.') +parser.add_argument('ground_truth_path', type=str, help='Path to the input groundtruth set.') +parser.add_argument('index_path_prefix', type=str, help='Path prefix for index files.') +parser.add_argument('K', type=int, help='k value for recall@K.') +parser.add_argument('W', type=int, help='Beamwidth for search.') +parser.add_argument('T', type=int, help='Number of threads to use for search.') + +args = parser.parse_args() +args = parser.parse_args() + +recall_at = args.K +W = args.W +# Use multi-threaded search only for batch mode. +num_threads = args.T +l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] + + +query_data = diskannpy.VectorFloat() +ground_truth_ids = diskannpy.VectorUnsigned() +ground_truth_dists = diskannpy.VectorFloat() + +num_queries, query_dims, query_aligned_dims = diskannpy.load_aligned_bin_float(args.query_path, query_data) +num_ground_truth, ground_truth_dims = diskannpy.load_truthset(args.ground_truth_path, ground_truth_ids, ground_truth_dists) + +query_data_numpy = np.zeros((num_queries,query_aligned_dims), dtype=np.float32) +for i in range(0, num_queries): + for d in range(0, query_dims): + query_data_numpy[i,d] = query_data[i * query_aligned_dims + d] + +index = diskannpy.DiskANNFloatIndex() +index.load_index(args.index_path_prefix, num_threads) +print("Index Loaded") + + +print("Ls QPS Recall@10") +print("========================") +for i, L in enumerate(l_search): + diskannpy.omp_set_num_threads(num_threads) + + qs = time.time() + ids, dists = index.batch_search_numpy_input(query_data_numpy, query_aligned_dims, + num_queries, recall_at, L, W, num_threads) + qe = time.time() + latency_stats = float((qe - qs) * 1000000) + qps = (num_queries / (qe - qs)) + + recall = diskannpy.calculate_recall_numpy_input(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + ids, recall_at, recall_at) + mean_latency = latency_stats / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + "{:>15}".format("{:.2f}".format(recall))) diff --git a/python/tests/test_search_memory_index.py b/python/tests/test_search_memory_index.py new file mode 100644 index 0000000000..e93e355540 --- /dev/null +++ b/python/tests/test_search_memory_index.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import argparse +import numpy as np +import diskannpy as vp + + +parser = argparse.ArgumentParser() +parser.add_argument('data_path', type=str, help='Path to the input base set of vectors.') +parser.add_argument('query_path', type=str, help='Path to the input query set of vectors.') +parser.add_argument('ground_truth_path', type=str, help='Path to the input groundtruth set.') +parser.add_argument('memory_index_path', type=str, help='Path to the built index.') +parser.add_argument('output_path_prefix', type=str, help='Prefix for the generated output files.') +args = parser.parse_args() + +recall_at = 10 +# Use multi-threaded search only for batch mode. +num_threads = 1 +single_query_mode = True +l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] + +query_data = vp.VectorFloat() +ground_truth_ids = vp.VectorUnsigned() +ground_truth_dists = vp.VectorFloat() + +num_queries, query_dims, query_aligned_dims = vp.load_aligned_bin_float(args.query_path, query_data) +num_ground_truth, ground_truth_dims = vp.load_truthset(args.ground_truth_path, ground_truth_ids, ground_truth_dists) + +index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, args.data_path) +index.load(file_name = args.memory_index_path) +print("Index Loaded") + +index.optimize_graph() +print("Graph Optimization Completed") + +if single_query_mode: + print("Ls QPS Mean Latency (mus) 99.9 Latency Recall@10") + print("================================================================") + for i, L in enumerate(l_search): + latency_stats = [] + query_result_ids = vp.VectorUnsigned() + s = time.time() + + for j in range(num_queries): + qs = time.time() + index.search_with_optimized_graph(query_data, j * query_aligned_dims, + recall_at, num_queries, L, + query_result_ids, + j * recall_at) + qe = time.time() + latency_stats.append(float((qe - qs) * 1000000)) + + e = time.time() + qps = (num_queries / (e - s)) + recall = vp.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + latency_stats.sort() + mean_latency = sum(latency_stats) / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>20}".format("{:.2f}".format(latency_stats[int((0.999 * num_queries))])) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = args.output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + vp.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) +else: + query_data = np.load('/mnt/SIFT1M/sift_query.npy') + print("Ls QPS Mean Latency (mus) Recall@10") + print("=============================================") + for i, L in enumerate(l_search): + vp.set_num_threads(num_threads) + + qs = time.time() + query_result_ids = index.batch_numpy_query(query_data, recall_at, num_queries, L) + qe = time.time() + latency_stats = float((qe - qs) * 1000000) + + query_result_ids = vp.VectorUnsigned(query_result_ids) + qps = (num_queries / (qe - qs)) + recall = vp.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + mean_latency = latency_stats / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = args.output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + vp.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) \ No newline at end of file diff --git a/python/tests/test_search_pq_memory_index.py b/python/tests/test_search_pq_memory_index.py new file mode 100644 index 0000000000..105b0de0dc --- /dev/null +++ b/python/tests/test_search_pq_memory_index.py @@ -0,0 +1,93 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import numpy as np +import diskannpy as vp + + +data_path = "/mnt/SIFT1M/sift_base.bin" +query_path = "/mnt/SIFT1M/sift_query.bin" +ground_truth_path = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/sift_groundtruth.bin" +memory_index_path = "/home/t-sjaiswal/diskann/build/tests/PQ_SIFT1M/test_build_pq_memory_index.bin" +pq_path = "/home/t-sjaiswal/diskann/build/tests/PQ_SIFT1M/test_pq_memory_index" +output_path_prefix = "/home/t-sjaiswal/diskann/build/tests/PQ_SIFT1M/test_search_pq_memory_index" + +recall_at = 10 +# Use multi-threaded search only for batch mode. +num_threads = 1 +single_query_mode = True +l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] + +query_data = vp.VectorFloat() +ground_truth_ids = vp.VectorUnsigned() +ground_truth_dists = vp.VectorFloat() + +num_queries, query_dims, query_aligned_dims = vp.load_aligned_bin_float(query_path, query_data) +num_ground_truth, ground_truth_dims = vp.load_truthset(ground_truth_path, ground_truth_ids, ground_truth_dists) + +index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, data_path) +index.load(file_name = memory_index_path) +print("Index Loaded") +index.pq_load(pq_prefix_path = pq_path) +print("PQ Data Loaded") + +index.optimize_graph() +print("Graph Optimization Completed") + +if single_query_mode: + print("Ls QPS Mean Latency (mus) 99.9 Latency Recall@10") + print("================================================================") + for i, L in enumerate(l_search): + latency_stats = [] + query_result_ids = vp.VectorUnsigned() + s = time.time() + + for j in range(num_queries): + qs = time.time() + index.pq_search(query_data, j * query_aligned_dims, recall_at, + num_queries, L, query_result_ids, j * recall_at) + qe = time.time() + latency_stats.append(float((qe - qs) * 1000000)) + + e = time.time() + qps = (num_queries / (e - s)) + recall = vp.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + latency_stats.sort() + mean_latency = sum(latency_stats) / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>20}".format("{:.2f}".format(latency_stats[int((0.999 * num_queries))])) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + vp.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) +else: + query_data = np.load('/mnt/SIFT1M/sift_query.npy') + print("Ls QPS Mean Latency (mus) Recall@10") + print("=============================================") + for i, L in enumerate(l_search): + vp.set_num_threads(num_threads) + + qs = time.time() + query_result_ids = index.pq_batch_numpy_query(query_data, recall_at, + num_queries, L) + qe = time.time() + latency_stats = float((qe - qs) * 1000000) + + query_result_ids = vp.VectorUnsigned(query_result_ids) + qps = (num_queries / (qe - qs)) + recall = vp.calculate_recall(num_queries, ground_truth_ids, + ground_truth_dists, ground_truth_dims, + query_result_ids, recall_at, + recall_at) + mean_latency = latency_stats / num_queries + print(str(L) + "{:>10}".format("{:.2f}".format(qps)) + + "{:>15}".format("{:.2f}".format(mean_latency)) + + "{:>15}".format("{:.2f}".format(recall))) + + result_path = output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + vp.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) diff --git a/src/aux_utils.cpp b/src/aux_utils.cpp index 6a2990ba66..59631d8834 100644 --- a/src/aux_utils.cpp +++ b/src/aux_utils.cpp @@ -565,7 +565,7 @@ namespace diskann { _u64 cur_node_id = 0; for (_u64 sector = 0; sector < n_sectors; sector++) { if (sector % 100000 == 0) { - diskann::cout << "Sector #" << sector << "written" << std::endl; + diskann::cout << "Sector #" << sector << " written" << std::endl; } memset(sector_buf.get(), 0, SECTOR_LEN); for (_u64 sector_node_id = 0; @@ -632,7 +632,7 @@ namespace diskann { } std::string index_prefix_path(indexFilePath); - std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; + std::string pq_pivots_path = index_prefix_path + "_pq_pivots"; std::string pq_compressed_vectors_path = index_prefix_path + "_pq_compressed.bin"; std::string mem_index_path = index_prefix_path + "_mem.index"; @@ -666,7 +666,7 @@ namespace diskann { diskann::cout << "Starting index build: R=" << R << " L=" << L << " Query RAM budget: " << final_index_ram_limit - << " Indexing ram budget: " << indexing_ram_budget + << " Indexing RAM budget: " << indexing_ram_budget << " T: " << num_threads << std::endl; auto s = std::chrono::high_resolution_clock::now(); diff --git a/src/index.cpp b/src/index.cpp index 2405878669..8bb5f3d343 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -31,6 +31,7 @@ #include "math_utils.h" #include "memory_mapper.h" #include "parameters.h" +#include "pq_flash_index.h" #include "partition_and_pq.h" #include "timer.h" #include "utils.h" @@ -47,12 +48,19 @@ namespace { template<> diskann::Distance *get_distance_function(diskann::Metric m) { if (m == diskann::Metric::FAST_L2) { - std::cout << "Here" << std::endl; + std::cout << "Using Fast L2 Distance Metric" << std::endl; return new diskann::DistanceFastL2(); + } else if (m == diskann::Metric::INNER_PRODUCT) { + std::cout << "Using Fast Inner Product Distance Metric" << std::endl; + return new diskann::DistanceFastInnerProduct(); } else if (m == diskann::Metric::L2) { - if (Avx2SupportedCPU) { - std::cout << "Using AVX2 distance computation" << std::endl; - return new diskann::DistanceL2(); + if (Avx512SupportedCPU) { + std::cout << "Using AVX512 distance computation" << std::endl; + return new diskann::AVX512DistanceL2Float(); + } else if (Avx2SupportedCPU) { + std::cout << "AVX512 not supported. Using AVX2 distance computation" + << std::endl; + return new diskann::AVX2DistanceL2Float(); } else if (AvxSupportedCPU) { std::cout << "AVX2 not supported. Using AVX distance computation" << std::endl; @@ -63,9 +71,8 @@ namespace { } } else { std::stringstream stream; - stream << "Only L2 metric supported as of now. Email " - "gopalsr@microsoft.com if you need cosine similarity or inner " - "product." + stream << "Only L2 and Inner Product metric supported as of now. Email " + "gopalsr@microsoft.com if you need support for other metrics." << std::endl; std::cerr << stream.str() << std::endl; throw diskann::ANNException(stream.str(), -1, __FUNCSIG__, __FILE__, @@ -180,6 +187,12 @@ namespace diskann { this->_distance = ::get_distance_function(m); _locks = std::vector(_max_points + _num_frozen_pts); + DistanceFastInnerProduct *cosine_distance = + dynamic_cast *>(_distance); + if (cosine_distance != nullptr) { + _normalize = true; + } + _width = 0; } @@ -187,6 +200,8 @@ namespace diskann { Index::~Index() { delete this->_distance; aligned_free(_data); + delete[] _pq_data; + aligned_free(_pq_table_dists); } template<> @@ -1008,6 +1023,227 @@ namespace diskann { _has_built = true; } + template + void Index::pq_build(const char *dataFilePath, + const char *indexFilePath, + Parameters ¶meters) { + std::string filename(dataFilePath); + std::string index_prefix_path(indexFilePath); + std::string save_path = index_prefix_path + "_normalized_data.bin"; + std::string pq_pivots_path = index_prefix_path + "_pq_pivots"; + std::string pq_compressed_vectors_path = + index_prefix_path + "_pq_compressed.bin"; + + const unsigned R = parameters.Get("R"); + const unsigned L = parameters.Get("L"); + const float alpha = parameters.Get("alpha"); + const unsigned num_chunks = parameters.Get("num_chunks"); + const unsigned num_threads = parameters.Get("num_threads"); + + if (num_threads != 0) { + omp_set_num_threads(num_threads); + mkl_set_num_threads(num_threads); + } + + _n_chunks = num_chunks; + diskann::cout << "Starting index build: R = " << R << " L = " << L + << " Alpha = " << alpha << " Threads = " << num_threads + << std::endl; + diskann::cout << "Compressing " << _dim << "-dimensional data into " + << _n_chunks << " bytes per vector." << std::endl; + diskann::cout << "Training data loaded of size " << _nd << std::endl; + + float *unaligned_data; + diskann::load_bin(filename, unaligned_data, _nd, _dim); + + if (_normalize) { + DistanceFastInnerProduct *base_norm = + dynamic_cast *>(_distance); + for (unsigned b = 0; b < _nd; b++) { + float norm = base_norm->norm(_data + b * _aligned_dim, _aligned_dim); + if (norm != std::numeric_limits::max()) { + for (unsigned i = 0; i < _dim; i++) { + unaligned_data[b * _dim + i] *= norm; + } + } + } + + diskann::save_bin(save_path.c_str(), unaligned_data, _nd, _dim); + filename = save_path; + } + + generate_pq_pivots(unaligned_data, _nd, (uint32_t) _dim, 256, + (uint32_t) _n_chunks, 15, pq_pivots_path); + generate_pq_data_from_pivots(filename, 256, (uint32_t) _n_chunks, + pq_pivots_path, pq_compressed_vectors_path); + delete[] unaligned_data; + + this->build(parameters); + } + + template + void Index::pq_load(const char *pq_prefix) { + std::string pq_table_path = std::string(pq_prefix) + "_pq_pivots"; + std::string pq_compressed_vectors = + std::string(pq_prefix) + "_pq_compressed.bin"; + + size_t pq_file_dim, pq_file_num_centroids; + get_bin_metadata(pq_table_path + ".bin", pq_file_num_centroids, + pq_file_dim); + + if (pq_file_num_centroids != 256) { + diskann::cout << "Error. Number of PQ centroids is not 256. Exiting." + << std::endl; + return; + } + + _u64 data_dim = pq_file_dim; + _u64 aligned_dim = ROUND_UP(pq_file_dim, 8); + + size_t npts_u64, nchunks_u64; + diskann::load_bin<_u8>(pq_compressed_vectors, _pq_data, npts_u64, + nchunks_u64); + + _n_chunks = nchunks_u64; + _pq_table.load_pq_centroid_bin(pq_table_path.c_str(), _n_chunks); + + if (_nd != npts_u64) { + diskann::cout << "Error. Mismatch of data points in the graph. Exiting." + << std::endl; + return; + } + + diskann::cout + << "Loaded PQ centroids and in-memory compressed vectors. #points: " + << npts_u64 << " #dim: " << data_dim << " #aligned_dim: " << aligned_dim + << " #chunks: " << _n_chunks << std::endl; + + diskann::alloc_aligned((void **) &_pq_table_dists, + 256 * _n_chunks * sizeof(float), 256); + diskann::cout << "Done.." << std::endl; + return; + } + + template + void Index::pq_search(T *query, size_t K, size_t L, + unsigned *indices) { + std::vector retset(L + 1); + std::vector init_ids(L); + float pq_coord_dists[256]; + unsigned v_neighbors[256]; + + DistanceInnerProduct *dist_fast = + dynamic_cast *>(_distance); + + if (_normalize) { + float norm = dist_fast->norm(query, _aligned_dim); + if (norm != std::numeric_limits::max()) { + for (unsigned i = 0; i < _dim; i++) { + query[i] *= norm; + } + } + } + + _pq_table.populate_chunk_distances(query, _pq_table_dists); + + boost::dynamic_bitset<> flags{_nd, 0}; + unsigned tmp_l = 0; + unsigned * neighbors = + (unsigned *) (_opt_graph + _node_size * _ep + _data_len); + unsigned MaxM_ep = *neighbors; + neighbors++; + + for (; tmp_l < L && tmp_l < MaxM_ep; tmp_l++) { + init_ids[tmp_l] = neighbors[tmp_l]; + flags[init_ids[tmp_l]] = true; + } + + while (tmp_l < L) { + unsigned id = rand() % _nd; + if (flags[id]) + continue; + flags[id] = true; + init_ids[tmp_l] = id; + tmp_l++; + } + + L = init_ids.size(); + diskann::pq_dist_fast(init_ids.data(), _pq_data, L, _n_chunks, + _pq_table_dists, pq_coord_dists); + + for (unsigned i = 0; i < init_ids.size(); i++) { + unsigned id = init_ids[i]; + retset[i] = Neighbor(id, pq_coord_dists[i], true); + flags[id] = true; + } + + std::sort(retset.begin(), retset.begin() + L); + int k = 0; + while (k < (int) L) { + int nk = L; + + if (retset[k].flag) { + retset[k].flag = false; + unsigned n = retset[k].id; + + unsigned *neighbors = + (unsigned *) (_opt_graph + _node_size * n + _data_len); + unsigned MaxM = *neighbors; + neighbors++; + + memset(v_neighbors, 0, MaxM); + unsigned visitable_neighbors = 0; + + for (unsigned m = 0; m < MaxM; ++m) { + unsigned id = neighbors[m]; + if (flags[id]) { + continue; + } + flags[id] = 1; + v_neighbors[visitable_neighbors] = id; + visitable_neighbors++; + } + + diskann::pq_dist_fast(v_neighbors, _pq_data, visitable_neighbors, + _n_chunks, _pq_table_dists, pq_coord_dists); + for (unsigned m = 0; m < visitable_neighbors; ++m) { + float dist = pq_coord_dists[m]; + if (dist >= retset[L - 1].distance) { + continue; + } + Neighbor nn(v_neighbors[m], dist, true); + int r = InsertIntoPool(retset.data(), L, nn); + + if (r < nk) { + nk = r; + } + } + } + + if (nk <= k) { + k = nk; + } else { + ++k; + } + } + + for (unsigned i = 0; i < L; i++) { + unsigned id = retset[i].id; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); + T * x = (T *) (_opt_graph + _node_size * id); + float norm_x = *x; + x++; + retset[i].distance = + dist_fast->compare(x, query, norm_x, (unsigned) _aligned_dim); + } + + std::sort(retset.begin(), retset.begin() + L); + + for (size_t i = 0; i < K; i++) { + indices[i] = retset[i].id; + } + } + template std::pair Index::search(const T *query, const size_t K, @@ -1081,7 +1317,8 @@ namespace diskann { _neighbor_len = (_width + 1) * sizeof(unsigned); _node_size = _data_len + _neighbor_len; _opt_graph = (char *) malloc(_node_size * _nd); - DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; + DistanceInnerProduct *dist_fast = + dynamic_cast *>(_distance); for (unsigned i = 0; i < _nd; i++) { char *cur_node_offset = _opt_graph + i * _node_size; float cur_norm = dist_fast->norm(_data + i * _aligned_dim, _aligned_dim); @@ -1103,12 +1340,11 @@ namespace diskann { template void Index::search_with_opt_graph(const T *query, size_t K, size_t L, unsigned *indices) { - DistanceFastL2 *dist_fast = (DistanceFastL2 *) _distance; + DistanceInnerProduct *dist_fast = + dynamic_cast *>(_distance); std::vector retset(L + 1); std::vector init_ids(L); - // std::mt19937 rng(rand()); - // GenRandom(rng, init_ids.data(), L, (unsigned) nd_); boost::dynamic_bitset<> flags{_nd, 0}; unsigned tmp_l = 0; @@ -1131,17 +1367,12 @@ namespace diskann { tmp_l++; } - for (unsigned i = 0; i < init_ids.size(); i++) { - unsigned id = init_ids[i]; - if (id >= _nd) - continue; - _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); - } L = 0; for (unsigned i = 0; i < init_ids.size(); i++) { unsigned id = init_ids[i]; if (id >= _nd) continue; + _mm_prefetch(_opt_graph + _node_size * id, _MM_HINT_T0); T * x = (T *) (_opt_graph + _node_size * id); float norm_x = *x; x++; @@ -1151,7 +1382,6 @@ namespace diskann { flags[id] = true; L++; } - // std::cout< full_pivot_data; - if (file_exists(pq_pivots_path)) { + std::string save_path = pq_pivots_path + ".bin"; + if (file_exists(save_path)) { size_t file_dim, file_num_centers; - diskann::load_bin(pq_pivots_path, full_pivot_data, file_num_centers, + diskann::load_bin(save_path, full_pivot_data, file_num_centers, file_dim); if (file_dim == dim && file_num_centers == num_centers) { diskann::cout << "PQ pivot file exists. Not generating again" @@ -338,7 +339,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, } } - diskann::save_bin(pq_pivots_path.c_str(), full_pivot_data.get(), + diskann::save_bin(save_path.c_str(), full_pivot_data.get(), (size_t) num_centers, dim); std::string centroids_path = pq_pivots_path + "_centroid.bin"; diskann::save_bin(centroids_path.c_str(), centroid.get(), (size_t) dim, @@ -355,7 +356,7 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, // streams the base file (data_file), and computes the closest centers in each // chunk to generate the compressed data_file and stores it in // pq_compressed_vectors_path. -// If the numbber of centers is < 256, it stores as byte vector, else as 4-byte +// If the number of centers is < 256, it stores as byte vector, else as 4-byte // vector in binary format. template int generate_pq_data_from_pivots(const std::string data_file, @@ -376,7 +377,7 @@ int generate_pq_data_from_pivots(const std::string data_file, std::unique_ptr rearrangement; std::unique_ptr chunk_offsets; - if (!file_exists(pq_pivots_path)) { + if (!file_exists(pq_pivots_path + ".bin")) { diskann::cout << "ERROR: PQ k-means pivot file not found" << std::endl; throw diskann::ANNException("PQ k-means pivot file not found", -1); } else { @@ -409,8 +410,8 @@ int generate_pq_data_from_pivots(const std::string data_file, size_t file_num_centers; size_t file_dim; - diskann::load_bin(pq_pivots_path, full_pivot_data, file_num_centers, - file_dim); + diskann::load_bin(pq_pivots_path + ".bin", full_pivot_data, + file_num_centers, file_dim); if (file_num_centers != num_centers) { std::stringstream stream; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 21da3521d6..d6fcd7c3ce 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -54,35 +54,6 @@ // returns region of `node_buf` containing [COORD(T)] #define OFFSET_TO_NODE_COORDS(node_buf) (T *) (node_buf) -namespace { - void aggregate_coords(const unsigned *ids, const _u64 n_ids, - const _u8 *all_coords, const _u64 ndims, _u8 *out) { - for (_u64 i = 0; i < n_ids; i++) { - memcpy(out + i * ndims, all_coords + ids[i] * ndims, ndims * sizeof(_u8)); - } - } - - void pq_dist_lookup(const _u8 *pq_ids, const _u64 n_pts, - const _u64 pq_nchunks, const float *pq_dists, - float *dists_out) { - _mm_prefetch((char *) dists_out, _MM_HINT_T0); - _mm_prefetch((char *) pq_ids, _MM_HINT_T0); - _mm_prefetch((char *) (pq_ids + 64), _MM_HINT_T0); - _mm_prefetch((char *) (pq_ids + 128), _MM_HINT_T0); - memset(dists_out, 0, n_pts * sizeof(float)); - for (_u64 chunk = 0; chunk < pq_nchunks; chunk++) { - const float *chunk_dists = pq_dists + 256 * chunk; - if (chunk < pq_nchunks - 1) { - _mm_prefetch((char *) (chunk_dists + 256), _MM_HINT_T0); - } - for (_u64 idx = 0; idx < n_pts; idx++) { - _u8 pq_centerid = pq_ids[pq_nchunks * idx + chunk]; - dists_out[idx] += chunk_dists[pq_centerid]; - } - } - } -} // namespace - namespace diskann { template<> PQFlashIndex<_u8>::PQFlashIndex( @@ -97,7 +68,7 @@ namespace diskann { this->dist_cmp = new DistanceL2UInt8(); if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 dist_cmp_float function." << std::endl; - this->dist_cmp_float = new DistanceL2(); + this->dist_cmp_float = new AVX2DistanceL2Float(); } else if (AvxSupportedCPU) { diskann::cout << "Using AVX dist_cmp_float function" << std::endl; this->dist_cmp_float = new AVXDistanceL2Float(); @@ -116,7 +87,7 @@ namespace diskann { diskann::cout << "Using AVX2 function for dist_cmp and dist_cmp_float" << std::endl; this->dist_cmp = new DistanceL2Int8(); - this->dist_cmp_float = new DistanceL2(); + this->dist_cmp_float = new AVX2DistanceL2Float(); } else if (AvxSupportedCPU) { diskann::cout << "No AVX2 support. Switching to AVX routines for " "dist_cmp, dist_cmp_float." @@ -139,8 +110,8 @@ namespace diskann { if (Avx2SupportedCPU) { diskann::cout << "Using AVX2 functions for dist_cmp and dist_cmp_float" << std::endl; - this->dist_cmp = new DistanceL2(); - this->dist_cmp_float = new DistanceL2(); + this->dist_cmp = new AVX2DistanceL2Float(); + this->dist_cmp_float = new AVX2DistanceL2Float(); } else if (AvxSupportedCPU) { diskann::cout << "No AVX2 support. Switching to AVX functions for " "dist_cmp and dist_cmp_float." @@ -151,8 +122,8 @@ namespace diskann { diskann::cout << "No AVX/AVX2 support. Switching to slow implementations " "for dist_cmp and dist_cmp_float" << std::endl; - this->dist_cmp = new AVXDistanceL2Float(); - this->dist_cmp_float = new AVXDistanceL2Float(); + this->dist_cmp = new SlowDistanceL2Float(); + this->dist_cmp_float = new SlowDistanceL2Float(); } } @@ -565,18 +536,20 @@ namespace diskann { int PQFlashIndex::load(uint32_t num_threads, const char *pq_prefix, const char *disk_index_file) { #endif - std::string pq_table_bin = std::string(pq_prefix) + "_pivots.bin"; + std::string pq_table_path = std::string(pq_prefix) + "_pq_pivots"; std::string pq_compressed_vectors = - std::string(pq_prefix) + "_compressed.bin"; + std::string(pq_prefix) + "_pq_compressed.bin"; std::string medoids_file = std::string(disk_index_file) + "_medoids.bin"; std::string centroids_file = std::string(disk_index_file) + "_centroids.bin"; size_t pq_file_dim, pq_file_num_centroids; #ifdef EXEC_ENV_OLS - get_bin_metadata(files, pq_table_bin, pq_file_num_centroids, pq_file_dim); + get_bin_metadata(files, pq_table_path + ".bin", pq_file_num_centroids, + pq_file_dim); #else - get_bin_metadata(pq_table_bin, pq_file_num_centroids, pq_file_dim); + get_bin_metadata(pq_table_path + ".bin", pq_file_num_centroids, + pq_file_dim); #endif this->disk_index_file = std::string(disk_index_file); @@ -603,9 +576,9 @@ namespace diskann { this->n_chunks = nchunks_u64; #ifdef EXEC_ENV_OLS - pq_table.load_pq_centroid_bin(files, pq_table_bin.c_str(), nchunks_u64); + pq_table.load_pq_centroid_bin(files, pq_table_path.c_str(), nchunks_u64); #else - pq_table.load_pq_centroid_bin(pq_table_bin.c_str(), nchunks_u64); + pq_table.load_pq_centroid_bin(pq_table_path.c_str(), nchunks_u64); #endif diskann::cout @@ -738,7 +711,7 @@ namespace diskann { use_medoids_data_as_centroids(); } - diskann::cout << "done.." << std::endl; + diskann::cout << "Index load complete." << std::endl; return 0; } @@ -810,10 +783,10 @@ namespace diskann { // lambda to batch compute query<-> node distances in PQ space auto compute_dists = [this, pq_coord_scratch, pq_dists]( const unsigned *ids, const _u64 n_ids, float *dists_out) { - ::aggregate_coords(ids, n_ids, this->data, this->n_chunks, - pq_coord_scratch); - ::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, - dists_out); + diskann::aggregate_coords(ids, n_ids, this->data, this->n_chunks, + pq_coord_scratch); + diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->n_chunks, pq_dists, + dists_out); }; Timer query_timer, io_timer, cpu_timer; std::vector retset(l_search + 1); diff --git a/src/utils.cpp b/src/utils.cpp index c3eec1d918..42b628fb02 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -49,12 +49,26 @@ bool cpuHasAvx2Support() { } return false; } + +bool cpuHasAvx512Support() { + int cpuInfo[4]; + __cpuid(cpuInfo, 0); + int n = cpuInfo[0]; + if (n >= 7) { + __cpuidex(cpuInfo, 7, 0); + static int avx512fMask = 0x10000; + return (cpuInfo[1] & avx512Mask) > 0; + } + return false; +} #endif #ifndef _WINDOWS bool AvxSupportedCPU = false; -bool Avx2SupportedCPU = true; +bool Avx2SupportedCPU = false; +bool Avx512SupportedCPU = true; #else bool AvxSupportedCPU = cpuHasAvxSupport(); bool Avx2SupportedCPU = cpuHasAvx2Support(); +bool Avx512SupportedCPU = cpuHasAvx512Support(); #endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fa6e8ea045..f7c7b03f81 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -12,6 +12,15 @@ else() target_link_libraries(build_memory_index ${PROJECT_NAME} -ltcmalloc) endif() +add_executable(build_pq_memory_index build_pq_memory_index.cpp ) +if(MSVC) + target_link_options(build_pq_memory_index PRIVATE /MACHINE:x64 /DEBUG:FULL "/INCLUDE:_tcmalloc") + target_link_libraries(build_pq_memory_index debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) + target_link_libraries(build_pq_memory_index optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib ${PROJECT_SOURCE_DIR}/dependencies/windows/tcmalloc/libtcmalloc_minimal.lib) +else() + target_link_libraries(build_pq_memory_index ${PROJECT_NAME} -ltcmalloc) +endif() + add_executable(search_memory_index search_memory_index.cpp ) if(MSVC) target_link_options(search_memory_index PRIVATE /MACHINE:x64 /DEBUG:FULL) @@ -21,6 +30,15 @@ else() target_link_libraries(search_memory_index ${PROJECT_NAME} aio -ltcmalloc) endif() +add_executable(search_pq_memory_index search_pq_memory_index.cpp ) +if(MSVC) + target_link_options(search_pq_memory_index PRIVATE /MACHINE:x64 /DEBUG:FULL) + target_link_libraries(search_pq_memory_index debug ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_DEBUG}/diskann_dll.lib) + target_link_libraries(search_pq_memory_index optimized ${CMAKE_LIBRARY_OUTPUT_DIRECTORY_RELEASE}/diskann_dll.lib) +else() + target_link_libraries(search_pq_memory_index ${PROJECT_NAME} aio -ltcmalloc) +endif() + add_executable(build_disk_index build_disk_index.cpp ) if(MSVC) target_link_options(build_disk_index PRIVATE /MACHINE:x64 /DEBUG:FULL "/INCLUDE:_tcmalloc") diff --git a/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..82afc7eb78 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -29,7 +29,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::FAST_L2, data_path.c_str()); auto s = std::chrono::high_resolution_clock::now(); index.build(paras); std::chrono::duration diff = @@ -45,7 +45,7 @@ int main(int argc, char** argv) { if (argc != 8) { std::cout << "Usage: " << argv[0] << " [data_type] [data_file.bin] " - "[output_index_file] " + "[output_index_file.bin] " << "[R] [L] [alpha]" << " [num_threads_to_use]. See README for more information on " "parameters." diff --git a/tests/build_pq_memory_index.cpp b/tests/build_pq_memory_index.cpp new file mode 100644 index 0000000000..d820034932 --- /dev/null +++ b/tests/build_pq_memory_index.cpp @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include "utils.h" + +#ifndef _WINDOWS +#include +#include +#else +#include +#endif + +#include "memory_mapper.h" + +template +int build_pq_in_memory_index(const std::string& data_path, + const std::string& pq_path, const unsigned R, + const unsigned L, const float alpha, + const unsigned num_chunks, + const std::string& save_path, + const unsigned num_threads) { + diskann::Parameters paras; + paras.Set("R", R); + paras.Set("L", L); + paras.Set("C", 750); + paras.Set("alpha", alpha); + paras.Set("saturate_graph", 0); + paras.Set("num_chunks", num_chunks); + paras.Set("num_threads", num_threads); + + diskann::Index index(diskann::FAST_L2, data_path.c_str()); + + auto s = std::chrono::high_resolution_clock::now(); + index.pq_build(data_path.c_str(), pq_path.c_str(), paras); + std::chrono::duration diff = + std::chrono::high_resolution_clock::now() - s; + + std::cout << "Indexing time: " << diff.count() << "\n"; + index.save(save_path.c_str()); + + return 0; +} + +int main(int argc, char** argv) { + if (argc != 10) { + std::cout << "Usage: " << argv[0] + << " [data_type] [data_file.bin] " + "[pq_output_prefix] [output_index_file.bin] " + << "[R] [L] [alpha] [num_chunks]" + << " [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 pq_path(argv[3]); + const std::string save_path(argv[4]); + const unsigned R = (unsigned) atoi(argv[5]); + const unsigned L = (unsigned) atoi(argv[6]); + const float alpha = (float) atof(argv[7]); + const unsigned num_chunks = (unsigned) atoi(argv[8]); + const unsigned num_threads = (unsigned) atoi(argv[9]); + + if (std::string(argv[1]) == std::string("int8")) + build_pq_in_memory_index(data_path, pq_path, R, L, alpha, + num_chunks, save_path, num_threads); + else if (std::string(argv[1]) == std::string("uint8")) + build_pq_in_memory_index(data_path, pq_path, R, L, alpha, + num_chunks, save_path, num_threads); + else if (std::string(argv[1]) == std::string("float")) + build_pq_in_memory_index(data_path, pq_path, R, L, alpha, num_chunks, + save_path, num_threads); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/search_pq_memory_index.cpp b/tests/search_pq_memory_index.cpp new file mode 100644 index 0000000000..357da1b9d7 --- /dev/null +++ b/tests/search_pq_memory_index.cpp @@ -0,0 +1,175 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include +#include +#include +#include + +#ifndef _WINDOWS +#include +#include +#include +#include +#endif + +#include "aux_utils.h" +#include "index.h" +#include "memory_mapper.h" +#include "utils.h" + +template +int search_pq_memory_index(int argc, char** argv) { + T* query = nullptr; + unsigned* gt_ids = nullptr; + float* gt_dists = nullptr; + size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim; + std::vector<_u64> Lvec; + + std::string data_file(argv[2]); + const std::string pq_path(argv[3]); + std::string memory_index_file(argv[4]); + _u64 num_threads = std::atoi(argv[5]); + std::string query_bin(argv[6]); + std::string truthset_bin(argv[7]); + _u64 recall_at = std::atoi(argv[8]); + std::string result_output_prefix(argv[9]); + bool use_optimized_search = std::atoi(argv[10]); + + if ((std::string(argv[1]) != std::string("float")) && + (use_optimized_search == true)) { + std::cout << "Error. Optimized search currently only supported for " + "floating point datatypes. Using un-optimized search." + << std::endl; + use_optimized_search = false; + } + + bool calc_recall_flag = false; + + for (int ctr = 10; ctr < argc; ctr++) { + _u64 curL = std::atoi(argv[ctr]); + if (curL >= recall_at) + Lvec.push_back(curL); + } + + if (Lvec.size() == 0) { + std::cout << "No valid Lsearch found. Lsearch must be at least recall_at." + << 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); + if (gt_num != query_num) { + std::cout << "Error. Mismatch in number of queries and ground truth data" + << std::endl; + } + calc_recall_flag = true; + } + + std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield); + std::cout.precision(2); + + auto metric = diskann::L2; + if (use_optimized_search) + metric = diskann::FAST_L2; + diskann::Index index(metric, data_file.c_str()); + index.load(memory_index_file.c_str()); // to load NSG + std::cout << "Index loaded" << std::endl; + index.pq_load(pq_path.c_str()); // to load NSG + std::cout << "PQ data loaded" << std::endl; + + if (use_optimized_search) + index.optimize_graph(); + + std::string recall_string = "Recall@" + std::to_string(recall_at); + std::cout << std::setw(4) << "Ls" << std::setw(12) << "QPS " << std::setw(18) + << "Mean Latency (mus)" << std::setw(15) << "99.9 Latency" + << std::setw(12) << recall_string << std::endl; + std::cout << "===============================================================" + "===============" + << std::endl; + + std::vector> query_result_ids(Lvec.size()); + std::vector latency_stats(query_num, 0); + + for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++) { + _u64 L = Lvec[test_id]; + query_result_ids[test_id].resize(recall_at * query_num); + omp_set_num_threads(num_threads); + + auto s = std::chrono::high_resolution_clock::now(); + //#pragma omp parallel for schedule(dynamic, 1) + for (int64_t i = 0; i < (int64_t) query_num; i++) { + auto qs = std::chrono::high_resolution_clock::now(); + index.pq_search(query + i * query_aligned_dim, recall_at, L, + query_result_ids[test_id].data() + i * recall_at); + auto qe = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = qe - qs; + latency_stats[i] = diff.count() * 1000000; + } + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + + float qps = (query_num / diff.count()); + + float recall = 0; + if (calc_recall_flag) { + recall = diskann::calculate_recall(query_num, gt_ids, gt_dists, gt_dim, + query_result_ids[test_id].data(), + recall_at, recall_at); + } + + std::sort(latency_stats.begin(), latency_stats.end()); + double mean_latency = 0; + for (uint64_t q = 0; q < query_num; q++) { + mean_latency += latency_stats[q]; + } + mean_latency /= query_num; + + std::cout << std::setw(4) << L << std::setw(12) << qps << std::setw(18) + << (float) mean_latency << std::setw(15) + << (float) latency_stats[(_u64)(0.999 * query_num)] + << std::setw(12) << recall << std::endl; + } + + std::cout << "Done searching. Now saving results " << std::endl; + _u64 test_id = 0; + for (auto L : Lvec) { + std::string cur_result_path = + result_output_prefix + "_" + std::to_string(L) + "_idx_uint32.bin"; + diskann::save_bin<_u32>(cur_result_path, query_result_ids[test_id].data(), + query_num, recall_at); + test_id++; + } + + diskann::aligned_free(query); + return 0; +} + +int main(int argc, char** argv) { + if (argc < 12) { + std::cout + << "Usage: " << argv[0] + << " [index_type] [data_file.bin] " + "[pq_path_prefix] [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_pq_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("uint8")) + search_pq_memory_index(argc, argv); + else if (std::string(argv[1]) == std::string("float")) + search_pq_memory_index(argc, argv); + else + std::cout << "Unsupported type. Use float/int8/uint8" << std::endl; +} diff --git a/tests/utils/bin_to_tsv.cpp b/tests/utils/bin_to_tsv.cpp new file mode 100644 index 0000000000..98e45d7a1a --- /dev/null +++ b/tests/utils/bin_to_tsv.cpp @@ -0,0 +1,69 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include +#include "utils.h" + +template +void block_convert(std::ofstream& writer, std::ifstream& reader, T* read_buf, + _u64 npts, _u64 ndims) { + reader.read((char*) read_buf, npts * ndims * sizeof(float)); + + for (_u64 i = 0; i < npts; i++) { + for (_u64 d = 0; d < ndims; d++) { + writer << read_buf[d + i * ndims]; + if (d < ndims - 1) + writer << "\t"; + else + writer << "\n"; + } + } +} + +int main(int argc, char** argv) { + if (argc != 4) { + std::cout << argv[0] << " input_bin output_tsv" << std::endl; + exit(-1); + } + std::string type_string(argv[1]); + if ((type_string != std::string("float")) && + (type_string != std::string("uint32")) && + (type_string != std::string("int8")) && + (type_string != std::string("uin8"))) { + std::cerr << "Error: type not supported. Use float/uint32/int8/uint8" << std::endl; + } + + std::ifstream reader(argv[2], std::ios::binary); + _u32 npts_u32; + _u32 ndims_u32; + reader.read((char*) &npts_u32, sizeof(_s32)); + reader.read((char*) &ndims_u32, sizeof(_s32)); + size_t npts = npts_u32; + size_t ndims = ndims_u32; + std::cout << "Dataset: #pts = " << npts << ", # dims = " << ndims + << std::endl; + + _u64 blk_size = 131072; + _u64 nblks = ROUND_UP(npts, blk_size) / blk_size; + + std::ofstream writer(argv[3]); + char* read_buf = new char[blk_size * ndims * 4]; + for (_u64 i = 0; i < nblks; i++) { + _u64 cblk_size = std::min(npts - i * blk_size, blk_size); + if (type_string == std::string("float")) + block_convert(writer, reader, (float*)read_buf, cblk_size, ndims); + else if (type_string == std::string("uint32")) + block_convert(writer, reader, (uint32_t*) read_buf, cblk_size, ndims); + else if (type_string == std::string("int8")) + block_convert(writer, reader, (int8_t*) read_buf, cblk_size, ndims); + else if (type_string == std::string("uint8")) + block_convert(writer, reader, (uint8_t*) read_buf, cblk_size, + ndims); + std::cout << "Block #" << i << " written" << std::endl; + } + + delete[] read_buf; + + writer.close(); + reader.close(); +} diff --git a/tests/utils/generate_pq.cpp b/tests/utils/generate_pq.cpp index 9e7a37c597..4d326a5aa6 100644 --- a/tests/utils/generate_pq.cpp +++ b/tests/utils/generate_pq.cpp @@ -11,7 +11,7 @@ bool generate_pq(const std::string& data_path, const std::string& index_prefix_path, const size_t num_pq_centers, const size_t num_pq_chunks, const float sampling_rate) { - std::string pq_pivots_path = index_prefix_path + "_pq_pivots.bin"; + std::string pq_pivots_path = index_prefix_path + "_pq_pivots"; std::string pq_compressed_vectors_path = index_prefix_path + "_compressed.bin";