From f09e6e0db2543e1527309092c5770eb8a40e0618 Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Mon, 1 Feb 2021 13:15:28 +0530 Subject: [PATCH 01/37] Vamana Python Bindings --- CMakeLists.txt | 2 +- include/parameters.h | 2 +- python/setup.py | 84 +++++++ python/src/vamana_bindings.cpp | 277 +++++++++++++++++++++++ python/tests/test_build_memory_index.py | 26 +++ python/tests/test_search_memory_index.py | 64 ++++++ 6 files changed, 453 insertions(+), 2 deletions(-) create mode 100644 python/setup.py create mode 100644 python/src/vamana_bindings.cpp create mode 100644 python/tests/test_build_memory_index.py create mode 100644 python/tests/test_search_memory_index.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd3886f5b..6074c03c87 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 -DUSE_AVX2 -fPIC) endif() add_subdirectory(src) 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/python/setup.py b/python/setup.py new file mode 100644 index 0000000000..9b0b8f3177 --- /dev/null +++ b/python/setup.py @@ -0,0 +1,84 @@ +# 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': ['-Ofast']} + arch_list = '-march -msse -msse2 -msse3 -mssse3 -msse4 -msse4a -msse4.1 -msse4.2 -mavx -mavx2'.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': []} + 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( + 'vamanapy', + ['src/vamana_bindings.cpp'], + include_dirs=["../include/", + pybind11.get_include(False), + pybind11.get_include(True)], + libraries=[], + language='c++', + extra_objects=['../build/src/libdiskann_s.a'], + ), +] + + +setup( + name="vamanapy", + version=__version__, + author="Shikhar Jaiswal", + author_email="t-sjaiswal@microsoft.com", + url="https://github.com/microsoft/diskann", + description="Vamana 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/vamana_bindings.cpp b/python/src/vamana_bindings.cpp new file mode 100644 index 0000000000..8f5913fb87 --- /dev/null +++ b/python/src/vamana_bindings.cpp @@ -0,0 +1,277 @@ +// 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("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("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); + 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")); +} diff --git a/python/tests/test_build_memory_index.py b/python/tests/test_build_memory_index.py new file mode 100644 index 0000000000..8d89e21c58 --- /dev/null +++ b/python/tests/test_build_memory_index.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +from vamanapy import Metric, Parameters, SinglePrecisionIndex + + +data_path = "/mnt/SIFT1M/sift_base.bin" +save_path = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/test_build_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_threads", 32) + +start = time.time() +index = SinglePrecisionIndex(Metric.L2, data_path) +index.build(params, []) +end = time.time() + +print("Indexing Time: " + str(end - start) + " seconds") + +index.save(save_path) diff --git a/python/tests/test_search_memory_index.py b/python/tests/test_search_memory_index.py new file mode 100644 index 0000000000..03af99d5e9 --- /dev/null +++ b/python/tests/test_search_memory_index.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +import time +import vamanapy 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/SIFT1M/test_build_memory_index.bin" +output_path_prefix = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/test_search_memory_index" + +recall_at = 10 +num_threads = 32 +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.optimize_graph() +print("Graph Optimization Completed") + +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() + vp.set_num_threads(num_threads) + 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 = output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + vp.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) From 5ae530c3d4e7cd5505c70e1251eb274bae120c58 Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Wed, 3 Feb 2021 20:16:01 +0530 Subject: [PATCH 02/37] AVX-512 support and Improved Optimized Search --- CMakeLists.txt | 2 +- include/distance.h | 90 +++++++++++++++++++++++++++++++++++++----- include/utils.h | 6 +-- python/setup.py | 2 +- src/index.cpp | 10 +++-- src/pq_flash_index.cpp | 8 ++-- src/utils.cpp | 16 +++++++- 7 files changed, 110 insertions(+), 24 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 6074c03c87..bfc1a15977 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 -fPIC) + 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_AVX512 -fPIC) endif() add_subdirectory(src) diff --git a/include/distance.h b/include/distance.h index a4f311e510..c76fe90702 100644 --- a/include/distance.h +++ b/include/distance.h @@ -174,7 +174,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 @@ -333,12 +380,39 @@ namespace diskann { float compare(const T *a, const T *b, unsigned size) const { float result = 0; #ifdef __GNUC__ -#ifdef __AVX__ +#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 *l = (float *) a; + const float *r = (float *) b; + 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); \ - tmp1 = _mm256_mul_ps(tmp1, tmp2); \ - dest = _mm256_add_ps(dest, tmp1); + dest = _mm256_fmadd_ps(tmp1, tmp2, dest); __m256 sum; __m256 l0, l1; @@ -350,9 +424,8 @@ namespace diskann { 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); } @@ -361,9 +434,8 @@ namespace diskann { 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__ diff --git a/include/utils.h b/include/utils.h index 6b9db5bf62..2dddfaf50d 100644 --- a/include/utils.h +++ b/include/utils.h @@ -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 index 9b0b8f3177..e0d82eeee8 100644 --- a/python/setup.py +++ b/python/setup.py @@ -15,7 +15,7 @@ class BuildExt(build_ext): """A custom build extension for adding compiler-specific options.""" c_opts = {'unix': ['-Ofast']} - arch_list = '-march -msse -msse2 -msse3 -mssse3 -msse4 -msse4a -msse4.1 -msse4.2 -mavx -mavx2'.split() + 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: diff --git a/src/index.cpp b/src/index.cpp index 2405878669..c0b6383bbe 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -50,9 +50,13 @@ namespace { std::cout << "Here" << std::endl; return new diskann::DistanceFastL2(); } 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; diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 21da3521d6..e5fa701623 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -97,7 +97,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 +116,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 +139,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." 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 From 91a624b617a7805f775832295287adf0bf05bc9f Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Fri, 12 Feb 2021 01:08:17 +0530 Subject: [PATCH 03/37] Add Support for Inner Product Metric --- include/distance.h | 234 +++++++++++++++++++++++++++------------------ include/neighbor.h | 10 +- src/index.cpp | 30 +++--- 3 files changed, 159 insertions(+), 115 deletions(-) diff --git a/include/distance.h b/include/distance.h index c76fe90702..2431f96b6a 100644 --- a/include/distance.h +++ b/include/distance.h @@ -377,227 +377,277 @@ 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 __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); +#define AVX512_L2NORM(addr, dest, tmp) \ + tmp = _mm512_loadu_ps(addr); \ + dest = _mm512_fmadd_ps(tmp, tmp, dest); __m512 sum; __m512 l0, l1; - __m512 r0, r1; unsigned D = (size + 15) & ~15U; unsigned DR = D % 32; 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; sum = _mm512_setzero_ps(); if (DR) { - AVX512_DOT(e_l, e_r, sum, l0, r0); + AVX512_L2NORM(e_l, sum, l0); } - 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); + 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_DOT(addr1, addr2, dest, tmp1, tmp2) \ - tmp1 = _mm256_loadu_ps(addr1); \ - tmp2 = _mm256_loadu_ps(addr2); \ - dest = _mm256_fmadd_ps(tmp1, tmp2, dest); +#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; 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); } 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 { + return DistanceInnerProduct::norm(a, size); + } + + 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 { + return 1 / std::sqrt(DistanceInnerProduct::norm(a, size)); + } + + 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)); + 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/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/src/index.cpp b/src/index.cpp index c0b6383bbe..07c26740f6 100644 --- a/src/index.cpp +++ b/src/index.cpp @@ -47,8 +47,11 @@ 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 (Avx512SupportedCPU) { std::cout << "Using AVX512 distance computation" << std::endl; @@ -67,9 +70,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__, @@ -1085,7 +1087,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); @@ -1107,12 +1110,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; @@ -1135,17 +1137,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++; @@ -1155,7 +1152,6 @@ namespace diskann { flags[id] = true; L++; } - // std::cout< Date: Tue, 16 Feb 2021 19:32:12 +0530 Subject: [PATCH 04/37] Add support for zero vectors --- include/distance.h | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index 2431f96b6a..0bed7363c4 100644 --- a/include/distance.h +++ b/include/distance.h @@ -1,6 +1,7 @@ #pragma once #include +#include #ifdef _WINDOWS #include #include @@ -635,12 +636,19 @@ namespace diskann { class DistanceFastInnerProduct : public DistanceInnerProduct { public: float norm(const T *a, unsigned size) const { - return 1 / std::sqrt(DistanceInnerProduct::norm(a, size)); + 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; From a6160a758b77441666998fee71433f5ad4935654 Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Thu, 18 Feb 2021 00:01:55 +0530 Subject: [PATCH 05/37] Fix Aligned Dimension Instruction --- include/utils.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/utils.h b/include/utils.h index 2dddfaf50d..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); From 0d91a56b4a14c9ba7a2cf9ce70bdf9770bc7856b Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Sat, 27 Feb 2021 19:12:30 +0530 Subject: [PATCH 06/37] Refactor L2 Distance for Non-Zero Vector Traversal --- include/distance.h | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/include/distance.h b/include/distance.h index 0bed7363c4..b037d1042a 100644 --- a/include/distance.h +++ b/include/distance.h @@ -623,7 +623,22 @@ namespace diskann { class DistanceFastL2 : public DistanceInnerProduct { public: float norm(const T *a, unsigned size) const { - return DistanceInnerProduct::norm(a, size); + 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 { From 36b542ce2434e4d607154e29a98d4813b8661b0f Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Sat, 6 Mar 2021 19:31:21 +0530 Subject: [PATCH 07/37] Add multi-threaded search support for Python bindings --- python/src/vamana_bindings.cpp | 1 + python/tests/test_build_memory_index.py | 2 +- python/tests/test_search_memory_index.py | 83 ++++++++++++++++-------- tests/build_memory_index.cpp | 2 +- 4 files changed, 58 insertions(+), 30 deletions(-) diff --git a/python/src/vamana_bindings.cpp b/python/src/vamana_bindings.cpp index 8f5913fb87..1afcc91807 100644 --- a/python/src/vamana_bindings.cpp +++ b/python/src/vamana_bindings.cpp @@ -267,6 +267,7 @@ PYBIND11_MODULE(vamanapy, m) { 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)); diff --git a/python/tests/test_build_memory_index.py b/python/tests/test_build_memory_index.py index 8d89e21c58..8aa01537f8 100644 --- a/python/tests/test_build_memory_index.py +++ b/python/tests/test_build_memory_index.py @@ -17,7 +17,7 @@ params.set("num_threads", 32) start = time.time() -index = SinglePrecisionIndex(Metric.L2, data_path) +index = SinglePrecisionIndex(Metric.FAST_L2, data_path) index.build(params, []) end = time.time() diff --git a/python/tests/test_search_memory_index.py b/python/tests/test_search_memory_index.py index 03af99d5e9..6518d2205d 100644 --- a/python/tests/test_search_memory_index.py +++ b/python/tests/test_search_memory_index.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import time +import numpy as np import vamanapy as vp @@ -12,7 +13,9 @@ output_path_prefix = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/test_search_memory_index" recall_at = 10 -num_threads = 32 +# 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() @@ -29,36 +32,60 @@ index.optimize_graph() print("Graph Optimization Completed") -print("Ls QPS Mean Latency (mus) 99.9 Latency Recall@10") -print("================================================================") +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 i, L in enumerate(l_search): - latency_stats = [] - query_result_ids = vp.VectorUnsigned() - vp.set_num_threads(num_threads) - 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 = 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) - 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) + query_result_ids = index.batch_numpy_query(query_data, recall_at, num_queries, L) qe = time.time() - latency_stats.append(float((qe - qs) * 1000000)) + latency_stats = 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))) + 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) + result_path = 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/tests/build_memory_index.cpp b/tests/build_memory_index.cpp index a5a13595e3..19a2bf45fa 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 = From 49560ec021704ff6cce75ed8bb4a8956bb00efc2 Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Mon, 15 Mar 2021 19:48:05 +0000 Subject: [PATCH 08/37] Add PQ Compression Support for Vamana --- include/index.h | 20 ++ include/pq_flash_index.h | 47 ++++ include/pq_table.h | 5 +- python/setup.py | 7 +- python/src/vamana_bindings.cpp | 48 ++++ python/tests/test_build_pq_memory_index.py | 28 +++ python/tests/test_search_pq_memory_index.py | 93 ++++++++ src/aux_utils.cpp | 6 +- src/index.cpp | 230 ++++++++++++++++++++ src/partition_and_pq.cpp | 15 +- src/pq_flash_index.cpp | 51 +---- tests/CMakeLists.txt | 18 ++ tests/build_memory_index.cpp | 2 +- tests/build_pq_memory_index.cpp | 79 +++++++ tests/search_pq_memory_index.cpp | 175 +++++++++++++++ tests/utils/generate_pq.cpp | 2 +- 16 files changed, 770 insertions(+), 56 deletions(-) create mode 100644 python/tests/test_build_pq_memory_index.py create mode 100644 python/tests/test_search_pq_memory_index.py create mode 100644 tests/build_pq_memory_index.cpp create mode 100644 tests/search_pq_memory_index.cpp 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/pq_flash_index.h b/include/pq_flash_index.h index c7601851c6..872475276c 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] 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/python/setup.py b/python/setup.py index e0d82eeee8..d6a72942e2 100644 --- a/python/setup.py +++ b/python/setup.py @@ -14,7 +14,7 @@ class BuildExt(build_ext): """A custom build extension for adding compiler-specific options.""" - c_opts = {'unix': ['-Ofast']} + c_opts = {'unix': ['-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 @@ -27,7 +27,7 @@ class BuildExt(build_ext): if no_arch_flag: c_opts['unix'].append('-march=native') - link_opts = {'unix': []} + 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']) @@ -58,7 +58,8 @@ def build_extensions(self): Extension( 'vamanapy', ['src/vamana_bindings.cpp'], - include_dirs=["../include/", + include_dirs=["../include/", + "/opt/intel/compilers_and_libraries/linux/mkl/include/", pybind11.get_include(False), pybind11.get_include(True)], libraries=[], diff --git a/python/src/vamana_bindings.cpp b/python/src/vamana_bindings.cpp index 1afcc91807..1961144c63 100644 --- a/python/src/vamana_bindings.cpp +++ b/python/src/vamana_bindings.cpp @@ -162,6 +162,10 @@ PYBIND11_MODULE(vamanapy, m) { } }, 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 == "") { @@ -178,6 +182,11 @@ PYBIND11_MODULE(vamanapy, m) { 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, @@ -274,5 +283,44 @@ PYBIND11_MODULE(vamanapy, m) { } 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_pq_memory_index.py b/python/tests/test_build_pq_memory_index.py new file mode 100644 index 0000000000..0cdb4b5e78 --- /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 vamanapy 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_pq_memory_index.py b/python/tests/test_search_pq_memory_index.py new file mode 100644 index 0000000000..7af286a490 --- /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 vamanapy 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 07c26740f6..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" @@ -186,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; } @@ -193,6 +200,8 @@ namespace diskann { Index::~Index() { delete this->_distance; aligned_free(_data); + delete[] _pq_data; + aligned_free(_pq_table_dists); } template<> @@ -1014,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, diff --git a/src/partition_and_pq.cpp b/src/partition_and_pq.cpp index 9da49b2203..f7ca8ac7fe 100644 --- a/src/partition_and_pq.cpp +++ b/src/partition_and_pq.cpp @@ -211,9 +211,10 @@ int generate_pq_pivots(const float *passed_train_data, size_t num_train, std::unique_ptr 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 e5fa701623..30b1f05a2a 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( @@ -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 @@ -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/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 19a2bf45fa..82afc7eb78 100644 --- a/tests/build_memory_index.cpp +++ b/tests/build_memory_index.cpp @@ -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/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"; From 2d32142d20d8a773436f26f7598e09d0694a15e0 Mon Sep 17 00:00:00 2001 From: ShikharJ Date: Sun, 2 May 2021 16:35:50 +0000 Subject: [PATCH 09/37] Update Paths in Python Tests --- python/tests/test_build_memory_index.py | 11 +++++++---- python/tests/test_search_memory_index.py | 25 +++++++++++++----------- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/python/tests/test_build_memory_index.py b/python/tests/test_build_memory_index.py index 8aa01537f8..5bc58bc3fa 100644 --- a/python/tests/test_build_memory_index.py +++ b/python/tests/test_build_memory_index.py @@ -2,11 +2,14 @@ # Licensed under the MIT license. import time +import argparse from vamanapy import Metric, Parameters, SinglePrecisionIndex -data_path = "/mnt/SIFT1M/sift_base.bin" -save_path = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/test_build_memory_index.bin" +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) @@ -17,10 +20,10 @@ params.set("num_threads", 32) start = time.time() -index = SinglePrecisionIndex(Metric.FAST_L2, data_path) +index = SinglePrecisionIndex(Metric.FAST_L2, args.data_path) index.build(params, []) end = time.time() print("Indexing Time: " + str(end - start) + " seconds") -index.save(save_path) +index.save(args.save_path) diff --git a/python/tests/test_search_memory_index.py b/python/tests/test_search_memory_index.py index 6518d2205d..24cb5c1088 100644 --- a/python/tests/test_search_memory_index.py +++ b/python/tests/test_search_memory_index.py @@ -2,15 +2,18 @@ # Licensed under the MIT license. import time +import argparse import numpy as np import vamanapy 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/SIFT1M/test_build_memory_index.bin" -output_path_prefix = "/home/t-sjaiswal/diskann/build/tests/SIFT1M/test_search_memory_index" +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. @@ -22,11 +25,11 @@ 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) +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, data_path) -index.load(file_name = memory_index_path) +index = vp.SinglePrecisionIndex(vp.Metric.FAST_L2, args.data_path) +index.load(file_name = args.memory_index_path) print("Index Loaded") index.optimize_graph() @@ -62,7 +65,7 @@ "{:>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" + 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') @@ -87,5 +90,5 @@ "{:>15}".format("{:.2f}".format(mean_latency)) + "{:>15}".format("{:.2f}".format(recall))) - result_path = output_path_prefix + "_" + str(L) + "_idx_uint32.bin" + 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 From f73a0b02bb320576a40fc61c178f947e27fa4a52 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 25 May 2021 12:36:51 -0700 Subject: [PATCH 10/37] remove comments --- include/pq_flash_index.h | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/include/pq_flash_index.h b/include/pq_flash_index.h index 872475276c..17f6478c06 100644 --- a/include/pq_flash_index.h +++ b/include/pq_flash_index.h @@ -147,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, From 4ecc22ef4f22ade69a0d04ab25fbbd6b06a6db7e Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 4 Jun 2021 14:50:41 -0700 Subject: [PATCH 11/37] setup.py file changed, diskann_bindings added --- CMakeLists.txt | 2 +- python/setup.py | 23 +- python/src/diskann_bindings.cpp | 231 ++++++++++++++++++++ python/tests/test_build_pq_memory_index.py | 2 +- python/tests/test_search_memory_index.py | 2 +- python/tests/test_search_pq_memory_index.py | 2 +- 6 files changed, 252 insertions(+), 10 deletions(-) create mode 100644 python/src/diskann_bindings.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index bfc1a15977..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_AVX512 -fPIC) + 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/python/setup.py b/python/setup.py index d6a72942e2..f53807d6c7 100644 --- a/python/setup.py +++ b/python/setup.py @@ -55,9 +55,20 @@ def build_extensions(self): ext_modules = [ + #Extension( + # 'vamanapy', + # ['src/vamana_bindings.cpp'], + # include_dirs=["../include/", + # "/opt/intel/compilers_and_libraries/linux/mkl/include/", + # pybind11.get_include(False), + # pybind11.get_include(True)], + # libraries=[], + # language='c++', + # extra_objects=['../build/src/libdiskann_s.a'], + #), Extension( - 'vamanapy', - ['src/vamana_bindings.cpp'], + 'diskannpy', + ['src/diskann_bindings.cpp'], include_dirs=["../include/", "/opt/intel/compilers_and_libraries/linux/mkl/include/", pybind11.get_include(False), @@ -70,12 +81,12 @@ def build_extensions(self): setup( - name="vamanapy", + name="diskannpy", version=__version__, - author="Shikhar Jaiswal", - author_email="t-sjaiswal@microsoft.com", + author="Shikhar Jaiswal, Harsha Vardhan Simhadri", + author_email="t-sjaiswal@microsoft.com, harshasi@microsoft.com", url="https://github.com/microsoft/diskann", - description="Vamana Bindings using PyBind11", + description="DiskANN Bindings using PyBind11", long_description="", ext_modules=ext_modules, install_requires=['numpy', 'pybind11'], diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp new file mode 100644 index 0000000000..988c501a4a --- /dev/null +++ b/python/src/diskann_bindings.cpp @@ -0,0 +1,231 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +#include + +#include +#include +#include +#include +#include + +#include "utils.h" +#include "memory_mapper.h" +#include "aligned_file_reader.h" +#include "linux_aligned_file_reader.h" +#include "pq_flash_index.h" + + +PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); + +namespace py = pybind11; +using namespace diskann; + +std::unique_ptr> FloatPQFlashIndexCreator() { + return new PQFlashIndex( + std::shared_ptr(new LinuxAlignedFileReader())); +} + +PYBIND11_MODULE(diskannpy, m) { + m.doc() = "DiskANN 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) + .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<>()); + // .def("get_ctx", &LinuxAlignedFileReader::get_ctx) + // .def("register_thread", &LinuxAlignedFileReader::register_thread) + // .def("open", &LinuxAlignedFileReader::open) + // .def("close", &LinuxAlignedFileReader::close) + // .def("read", &LinuxAlignedFileReader::read); + + 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, "DiskANNFloatIndex") + .def(py::init(&FloatPQFlashIndexCreator); + // .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( +// "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( +// "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")); +} diff --git a/python/tests/test_build_pq_memory_index.py b/python/tests/test_build_pq_memory_index.py index 0cdb4b5e78..f076fd43f3 100644 --- a/python/tests/test_build_pq_memory_index.py +++ b/python/tests/test_build_pq_memory_index.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. import time -from vamanapy import Metric, Parameters, SinglePrecisionIndex +from diskannpy import Metric, Parameters, SinglePrecisionIndex data_path = "/mnt/SIFT1M/sift_base.bin" diff --git a/python/tests/test_search_memory_index.py b/python/tests/test_search_memory_index.py index 24cb5c1088..e93e355540 100644 --- a/python/tests/test_search_memory_index.py +++ b/python/tests/test_search_memory_index.py @@ -4,7 +4,7 @@ import time import argparse import numpy as np -import vamanapy as vp +import diskannpy as vp parser = argparse.ArgumentParser() diff --git a/python/tests/test_search_pq_memory_index.py b/python/tests/test_search_pq_memory_index.py index 7af286a490..105b0de0dc 100644 --- a/python/tests/test_search_pq_memory_index.py +++ b/python/tests/test_search_pq_memory_index.py @@ -3,7 +3,7 @@ import time import numpy as np -import vamanapy as vp +import diskannpy as vp data_path = "/mnt/SIFT1M/sift_base.bin" From efc9a0affa84a4206202e5fa1e9b8192d0cb1113 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 8 Jun 2021 01:16:54 -0700 Subject: [PATCH 12/37] added pybind for diskann search --- python/src/diskann_bindings.cpp | 109 +++++++++++-------------- python/tests/test_search_disk_index.py | 93 +++++++++++++++++++++ 2 files changed, 140 insertions(+), 62 deletions(-) create mode 100644 python/tests/test_search_disk_index.py diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 988c501a4a..255c65b15a 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -23,8 +23,10 @@ namespace py = pybind11; using namespace diskann; std::unique_ptr> FloatPQFlashIndexCreator() { - return new PQFlashIndex( - std::shared_ptr(new LinuxAlignedFileReader())); + std::shared_ptr reader(new LinuxAlignedFileReader()); + auto index = new PQFlashIndex(reader); + std::unique_ptr> unique_ptr_index(index); + return unique_ptr_index; } PYBIND11_MODULE(diskannpy, m) { @@ -168,64 +170,47 @@ PYBIND11_MODULE(diskannpy, m) { py::class_>(m, "DiskANNFloatIndex") - .def(py::init(&FloatPQFlashIndexCreator); - // .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( -// "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( -// "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(py::init(&FloatPQFlashIndexCreator)) + .def( + "load", + [](PQFlashIndex &self, const std::string& index_path_prefix) { + + const std::string pq_path = index_path_prefix + std::string("_pq"); + const std::string index_path = + index_path_prefix + std::string("_disk.index"); + self.load(1, pq_path.c_str(), index_path.c_str()); + std::vector node_list; + _u64 num_nodes_to_cache = 100000; + self.cache_bfs_levels(num_nodes_to_cache, node_list); + std::cout << "loaded index, cached " << node_list.size() + << " nodes based on BFS" << std::endl; + }, + py::arg("index_path_prefix")) + .def( + "search", + [](PQFlashIndex &self, const float *query, const _u64 dim, + const _u64 knn, const _u64 l_search, const _u64 beam_width, + _u64 *ids, float *dists) { + QueryStats stats; + self.cached_beam_search(query, knn, l_search, ids, dists, + beam_width, &stats); + }, + py::arg("query"), py::arg("dim"), py::arg("knn") = 10, + py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), + py::arg("dists")) + .def( + "batch_search", + [](PQFlashIndex &self, const float *query_data, + const _u64 nqueries, const _u64 dim, const _u64 knn, + const _u64 l_search, const _u64 beam_width, _u64 *ids, + float *dists) { +#pragma omp parallel for schedule(dynamic, 1) + for (_u64 i = 0; i < nqueries; ++i) + self.cached_beam_search(query_data + i * dim, knn, l_search, + ids + i * knn, dists + i * knn, + beam_width); + }, + py::arg("query_data"), py::arg("nqueries"), py::arg("dim"), + py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), + py::arg("ids"), py::arg("dists")); } diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py new file mode 100644 index 0000000000..8df64ff87e --- /dev/null +++ b/python/tests/test_search_disk_index.py @@ -0,0 +1,93 @@ +# 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('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('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.') +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] +W = 4 + +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(file_name = args.index_path_prefix) +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.VectorUnsigned() + s = time.time() + + for j in range(num_queries): + qs = time.time() + index.search(query_data, query_aligned_dims, recall_t, L, W, + query_result_ids + j * recall_at, query_result_dists + j * recall_at) + 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.set_num_threads(num_threads) + + qs = time.time() + query_result_ids = index.batch_search(query_data, num_queries, query_aligned_dims, recall_t, L, W, query_result_ids, query_result_dists) + qe = time.time() + latency_stats = float((qe - qs) * 1000000) + + query_result_ids = diskannpy.VectorUnsigned(query_result_ids) + 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) \ No newline at end of file From ae2071c2fc0523500a58e09dc300c73446b4ba16 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 29 Jun 2021 21:40:45 -0700 Subject: [PATCH 13/37] added build code to python bindings --- python/src/diskann_bindings.cpp | 96 +++++++++++++++++++++++++++++---- 1 file changed, 85 insertions(+), 11 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 255c65b15a..76f46ee599 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include #include #include @@ -36,9 +37,7 @@ PYBIND11_MODULE(diskannpy, m) { py::bind_vector>(m, "VectorUnsigned"); py::bind_vector>(m, "VectorFloat"); - py::enum_(m, "Metric") - .value("L2", Metric::L2) - .export_values(); + py::enum_(m, "Metric").value("L2", Metric::L2).export_values(); py::class_(m, "Parameters") .def(py::init<>()) @@ -56,7 +55,7 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("name"), py::arg("value")); py::class_(m, "Neighbor") - .def(py::init<>()) + .def(py::init<>()) .def(py::init()) .def(py::self < py::self) .def(py::self == py::self); @@ -68,7 +67,7 @@ PYBIND11_MODULE(diskannpy, m) { .def(py::self == py::self); py::class_(m, "AlignedFileReader"); - + py::class_(m, "LinuxAlignedFileReader") .def(py::init<>()); // .def("get_ctx", &LinuxAlignedFileReader::get_ctx) @@ -167,14 +166,11 @@ PYBIND11_MODULE(diskannpy, m) { 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(&FloatPQFlashIndexCreator)) .def( - "load", - [](PQFlashIndex &self, const std::string& index_path_prefix) { - + "load_index", + [](PQFlashIndex &self, const std::string &index_path_prefix) { const std::string pq_path = index_path_prefix + std::string("_pq"); const std::string index_path = index_path_prefix + std::string("_disk.index"); @@ -212,5 +208,83 @@ PYBIND11_MODULE(diskannpy, m) { }, py::arg("query_data"), py::arg("nqueries"), py::arg("dim"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), - py::arg("ids"), py::arg("dists")); + py::arg("ids"), py::arg("dists")) + .def( + "build", + [](PQFlashIndex & self, + const char *dataFilePath, const std::string &index_prefix_path, + unsigned R, unsigned L, double final_index_ram_limit, + double indexing_ram_budget, unsigned num_threads) { + 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"; + std::string disk_index_path = index_prefix_path + "_disk.index"; + std::string medoids_path = disk_index_path + "_medoids.bin"; + std::string centroids_path = disk_index_path + "_centroids.bin"; + std::string sample_base_prefix = index_prefix_path + "_sample"; + + if (num_threads != 0) { + omp_set_num_threads(num_threads); + mkl_set_num_threads(num_threads); + } + + cout << "Starting index build: R=" << R << " L=" << L + << " Query RAM budget: " << final_index_ram_limit + << " Indexing RAM budget: " << indexing_ram_budget + << " T: " << num_threads << std::endl; + + auto s = std::chrono::high_resolution_clock::now(); + + size_t points_num, dim; + + get_bin_metadata(dataFilePath, points_num, dim); + + size_t num_pq_chunks = + (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); + + num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; + num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; + num_pq_chunks = + num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; + + cout << "Compressing " << dim << "-dimensional data into " + << num_pq_chunks << " bytes per vector." << std::endl; + + size_t train_size, train_dim; + float *train_data; + + double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); + // generates random sample and sets it to train_data and updates + // train_size + gen_random_slice(dataFilePath, p_val, train_data, train_size, + train_dim); + + cout << "Training data loaded of size " << train_size << std::endl; + + generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, + (uint32_t) num_pq_chunks, 15, pq_pivots_path); + generate_pq_data_from_pivots( + dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, + pq_compressed_vectors_path); + + delete[] train_data; + + build_merged_vamana_index( + dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, + mem_index_path, medoids_path, centroids_path); + + create_disk_layout(dataFilePath, mem_index_path, + disk_index_path); + + double sample_sampling_rate = (150000.0 / points_num); + gen_random_slice(dataFilePath, sample_base_prefix, + sample_sampling_rate); + + std::remove(mem_index_path.c_str()); + + auto e = std::chrono::high_resolution_clock::now(); + std::chrono::duration diff = e - s; + cout << "Indexing time: " << diff.count() << std::endl; + }); } From f1ea75ec34d47b8ad6baa0548e7cf9c7b1e79791 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 6 Jul 2021 17:54:54 -0700 Subject: [PATCH 14/37] libaio link in pybind works --- python/setup.py | 22 +--- python/src/diskann_bindings.cpp | 162 ++++++++++++------------- python/tests/test_search_disk_index.py | 1 - 3 files changed, 85 insertions(+), 100 deletions(-) diff --git a/python/setup.py b/python/setup.py index f53807d6c7..592937e089 100644 --- a/python/setup.py +++ b/python/setup.py @@ -55,28 +55,18 @@ def build_extensions(self): ext_modules = [ - #Extension( - # 'vamanapy', - # ['src/vamana_bindings.cpp'], - # include_dirs=["../include/", - # "/opt/intel/compilers_and_libraries/linux/mkl/include/", - # pybind11.get_include(False), - # pybind11.get_include(True)], - # libraries=[], - # language='c++', - # extra_objects=['../build/src/libdiskann_s.a'], - #), - Extension( + Extension( 'diskannpy', ['src/diskann_bindings.cpp'], - include_dirs=["../include/", - "/opt/intel/compilers_and_libraries/linux/mkl/include/", + include_dirs=['../include/', + '/opt/intel/compilers_and_libraries/linux/mkl/include/', + '/usr/include', pybind11.get_include(False), pybind11.get_include(True)], - libraries=[], + libraries=['aio'], language='c++', extra_objects=['../build/src/libdiskann_s.a'], - ), + ) ] diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 76f46ee599..09518f55cc 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -2,7 +2,6 @@ // Licensed under the MIT license. #include -#include #include #include @@ -10,9 +9,6 @@ #include #include -#include "utils.h" -#include "memory_mapper.h" -#include "aligned_file_reader.h" #include "linux_aligned_file_reader.h" #include "pq_flash_index.h" @@ -208,83 +204,83 @@ PYBIND11_MODULE(diskannpy, m) { }, py::arg("query_data"), py::arg("nqueries"), py::arg("dim"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), - py::arg("ids"), py::arg("dists")) - .def( - "build", - [](PQFlashIndex & self, - const char *dataFilePath, const std::string &index_prefix_path, - unsigned R, unsigned L, double final_index_ram_limit, - double indexing_ram_budget, unsigned num_threads) { - 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"; - std::string disk_index_path = index_prefix_path + "_disk.index"; - std::string medoids_path = disk_index_path + "_medoids.bin"; - std::string centroids_path = disk_index_path + "_centroids.bin"; - std::string sample_base_prefix = index_prefix_path + "_sample"; - - if (num_threads != 0) { - omp_set_num_threads(num_threads); - mkl_set_num_threads(num_threads); - } - - cout << "Starting index build: R=" << R << " L=" << L - << " Query RAM budget: " << final_index_ram_limit - << " Indexing RAM budget: " << indexing_ram_budget - << " T: " << num_threads << std::endl; - - auto s = std::chrono::high_resolution_clock::now(); - - size_t points_num, dim; - - get_bin_metadata(dataFilePath, points_num, dim); - - size_t num_pq_chunks = - (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); - - num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; - num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - num_pq_chunks = - num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; - - cout << "Compressing " << dim << "-dimensional data into " - << num_pq_chunks << " bytes per vector." << std::endl; - - size_t train_size, train_dim; - float *train_data; - - double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); - // generates random sample and sets it to train_data and updates - // train_size - gen_random_slice(dataFilePath, p_val, train_data, train_size, - train_dim); - - cout << "Training data loaded of size " << train_size << std::endl; - - generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - (uint32_t) num_pq_chunks, 15, pq_pivots_path); - generate_pq_data_from_pivots( - dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, - pq_compressed_vectors_path); - - delete[] train_data; - - build_merged_vamana_index( - dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, - mem_index_path, medoids_path, centroids_path); - - create_disk_layout(dataFilePath, mem_index_path, - disk_index_path); - - double sample_sampling_rate = (150000.0 / points_num); - gen_random_slice(dataFilePath, sample_base_prefix, - sample_sampling_rate); - - std::remove(mem_index_path.c_str()); - - auto e = std::chrono::high_resolution_clock::now(); - std::chrono::duration diff = e - s; - cout << "Indexing time: " << diff.count() << std::endl; - }); + py::arg("ids"), py::arg("dists")); + //.def( + // "build", + // [](PQFlashIndex & self, + // const char *dataFilePath, const std::string &index_prefix_path, + // unsigned R, unsigned L, double final_index_ram_limit, + // double indexing_ram_budget, unsigned num_threads) { + // 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"; + // std::string disk_index_path = index_prefix_path + "_disk.index"; + // std::string medoids_path = disk_index_path + "_medoids.bin"; + // std::string centroids_path = disk_index_path + "_centroids.bin"; + // std::string sample_base_prefix = index_prefix_path + "_sample"; + + // if (num_threads != 0) { + // omp_set_num_threads(num_threads); + // mkl_set_num_threads(num_threads); + // } + + // cout << "Starting index build: R=" << R << " L=" << L + // << " Query RAM budget: " << final_index_ram_limit + // << " Indexing RAM budget: " << indexing_ram_budget + // << " T: " << num_threads << std::endl; + + // auto s = std::chrono::high_resolution_clock::now(); + + // size_t points_num, dim; + + // get_bin_metadata(dataFilePath, points_num, dim); + + // size_t num_pq_chunks = + // (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); + + // num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; + // num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; + // num_pq_chunks = + // num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; + + // cout << "Compressing " << dim << "-dimensional data into " + // << num_pq_chunks << " bytes per vector." << std::endl; + + // size_t train_size, train_dim; + // float *train_data; + + // double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); + // // generates random sample and sets it to train_data and updates + // // train_size + // gen_random_slice(dataFilePath, p_val, train_data, train_size, + // train_dim); + + // cout << "Training data loaded of size " << train_size << std::endl; + + // generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, + // (uint32_t) num_pq_chunks, 15, pq_pivots_path); + // generate_pq_data_from_pivots( + // dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, + // pq_compressed_vectors_path); + + // delete[] train_data; + + // build_merged_vamana_index( + // dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, + // mem_index_path, medoids_path, centroids_path); + + // create_disk_layout(dataFilePath, mem_index_path, + // disk_index_path); + + // double sample_sampling_rate = (150000.0 / points_num); + // gen_random_slice(dataFilePath, sample_base_prefix, + // sample_sampling_rate); + + // std::remove(mem_index_path.c_str()); + + // auto e = std::chrono::high_resolution_clock::now(); + // std::chrono::duration diff = e - s; + // cout << "Indexing time: " << diff.count() << std::endl; + // }); } diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index 8df64ff87e..8d56b418df 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -8,7 +8,6 @@ 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('index_path_prefix', type=str, help='Path prefix for index files.') From 3c650442ed38d454206dff581296ef1e12125730 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Sun, 11 Jul 2021 01:23:12 -0700 Subject: [PATCH 15/37] added build code to python bindings --- python/setup.py | 2 +- python/src/diskann_bindings.cpp | 112 +++++-------------------- python/tests/test_search_disk_index.py | 2 +- tests/utils/bin_to_tsv.cpp | 69 +++++++++++++++ 4 files changed, 94 insertions(+), 91 deletions(-) create mode 100644 tests/utils/bin_to_tsv.cpp diff --git a/python/setup.py b/python/setup.py index 592937e089..0a59ae9961 100644 --- a/python/setup.py +++ b/python/setup.py @@ -14,7 +14,7 @@ class BuildExt(build_ext): """A custom build extension for adding compiler-specific options.""" - c_opts = {'unix': ['-Ofast', '-DMKL_ILP64', '-m64', '-Wl,--no-as-needed']} + 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 diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 09518f55cc..1fc94d23a4 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include +#include #include #include @@ -10,6 +11,7 @@ #include #include "linux_aligned_file_reader.h" +#include "aux_utils.h" #include "pq_flash_index.h" @@ -19,13 +21,6 @@ PYBIND11_MAKE_OPAQUE(std::vector); namespace py = pybind11; using namespace diskann; -std::unique_ptr> FloatPQFlashIndexCreator() { - std::shared_ptr reader(new LinuxAlignedFileReader()); - auto index = new PQFlashIndex(reader); - std::unique_ptr> unique_ptr_index(index); - return unique_ptr_index; -} - PYBIND11_MODULE(diskannpy, m) { m.doc() = "DiskANN Python Bindings"; m.attr("__version__") = "0.1.0"; @@ -163,11 +158,15 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("file_name"), py::arg("data"), py::arg("npts"), py::arg("dims")); py::class_>(m, "DiskANNFloatIndex") - .def(py::init(&FloatPQFlashIndexCreator)) + .def(py::init([]() { + std::shared_ptr reader(new LinuxAlignedFileReader()); + auto index = new PQFlashIndex(reader); + return index; + })) .def( "load_index", [](PQFlashIndex &self, const std::string &index_path_prefix) { - const std::string pq_path = index_path_prefix + std::string("_pq"); + const std::string pq_path = index_path_prefix; const std::string index_path = index_path_prefix + std::string("_disk.index"); self.load(1, pq_path.c_str(), index_path.c_str()); @@ -204,83 +203,18 @@ PYBIND11_MODULE(diskannpy, m) { }, py::arg("query_data"), py::arg("nqueries"), py::arg("dim"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), - py::arg("ids"), py::arg("dists")); - //.def( - // "build", - // [](PQFlashIndex & self, - // const char *dataFilePath, const std::string &index_prefix_path, - // unsigned R, unsigned L, double final_index_ram_limit, - // double indexing_ram_budget, unsigned num_threads) { - // 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"; - // std::string disk_index_path = index_prefix_path + "_disk.index"; - // std::string medoids_path = disk_index_path + "_medoids.bin"; - // std::string centroids_path = disk_index_path + "_centroids.bin"; - // std::string sample_base_prefix = index_prefix_path + "_sample"; - - // if (num_threads != 0) { - // omp_set_num_threads(num_threads); - // mkl_set_num_threads(num_threads); - // } - - // cout << "Starting index build: R=" << R << " L=" << L - // << " Query RAM budget: " << final_index_ram_limit - // << " Indexing RAM budget: " << indexing_ram_budget - // << " T: " << num_threads << std::endl; - - // auto s = std::chrono::high_resolution_clock::now(); - - // size_t points_num, dim; - - // get_bin_metadata(dataFilePath, points_num, dim); - - // size_t num_pq_chunks = - // (size_t)(std::floor)(_u64(final_index_ram_limit / points_num)); - - // num_pq_chunks = num_pq_chunks <= 0 ? 1 : num_pq_chunks; - // num_pq_chunks = num_pq_chunks > dim ? dim : num_pq_chunks; - // num_pq_chunks = - // num_pq_chunks > MAX_PQ_CHUNKS ? MAX_PQ_CHUNKS : num_pq_chunks; - - // cout << "Compressing " << dim << "-dimensional data into " - // << num_pq_chunks << " bytes per vector." << std::endl; - - // size_t train_size, train_dim; - // float *train_data; - - // double p_val = ((double) TRAINING_SET_SIZE / (double) points_num); - // // generates random sample and sets it to train_data and updates - // // train_size - // gen_random_slice(dataFilePath, p_val, train_data, train_size, - // train_dim); - - // cout << "Training data loaded of size " << train_size << std::endl; - - // generate_pq_pivots(train_data, train_size, (uint32_t) dim, 256, - // (uint32_t) num_pq_chunks, 15, pq_pivots_path); - // generate_pq_data_from_pivots( - // dataFilePath, 256, (uint32_t) num_pq_chunks, pq_pivots_path, - // pq_compressed_vectors_path); - - // delete[] train_data; - - // build_merged_vamana_index( - // dataFilePath, _compareMetric, L, R, p_val, indexing_ram_budget, - // mem_index_path, medoids_path, centroids_path); - - // create_disk_layout(dataFilePath, mem_index_path, - // disk_index_path); - - // double sample_sampling_rate = (150000.0 / points_num); - // gen_random_slice(dataFilePath, sample_base_prefix, - // sample_sampling_rate); - - // std::remove(mem_index_path.c_str()); - - // auto e = std::chrono::high_resolution_clock::now(); - // std::chrono::duration diff = e - s; - // cout << "Indexing time: " << diff.count() << std::endl; - // }); -} + py::arg("ids"), py::arg("dists")) + .def( + "build", + [](PQFlashIndex &self, const char *dataFilePath, + 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(dataFilePath, index_prefix_path, + params.c_str(), diskann::Metric::L2); + }); +} \ No newline at end of file diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index 8d56b418df..8ded38e1fb 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -29,7 +29,7 @@ num_ground_truth, ground_truth_dims = diskannpy.load_truthset(args.ground_truth_path, ground_truth_ids, ground_truth_dists) index = diskannpy.DiskANNFloatIndex() -index.load(file_name = args.index_path_prefix) +index.load_index(index_path_prefix = args.index_path_prefix) print("Index Loaded") #index.optimize_graph() 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(); +} From 7c0cd60948824b0343b0bedf14f2c112f8c4970e Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Sun, 11 Jul 2021 17:38:47 -0700 Subject: [PATCH 16/37] created diskann wrapper class in pybind to avoid shared_ptr deallocation --- python/src/diskann_bindings.cpp | 41 ++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 1fc94d23a4..9c6b01597a 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -21,6 +22,23 @@ 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; + } +}; +#endif + PYBIND11_MODULE(diskannpy, m) { m.doc() = "DiskANN Python Bindings"; m.attr("__version__") = "0.1.0"; @@ -157,33 +175,30 @@ PYBIND11_MODULE(diskannpy, m) { 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([]() { - std::shared_ptr reader(new LinuxAlignedFileReader()); - auto index = new PQFlashIndex(reader); - return index; + py::class_>(m, "DiskANNFloatIndex") + .def(py::init([]() { return new DiskANNIndex(); })) .def( "load_index", - [](PQFlashIndex &self, const std::string &index_path_prefix) { + [](DiskANNIndex &self, const std::string &index_path_prefix) { const std::string pq_path = index_path_prefix; const std::string index_path = index_path_prefix + std::string("_disk.index"); - self.load(1, pq_path.c_str(), index_path.c_str()); + self.pq_flash_index->load(1, pq_path.c_str(), index_path.c_str()); std::vector node_list; _u64 num_nodes_to_cache = 100000; - self.cache_bfs_levels(num_nodes_to_cache, node_list); + self.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; }, py::arg("index_path_prefix")) .def( "search", - [](PQFlashIndex &self, const float *query, const _u64 dim, + [](DiskANNIndex &self, const float *query, const _u64 dim, const _u64 knn, const _u64 l_search, const _u64 beam_width, _u64 *ids, float *dists) { QueryStats stats; - self.cached_beam_search(query, knn, l_search, ids, dists, + self.pq_flash_index->cached_beam_search(query, knn, l_search, ids, dists, beam_width, &stats); }, py::arg("query"), py::arg("dim"), py::arg("knn") = 10, @@ -191,13 +206,13 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("dists")) .def( "batch_search", - [](PQFlashIndex &self, const float *query_data, + [](DiskANNIndex &self, const float *query_data, const _u64 nqueries, const _u64 dim, const _u64 knn, const _u64 l_search, const _u64 beam_width, _u64 *ids, float *dists) { #pragma omp parallel for schedule(dynamic, 1) for (_u64 i = 0; i < nqueries; ++i) - self.cached_beam_search(query_data + i * dim, knn, l_search, + self.pq_flash_index->cached_beam_search(query_data + i * dim, knn, l_search, ids + i * knn, dists + i * knn, beam_width); }, @@ -206,7 +221,7 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("ids"), py::arg("dists")) .def( "build", - [](PQFlashIndex &self, const char *dataFilePath, + [](DiskANNIndex &self, const char *dataFilePath, const char *index_prefix_path, unsigned R, unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads) { From 59eebe0382395aa68ec039dd9ec3efe39065b2ca Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Sun, 11 Jul 2021 18:13:39 -0700 Subject: [PATCH 17/37] bug fix in python test search disk index --- python/src/diskann_bindings.cpp | 2 +- python/tests/test_search_disk_index.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 9c6b01597a..a89066116f 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -186,7 +186,7 @@ PYBIND11_MODULE(diskannpy, m) { index_path_prefix + std::string("_disk.index"); self.pq_flash_index->load(1, pq_path.c_str(), index_path.c_str()); std::vector node_list; - _u64 num_nodes_to_cache = 100000; + _u64 num_nodes_to_cache = 1000; self.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; diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index 8ded38e1fb..d91814139c 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -46,7 +46,7 @@ for j in range(num_queries): qs = time.time() - index.search(query_data, query_aligned_dims, recall_t, L, W, + index.search(query_data, query_aligned_dims, recall_at, L, W, query_result_ids + j * recall_at, query_result_dists + j * recall_at) qe = time.time() latency_stats.append(float((qe - qs) * 1000000)) @@ -73,7 +73,7 @@ diskannpy.set_num_threads(num_threads) qs = time.time() - query_result_ids = index.batch_search(query_data, num_queries, query_aligned_dims, recall_t, L, W, query_result_ids, query_result_dists) + query_result_ids = index.batch_search(query_data, num_queries, query_aligned_dims, recall_at, L, W, query_result_ids, query_result_dists) qe = time.time() latency_stats = float((qe - qs) * 1000000) From 79d93f6a9f968bf190c3e9f039d3c018875fe1da Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Sun, 11 Jul 2021 18:22:38 -0700 Subject: [PATCH 18/37] bug fix in python test search disk index --- src/pq_flash_index.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 30b1f05a2a..36752d17a6 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -711,7 +711,7 @@ namespace diskann { use_medoids_data_as_centroids(); } - diskann::cout << "done.." << std::endl; + diskann::cout << "Index load complete." << std::endl; return 0; } From 97e9c468721016f16ab37e5a6de1bae02a73bd13 Mon Sep 17 00:00:00 2001 From: Suhas Jayaram Subramanya Date: Tue, 13 Jul 2021 01:00:57 +0000 Subject: [PATCH 19/37] working test search disk index pybind --- python/src/diskann_bindings.cpp | 21 ++++++++++++++------- python/tests/test_search_disk_index.py | 8 ++++---- src/pq_flash_index.cpp | 4 ++-- 3 files changed, 20 insertions(+), 13 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index a89066116f..a438c1c29a 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -194,14 +194,21 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("index_path_prefix")) .def( "search", - [](DiskANNIndex &self, const float *query, const _u64 dim, - const _u64 knn, const _u64 l_search, const _u64 beam_width, - _u64 *ids, float *dists) { + [](DiskANNIndex &self, 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; - self.pq_flash_index->cached_beam_search(query, knn, l_search, ids, dists, - beam_width, &stats); + if (ids.size() < knn * num_queries) { + ids.resize(knn * num_queries); + dists.resize(knn * num_queries); + } + std::vector<_u64> _u64_ids(knn); + self.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(dataFilePath, index_prefix_path, params.c_str(), diskann::Metric::L2); }); -} \ No newline at end of file +} diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index d91814139c..56d0bd0eaa 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -41,13 +41,13 @@ for i, L in enumerate(l_search): latency_stats = [] query_result_ids = diskannpy.VectorUnsigned() - query_result_dists = diskannpy.VectorUnsigned() + query_result_dists = diskannpy.VectorFloat() s = time.time() for j in range(num_queries): qs = time.time() - index.search(query_data, query_aligned_dims, recall_at, L, W, - query_result_ids + j * recall_at, query_result_dists + j * recall_at) + 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)) @@ -89,4 +89,4 @@ "{:>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) \ No newline at end of file + diskannpy.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp index 36752d17a6..d6fcd7c3ce 100644 --- a/src/pq_flash_index.cpp +++ b/src/pq_flash_index.cpp @@ -122,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(); } } From 03151f0c24311aa09ea008716790f23f80e70e13 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 29 Jul 2021 14:14:10 -0700 Subject: [PATCH 20/37] python interface test build disk index --- python/tests/test_build_disk_index.py | 25 +++++++++++++++++++++++++ python/tests/test_search_disk_index.py | 8 ++++++-- 2 files changed, 31 insertions(+), 2 deletions(-) create mode 100644 python/tests/test_build_disk_index.py 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_search_disk_index.py b/python/tests/test_search_disk_index.py index 56d0bd0eaa..c8d5c966a2 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -72,12 +72,16 @@ for i, L in enumerate(l_search): diskannpy.set_num_threads(num_threads) + query_result_ids = diskannpy.VectorUnsigned() + query_result_dists = diskannpy.VectorFloat() + qs = time.time() - query_result_ids = index.batch_search(query_data, num_queries, query_aligned_dims, recall_at, L, W, query_result_ids, query_result_dists) + query_result_ids = index.batch_search(query_data, num_queries, + query_aligned_dims, recall_at, L, W, + query_result_ids, query_result_dists) qe = time.time() latency_stats = float((qe - qs) * 1000000) - query_result_ids = diskannpy.VectorUnsigned(query_result_ids) qps = (num_queries / (qe - qs)) recall = diskannpy.calculate_recall(num_queries, ground_truth_ids, ground_truth_dists, ground_truth_dims, From cdcf508620897d9156a26d44f4d458e0fed31b08 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 29 Jul 2021 19:33:33 -0700 Subject: [PATCH 21/37] fixes to batch query mode --- python/src/diskann_bindings.cpp | 54 ++++++++++++++++---------- python/tests/test_search_disk_index.py | 12 ++++-- 2 files changed, 42 insertions(+), 24 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index a438c1c29a..da0529d7c2 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -194,36 +194,50 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("index_path_prefix")) .def( "search", - [](DiskANNIndex &self, 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, + [](DiskANNIndex &self, 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); - self.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 _u64_ids(knn); + self.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]; }, py::arg("query"), py::arg("query_idx"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), py::arg("dists")) .def( "batch_search", - [](DiskANNIndex &self, const float *query_data, - const _u64 nqueries, const _u64 dim, const _u64 knn, - const _u64 l_search, const _u64 beam_width, _u64 *ids, - float *dists) { + [](DiskANNIndex &self, 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) { + + if (ids.size() < knn * num_queries) { + ids.resize(knn * num_queries); + dists.resize(knn * num_queries); + } #pragma omp parallel for schedule(dynamic, 1) - for (_u64 i = 0; i < nqueries; ++i) - self.pq_flash_index->cached_beam_search(query_data + i * dim, knn, l_search, - ids + i * knn, dists + i * knn, - beam_width); + for (_u64 q = 0; q < num_queries; ++q) { + std::vector<_u64> u64_ids(knn); + + self.pq_flash_index->cached_beam_search( + queries.data() + q * dim, knn, l_search, u64_ids.data() + q * knn, + dists.data() + q * knn, beam_width); + for (_u64 i = 0; i < knn; i++) + ids[(q * knn) + i] = u64_ids[i]; + } }, - py::arg("query_data"), py::arg("nqueries"), py::arg("dim"), + py::arg("queries"), py::arg("nqueries"), py::arg("dim"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), py::arg("dists")) .def( diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index c8d5c966a2..7c4b8f2440 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -17,7 +17,7 @@ recall_at = 10 # Use multi-threaded search only for batch mode. num_threads = 1 -single_query_mode = True +single_query_mode = False l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] W = 4 @@ -42,6 +42,10 @@ latency_stats = [] query_result_ids = diskannpy.VectorUnsigned() query_result_dists = diskannpy.VectorFloat() + + ids.resize(knn * num_queries) + dists.resize(knn * num_queries) + s = time.time() for j in range(num_queries): @@ -76,9 +80,9 @@ query_result_dists = diskannpy.VectorFloat() qs = time.time() - query_result_ids = index.batch_search(query_data, num_queries, - query_aligned_dims, recall_at, L, W, - query_result_ids, query_result_dists) + index.batch_search(query_data, query_aligned_dims, num_queries, + recall_at, L, W, + query_result_ids, query_result_dists) qe = time.time() latency_stats = float((qe - qs) * 1000000) From 59a8467d2b0574d774d60fde266f760d1be4fd80 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 30 Jul 2021 00:29:36 -0700 Subject: [PATCH 22/37] batch search works, but with 1 thread --- python/src/diskann_bindings.cpp | 16 ++++++++++++---- python/tests/test_search_disk_index.py | 8 +++----- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index da0529d7c2..b57998d445 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include "linux_aligned_file_reader.h" #include "aux_utils.h" @@ -220,26 +221,33 @@ PYBIND11_MODULE(diskannpy, m) { [](DiskANNIndex &self, 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) { + 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); } + + py::gil_scoped_release release; + + 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); self.pq_flash_index->cached_beam_search( - queries.data() + q * dim, knn, l_search, u64_ids.data() + q * knn, + 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]; } + + py::gil_scoped_acquire acquire; }, - py::arg("queries"), py::arg("nqueries"), py::arg("dim"), + py::arg("queries"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), - py::arg("ids"), py::arg("dists")) + py::arg("ids"), py::arg("dists"), py::arg("num_threads")) .def( "build", [](DiskANNIndex &self, const char *dataFilePath, diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index 7c4b8f2440..516c3632ac 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -16,7 +16,7 @@ recall_at = 10 # Use multi-threaded search only for batch mode. -num_threads = 1 +num_threads = 16 single_query_mode = False l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] W = 4 @@ -43,9 +43,6 @@ query_result_ids = diskannpy.VectorUnsigned() query_result_dists = diskannpy.VectorFloat() - ids.resize(knn * num_queries) - dists.resize(knn * num_queries) - s = time.time() for j in range(num_queries): @@ -82,7 +79,8 @@ qs = time.time() index.batch_search(query_data, query_aligned_dims, num_queries, recall_at, L, W, - query_result_ids, query_result_dists) + query_result_ids, query_result_dists, + num_threads) qe = time.time() latency_stats = float((qe - qs) * 1000000) From 6a3011b5f3a91ae0ae960c5628314467b787f348 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 30 Jul 2021 11:32:13 -0700 Subject: [PATCH 23/37] added num_thread arg to load function in diskann bindings --- python/src/diskann_bindings.cpp | 26 ++++++++++++-------------- python/tests/test_search_disk_index.py | 13 +++++++++---- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index b57998d445..a199dc21e9 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -181,18 +181,19 @@ PYBIND11_MODULE(diskannpy, m) { })) .def( "load_index", - [](DiskANNIndex &self, const std::string &index_path_prefix) { + [](DiskANNIndex &self, 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"); - self.pq_flash_index->load(1, pq_path.c_str(), index_path.c_str()); + self.pq_flash_index->load(num_threads, pq_path.c_str(), index_path.c_str()); std::vector node_list; _u64 num_nodes_to_cache = 1000; self.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; }, - py::arg("index_path_prefix")) + py::arg("index_path_prefix"), py::arg("num_threads")) .def( "search", [](DiskANNIndex &self, std::vector &query, @@ -220,18 +221,17 @@ PYBIND11_MODULE(diskannpy, m) { "batch_search", [](DiskANNIndex &self, 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) { + 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); } - py::gil_scoped_release release; + py::gil_scoped_release release; - omp_set_num_threads(num_threads); + 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); @@ -243,14 +243,12 @@ PYBIND11_MODULE(diskannpy, m) { ids[(q * knn) + i] = u64_ids[i]; } - py::gil_scoped_acquire acquire; + py::gil_scoped_acquire acquire; }, py::arg("queries"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), py::arg("dists"), py::arg("num_threads")) - .def( - "build", - [](DiskANNIndex &self, const char *dataFilePath, + .def("build", [](DiskANNIndex &self, const char *dataFilePath, const char *index_prefix_path, unsigned R, unsigned L, double final_index_ram_limit, double indexing_ram_budget, unsigned num_threads) { diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index 516c3632ac..a01a5dd87c 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -12,14 +12,19 @@ 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 = 10 +recall_at = args.K +W = args.W # Use multi-threaded search only for batch mode. -num_threads = 16 +num_threads = args.T single_query_mode = False l_search = [40, 50, 60, 70, 80, 90, 100, 110, 120] -W = 4 + query_data = diskannpy.VectorFloat() ground_truth_ids = diskannpy.VectorUnsigned() @@ -29,7 +34,7 @@ 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(index_path_prefix = args.index_path_prefix) +index.load_index(index_path_prefix = args.index_path_prefix, num_threads) print("Index Loaded") #index.optimize_graph() From b76581fea1073c5094b2f0ebe7c158f02dc43849 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 30 Jul 2021 11:34:37 -0700 Subject: [PATCH 24/37] added num_thread arg to load function in diskann bindings --- python/tests/test_search_disk_index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index a01a5dd87c..b1957ce25f 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -34,7 +34,7 @@ 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(index_path_prefix = args.index_path_prefix, num_threads) +index.load_index(args.index_path_prefix, num_threads) print("Index Loaded") #index.optimize_graph() From ea2376bf06c9073ff4069c4297e32df3d99bb8e8 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 30 Jul 2021 12:43:07 -0700 Subject: [PATCH 25/37] remove gil lock/unlock --- python/src/diskann_bindings.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index a199dc21e9..79e8e2abe2 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -10,7 +10,6 @@ #include #include #include -#include #include "linux_aligned_file_reader.h" #include "aux_utils.h" @@ -228,9 +227,6 @@ PYBIND11_MODULE(diskannpy, m) { ids.resize(knn * num_queries); dists.resize(knn * num_queries); } - - py::gil_scoped_release release; - omp_set_num_threads(num_threads); #pragma omp parallel for schedule(dynamic, 1) for (_u64 q = 0; q < num_queries; ++q) { @@ -242,8 +238,6 @@ PYBIND11_MODULE(diskannpy, m) { for (_u64 i = 0; i < knn; i++) ids[(q * knn) + i] = u64_ids[i]; } - - py::gil_scoped_acquire acquire; }, py::arg("queries"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), From f7fe1d94f4492600b6e4543ea35367a48b6b40fb Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Mon, 2 Aug 2021 19:21:00 -0700 Subject: [PATCH 26/37] add omp get max threads to python wrapper --- python/src/diskann_bindings.cpp | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 79e8e2abe2..90c19cf87b 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -79,17 +79,14 @@ PYBIND11_MODULE(diskannpy, m) { py::class_(m, "LinuxAlignedFileReader") .def(py::init<>()); - // .def("get_ctx", &LinuxAlignedFileReader::get_ctx) - // .def("register_thread", &LinuxAlignedFileReader::register_thread) - // .def("open", &LinuxAlignedFileReader::open) - // .def("close", &LinuxAlignedFileReader::close) - // .def("read", &LinuxAlignedFileReader::read); m.def( - "set_num_threads", + "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) { From bfd3f085b63b258017158af71a8b9a87a6de006d Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 3 Aug 2021 10:30:20 -0700 Subject: [PATCH 27/37] added param names to build function in python wrapper --- python/src/diskann_bindings.cpp | 28 +++++++++++++++++----------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 90c19cf87b..e4c29b3ca9 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -239,15 +239,21 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("queries"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), py::arg("dists"), py::arg("num_threads")) - .def("build", [](DiskANNIndex &self, const char *dataFilePath, - 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(dataFilePath, index_prefix_path, - params.c_str(), diskann::Metric::L2); - }); + .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") ); } From 3f045a6e97189ec4dfcbf0998bdc205b5d06bbda Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 3 Aug 2021 15:47:05 -0700 Subject: [PATCH 28/37] added return value for index_load success --- python/src/diskann_bindings.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index e4c29b3ca9..78b0d62ecc 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -182,12 +182,18 @@ PYBIND11_MODULE(diskannpy, m) { const std::string pq_path = index_path_prefix; const std::string index_path = index_path_prefix + std::string("_disk.index"); - self.pq_flash_index->load(num_threads, pq_path.c_str(), index_path.c_str()); + int load_success = + self.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; self.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; }, py::arg("index_path_prefix"), py::arg("num_threads")) .def( From dd75524e527ede77d667cf78b93039ce645e3241 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Wed, 4 Aug 2021 23:05:17 -0700 Subject: [PATCH 29/37] search and batch search API with py::array query --- python/src/diskann_bindings.cpp | 85 +++++++++++++++++++++++++++------ 1 file changed, 71 insertions(+), 14 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 78b0d62ecc..3e42258da5 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -173,24 +173,24 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("file_name"), py::arg("data"), py::arg("npts"), py::arg("dims")); py::class_>(m, "DiskANNFloatIndex") - .def(py::init([]() { return new DiskANNIndex(); - })) + .def(py::init([]() { return new DiskANNIndex(); })) .def( "load_index", [](DiskANNIndex &self, const std::string &index_path_prefix, - const int num_threads) { + 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 = - self.pq_flash_index->load(num_threads, pq_path.c_str(), index_path.c_str()); + int load_success = self.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; - self.pq_flash_index->cache_bfs_levels(num_nodes_to_cache, node_list); + self.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; @@ -202,7 +202,6 @@ PYBIND11_MODULE(diskannpy, m) { 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); @@ -216,9 +215,9 @@ PYBIND11_MODULE(diskannpy, m) { for (_u64 i = 0; i < knn; i++) ids[(query_idx * knn) + i] = _u64_ids[i]; }, - py::arg("query"), py::arg("query_idx"), py::arg("dim"), py::arg("num_queries"), py::arg("knn") = 10, - py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), - py::arg("dists")) + 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 &self, std::vector &queries, @@ -243,7 +242,7 @@ PYBIND11_MODULE(diskannpy, m) { } }, py::arg("queries"), py::arg("dim"), py::arg("num_queries"), - py::arg("knn") = 10, py::arg("l_search"), py::arg("beam_width"), + py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), py::arg("ids"), py::arg("dists"), py::arg("num_threads")) .def( "build", @@ -259,7 +258,65 @@ PYBIND11_MODULE(diskannpy, m) { 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::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")) + .def( + "pq_single_numpy_query", + [](DiskANNIndex &self, + py::array_t + & query, + const _u64 dim, const _u64 knn, const _u64 l_search, + const _u64 beam_width) { + py::array_t ids(knn); + + std::vector u32_ids(knn); + std::vector<_u64> u64_ids(knn); + std::vector dists(knn); + QueryStats stats; + + self.pq_flash_index->cached_beam_search( + query.mutable_data(), knn, l_search, u64_ids.data(), + dists.data(), beam_width, &stats); + + auto r = ids.mutable_unchecked<1>(); + for (_u64 i = 0; i < knn; ++i) + r(i) = (unsigned) u64_ids[i]; + + return ids; + }, + py::arg("query"), py::arg("dim"), py::arg("knn"), py::arg("l_search"), + py::arg("beam_width")) + .def( + "pq_batch_numpy_query", + [](DiskANNIndex &self, + 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(knn * num_queries); + + std::vector u32_ids(knn * num_queries); + std::vector<_u64> u64_ids(knn * num_queries); + std::vector dists(knn * num_queries); + QueryStats stats; + +#pragma omp parallel for schedule(dynamic, 1) + for (_u64 i = 0; i < num_queries; i++) { + self.pq_flash_index->cached_beam_search( + queries.mutable_data(i), knn, l_search, + u64_ids.data() + i * knn, dists.data() + i * knn, beam_width, + &stats); + } + + auto r = ids.mutable_unchecked<1>(); + for (_u64 i = 0; i < knn * num_queries; ++i) + r(i) = (unsigned) u64_ids[i]; + + return ids; + }, + 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")); } From b232eaf5882c954bde64959419d9fcfb582fa1db Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 5 Aug 2021 01:12:52 -0700 Subject: [PATCH 30/37] added a test file for numpy interface search --- python/src/diskann_bindings.cpp | 18 ++--- python/tests/test_search_disk_index_numpy.py | 76 ++++++++++++++++++++ 2 files changed, 85 insertions(+), 9 deletions(-) create mode 100644 python/tests/test_search_disk_index_numpy.py diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 3e42258da5..328b282a87 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -262,33 +262,33 @@ PYBIND11_MODULE(diskannpy, m) { py::arg("L"), py::arg("final_index_ram_limit"), py::arg("indexing_ram_limit"), py::arg("num_threads")) .def( - "pq_single_numpy_query", + "search_numpy_input", [](DiskANNIndex &self, 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); - std::vector dists(knn); QueryStats stats; self.pq_flash_index->cached_beam_search( query.mutable_data(), knn, l_search, u64_ids.data(), - dists.data(), beam_width, &stats); + 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 ids; + return std::make_pair(ids, dists); }, py::arg("query"), py::arg("dim"), py::arg("knn"), py::arg("l_search"), py::arg("beam_width")) .def( - "pq_batch_numpy_query", + "batch_search_numpy_input", [](DiskANNIndex &self, py::array_t & queries, @@ -296,25 +296,25 @@ PYBIND11_MODULE(diskannpy, m) { const _u64 l_search, const _u64 beam_width, const int num_threads) { py::array_t ids(knn * num_queries); + py::array_t dists(knn * num_queries); std::vector u32_ids(knn * num_queries); std::vector<_u64> u64_ids(knn * num_queries); - std::vector dists(knn * num_queries); QueryStats stats; #pragma omp parallel for schedule(dynamic, 1) for (_u64 i = 0; i < num_queries; i++) { self.pq_flash_index->cached_beam_search( queries.mutable_data(i), knn, l_search, - u64_ids.data() + i * knn, dists.data() + i * knn, beam_width, - &stats); + u64_ids.data() + i * knn, dists.mutable_data(i * knn), + beam_width, &stats); } auto r = ids.mutable_unchecked<1>(); for (_u64 i = 0; i < knn * num_queries; ++i) r(i) = (unsigned) u64_ids[i]; - return ids; + return std::make_pair(ids, dists); }, py::arg("queries"), py::arg("dim"), py::arg("num_queries"), py::arg("knn"), py::arg("l_search"), py::arg("beam_width"), 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..0c455890ea --- /dev/null +++ b/python/tests/test_search_disk_index_numpy.py @@ -0,0 +1,76 @@ +# 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) + +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") + +#index.optimize_graph() +#print("Graph Optimization Completed") +if single_query_mode: + pass +else: + print("Ls QPS Mean Latency (mus) Recall@10") + print("=============================================") + for i, L in enumerate(l_search): + diskannpy.set_num_threads(num_threads) + + qs = time.time() + ids, dists = index.batch_search_numpy_input(query_data, 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)) + + query_result_ids = diskannpy.VectorUnsigned(num_queries * query_aligned_dims) + query_result_dists = diskannpy.VectorFloat(num_queries * query_aligned_dims) + for q in range(0,num_queries): + for r in range(0,knn): + query_result_ids[q*num_queries + r] = ids[q*num_queries + r] + query_result_dists[q*num_queries + r] = dists[q*num_queries + r] + + 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) From 0440668535a3c965f99bea53112657831e3b3263 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 5 Aug 2021 01:48:08 -0700 Subject: [PATCH 31/37] calculate recall with numpy results input --- python/src/diskann_bindings.cpp | 53 ++++++++++++++++++-- python/tests/test_search_disk_index.py | 2 +- python/tests/test_search_disk_index_numpy.py | 14 ++---- 3 files changed, 52 insertions(+), 17 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 328b282a87..fa6eaa9c34 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -15,7 +15,6 @@ #include "aux_utils.h" #include "pq_flash_index.h" - PYBIND11_MAKE_OPAQUE(std::vector); PYBIND11_MAKE_OPAQUE(std::vector); @@ -25,10 +24,10 @@ using namespace diskann; #ifdef __linux__ template struct DiskANNIndex { - PQFlashIndex* pq_flash_index; + PQFlashIndex * pq_flash_index; std::shared_ptr reader; - - DiskANNIndex(){ + + DiskANNIndex() { reader = std::make_shared(); pq_flash_index = new PQFlashIndex(reader); } @@ -166,6 +165,50 @@ PYBIND11_MODULE(diskannpy, m) { 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, @@ -269,7 +312,7 @@ PYBIND11_MODULE(diskannpy, m) { const _u64 dim, const _u64 knn, const _u64 l_search, const _u64 beam_width) { py::array_t ids(knn); - py::array_t dists(knn); + py::array_t dists(knn); std::vector u32_ids(knn); std::vector<_u64> u64_ids(knn); diff --git a/python/tests/test_search_disk_index.py b/python/tests/test_search_disk_index.py index b1957ce25f..5d44a83476 100644 --- a/python/tests/test_search_disk_index.py +++ b/python/tests/test_search_disk_index.py @@ -76,7 +76,7 @@ print("Ls QPS Mean Latency (mus) Recall@10") print("=============================================") for i, L in enumerate(l_search): - diskannpy.set_num_threads(num_threads) + diskannpy.omp_set_num_threads(num_threads) query_result_ids = diskannpy.VectorUnsigned() query_result_dists = diskannpy.VectorFloat() diff --git a/python/tests/test_search_disk_index_numpy.py b/python/tests/test_search_disk_index_numpy.py index 0c455890ea..65980a6b04 100644 --- a/python/tests/test_search_disk_index_numpy.py +++ b/python/tests/test_search_disk_index_numpy.py @@ -49,7 +49,7 @@ print("Ls QPS Mean Latency (mus) Recall@10") print("=============================================") for i, L in enumerate(l_search): - diskannpy.set_num_threads(num_threads) + diskannpy.omp_set_num_threads(num_threads) qs = time.time() ids, dists = index.batch_search_numpy_input(query_data, query_aligned_dims, @@ -58,17 +58,9 @@ latency_stats = float((qe - qs) * 1000000) qps = (num_queries / (qe - qs)) - query_result_ids = diskannpy.VectorUnsigned(num_queries * query_aligned_dims) - query_result_dists = diskannpy.VectorFloat(num_queries * query_aligned_dims) - for q in range(0,num_queries): - for r in range(0,knn): - query_result_ids[q*num_queries + r] = ids[q*num_queries + r] - query_result_dists[q*num_queries + r] = dists[q*num_queries + r] - - recall = diskannpy.calculate_recall(num_queries, ground_truth_ids, + recall = diskannpy.calculate_recall_numpy(num_queries, ground_truth_ids, ground_truth_dists, ground_truth_dims, - query_result_ids, recall_at, - recall_at) + 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))) From 29c3caf286a79611a317903c38a0f80895ddc84e Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 5 Aug 2021 01:50:18 -0700 Subject: [PATCH 32/37] calculate recall with numpy results input --- python/tests/test_search_disk_index_numpy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tests/test_search_disk_index_numpy.py b/python/tests/test_search_disk_index_numpy.py index 65980a6b04..03b60c3268 100644 --- a/python/tests/test_search_disk_index_numpy.py +++ b/python/tests/test_search_disk_index_numpy.py @@ -15,6 +15,7 @@ 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 @@ -58,7 +59,7 @@ latency_stats = float((qe - qs) * 1000000) qps = (num_queries / (qe - qs)) - recall = diskannpy.calculate_recall_numpy(num_queries, ground_truth_ids, + 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 From 80374e3fabd0e82a6784183bca6ef53ec2008a71 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 5 Aug 2021 16:27:34 -0700 Subject: [PATCH 33/37] tested numpy interface to batch query --- python/src/diskann_bindings.cpp | 6 +-- python/tests/test_search_disk_index_numpy.py | 43 ++++++++------------ 2 files changed, 19 insertions(+), 30 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index fa6eaa9c34..f1db217fa2 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -341,21 +341,19 @@ PYBIND11_MODULE(diskannpy, m) { py::array_t ids(knn * num_queries); py::array_t dists(knn * num_queries); - std::vector u32_ids(knn * num_queries); std::vector<_u64> u64_ids(knn * num_queries); - QueryStats stats; #pragma omp parallel for schedule(dynamic, 1) for (_u64 i = 0; i < num_queries; i++) { self.pq_flash_index->cached_beam_search( queries.mutable_data(i), knn, l_search, u64_ids.data() + i * knn, dists.mutable_data(i * knn), - beam_width, &stats); + beam_width); } auto r = ids.mutable_unchecked<1>(); for (_u64 i = 0; i < knn * num_queries; ++i) - r(i) = (unsigned) u64_ids[i]; + r(i) = (unsigned) u64_ids[i]; return std::make_pair(ids, dists); }, diff --git a/python/tests/test_search_disk_index_numpy.py b/python/tests/test_search_disk_index_numpy.py index 03b60c3268..196541e713 100644 --- a/python/tests/test_search_disk_index_numpy.py +++ b/python/tests/test_search_disk_index_numpy.py @@ -10,7 +10,6 @@ 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.') @@ -22,7 +21,6 @@ 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] @@ -37,33 +35,26 @@ 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") -#index.optimize_graph() -#print("Graph Optimization Completed") -if single_query_mode: - pass -else: - print("Ls QPS Mean Latency (mus) 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, 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(mean_latency)) + "{:>15}".format("{:.2f}".format(recall))) +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)) - result_path = args.output_path_prefix + "_" + str(L) + "_idx_uint32.bin" - diskannpy.save_bin_u32(result_path, query_result_ids, num_queries, recall_at) + 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))) From 534c6f6b69ccd2a0ecd99a97ff54dacbc83c4dc5 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Thu, 5 Aug 2021 19:34:39 -0700 Subject: [PATCH 34/37] query np array accessed via non-mutable interface --- python/src/diskann_bindings.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index f1db217fa2..655b43a8c5 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -319,7 +319,7 @@ PYBIND11_MODULE(diskannpy, m) { QueryStats stats; self.pq_flash_index->cached_beam_search( - query.mutable_data(), knn, l_search, u64_ids.data(), + query.data(), knn, l_search, u64_ids.data(), dists.mutable_data(), beam_width, &stats); auto r = ids.mutable_unchecked<1>(); @@ -346,7 +346,7 @@ PYBIND11_MODULE(diskannpy, m) { #pragma omp parallel for schedule(dynamic, 1) for (_u64 i = 0; i < num_queries; i++) { self.pq_flash_index->cached_beam_search( - queries.mutable_data(i), knn, l_search, + queries.data(i), knn, l_search, u64_ids.data() + i * knn, dists.mutable_data(i * knn), beam_width); } From 8b8b31a32634799fdb6e5ef23950ae08c591b617 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Fri, 6 Aug 2021 12:46:18 -0700 Subject: [PATCH 35/37] reshape numpy interface output to {nq,knn} --- python/src/diskann_bindings.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 655b43a8c5..8e9f1d232c 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -338,8 +338,8 @@ PYBIND11_MODULE(diskannpy, m) { 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(knn * num_queries); - py::array_t dists(knn * num_queries); + py::array_t ids({num_queries, knn}); + py::array_t dists({num_queries, knn}); std::vector<_u64> u64_ids(knn * num_queries); @@ -347,13 +347,14 @@ PYBIND11_MODULE(diskannpy, m) { for (_u64 i = 0; i < num_queries; i++) { self.pq_flash_index->cached_beam_search( queries.data(i), knn, l_search, - u64_ids.data() + i * knn, dists.mutable_data(i * knn), + u64_ids.data() + i * knn, dists.mutable_data(i), beam_width); } - auto r = ids.mutable_unchecked<1>(); - for (_u64 i = 0; i < knn * num_queries; ++i) - r(i) = (unsigned) u64_ids[i]; + 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); }, From cc63fe9ebffb6f409cab54f12a3bd56299d08d75 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Mon, 9 Aug 2021 18:50:08 -0700 Subject: [PATCH 36/37] moved diskann bindings methods to a teampled c++ class --- python/src/diskann_bindings.cpp | 251 ++++++++++++++++---------------- 1 file changed, 122 insertions(+), 129 deletions(-) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 8e9f1d232c..1676f118b0 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -35,12 +35,115 @@ struct DiskANNIndex { ~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.0"; + m.attr("__version__") = "0.1.1"; py::bind_vector>(m, "VectorUnsigned"); py::bind_vector>(m, "VectorFloat"); @@ -217,76 +320,23 @@ PYBIND11_MODULE(diskannpy, m) { py::class_>(m, "DiskANNFloatIndex") .def(py::init([]() { return new DiskANNIndex(); })) - .def( - "load_index", - [](DiskANNIndex &self, 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 = self.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; - self.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; - }, - py::arg("index_path_prefix"), py::arg("num_threads")) - .def( - "search", - [](DiskANNIndex &self, 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); - self.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]; - }, - 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 &self, 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); - - self.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]; - } - }, - 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("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, @@ -303,62 +353,5 @@ PYBIND11_MODULE(diskannpy, m) { }, 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")) - .def( - "search_numpy_input", - [](DiskANNIndex &self, - 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; - - self.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); - }, - py::arg("query"), py::arg("dim"), py::arg("knn"), py::arg("l_search"), - py::arg("beam_width")) - .def( - "batch_search_numpy_input", - [](DiskANNIndex &self, - 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++) { - self.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); - }, - 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")); + py::arg("indexing_ram_limit"), py::arg("num_threads")); } From 2a83592b9acfbb4a7b7a1022bfd920eb9869b750 Mon Sep 17 00:00:00 2001 From: Harsha Vardhan Simhadri Date: Tue, 10 Aug 2021 10:38:07 -0700 Subject: [PATCH 37/37] added int8 and unit8 index interface --- python/src/diskann_bindings.cpp | 83 +++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/python/src/diskann_bindings.cpp b/python/src/diskann_bindings.cpp index 1676f118b0..3147b79130 100644 --- a/python/src/diskann_bindings.cpp +++ b/python/src/diskann_bindings.cpp @@ -17,6 +17,9 @@ 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; @@ -147,6 +150,9 @@ PYBIND11_MODULE(diskannpy, m) { 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(); @@ -354,4 +360,81 @@ PYBIND11_MODULE(diskannpy, m) { 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")); }