├── evaluation.png
├── output
    └── mnist
    │   └── sn.bin
├── src
    ├── CMakeLists.txt
    ├── index.cpp
    ├── index_random.cpp
    ├── util.cpp
    └── index_mips.cpp
├── test
    ├── CMakeLists.txt
    ├── test_mips_index.cpp
    └── test_mips_search.cpp
├── include
    ├── ssg
    │   ├── util.h
    │   ├── index_random.h
    │   ├── index.h
    │   ├── parameters.h
    │   ├── index_ssg.h
    │   ├── neighbor.h
    │   └── distance.h
    └── mips
    │   └── index_mips.h
├── scripts
    └── run.sh
├── CMakeLists.txt
├── .gitignore
└── README.md


/evaluation.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ZJU-DAILY/PSP/HEAD/evaluation.png


--------------------------------------------------------------------------------
/output/mnist/sn.bin:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ZJU-DAILY/PSP/HEAD/output/mnist/sn.bin


--------------------------------------------------------------------------------
/src/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | set(CMAKE_CXX_STANDARD 17)
 2 | 
 3 | list(
 4 |     APPEND CPP_SOURCES
 5 |     index.cpp
 6 |     index_random.cpp
 7 |     index_mips.cpp
 8 |     util.cpp
 9 | )
10 | 
11 | add_library(${PROJECT_NAME} ${CPP_SOURCES})
12 | add_library(${PROJECT_NAME}_s STATIC ${CPP_SOURCES})


--------------------------------------------------------------------------------
/test/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | set(CMAKE_CXX_STANDARD 17)
 2 | 
 3 | add_executable(test_mips_index test_mips_index.cpp)
 4 | target_link_libraries(test_mips_index ${PROJECT_NAME})
 5 | 
 6 | add_executable(test_mips_search test_mips_search.cpp)
 7 | target_link_libraries(test_mips_search ${PROJECT_NAME})
 8 | 
 9 | 
10 | 
11 | 
12 | 
13 | 


--------------------------------------------------------------------------------
/include/ssg/util.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <random>
 4 | 
 5 | namespace efanna2e {
 6 | 
 7 | void GenRandom(std::mt19937& rng, unsigned* addr, unsigned size, unsigned N);
 8 | 
 9 | float* load_data(const char* filename, unsigned& num, unsigned& dim);
10 | 
11 | float* data_align(float* data_ori, unsigned point_num, unsigned& dim);
12 | 
13 | }  // namespace efanna2e
14 | 
15 | 


--------------------------------------------------------------------------------
/src/index.cpp:
--------------------------------------------------------------------------------
 1 | #include <ssg/index.h>
 2 | namespace efanna2e {
 3 | Index::Index(const size_t dimension, const size_t n, Metric metric = L2)
 4 |     : dimension_(dimension), nd_(n), has_built(false) {
 5 |   switch (metric) {
 6 |     case L2:
 7 |       distance_ = new DistanceL2();
 8 |       break;
 9 |     default:
10 |       distance_ = new DistanceL2();
11 |       break;
12 |   }
13 | }
14 | Index::~Index() {}
15 | } 


--------------------------------------------------------------------------------
/scripts/run.sh:
--------------------------------------------------------------------------------
 1 | #! /bin/bash
 2 | 
 3 | mkdir build
 4 | cd build
 5 | cmake ..
 6 | make
 7 | 
 8 | if [ $? -eq 0 ]; then
 9 |   echo "Build successful"
10 |   # ./test/test_mips_index "../datasets/mnist/database_mnist.bin" "../output/mnist/database_mnist.knng" 512 40 60 5 "../output/mnist/mnist.mips" 784
11 |   # ./test/test_mips_search "../datasets/mnist/database_mnist.bin" "../datasets/mnist/query_mnist.bin" "../output/mnist/mnist.mips" 150 100 "../output/mnist/result.txt" 784
12 | else
13 |   echo "Build failed"
14 | fi


--------------------------------------------------------------------------------
/src/index_random.cpp:
--------------------------------------------------------------------------------
 1 | #include <ssg/index_random.h>
 2 | 
 3 | namespace efanna2e {
 4 | 
 5 | IndexRandom::IndexRandom(const size_t dimension, const size_t n)
 6 |     : Index(dimension, n, L2) {
 7 |   has_built = true;
 8 | }
 9 | IndexRandom::~IndexRandom() {}
10 | void IndexRandom::Build(size_t n, const float *data,
11 |                         const Parameters &parameters) {
12 |   data_ = data;
13 |   nd_ = n;
14 | 
15 |   // Do Nothing
16 | 
17 |   has_built = true;
18 | }
19 | void IndexRandom::Search(const float *query, const float *x, size_t k,
20 |                          const Parameters &parameters, unsigned *indices) {
21 |   GenRandom(rng, indices, k, nd_);
22 | }
23 | 
24 | } 


--------------------------------------------------------------------------------
/include/ssg/index_random.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include "index.h"
 4 | #include "util.h"
 5 | 
 6 | namespace efanna2e {
 7 | 
 8 | class IndexRandom : public Index {
 9 |  public:
10 |   IndexRandom(const size_t dimension, const size_t n);
11 |   virtual ~IndexRandom();
12 |   std::mt19937 rng;
13 |   void Save(const char *filename) override {}
14 |   void Load(const char *filename) override {}
15 |   virtual void Build(size_t n, const float *data,
16 |                      const Parameters &parameters) override;
17 | 
18 |   virtual void Search(const float *query, const float *x, size_t k,
19 |                       const Parameters &parameters, unsigned *indices) override;
20 | };
21 | 
22 | }  // namespace efanna2e


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | cmake_minimum_required(VERSION 3.21)
 2 | project(MIPS)
 3 | 
 4 | set(CMAKE_CXX_STANDARD 17)
 5 | 
 6 | include_directories(${PROJECT_SOURCE_DIR}/include)
 7 | 
 8 | 
 9 | find_package(Boost REQUIRED)
10 | if (Boost_FOUND)
11 |     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${Boost_C_FLAGS}")
12 |     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${Boost_CXX_FLAGS}")
13 |     include_directories(${Boost_INCLUDE_DIRS})
14 | else()
15 |     message(FATAL_ERROR "Boost dynamic-bitset is required")
16 | endif()
17 | 
18 | # OpenMP
19 | find_package(OpenMP REQUIRED)
20 | if (OPENMP_FOUND)
21 |     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
22 |     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
23 | else()
24 |     message(FATAL_ERROR "OpenMP is required")
25 | endif()
26 | 
27 | # Compile flags
28 | 
29 | add_subdirectory(src)
30 | add_subdirectory(test)
31 | 
32 | add_definitions(-std=c++17 -O3 -lboost -ltcmalloc_minimal -lprofiler -march=native -Wall -DINFO -mavx)
33 | 
34 | 
35 | 


--------------------------------------------------------------------------------
/include/ssg/index.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <cstddef>
 4 | #include <fstream>
 5 | #include <string>
 6 | #include <vector>
 7 | 
 8 | #include "distance.h"
 9 | #include "parameters.h"
10 | 
11 | namespace efanna2e {
12 | 
13 | class Index {
14 |  public:
15 |   explicit Index(const size_t dimension, const size_t n, Metric metric);
16 | 
17 |   virtual ~Index();
18 | 
19 |   virtual void Build(size_t n, const float *data,
20 |                      const Parameters &parameters) = 0;
21 | 
22 |   virtual void Search(const float *query, const float *x, size_t k,
23 |                       const Parameters &parameters, unsigned *indices) = 0;
24 | 
25 |   virtual void Save(const char *filename) = 0;
26 | 
27 |   virtual void Load(const char *filename) = 0;
28 | 
29 |   inline bool HasBuilt() const { return has_built; }
30 | 
31 |   inline size_t GetDimension() const { return dimension_; };
32 | 
33 |   inline size_t GetSizeOfDataset() const { return nd_; }
34 | 
35 |   inline const float *GetDataset() const { return data_; }
36 | 
37 |  protected:
38 |   const size_t dimension_;
39 |   const float *data_;
40 |   size_t nd_;
41 |   bool has_built;
42 |   Distance *distance_;
43 | };
44 | 
45 | }  // namespace efanna2e


--------------------------------------------------------------------------------
/test/test_mips_index.cpp:
--------------------------------------------------------------------------------
 1 | #include "ssg/index_random.h"
 2 | #include "ssg/util.h"
 3 | #include "mips/index_mips.h"
 4 | 
 5 | 
 6 | int main(int argc, char** argv) {
 7 | 
 8 |   std::cerr << "Data Path: " << argv[1] << std::endl;
 9 | 
10 |   unsigned points_num, dim = (unsigned)atoi(argv[8]);
11 |   float* data_load = nullptr;
12 |   data_load = efanna2e::load_data(argv[1], points_num, dim);
13 |   data_load = efanna2e::data_align(data_load, points_num, dim);
14 | 
15 |   std::string nn_graph_path(argv[2]);
16 |   unsigned L = (unsigned)atoi(argv[3]);
17 |   unsigned R = (unsigned)atoi(argv[4]);
18 |   float A = (float)atof(argv[5]);
19 |   unsigned M = (unsigned)atoi(argv[6]);
20 | 
21 |   std::cout << "L = " << L << ", ";
22 |   std::cout << "R = " << R << ", ";
23 |   std::cout << "Angle = " << A << std::endl;
24 |   std::cout << "KNNG = " << nn_graph_path << std::endl;
25 |   efanna2e::IndexRandom init_index(dim, points_num);
26 |   efanna2e::IndexMips index(dim, points_num, efanna2e::L2,
27 |                            (efanna2e::Index*)(&init_index));
28 |   efanna2e::Parameters paras;
29 |   paras.Set<unsigned>("L", L);
30 |   paras.Set<unsigned>("R", R);
31 |   paras.Set<float>("A", A);
32 |   paras.Set<unsigned>("n_try", 10);
33 |   paras.Set<unsigned>("M", M);
34 |   paras.Set<std::string>("nn_graph_path", nn_graph_path);
35 | 
36 |   std::cerr << "Output MIPS-SSG Path: " << argv[7] << std::endl;
37 | 
38 |   auto s = std::chrono::high_resolution_clock::now();
39 |   index.Build(points_num, data_load, paras);
40 |   auto e = std::chrono::high_resolution_clock::now();
41 |   std::chrono::duration<double> diff = e - s;
42 |   std::cout << "Build Time: " << diff.count() << "\n";
43 | 
44 |   index.Save(argv[7]);
45 | 
46 |   return 0;
47 | }


--------------------------------------------------------------------------------
/include/ssg/parameters.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <sstream>
 4 | #include <typeinfo>
 5 | #include <unordered_map>
 6 | namespace efanna2e {
 7 | 
 8 | class Parameters {
 9 |  public:
10 |   template <typename ParamType>
11 |   inline void Set(const std::string &name, const ParamType &value) {
12 |     std::stringstream sstream;
13 |     sstream << value;
14 |     params[name] = sstream.str();
15 |   }
16 | 
17 |   inline std::string GetRaw(const std::string &name) const {
18 |     auto item = params.find(name);
19 |     if (item == params.end()) {
20 |       throw std::invalid_argument("Invalid parameter name.");
21 |     } else {
22 |       return item->second;
23 |     }
24 |   }
25 | 
26 |   template <typename ParamType>
27 |   inline ParamType Get(const std::string &name) const {
28 |     auto item = params.find(name);
29 |     if (item == params.end()) {
30 |       throw std::invalid_argument("Invalid parameter name.");
31 |     } else {
32 |       return ConvertStrToValue<ParamType>(item->second);
33 |     }
34 |   }
35 | 
36 |   template <typename ParamType>
37 |   inline ParamType Get(const std::string &name,
38 |                        const ParamType &default_value) {
39 |     try {
40 |       return Get<ParamType>(name);
41 |     } catch (std::invalid_argument e) {
42 |       return default_value;
43 |     }
44 |   }
45 | 
46 |  private:
47 |   std::unordered_map<std::string, std::string> params;
48 | 
49 |   template <typename ParamType>
50 |   inline ParamType ConvertStrToValue(const std::string &str) const {
51 |     std::stringstream sstream(str);
52 |     ParamType value;
53 |     if (!(sstream >> value) || !sstream.eof()) {
54 |       std::stringstream err;
55 |       err << "Failed to convert value '" << str
56 |           << "' to type: " << typeid(value).name();
57 |       throw std::runtime_error(err.str());
58 |     }
59 |     return value;
60 |   }
61 | };
62 | 
63 | } 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | pip-wheel-metadata/
 24 | share/python-wheels/
 25 | *.egg-info/
 26 | .installed.cfg
 27 | *.egg
 28 | MANIFEST
 29 | 
 30 | # PyInstaller
 31 | #  Usually these files are written by a python script from a template
 32 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 33 | *.manifest
 34 | *.spec
 35 | 
 36 | # Installer logs
 37 | pip-log.txt
 38 | pip-delete-this-directory.txt
 39 | 
 40 | # Unit test / coverage reports
 41 | htmlcov/
 42 | .tox/
 43 | .nox/
 44 | .coverage
 45 | .coverage.*
 46 | .cache
 47 | nosetests.xml
 48 | coverage.xml
 49 | *.cover
 50 | *.py,cover
 51 | .hypothesis/
 52 | .pytest_cache/
 53 | 
 54 | # Translations
 55 | *.mo
 56 | *.pot
 57 | 
 58 | # Django stuff:
 59 | *.log
 60 | local_settings.py
 61 | db.sqlite3
 62 | db.sqlite3-journal
 63 | 
 64 | # Flask stuff:
 65 | instance/
 66 | .webassets-cache
 67 | 
 68 | # Scrapy stuff:
 69 | .scrapy
 70 | 
 71 | # Sphinx documentation
 72 | docs/_build/
 73 | 
 74 | # PyBuilder
 75 | target/
 76 | 
 77 | # Jupyter Notebook
 78 | .ipynb_checkpoints
 79 | 
 80 | # IPython
 81 | profile_default/
 82 | ipython_config.py
 83 | 
 84 | # pyenv
 85 | .python-version
 86 | 
 87 | # pipenv
 88 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 89 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 90 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 91 | #   install all needed dependencies.
 92 | #Pipfile.lock
 93 | 
 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 95 | __pypackages__/
 96 | 
 97 | # Celery stuff
 98 | celerybeat-schedule
 99 | celerybeat.pid
100 | 
101 | # SageMath parsed files
102 | *.sage.py
103 | 
104 | # Environments
105 | .env
106 | .venv
107 | env/
108 | venv/
109 | ENV/
110 | env.bak/
111 | venv.bak/
112 | 
113 | # Spyder project settings
114 | .spyderproject
115 | .spyproject
116 | 
117 | # Rope project settings
118 | .ropeproject
119 | 
120 | # mkdocs documentation
121 | /site
122 | 
123 | # mypy
124 | .mypy_cache/
125 | .dmypy.json
126 | dmypy.json
127 | 
128 | # Pyre type checker
129 | .pyre/
130 | 
131 | # Personal config
132 | 
133 | **/.idea/
134 | */.DS_Store


--------------------------------------------------------------------------------
/include/ssg/index_ssg.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <boost/dynamic_bitset.hpp>
 4 | #include <cassert>
 5 | #include <sstream>
 6 | #include <stack>
 7 | #include <string>
 8 | #include <unordered_map>
 9 | 
10 | #include "index.h"
11 | #include "neighbor.h"
12 | #include "parameters.h"
13 | #include "util.h"
14 | 
15 | namespace efanna2e {
16 | 
17 | class IndexSSG : public Index {
18 |  public:
19 |   explicit IndexSSG(const size_t dimension, const size_t n, Metric m,
20 |                     Index *initializer);
21 | 
22 |   virtual ~IndexSSG();
23 | 
24 |   virtual void Save(const char *filename) override;
25 |   virtual void Load(const char *filename) override;
26 | 
27 |   virtual void Build(size_t n, const float *data,
28 |                      const Parameters &parameters) override;
29 | 
30 |   virtual void Search(const float *query, const float *x, size_t k,
31 |                       const Parameters &parameters, unsigned *indices) override;
32 |   void SearchWithOptGraph(const float *query, size_t K,
33 |                           const Parameters &parameters, unsigned *indices);
34 |   void OptimizeGraph(const float *data);
35 | 
36 |  protected:
37 |   typedef std::vector<std::vector<unsigned>> CompactGraph;
38 |   typedef std::vector<SimpleNeighbors> LockGraph;
39 |   typedef std::vector<nhood> KNNGraph;
40 | 
41 |   CompactGraph final_graph_;
42 |   Index *initializer_;
43 | 
44 |   void init_graph(const Parameters &parameters);
45 |   void get_neighbors(const float *query, const Parameters &parameter,
46 |                      std::vector<Neighbor> &retset,
47 |                      std::vector<Neighbor> &fullset);
48 |   void get_neighbors(const unsigned q, const Parameters &parameter,
49 |                      std::vector<Neighbor> &pool);
50 |   void sync_prune(unsigned q, std::vector<Neighbor> &pool,
51 |                   const Parameters &parameter, float threshold,
52 |                   SimpleNeighbor *cut_graph_);
53 |   void Link(const Parameters &parameters, SimpleNeighbor *cut_graph_);
54 |   void InterInsert(unsigned n, unsigned range, float threshold,
55 |                    std::vector<std::mutex> &locks, SimpleNeighbor *cut_graph_);
56 |   void Load_nn_graph(const char *filename);
57 |   void strong_connect(const Parameters &parameter);
58 | 
59 |   void DFS(boost::dynamic_bitset<> &flag,
60 |            std::vector<std::pair<unsigned, unsigned>> &edges, unsigned root,
61 |            unsigned &cnt);
62 |   bool check_edge(unsigned u, unsigned t);
63 |   void findroot(boost::dynamic_bitset<> &flag, unsigned &root,
64 |                 const Parameters &parameter);
65 |   void DFS_expand(const Parameters &parameter);
66 | 
67 |  private:
68 |   unsigned width;
69 |   unsigned ep_; //not in use
70 |   std::vector<unsigned> eps_;
71 |   std::vector<std::mutex> locks;
72 |   char *opt_graph_;
73 |   size_t node_size;
74 |   size_t data_len;
75 |   size_t neighbor_len;
76 |   KNNGraph nnd_graph;
77 | };
78 | 
79 | }  // namespace efanna2e


--------------------------------------------------------------------------------
/src/util.cpp:
--------------------------------------------------------------------------------
  1 | #include "ssg/util.h"
  2 | 
  3 | #include <malloc.h>
  4 | #include <algorithm>
  5 | #include <cstdlib>
  6 | #include <cstring>
  7 | #include <fstream>
  8 | #include <iostream>
  9 | #include <random>
 10 | 
 11 | namespace efanna2e {
 12 | 
 13 | void GenRandom(std::mt19937& rng, unsigned* addr, unsigned size, unsigned N) {
 14 |   for (unsigned i = 0; i < size; ++i) {
 15 |     addr[i] = rng() % (N - size);
 16 |   }
 17 |   std::sort(addr, addr + size);
 18 |   for (unsigned i = 1; i < size; ++i) {
 19 |     if (addr[i] <= addr[i - 1]) {
 20 |       addr[i] = addr[i - 1] + 1;
 21 |     }
 22 |   }
 23 |   unsigned off = rng() % N;
 24 |   for (unsigned i = 0; i < size; ++i) {
 25 |     addr[i] = (addr[i] + off) % N;
 26 |   }
 27 | }
 28 | 
 29 | float* load_only_data(const char* filename, unsigned num, unsigned dim) {
 30 |   std::ifstream in(filename, std::ios::binary);
 31 |   if (!in.is_open()) {
 32 |     std::cerr << "Open file error" << std::endl;
 33 |     exit(-1);
 34 |   }
 35 | 
 36 |   std::cout << "Dim = " << dim << std::endl;
 37 |   float* data = new float[(size_t)num * (size_t)dim];
 38 | 
 39 |   in.seekg(0, std::ios::beg);
 40 |   for (size_t i = 0; i < num; i++) {
 41 |     in.read((char*)(data + i * dim), dim * sizeof(float));
 42 |   }
 43 |   in.close();
 44 | 
 45 |   return data;
 46 | }
 47 | 
 48 | float* load_data(const char* filename, unsigned& num, unsigned& dim) {
 49 |   std::ifstream in(filename, std::ios::binary);
 50 |   if (!in.is_open()) {
 51 |     std::cerr << "Open file error" << std::endl;
 52 |     exit(-1);
 53 |   }
 54 | 
 55 |   std::cout << "Dim = " << dim << std::endl;
 56 |   in.seekg(0, std::ios::end);
 57 |   std::ios::pos_type ss = in.tellg();
 58 |   size_t fsize = (size_t)ss;
 59 |   num = (unsigned)(fsize / (dim) / 4);
 60 |   std::cout<< "Num = " << num << std::endl;
 61 |   float* data = new float[(size_t)num * (size_t)dim];
 62 |   in.seekg(0, std::ios::beg);
 63 |   for (size_t i = 0; i < num; i++) {
 64 |     in.read((char*)(data + i * dim), dim * sizeof(float));
 65 |   }
 66 |   in.close();
 67 | 
 68 |   return data;
 69 | }
 70 | 
 71 | float* load_bin_data(const char* filename, unsigned& num, unsigned& dim) {
 72 |   std::ifstream in(filename, std::ios::binary);
 73 |   if (!in.is_open()) {
 74 |     std::cerr << "Open file error" << std::endl;
 75 |     exit(-1);
 76 |   }
 77 | 
 78 |   in.read((char*)&num, 4);
 79 |   std::cout<< "Num = " << num << std::endl;
 80 |   
 81 |   // in.seekg(0, std::ios::end);
 82 |   // std::ios::pos_type ss = in.tellg();
 83 |   // size_t fsize = (size_t)ss;
 84 |   // num = (unsigned)(fsize / (dim + 1) / 4);
 85 |   in.read((char*)&dim, 4);
 86 |   std::cout << "Dim = " << dim << std::endl;
 87 |   float* data = new float[(size_t)num * (size_t)dim];
 88 | 
 89 |   in.seekg(0, std::ios::beg);
 90 |   for (size_t i = 0; i < num; i++) {
 91 |     // in.seekg(4, std::ios::cur);
 92 |     in.read((char*)(data + i * dim), dim * sizeof(float));
 93 |     // std::cout << "data[" << i << "]: " << data[i] << std::endl;
 94 |   }
 95 |   in.close();
 96 | 
 97 |   return data;
 98 | }
 99 | 
100 | float* data_align(float* data_ori, unsigned point_num, unsigned& dim) {
101 | #ifdef __GNUC__
102 | #ifdef __AVX__
103 | #define DATA_ALIGN_FACTOR 8
104 | #else
105 | #ifdef __SSE2__
106 | #define DATA_ALIGN_FACTOR 4
107 | #else
108 | #define DATA_ALIGN_FACTOR 1
109 | #endif
110 | #endif
111 | #endif
112 |   float* data_new = 0;
113 |   unsigned new_dim =
114 |       (dim + DATA_ALIGN_FACTOR - 1) / DATA_ALIGN_FACTOR * DATA_ALIGN_FACTOR;
115 | #ifdef __APPLE__
116 |   data_new = new float[(size_t)new_dim * (size_t)point_num];
117 | #else
118 |   data_new =
119 |       (float*)memalign(DATA_ALIGN_FACTOR * 4,
120 |                        (size_t)point_num * (size_t)new_dim * sizeof(float));
121 | #endif
122 | 
123 |   for (size_t i = 0; i < point_num; i++) {
124 |     memcpy(data_new + i * new_dim, data_ori + i * dim, dim * sizeof(float));
125 |     memset(data_new + i * new_dim + dim, 0, (new_dim - dim) * sizeof(float));
126 |   }
127 | 
128 |   dim = new_dim;
129 | 
130 | #ifdef __APPLE__
131 |   delete[] data_ori;
132 | #else
133 |   free(data_ori);
134 | #endif
135 | 
136 |   return data_new;
137 | }
138 | 
139 | }


--------------------------------------------------------------------------------
/include/mips/index_mips.h:
--------------------------------------------------------------------------------
  1 | #pragma once
  2 | 
  3 | #include <boost/dynamic_bitset.hpp>
  4 | #include <cassert>
  5 | #include <sstream>
  6 | #include <stack>
  7 | #include <string>
  8 | #include <unordered_map>
  9 | 
 10 | #include "../ssg/index.h"
 11 | #include "../ssg/neighbor.h"
 12 | #include "../ssg/parameters.h"
 13 | #include "../ssg/util.h"
 14 | #include "../ssg/distance.h"
 15 | 
 16 | 
 17 | namespace efanna2e {
 18 | 
 19 |   class IndexMips : public Index {
 20 |     public:
 21 |       explicit IndexMips(const size_t dimension, const size_t n, Metric m,
 22 |                         Index *initializer);
 23 | 
 24 |       virtual ~IndexMips();
 25 | 
 26 |       virtual void Save(const char *filename) override;
 27 |       virtual void Load(const char *filename) override;
 28 |       virtual void Search(const float *query, const float *x, size_t k,
 29 |                           const Parameters &parameters, unsigned *indices) override;
 30 |       virtual void Build(size_t n, const float *data,
 31 |                         const Parameters &parameters) override;
 32 | 
 33 |       int Search_Mips_IP_Cal(const float *query, const float *x, size_t k,
 34 |                           const Parameters &parameters, unsigned *indices, std::vector<int> &init_nodes);
 35 | 
 36 |       int Search_Mips_IP_Cal_with_No_SN(const float *query, const float *x, size_t k, const Parameters &parameters, unsigned *indices);
 37 | 
 38 |       void init_eps(std::vector<int> &nodes) {
 39 |         eps_.resize(nodes.size());
 40 |         for (int i = 0; i < nodes.size(); i++) {
 41 |           eps_[i] = nodes[i];
 42 |         }
 43 |       }
 44 | 
 45 |       int check_connected_component();
 46 | 
 47 |       void SaveData(float *data) {
 48 |         data_ = data;
 49 |       }
 50 | 
 51 |     protected:
 52 |       typedef std::vector<std::vector<unsigned>> CompactGraph;
 53 |       typedef std::vector<SimpleNeighbors> LockGraph;
 54 |       typedef std::vector<nhood> KNNGraph;
 55 | 
 56 |       CompactGraph final_graph_;
 57 |       Index *initializer_;
 58 | 
 59 |       void init_graph(const Parameters &parameters);
 60 |       void get_neighbors(const float *query, const Parameters &parameter,
 61 |                         std::vector<Neighbor> &retset,
 62 |                         std::vector<Neighbor> &fullset);
 63 |       void get_neighbors(const unsigned q, const Parameters &parameter,
 64 |                         std::vector<Neighbor> &pool, boost::dynamic_bitset<> &flags);
 65 |       void sync_prune(unsigned q, std::vector<Neighbor> &pool,
 66 |                       const Parameters &parameter, float threshold,
 67 |                       SimpleNeighbor *cut_graph_);
 68 | 
 69 |       void Link(const Parameters &parameters, SimpleNeighbor *cut_graph_);
 70 | 
 71 |       void InterInsert(unsigned n, unsigned range, float threshold,
 72 |                       std::vector<std::mutex> &locks, SimpleNeighbor *cut_graph_);
 73 | 
 74 |       void Load_nn_graph(const char *filename);
 75 | 
 76 |       void strong_connect(const Parameters &parameter);
 77 | 
 78 |       void get_refine_neighbors(const unsigned ep, const float *query, const Parameters &parameter,
 79 |                         std::vector<Neighbor> &retset,
 80 |                         std::vector<Neighbor> &fullset);
 81 | 
 82 |       void add_mips_neighbors(const unsigned n, const Parameters &parameter,
 83 |                         std::vector<IpNeighbor> &retset,
 84 |                         SimpleNeighbor *cut_graph_);
 85 |       void get_mips_neighbors(const unsigned ep, const float *query, const Parameters &parameter,
 86 |                         std::vector<IpNeighbor> &retset); 
 87 |       
 88 |       void DFS(boost::dynamic_bitset<> &flag,
 89 |               std::vector<std::pair<unsigned, unsigned>> &edges, unsigned root,
 90 |               unsigned &cnt);
 91 |       bool check_edge(unsigned u, unsigned t);
 92 |       void findroot(boost::dynamic_bitset<> &flag, unsigned &root,
 93 |                     const Parameters &parameter);
 94 |       void DFS_expand(const Parameters &parameter);
 95 | 
 96 |     private:
 97 |       unsigned width;
 98 |       unsigned ep_; //not in use
 99 |       std::vector<unsigned> eps_;
100 |       std::vector<std::mutex> locks;
101 |       char *opt_graph_;
102 |       size_t node_size;
103 |       size_t data_len;
104 |       size_t neighbor_len;
105 |       KNNGraph nnd_graph;
106 |       std::vector<float> norms_;
107 |       std::vector<int> init_points_;
108 |   };
109 | 
110 | }
111 | 
112 | 


--------------------------------------------------------------------------------
/test/test_mips_search.cpp:
--------------------------------------------------------------------------------
  1 | #include <chrono>
  2 | 
  3 | #include "ssg/index_random.h"
  4 | #include "mips/index_mips.h"
  5 | #include "ssg/util.h"
  6 | 
  7 | void save_result(char* filename, std::vector<std::vector<unsigned> >& results) {
  8 |   std::ofstream out(filename, std::ios::binary | std::ios::out);
  9 | 
 10 |   for (unsigned i = 0; i < results.size(); i++) {
 11 |     unsigned GK = (unsigned)results[i].size();
 12 |     out.write((char*)&GK, sizeof(unsigned));
 13 |     out.write((char*)results[i].data(), GK * sizeof(unsigned));
 14 |   }
 15 |   out.close();
 16 | }
 17 | 
 18 | void read_kmeans(const std::string& filename, int query_count, std::vector<int>& query_cluster, std::vector<std::vector<int>>& init_nodes_clusters, int number_of_clusters) {
 19 |     std::ifstream file(filename, std::ios::binary);
 20 |     if (!file) {
 21 |         std::cerr << "Cannot open file: " << filename << std::endl;
 22 |         return;
 23 |     }
 24 | 
 25 |     query_cluster.resize(query_count);
 26 |     file.read(reinterpret_cast<char*>(query_cluster.data()), query_count * sizeof(int));
 27 | 
 28 |     std::vector<int> init_nodes;
 29 |     file.seekg(0, std::ios::end);
 30 |     size_t file_size = file.tellg();
 31 |     size_t init_nodes_count = (file_size - query_count * sizeof(int)) / sizeof(int);
 32 |     file.seekg(query_count * sizeof(int), std::ios::beg);
 33 | 
 34 |     init_nodes.resize(init_nodes_count);
 35 |     file.read(reinterpret_cast<char*>(init_nodes.data()), init_nodes_count * sizeof(int));
 36 | 
 37 |     file.close();
 38 | 
 39 |     size_t cluster_size = number_of_clusters;
 40 |     for (size_t i = 0; i < init_nodes.size(); i += cluster_size) {
 41 |         std::vector<int> cluster(init_nodes.begin() + i, init_nodes.begin() + i + cluster_size);
 42 |         init_nodes_clusters.push_back(cluster);
 43 |     }
 44 | }
 45 | 
 46 | int main(int argc, char** argv) {
 47 |   if (argc < 8) {
 48 |     std::cout << "./run data_file query_file ssg_path L K result_path dim"
 49 |               << std::endl;
 50 |     exit(-1);
 51 |   }
 52 | 
 53 |   std::cerr << "Data Path: " << argv[1] << std::endl;
 54 | 
 55 |   unsigned points_num, dim = (unsigned)atoi(argv[7]);
 56 |   float* data_load = nullptr;
 57 |   data_load = efanna2e::load_data(argv[1], points_num, dim);
 58 |   data_load = efanna2e::data_align(data_load, points_num, dim);
 59 | 
 60 |   std::cerr << "Query Path: " << argv[2] << std::endl;
 61 | 
 62 |   unsigned query_num, query_dim = (unsigned)atoi(argv[7]);
 63 |   float* query_load = nullptr;
 64 |   query_load = efanna2e::load_data(argv[2], query_num, query_dim);
 65 |   query_load = efanna2e::data_align(query_load, query_num, query_dim);
 66 | 
 67 |   assert(dim == query_dim);
 68 | 
 69 |   efanna2e::IndexRandom init_index(dim, points_num);
 70 |   efanna2e::IndexMips index(dim, points_num, efanna2e::FAST_L2,
 71 |                            (efanna2e::Index*)(&init_index));
 72 | 
 73 |   std::cerr << "SSG Path: " << argv[3] << std::endl;
 74 |   std::cerr << "Result Path: " << argv[6] << std::endl;
 75 |   index.SaveData(data_load);
 76 |   index.Load(argv[3]);
 77 | 
 78 |   unsigned L = (unsigned)atoi(argv[4]);
 79 |   unsigned K = (unsigned)atoi(argv[5]);
 80 | 
 81 |   std::cerr << "L = " << L << ", ";
 82 |   std::cerr << "K = " << K << std::endl;
 83 | 
 84 |   efanna2e::Parameters paras;
 85 |   paras.Set<unsigned>("L_search", L);
 86 |   std::vector<std::vector<unsigned> > res(query_num);
 87 |   for (unsigned i = 0; i < query_num; i++) res[i].resize(K);
 88 | 
 89 |   /* optional entry points initialization */
 90 |   // std::string filename = "../output/mnist/sn.bin";
 91 |   // std::vector<int> query_cluster;
 92 |   // std::vector<std::vector<int>> init_nodes_clusters;
 93 |   // read_kmeans(filename, query_num, query_cluster, init_nodes_clusters, 100);
 94 | 
 95 |   auto num = 0.0;
 96 | 
 97 |   std::vector<std::pair<float,float>> cal_pair;
 98 | 
 99 | 
100 |   auto start = std::chrono::high_resolution_clock::now();
101 | 
102 |   #pragma omp parallel for
103 |   for (unsigned i = 0; i < (int)query_num; i++) {
104 |     // int dis_cal = index.Search_Mips_IP_Cal(query_load + i * dim, data_load, K, paras, res[i].data(), init_nodes_clusters[query_cluster[i]]);
105 |     int dis_cal = index.Search_Mips_IP_Cal_with_No_SN(query_load + i * dim, data_load, K, paras, res[i].data());
106 |     num += dis_cal;
107 |   }
108 | 
109 |   auto end = std::chrono::high_resolution_clock::now();
110 | 
111 |   std::cout << "Average Distance Computation: " << num / (int)query_num << std::endl;
112 |   std::cout << "Average Query Time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() / static_cast<double>(query_num) << "ms" << std::endl;
113 |   save_result(argv[6], res);
114 | 
115 |   return 0;
116 | }


--------------------------------------------------------------------------------
/include/ssg/neighbor.h:
--------------------------------------------------------------------------------
  1 | #pragma once
  2 | 
  3 | #include <pthread.h>
  4 | #include <cstddef>
  5 | #include <cstring>
  6 | #include <mutex>
  7 | #include <random>
  8 | #include <vector>
  9 | 
 10 | #include "util.h"
 11 | 
 12 | namespace efanna2e {
 13 | 
 14 | struct Neighbor {
 15 |   unsigned id;
 16 |   float distance;
 17 |   bool flag;
 18 | 
 19 |   Neighbor() = default;
 20 |   Neighbor(unsigned id, float distance, bool f)
 21 |       : id{id}, distance{distance}, flag(f) {}
 22 | 
 23 |   inline bool operator<(const Neighbor &other) const {
 24 |     return distance < other.distance;
 25 |   }
 26 | };
 27 | 
 28 | struct IpNeighbor {
 29 |   unsigned id;
 30 |   float distance;
 31 |   bool flag;
 32 | 
 33 |   IpNeighbor() = default;
 34 |   IpNeighbor(unsigned id, float distance, bool f)
 35 |       : id{id}, distance{distance}, flag(f) {}
 36 | 
 37 |   inline bool operator<(const IpNeighbor &other) const {
 38 |     return distance > other.distance;
 39 |   }
 40 | };
 41 | 
 42 | typedef std::lock_guard<std::mutex> LockGuard;
 43 | struct nhood {
 44 |   // std::mutex lock;
 45 |   pthread_mutex_t lock;
 46 |   std::vector<Neighbor> pool;
 47 |   unsigned M;
 48 | 
 49 |   std::vector<unsigned> nn_old;
 50 |   std::vector<unsigned> nn_new;
 51 |   std::vector<unsigned> rnn_old;
 52 |   std::vector<unsigned> rnn_new;
 53 | 
 54 |   nhood() {}
 55 |   nhood(unsigned l, unsigned s, std::mt19937 &rng, unsigned N) {
 56 |     M = s;
 57 |     nn_new.resize(s * 2);
 58 |     GenRandom(rng, &nn_new[0], (unsigned)nn_new.size(), N);
 59 |     nn_new.reserve(s * 2);
 60 |     pool.reserve(l);
 61 |     pthread_mutex_init(&lock, NULL);
 62 |   }
 63 | 
 64 |   nhood(const nhood &other) {
 65 |     M = other.M;
 66 |     std::copy(other.nn_new.begin(), other.nn_new.end(),
 67 |               std::back_inserter(nn_new));
 68 |     nn_new.reserve(other.nn_new.capacity());
 69 |     pool.reserve(other.pool.capacity());
 70 |   }
 71 |   void insert(unsigned id, float dist) {
 72 |     // LockGuard guard(lock);
 73 |     pthread_mutex_lock(&lock);
 74 |     if (dist > pool.front().distance) {
 75 |       pthread_mutex_unlock(&lock);
 76 |       return;
 77 |     }
 78 |     for (unsigned i = 0; i < pool.size(); i++) {
 79 |       if (id == pool[i].id) {
 80 |         pthread_mutex_unlock(&lock);
 81 |         return;
 82 |       }
 83 |     }
 84 |     if (pool.size() < pool.capacity()) {
 85 |       pool.push_back(Neighbor(id, dist, true));
 86 |       std::push_heap(pool.begin(), pool.end());
 87 |     } else {
 88 |       std::pop_heap(pool.begin(), pool.end());
 89 |       pool[pool.size() - 1] = Neighbor(id, dist, true);
 90 |       std::push_heap(pool.begin(), pool.end());
 91 |     }
 92 |     pthread_mutex_unlock(&lock);
 93 |   }
 94 | 
 95 |   template <typename C>
 96 |   void join(C callback) const {
 97 |     for (unsigned const i : nn_new) {
 98 |       for (unsigned const j : nn_new) {
 99 |         if (i < j) {
100 |           callback(i, j);
101 |         }
102 |       }
103 |       for (unsigned j : nn_old) {
104 |         callback(i, j);
105 |       }
106 |     }
107 |   }
108 | };
109 | 
110 | struct LockNeighbor {
111 |   // std::mutex lock;
112 |   pthread_mutex_t lock;
113 |   std::vector<Neighbor> pool;
114 | };
115 | 
116 | struct SimpleNeighbor {
117 |   unsigned id;
118 |   float distance;
119 | 
120 |   SimpleNeighbor() = default;
121 |   SimpleNeighbor(unsigned id_, float distance_)
122 |       : id(id_), distance(distance_) {}
123 | 
124 |   inline bool operator<(const SimpleNeighbor &other) const {
125 |     return distance < other.distance;
126 |   }
127 | };
128 | 
129 | struct SimpleIpNeighbor {
130 |   unsigned id;
131 |   float distance;
132 | 
133 |   SimpleIpNeighbor() = default;
134 |   SimpleIpNeighbor(unsigned id_, float distance_)
135 |       : id(id_), distance(distance_) {}
136 | 
137 |   inline bool operator<(const SimpleIpNeighbor &other) const {
138 |     return distance > other.distance;
139 |   }
140 | };
141 | 
142 | struct SimpleNeighbors {
143 |   std::vector<SimpleNeighbor> pool;
144 | };
145 | 
146 | struct SimpleIpNeighbors {
147 |   std::vector<SimpleIpNeighbor> pool;
148 | };
149 | 
150 | static inline int InsertIntoPool(Neighbor *addr, unsigned K, Neighbor nn) {
151 |   // find the location to insert
152 |   int left = 0, right = K - 1;
153 |   if (addr[left].distance > nn.distance) {
154 |     memmove((char *)&addr[left + 1], &addr[left], K * sizeof(Neighbor));
155 |     addr[left] = nn;
156 |     return left;
157 |   }
158 |   if (addr[right].distance < nn.distance) {
159 |     addr[K] = nn;
160 |     return K;
161 |   }
162 |   while (left < right - 1) {
163 |     int mid = (left + right) / 2;
164 |     if (addr[mid].distance > nn.distance)
165 |       right = mid;
166 |     else
167 |       left = mid;
168 |   }
169 |   // check equal ID
170 | 
171 |   while (left > 0) {
172 |     if (addr[left].distance < nn.distance) break;
173 |     if (addr[left].id == nn.id) return K + 1;
174 |     left--;
175 |   }
176 |   if (addr[left].id == nn.id || addr[right].id == nn.id) return K + 1;
177 |   memmove((char *)&addr[right + 1], &addr[right],
178 |           (K - right) * sizeof(Neighbor));
179 |   addr[right] = nn;
180 |   return right;
181 | }
182 | 
183 | static inline int InsertIntoIpPool(IpNeighbor *addr, unsigned K, IpNeighbor nn) {
184 |   // find the location to insert
185 |   int left = 0, right = K - 1;
186 |   if (addr[left].distance < nn.distance) {
187 |     memmove((char *)&addr[left + 1], &addr[left], K * sizeof(IpNeighbor));
188 |     addr[left] = nn;
189 |     return left;
190 |   }
191 |   if (addr[right].distance > nn.distance) {
192 |     addr[K] = nn;
193 |     return K;
194 |   }
195 |   while (left < right - 1) {
196 |     int mid = (left + right) / 2;
197 |     if (addr[mid].distance < nn.distance)
198 |       right = mid;
199 |     else
200 |       left = mid;
201 |   }
202 |   // check equal ID
203 | 
204 |   while (left > 0) {
205 |     if (addr[left].distance < nn.distance) break;
206 |     if (addr[left].id == nn.id) return K + 1;
207 |     left--;
208 |   }
209 |   if (addr[left].id == nn.id || addr[right].id == nn.id) return K + 1;
210 |   memmove((char *)&addr[right + 1], &addr[right],
211 |           (K - right) * sizeof(IpNeighbor));
212 |   addr[right] = nn;
213 |   return right;
214 | }
215 | 
216 | } 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | <h1 align="center">PSP - Proximity Graph with Spherical Pathway for MIPS </h1> 
  2 | 
  3 | This repository contains the source code for our paper: **Maximum Inner Product is Query-Scaled Nearest Neighbor** (VLDB25).
  4 | 
  5 | **If you're interested, you can also check out our further extended work: [SIGIR2025 MAG](https://github.com/ZJU-DAILY/MAG)**
  6 | ## 1 Abstract
  7 | 
  8 | This paper presents a novel theoretical framework that equates MIPS with NNS without requiring space transformation, thereby allowing us to leverage advanced graph-based indices for NNS and efficient edge pruning strategies, significantly reducing unnecessary computations. 
  9 | 
 10 | ## 2 Competitors
 11 | > The initial code for ip-nsw and ip-nsw+ came from the original papers, and we reconstructed the relevant code using hnswlib 0.8.0. For Möbius-graph and NAPG, we did not find available open-source code, so we implemented our own version, which will also be open-sourced for relevant evaluations. ScaNN is a commercial product, so we used the open-source version of the code, it achieves excellent results through careful parameter tuning. Its official examples do not enable multithreading. If you are conducting evaluations using multithreading, please invoke the relevant functions. 
 12 | 
 13 | * ip-NSW ([Paper](https://proceedings.neurips.cc/paper_files/paper/2018/file/229754d7799160502a143a72f6789927-Paper.pdf)): A graph based method using inner product navigable small world graph. 
 14 | 
 15 | * ip-NSW+ ([Paper](https://aaai.org/ojs/index.php/AAAI/article/view/5344/5200)): An enhancement of ip-NSW that introduces an additional angular proximity graph.
 16 | * Möbius-Graph ([Paper](https://proceedings.neurips.cc/paper/2019/file/0fd7e4f42a8b4b4ef33394d35212b13e-Paper.pdf)): A graph based method that reduces the MIPS problem to an NNS problem using Möbius transformation. Since the original code is not available, we implemented a version based on the paper.
 17 | * NAPG ([Paper](https://dl.acm.org/doi/abs/10.1145/3447548.3467412)): A recent graph-based method claiming state-of-the-art performance by improving ip-NSW with a specialized metric, using an adaptive $\alpha$ for different norm distributions.
 18 | * Fargo ([Paper](https://www.vldb.org/pvldb/vol16/p1100-zheng.pdf)): The latest state-of-the-art LSH based method with theoretical guarantees.
 19 | * ScaNN ([Paper](http://proceedings.mlr.press/v119/guo20h/guo20h.pdf)): The state-of-the-art quantization method.
 20 | ## 3 Datasets
 21 | 
 22 | The data format is: Number of vector (n) * Dimension (d).
 23 | 
 24 | *: Data source will be released upon publication.
 25 | 
 26 | | Dataset                                                      | Base Size   | Dim  | Query Size | Modality   |
 27 | | ------------------------------------------------------------ | ----------- | ---- | ---------- | ---------- |
 28 | | MNIST ([link](https://yann.lecun.com/exdb/mnist/index.html)) | 60,000      | 784  | 10,000     | Image      |
 29 | | DBpedia100K ([link](https://huggingface.co/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-3072-100K)) | 100,000     | 3072 | 1,000      | Text       |
 30 | | DBpedia1M ([link](https://huggingface.co/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M)) | 1,000,000   | 1536 | 1,000      | Text       |
 31 | | Music100 ([link](https://github.com/stanis-morozov/ip-nsw))  | 1,000,000   | 100  | 10,000     | Audio      |
 32 | | Text2Image1M ([link](https://research.yandex.com/blog/benchmarks-for-billion-scale-similarity-search)) | 1,000,000   | 200  | 100,000    | Multi      |
 33 | | Text2Image10M ([link](https://research.yandex.com/blog/benchmarks-for-billion-scale-similarity-search)) | 10,000,000  | 200  | 100,000    | Multi      |
 34 | | Laion10M ([link](https://arxiv.org/abs/2210.08402)) | 12,244,692      | 512  | 1,000      | Multi      |
 35 | | Commerce100M*                                                | 100,279,529 | 48   | 64,111     | E-commerce |
 36 | 
 37 | ## 4 Building Instruction
 38 | 
 39 | ### Prerequisites
 40 | 
 41 | - GCC 4.9+ with OpenMP
 42 | - CMake 2.8+
 43 | - Boost 1.55+
 44 | - Faiss (optional)
 45 | 
 46 | ### Compile On Linux
 47 | 
 48 | ```shell
 49 | $ mkdir build/ && cd build/
 50 | $ cmake ..
 51 | $ make -j
 52 | ```
 53 | 
 54 | ## 5 Usage
 55 | 
 56 | ### Code Structure
 57 | 
 58 | - **datasets**: datasets
 59 | - **include**: C++ class interface
 60 | - **output**: PSP index, k-MIP result file
 61 | - **script**: some scripts to run the experiments
 62 | - **src**: main function implementation
 63 | - **test**: test codes
 64 | 
 65 | ### How to use
 66 | 
 67 | #### Step 1. Build kNN Graph
 68 | 
 69 | Firstly, we need to prepare a kNN graph.  You can use Faiss and other libs.
 70 | 
 71 | #### Step 2. PSP indexing
 72 | 
 73 | ```shell
 74 | ./test/test_mips_index DATA_PATH KNNG_PATH L R Angle M PSP_PATH DIM
 75 | ```
 76 | 
 77 | - `DATA_PATH` is the path of the base data in `bin` format.
 78 | - `KNNG_PATH` is the path of the pre-built kNN graph in *Step 1.*.
 79 | - `L` candidate pool size.
 80 | - `R`maximum out-degree.
 81 | - `Angle` minimal angle between edges.
 82 | - `M` IP neighbor.
 83 | - `PSP_PATH` is the path of the generated PSP index.
 84 | - `DIM` dimension of dataset.
 85 | 
 86 | #### Step 3. PSP searching
 87 | 
 88 | ```shell
 89 | ./test/test_mips_search DATA_PATH QUERY_PATH PSP_PATH searh_L K RESULT_PATH DIM
 90 | ```
 91 | 
 92 | - `DATA_PATH` is the path of the base data in `bin` format.
 93 | - `QUERY_PATH` is the path of the query data in `bin` format.
 94 | - `PSP_PATH` is the path of the generated PSP index.
 95 | - `search_L` search pool size, the larger the better but slower (must larger than K).
 96 | - `K` the result size.
 97 | - `PSP_PATH` is the path of the result neighbors.
 98 | - `DIM` dimension of dataset.
 99 | 
100 | ## 6 To-Do Lists
101 | - ✅ Open-source code is available for the major components of the original paper.
102 | - ✅ Real-world evaluation scenarios of Commerce100M on the Shopee platform.
103 | - 🔄 More automated entry point selection and early stopping techniques.
104 | - 🔄 Improve compatibility of SIMD-related PQ codes.
105 | - 🔄 Python wrapper.
106 | - 🔄 A wider range of datasets with diverse norm distributions.
107 | ## 7 Performance
108 | 
109 | #### Evaluation Metric
110 | 
111 | - QPS, Distance computation (for graph-based method)
112 | 
113 | ![evaluation](./evaluation.png)
114 | 


--------------------------------------------------------------------------------
/include/ssg/distance.h:
--------------------------------------------------------------------------------
  1 | #pragma once
  2 | 
  3 | #include <x86intrin.h>
  4 | #include <iostream>
  5 | namespace efanna2e {
  6 | enum Metric { L2 = 0, INNER_PRODUCT = 1, FAST_L2 = 2, PQ = 3 };
  7 | class Distance {
  8 |  public:
  9 |   virtual float compare(const float *a, const float *b,
 10 |                         unsigned length) const = 0;
 11 |   virtual ~Distance() {}
 12 | };
 13 | 
 14 | class DistanceL2 : public Distance {
 15 |  public:
 16 |   float compare(const float *a, const float *b, unsigned size) const {
 17 |     float result = 0;
 18 | 
 19 | #ifdef __GNUC__
 20 | #ifdef __AVX__
 21 | 
 22 | #define AVX_L2SQR(addr1, addr2, dest, tmp1, tmp2) \
 23 |   tmp1 = _mm256_loadu_ps(addr1);                  \
 24 |   tmp2 = _mm256_loadu_ps(addr2);                  \
 25 |   tmp1 = _mm256_sub_ps(tmp1, tmp2);               \
 26 |   tmp1 = _mm256_mul_ps(tmp1, tmp1);               \
 27 |   dest = _mm256_add_ps(dest, tmp1);
 28 | 
 29 |     __m256 sum;
 30 |     __m256 l0, l1;
 31 |     __m256 r0, r1;
 32 |     unsigned D = (size + 7) & ~7U;
 33 |     unsigned DR = D % 16;
 34 |     unsigned DD = D - DR;
 35 |     const float *l = a;
 36 |     const float *r = b;
 37 |     const float *e_l = l + DD;
 38 |     const float *e_r = r + DD;
 39 |     float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0};
 40 | 
 41 |     sum = _mm256_loadu_ps(unpack);
 42 |     if (DR) {
 43 |       AVX_L2SQR(e_l, e_r, sum, l0, r0);
 44 |     }
 45 | 
 46 |     for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) {
 47 |       AVX_L2SQR(l, r, sum, l0, r0);
 48 |       AVX_L2SQR(l + 8, r + 8, sum, l1, r1);
 49 |     }
 50 |     _mm256_storeu_ps(unpack, sum);
 51 |     result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] +
 52 |              unpack[5] + unpack[6] + unpack[7];
 53 | 
 54 | #else
 55 | #ifdef __SSE2__
 56 | #define SSE_L2SQR(addr1, addr2, dest, tmp1, tmp2) \
 57 |   tmp1 = _mm_load_ps(addr1);                      \
 58 |   tmp2 = _mm_load_ps(addr2);                      \
 59 |   tmp1 = _mm_sub_ps(tmp1, tmp2);                  \
 60 |   tmp1 = _mm_mul_ps(tmp1, tmp1);                  \
 61 |   dest = _mm_add_ps(dest, tmp1);
 62 | 
 63 |     __m128 sum;
 64 |     __m128 l0, l1, l2, l3;
 65 |     __m128 r0, r1, r2, r3;
 66 |     unsigned D = (size + 3) & ~3U;
 67 |     unsigned DR = D % 16;
 68 |     unsigned DD = D - DR;
 69 |     const float *l = a;
 70 |     const float *r = b;
 71 |     const float *e_l = l + DD;
 72 |     const float *e_r = r + DD;
 73 |     float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0};
 74 | 
 75 |     sum = _mm_load_ps(unpack);
 76 |     switch (DR) {
 77 |       case 12:
 78 |         SSE_L2SQR(e_l + 8, e_r + 8, sum, l2, r2);
 79 |       case 8:
 80 |         SSE_L2SQR(e_l + 4, e_r + 4, sum, l1, r1);
 81 |       case 4:
 82 |         SSE_L2SQR(e_l, e_r, sum, l0, r0);
 83 |       default:
 84 |         break;
 85 |     }
 86 |     for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) {
 87 |       SSE_L2SQR(l, r, sum, l0, r0);
 88 |       SSE_L2SQR(l + 4, r + 4, sum, l1, r1);
 89 |       SSE_L2SQR(l + 8, r + 8, sum, l2, r2);
 90 |       SSE_L2SQR(l + 12, r + 12, sum, l3, r3);
 91 |     }
 92 |     _mm_storeu_ps(unpack, sum);
 93 |     result += unpack[0] + unpack[1] + unpack[2] + unpack[3];
 94 | 
 95 | // nomal distance
 96 | #else
 97 | 
 98 |     float diff0, diff1, diff2, diff3;
 99 |     const float* last = a + size;
100 |     const float* unroll_group = last - 3;
101 | 
102 |     /* Process 4 items with each loop for efficiency. */
103 |     while (a < unroll_group) {
104 |       diff0 = a[0] - b[0];
105 |       diff1 = a[1] - b[1];
106 |       diff2 = a[2] - b[2];
107 |       diff3 = a[3] - b[3];
108 |       result += diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;
109 |       a += 4;
110 |       b += 4;
111 |     }
112 |     /* Process last 0-3 pixels.  Not needed for standard vector lengths. */
113 |     while (a < last) {
114 |       diff0 = *a++ - *b++;
115 |       result += diff0 * diff0;
116 |     }
117 | #endif
118 | #endif
119 | #endif
120 | 
121 |     return result;
122 |   }
123 | };
124 | 
125 | class DistanceInnerProduct : public Distance {
126 |  public:
127 |   float compare(const float *a, const float *b, unsigned size) const {
128 |     float result = 0;
129 | #ifdef __GNUC__
130 | #ifdef __AVX__
131 | #define AVX_DOT(addr1, addr2, dest, tmp1, tmp2) \
132 |   tmp1 = _mm256_loadu_ps(addr1);                \
133 |   tmp2 = _mm256_loadu_ps(addr2);                \
134 |   tmp1 = _mm256_mul_ps(tmp1, tmp2);             \
135 |   dest = _mm256_add_ps(dest, tmp1);
136 | 
137 |     __m256 sum;
138 |     __m256 l0, l1;
139 |     __m256 r0, r1;
140 |     unsigned D = (size + 7) & ~7U;
141 |     unsigned DR = D % 16;
142 |     unsigned DD = D - DR;
143 |     const float *l = a;
144 |     const float *r = b;
145 |     const float *e_l = l + DD;
146 |     const float *e_r = r + DD;
147 |     float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0};
148 | 
149 |     sum = _mm256_loadu_ps(unpack);
150 |     if (DR) {
151 |       AVX_DOT(e_l, e_r, sum, l0, r0);
152 |     }
153 | 
154 |     for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) {
155 |       AVX_DOT(l, r, sum, l0, r0);
156 |       AVX_DOT(l + 8, r + 8, sum, l1, r1);
157 |     }
158 |     _mm256_storeu_ps(unpack, sum);
159 |     result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] +
160 |              unpack[5] + unpack[6] + unpack[7];
161 | 
162 | #else
163 | #ifdef __SSE2__
164 | #define SSE_DOT(addr1, addr2, dest, tmp1, tmp2) \
165 |   tmp1 = _mm_loadu_ps(addr1);                \
166 |   tmp2 = _mm_loadu_ps(addr2);                \
167 |   tmp1 = _mm_mul_ps(tmp1, tmp2);             \
168 |   dest = _mm_add_ps(dest, tmp1);
169 |     __m128 sum;
170 |     __m128 l0, l1, l2, l3;
171 |     __m128 r0, r1, r2, r3;
172 |     unsigned D = (size + 3) & ~3U;
173 |     unsigned DR = D % 16;
174 |     unsigned DD = D - DR;
175 |     const float *l = a;
176 |     const float *r = b;
177 |     const float *e_l = l + DD;
178 |     const float *e_r = r + DD;
179 |     float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0};
180 | 
181 |     sum = _mm_load_ps(unpack);
182 |     switch (DR) {
183 |       case 12:
184 |         SSE_DOT(e_l + 8, e_r + 8, sum, l2, r2);
185 |       case 8:
186 |         SSE_DOT(e_l + 4, e_r + 4, sum, l1, r1);
187 |       case 4:
188 |         SSE_DOT(e_l, e_r, sum, l0, r0);
189 |       default:
190 |         break;
191 |     }
192 |     for (unsigned i = 0; i < DD; i += 16, l += 16, r += 16) {
193 |       SSE_DOT(l, r, sum, l0, r0);
194 |       SSE_DOT(l + 4, r + 4, sum, l1, r1);
195 |       SSE_DOT(l + 8, r + 8, sum, l2, r2);
196 |       SSE_DOT(l + 12, r + 12, sum, l3, r3);
197 |     }
198 |     _mm_storeu_ps(unpack, sum);
199 |     result += unpack[0] + unpack[1] + unpack[2] + unpack[3];
200 | #else
201 | 
202 |     float dot0, dot1, dot2, dot3;
203 |     const float* last = a + size;
204 |     const float* unroll_group = last - 3;
205 | 
206 |     /* Process 4 items with each loop for efficiency. */
207 |     while (a < unroll_group) {
208 |       dot0 = a[0] * b[0];
209 |       dot1 = a[1] * b[1];
210 |       dot2 = a[2] * b[2];
211 |       dot3 = a[3] * b[3];
212 |       result += dot0 + dot1 + dot2 + dot3;
213 |       a += 4;
214 |       b += 4;
215 |     }
216 |     /* Process last 0-3 pixels.  Not needed for standard vector lengths. */
217 |     while (a < last) {
218 |       result += *a++ * *b++;
219 |     }
220 | #endif
221 | #endif
222 | #endif
223 |     return result;
224 |   }
225 | };
226 | class DistanceFastL2 : public DistanceInnerProduct {
227 |  public:
228 |   float norm(const float *a, unsigned size) const {
229 |     float result = 0;
230 | #ifdef __GNUC__
231 | #ifdef __AVX__
232 | #define AVX_L2NORM(addr, dest, tmp) \
233 |   tmp = _mm256_loadu_ps(addr);      \
234 |   tmp = _mm256_mul_ps(tmp, tmp);    \
235 |   dest = _mm256_add_ps(dest, tmp);
236 | 
237 |     __m256 sum;
238 |     __m256 l0, l1;
239 |     unsigned D = (size + 7) & ~7U;
240 |     unsigned DR = D % 16;
241 |     unsigned DD = D - DR;
242 |     const float *l = a;
243 |     const float *e_l = l + DD;
244 |     float unpack[8] __attribute__((aligned(32))) = {0, 0, 0, 0, 0, 0, 0, 0};
245 | 
246 |     sum = _mm256_loadu_ps(unpack);
247 |     if (DR) {
248 |       AVX_L2NORM(e_l, sum, l0);
249 |     }
250 |     for (unsigned i = 0; i < DD; i += 16, l += 16) {
251 |       AVX_L2NORM(l, sum, l0);
252 |       AVX_L2NORM(l + 8, sum, l1);
253 |     }
254 |     _mm256_storeu_ps(unpack, sum);
255 |     result = unpack[0] + unpack[1] + unpack[2] + unpack[3] + unpack[4] +
256 |              unpack[5] + unpack[6] + unpack[7];
257 | #else
258 | #ifdef __SSE2__
259 | #define SSE_L2NORM(addr, dest, tmp) \
260 |   tmp = _mm_loadu_ps(addr);      \
261 |   tmp = _mm_mul_ps(tmp, tmp);    \
262 |   dest = _mm_add_ps(dest, tmp);
263 | 
264 |     __m128 sum;
265 |     __m128 l0, l1, l2, l3;
266 |     unsigned D = (size + 3) & ~3U;
267 |     unsigned DR = D % 16;
268 |     unsigned DD = D - DR;
269 |     const float *l = a;
270 |     const float *e_l = l + DD;
271 |     float unpack[4] __attribute__((aligned(16))) = {0, 0, 0, 0};
272 | 
273 |     sum = _mm_load_ps(unpack);
274 |     switch (DR) {
275 |       case 12:
276 |         SSE_L2NORM(e_l + 8, sum, l2);
277 |       case 8:
278 |         SSE_L2NORM(e_l + 4, sum, l1);
279 |       case 4:
280 |         SSE_L2NORM(e_l, sum, l0);
281 |       default:
282 |         break;
283 |     }
284 |     for (unsigned i = 0; i < DD; i += 16, l += 16) {
285 |       SSE_L2NORM(l, sum, l0);
286 |       SSE_L2NORM(l + 4, sum, l1);
287 |       SSE_L2NORM(l + 8, sum, l2);
288 |       SSE_L2NORM(l + 12, sum, l3);
289 |     }
290 |     _mm_storeu_ps(unpack, sum);
291 |     result += unpack[0] + unpack[1] + unpack[2] + unpack[3];
292 | #else
293 |     float dot0, dot1, dot2, dot3;
294 |     const float* last = a + size;
295 |     const float* unroll_group = last - 3;
296 | 
297 |     /* Process 4 items with each loop for efficiency. */
298 |     while (a < unroll_group) {
299 |       dot0 = a[0] * a[0];
300 |       dot1 = a[1] * a[1];
301 |       dot2 = a[2] * a[2];
302 |       dot3 = a[3] * a[3];
303 |       result += dot0 + dot1 + dot2 + dot3;
304 |       a += 4;
305 |     }
306 |     /* Process last 0-3 pixels.  Not needed for standard vector lengths. */
307 |     while (a < last) {
308 |       result += (*a) * (*a);
309 |       a++;
310 |     }
311 | #endif
312 | #endif
313 | #endif
314 |     return result;
315 |   }
316 |   using DistanceInnerProduct::compare;
317 |   float compare(const float *a, const float *b, float norm,
318 |                 unsigned size) const {  // not implement
319 |     float result = -2 * DistanceInnerProduct::compare(a, b, size);
320 |     result += norm;
321 |     return result;
322 |   }
323 | };
324 | }


--------------------------------------------------------------------------------
/src/index_mips.cpp:
--------------------------------------------------------------------------------
  1 | #include "mips/index_mips.h"
  2 | 
  3 | #include <omp.h>
  4 | #include <bitset>
  5 | #include <chrono>
  6 | #include <cmath>
  7 | #include <queue>
  8 | #include <boost/dynamic_bitset.hpp>
  9 | 
 10 | #include "ssg/parameters.h"
 11 | 
 12 | constexpr double kPi = 3.14159265358979323846264;
 13 | 
 14 | namespace efanna2e {
 15 | #define _CONTROL_NUM 100
 16 | 
 17 | IndexMips::IndexMips(const size_t dimension, const size_t n, Metric m,
 18 |                      Index *initializer)
 19 |     : Index(dimension, n, m), initializer_(initializer) {
 20 |     }
 21 | 
 22 | IndexMips::~IndexMips() {}
 23 | 
 24 | void IndexMips::Save(const char *filename) {
 25 |   std::ofstream out(filename, std::ios::binary | std::ios::out);
 26 |   assert(final_graph_.size() == nd_);
 27 |   out.write((char *)&width, sizeof(unsigned));
 28 |   unsigned n_ep=eps_.size();
 29 |   out.write((char *)&n_ep, sizeof(unsigned));
 30 |   out.write((char *)eps_.data(), n_ep*sizeof(unsigned));
 31 |   for (unsigned i = 0; i < nd_; i++) {
 32 |     unsigned GK = (unsigned)final_graph_[i].size();
 33 |     // std::cout << GK << std::endl;
 34 |     out.write((char *)&GK, sizeof(unsigned));
 35 |     out.write((char *)final_graph_[i].data(), GK * sizeof(unsigned));
 36 |   }
 37 |   out.close();
 38 | }
 39 | 
 40 | void IndexMips::Load(const char *filename) {
 41 |   /* width | eps_(vector<unsigned) entries | k (neighbor size) | k * unsigned (neighbors)*/
 42 |   std::ifstream in(filename, std::ios::binary);
 43 |   in.read((char *)&width, sizeof(unsigned));
 44 |   unsigned n_ep=0;
 45 |   in.read((char *)&n_ep, sizeof(unsigned));
 46 |   eps_.resize(n_ep);
 47 |   in.read((char *)eps_.data(), n_ep*sizeof(unsigned));
 48 |   unsigned cc = 0;
 49 |   while (!in.eof()) {
 50 |     unsigned k;
 51 |     in.read((char *)&k, sizeof(unsigned));
 52 |     if (in.eof()) break;
 53 |     cc += k;
 54 |     std::vector<unsigned> tmp(k);
 55 |     in.read((char *)tmp.data(), k * sizeof(unsigned));
 56 |     final_graph_.push_back(tmp);
 57 |   }
 58 |   cc /= nd_;
 59 |   DistanceFastL2 *norm_dis = (DistanceFastL2 *)distance_;
 60 |   norms_.resize(nd_);
 61 |   for (size_t i = 0; i < nd_; i++) {
 62 |     norms_[i] = std::sqrt(norm_dis->norm(data_ + dimension_ * i, dimension_)) ;
 63 |   }
 64 | 
 65 |   std::cerr << "Average Degree = " << cc << std::endl;
 66 | }
 67 | 
 68 | void IndexMips::Load_nn_graph(const char *filename) {
 69 |   /* k (k neighbor)| num * ( id + k * id)*/
 70 |   std::ifstream in(filename, std::ios::binary);
 71 |   unsigned k;
 72 |   std::cout << std::string(filename) << std::endl;
 73 |   in.read((char *)&k, sizeof(unsigned));
 74 |   std::cout << k << std::endl;
 75 |   in.seekg(0, std::ios::end);
 76 |   std::ios::pos_type ss = in.tellg();
 77 |   size_t fsize = (size_t)ss;
 78 |   size_t num = (unsigned)(fsize / (k + 1) / 4);
 79 |   in.seekg(0, std::ios::beg);
 80 |   std::cout << num << std::endl;
 81 |   final_graph_.resize(num);
 82 |   final_graph_.reserve(num);
 83 |   unsigned kk = (k + 3) / 4 * 4;
 84 |   for (size_t i = 0; i < num; i++) {
 85 |     in.seekg(4, std::ios::cur);
 86 |     final_graph_[i].resize(k);
 87 |     final_graph_[i].reserve(kk);
 88 |     in.read((char *)final_graph_[i].data(), k * sizeof(unsigned));
 89 |   }
 90 |   in.close();
 91 | }
 92 | 
 93 | /* select l neighbors candidate from 2-hop*/
 94 | void IndexMips::get_neighbors(const unsigned q, const Parameters &parameter,
 95 |                              std::vector<Neighbor> &pool, boost::dynamic_bitset<> &flags) {
 96 |   // boost::dynamic_bitset<> flags{nd_, 0};
 97 |   unsigned L = parameter.Get<unsigned>("L");
 98 |   flags[q] = true;
 99 |   for (unsigned i = 0; i < final_graph_[q].size(); i++) {
100 |     unsigned nid = final_graph_[q][i];
101 |     for (unsigned nn = 0; nn < final_graph_[nid].size(); nn++) {
102 |       unsigned nnid = final_graph_[nid][nn];
103 |       if (flags[nnid]) continue;
104 |       flags[nnid] = true;
105 |       // float dist = 0;
106 |       float dist = distance_->compare(data_ + dimension_ * q,
107 |                                       data_ + dimension_ * nnid, dimension_);
108 |       pool.push_back(Neighbor(nnid, dist, true));
109 |       if (pool.size() >= L) break;
110 |     }
111 |     if (pool.size() >= L) break;
112 |   }
113 | }
114 | 
115 | void IndexMips::get_neighbors(const float *query, const Parameters &parameter,
116 |                              std::vector<Neighbor> &retset,
117 |                              std::vector<Neighbor> &fullset) {
118 |   unsigned L = parameter.Get<unsigned>("L");
119 | 
120 |   retset.resize(L + 1);
121 |   std::vector<unsigned> init_ids(L);
122 |   std::mt19937 rng(rand());
123 |   GenRandom(rng, init_ids.data(), L, (unsigned)nd_);
124 | 
125 |   boost::dynamic_bitset<> flags{nd_, 0};
126 |   L = 0;
127 |   for (unsigned i = 0; i < init_ids.size(); i++) {
128 |     unsigned id = init_ids[i];
129 |     if (id >= nd_) continue;
130 |     float dist = distance_->compare(data_ + dimension_ * (size_t)id, query,
131 |                                     (unsigned)dimension_);
132 |     retset[i] = Neighbor(id, dist, true);
133 |     flags[id] = 1;
134 |     L++;
135 |   }
136 |   /* sort the inital return set*/
137 |   std::sort(retset.begin(), retset.begin() + L);
138 |   int k = 0;
139 |   while (k < (int)L) { /* until no new node*/
140 |     int nk = L;
141 | 
142 |     if (retset[k].flag) {
143 |       retset[k].flag = false;
144 |       unsigned n = retset[k].id;
145 | 
146 |       for (unsigned m = 0; m < final_graph_[n].size(); ++m) {
147 |         unsigned id = final_graph_[n][m];
148 |         if (flags[id]) continue;
149 |         flags[id] = 1;
150 | 
151 |         float dist = distance_->compare(query, data_ + dimension_ * (size_t)id,
152 |                                         (unsigned)dimension_);
153 |         Neighbor nn(id, dist, true);
154 |         fullset.push_back(nn);
155 |         if (dist >= retset[L - 1].distance) continue;
156 |         int r = InsertIntoPool(retset.data(), L, nn);
157 | 
158 |         if (L + 1 < retset.size()) ++L;
159 |         if (r < nk) nk = r;
160 |       }
161 |     }
162 |     if (nk <= k) /* back to new node's neighbors*/
163 |       k = nk;
164 |     else
165 |       ++k;
166 |   }
167 | }
168 | 
169 | void IndexMips::get_refine_neighbors(const unsigned ep, const float *query, const Parameters &parameter,
170 |                              std::vector<Neighbor> &retset,
171 |                              std::vector<Neighbor> &fullset) {
172 |   unsigned L = parameter.Get<unsigned>("M");
173 | 
174 |   retset.resize(L + 1);
175 |   std::vector<unsigned> init_ids(L);
176 |   // initializer_->Search(query, nullptr, L, parameter, init_ids.data());
177 | 
178 |   boost::dynamic_bitset<> flags{nd_, 0};
179 |   boost::dynamic_bitset<> visit_flags{nd_, 0};
180 | 
181 |   L = 0;
182 | 
183 |   for (unsigned i = 0; i < final_graph_[ep].size(); i++) {
184 |     visit_flags[final_graph_[ep][i]] = true;
185 |     float dist = distance_->compare(data_ + dimension_ * (size_t)final_graph_[ep][i], query,
186 |                                     (unsigned)dimension_);
187 |     fullset.push_back(Neighbor(final_graph_[ep][i], dist, true));
188 |   }
189 | 
190 |   for (unsigned i = 0; i < init_ids.size() && i < final_graph_[ep].size(); i++) {
191 |     init_ids[i] = final_graph_[ep][i];
192 |     flags[init_ids[i]] = true;
193 |     L++;
194 |   }
195 | 
196 |   while (L < init_ids.size()) {
197 |     unsigned id = rand() % nd_;
198 |     if (flags[id]) continue;
199 |     init_ids[L] = id;
200 |     L++;
201 |     flags[id] = true;
202 |   }
203 | 
204 |   for (unsigned i = 0; i < init_ids.size(); i++) {
205 |     unsigned id = init_ids[i];
206 |     if (id >= nd_) continue;
207 |     // std::cout<<id<<std::endl;
208 |     float dist = distance_->compare(data_ + dimension_ * (size_t)id, query,
209 |                                     (unsigned)dimension_);
210 |     retset[i] = Neighbor(id, dist, true);
211 |   }
212 |   /* sort the inital return set*/
213 |   std::sort(retset.begin(), retset.begin() + L);
214 |   int k = 0;
215 |   while (k < (int)L) {
216 |     int nk = L;
217 | 
218 |     if (retset[k].flag) {
219 |       retset[k].flag = false;
220 |       unsigned n = retset[k].id;
221 | 
222 |       for (unsigned m = 0; m < final_graph_[n].size(); ++m) {
223 |         unsigned id = final_graph_[n][m];
224 |         if (flags[id]) continue;
225 |         flags[id] = 1;
226 | 
227 |         float dist = distance_->compare(query, data_ + dimension_ * (size_t)id,
228 |                                         (unsigned)dimension_);
229 |         Neighbor nn(id, dist, true);
230 |         if (dist >= retset[L - 1].distance) continue;
231 | 
232 |         if (!visit_flags[id]) {
233 |             fullset.push_back(nn);
234 |         }
235 |         int r = InsertIntoPool(retset.data(), L, nn);
236 | 
237 |         if (L + 1 < retset.size()) ++L;
238 |         if (r < nk) nk = r;
239 |       }
240 |     }
241 |     if (nk <= k)
242 |       k = nk;
243 |     else
244 |       ++k;
245 |   }
246 | }
247 | 
248 | void IndexMips::get_mips_neighbors(const unsigned ep, const float *query, const Parameters &parameter,
249 |                              std::vector<IpNeighbor> &retset) {
250 |   unsigned L = parameter.Get<unsigned>("M");
251 |   float threshold = std::cos(40 / 180 * kPi);
252 |   retset.resize(L + 1);
253 |   std::vector<IpNeighbor> ipset;
254 |   ipset.resize(L + 1);
255 |   
256 |   std::vector<unsigned> init_ids(L);
257 |   boost::dynamic_bitset<> flags{nd_, 0};
258 | 
259 |   L = 0;
260 | 
261 |   DistanceInnerProduct *dis_inner = new DistanceInnerProduct();
262 | 
263 |   for (unsigned i = 0; i < init_ids.size() && i < final_graph_[ep].size(); i++) {
264 |     init_ids[i] = final_graph_[ep][i];
265 |     flags[init_ids[i]] = true;
266 |     L++;
267 |   }
268 | 
269 |   while (L < init_ids.size()) {
270 |     unsigned id = rand() % nd_;
271 |     if (flags[id]) continue;
272 |     init_ids[L] = id;
273 |     L++;
274 |     flags[id] = true;
275 |   }
276 |   int count = 0;
277 |   for (unsigned i = 0; i < init_ids.size(); i++) {
278 |     unsigned id = init_ids[i];
279 |     if (id >= nd_) continue;
280 |     // std::cout<<id<<std::endl;
281 |     float dist = dis_inner->compare(data_ + dimension_ * (size_t)id, query,
282 |                                     (unsigned)dimension_);
283 |     retset[i] = IpNeighbor(id, dist, true);
284 |     auto cos = dist / norms_[id] / norms_[ep];
285 |     if (cos >= threshold) {
286 |       ipset[count] = IpNeighbor(id, dist, true);
287 |       count++;
288 |     }
289 | 
290 |   }
291 |   /* sort the inital return set*/
292 |   std::sort(retset.begin(), retset.begin() + L);
293 |   std::sort(ipset.begin(), ipset.begin() + count);
294 |   int k = 0;
295 |   while (k < (int)L) {
296 |     int nk = L;
297 | 
298 |     if (retset[k].flag) {
299 |       retset[k].flag = false;
300 |       unsigned n = retset[k].id;
301 | 
302 |       for (unsigned m = 0; m < final_graph_[n].size(); ++m) {
303 |         unsigned id = final_graph_[n][m];
304 |         if (flags[id]) continue;
305 |         flags[id] = 1;
306 |     
307 |         float dist = dis_inner->compare(query, data_ + dimension_ * (size_t)id,
308 |                                         (unsigned)dimension_);
309 |         
310 |         IpNeighbor nn(id, dist, true);
311 |         if (dist <= retset[L - 1].distance) continue;
312 | 
313 |         int r = InsertIntoIpPool(retset.data(), L, nn);
314 | 
315 |         auto cos = dist / norms_[id] / norms_[ep];
316 |         if (cos >= threshold) {
317 |           int ir = InsertIntoIpPool(ipset.data(), count, nn);
318 |           if (count + 1 < ipset.size()) ++count;
319 |         }
320 | 
321 |         if (L + 1 < retset.size()) ++L;
322 |         if (r < nk) nk = r;
323 |       }
324 |     }
325 |     if (nk <= k)
326 |       k = nk;
327 |     else
328 |       ++k;
329 |   }
330 | }
331 | 
332 | 
333 | 
334 | /* find the entry node */
335 | void IndexMips::init_graph(const Parameters &parameters) {
336 |   float *center = new float[dimension_];
337 |   for (unsigned j = 0; j < dimension_; j++) center[j] = 0;
338 |   for (unsigned i = 0; i < nd_; i++) {
339 |     for (unsigned j = 0; j < dimension_; j++) {
340 |       center[j] += data_[i * dimension_ + j];
341 |     }
342 |   }
343 |   for (unsigned j = 0; j < dimension_; j++) {
344 |     center[j] /= nd_;
345 |   }
346 |   std::vector<Neighbor> tmp, pool;
347 |   get_neighbors(center, parameters, tmp, pool);
348 |   ep_ = tmp[0].id;  
349 |   std::cout << ep_ << std::endl;
350 |   DistanceFastL2 *norm_dis = (DistanceFastL2 *)distance_;
351 |   norms_.resize(nd_);
352 |   for (size_t i = 0; i < nd_; i++) {
353 |     norms_[i] = std::sqrt(norm_dis->norm(data_ + dimension_ * i, dimension_)) ;
354 |   }
355 | }
356 | 
357 | 
358 | void IndexMips::sync_prune(unsigned q, std::vector<Neighbor> &pool,
359 |                           const Parameters &parameters, float threshold,
360 |                           SimpleNeighbor *cut_graph_) {
361 |   unsigned range = parameters.Get<unsigned>("R"); /* max neighbor size (after pruning)*/
362 |   width = range;
363 |   unsigned start = 0;
364 | 
365 |   boost::dynamic_bitset<> flags{nd_, 0};
366 |   for (unsigned i = 0; i < pool.size(); ++i) {
367 |     flags[pool[i].id] = 1;
368 |   }
369 | 
370 |   // std::cout << pool.size() << std::endl;
371 |   for (unsigned nn = 0; nn < final_graph_[q].size(); nn++) {
372 |     unsigned id = final_graph_[q][nn];
373 |     if (flags[id]) continue;
374 |     float dist = distance_->compare(data_ + dimension_ * (size_t)q,
375 |                                     data_ + dimension_ * (size_t)id,
376 |                                     (unsigned)dimension_);
377 |     pool.push_back(Neighbor(id, dist, true));
378 |   }
379 | 
380 |   std::sort(pool.begin(), pool.end());
381 |   std::vector<Neighbor> result;
382 |   if (pool[start].id == q) start++;
383 |   result.push_back(pool[start]);
384 | 
385 |   auto norm = norms_[q];
386 |   while (result.size() < range && (++start) < pool.size()) {
387 |     auto &p = pool[start];
388 |     bool occlude = false;
389 |     auto p_norm = norms_[p.id];
390 |     for (unsigned t = 0; t < result.size(); t++) {
391 |       if (p.id == result[t].id) {
392 |         occlude = true;
393 |         break;
394 |       }
395 |       float djk = distance_->compare(data_ + dimension_ * (size_t)result[t].id,
396 |                                      data_ + dimension_ * (size_t)p.id,
397 |                                      (unsigned)dimension_);
398 |       float cos_ij = (p.distance + result[t].distance - djk) / 2 /
399 |                      sqrt(p.distance * result[t].distance);
400 |       if (cos_ij > threshold) {
401 |         occlude = true;
402 |         break;
403 |       }
404 |     }
405 |     if (!occlude) result.push_back(p);
406 |   }
407 | 
408 |   SimpleNeighbor *des_pool = cut_graph_ + (size_t)q * (size_t)range;
409 |   for (size_t t = 0; t < result.size(); t++) {
410 |     des_pool[t].id = result[t].id;
411 |     des_pool[t].distance = result[t].distance;
412 |   }
413 |   if (result.size() < range) {
414 |     des_pool[result.size()].distance = -1;
415 |   }
416 | }
417 | 
418 | void IndexMips::InterInsert(unsigned n, unsigned range, float threshold,
419 |                            std::vector<std::mutex> &locks,
420 |                            SimpleNeighbor *cut_graph_) {
421 |   SimpleNeighbor *src_pool = cut_graph_ + (size_t)n * (size_t)range;
422 |   for (size_t i = 0; i < range; i++) {
423 |     if (src_pool[i].distance == -1) break;
424 | 
425 |     SimpleNeighbor sn(n, src_pool[i].distance);
426 |     size_t des = src_pool[i].id;
427 |     SimpleNeighbor *des_pool = cut_graph_ + des * (size_t)range;
428 | 
429 |     std::vector<SimpleNeighbor> temp_pool;
430 |     int dup = 0;
431 |     {
432 |       LockGuard guard(locks[des]);
433 |       for (size_t j = 0; j < range; j++) {
434 |         if (des_pool[j].distance == -1) break;
435 |         if (n == des_pool[j].id) {
436 |           dup = 1;
437 |           break;
438 |         }
439 |         temp_pool.push_back(des_pool[j]);
440 |       }
441 |     }
442 |     if (dup) continue;
443 | 
444 |     temp_pool.push_back(sn);
445 | 
446 |     if (temp_pool.size() > range) {
447 |       std::vector<SimpleNeighbor> result;
448 |       unsigned start = 0;
449 |       std::sort(temp_pool.begin(), temp_pool.end());
450 |       result.push_back(temp_pool[start]);
451 |       while (result.size() < range && (++start) < temp_pool.size()) {
452 |         auto &p = temp_pool[start];
453 |         bool occlude = false;
454 |         for (unsigned t = 0; t < result.size(); t++) {
455 |           if (p.id == result[t].id) {
456 |             occlude = true;
457 |             break;
458 |           }
459 |           float djk = distance_->compare(
460 |               data_ + dimension_ * (size_t)result[t].id,
461 |               data_ + dimension_ * (size_t)p.id, (unsigned)dimension_);
462 |           float cos_ij = (p.distance + result[t].distance - djk) / 2 /
463 |                          sqrt(p.distance * result[t].distance);
464 |           if (cos_ij > threshold) {
465 |             occlude = true;
466 |             break;
467 |           }
468 |         }
469 |         if (!occlude) result.push_back(p);
470 |       }
471 |       {
472 |         LockGuard guard(locks[des]);
473 |         for (unsigned t = 0; t < result.size(); t++) {
474 |           des_pool[t] = result[t];
475 |         }
476 |         if (result.size() < range) {
477 |           des_pool[result.size()].distance = -1;
478 |         }
479 |       }
480 |     } else {
481 |       LockGuard guard(locks[des]);
482 |       for (unsigned t = 0; t < range; t++) {
483 |         if (des_pool[t].distance == -1) {
484 |           des_pool[t] = sn;
485 |           if (t + 1 < range) des_pool[t + 1].distance = -1;
486 |           break;
487 |         }
488 |       }
489 |     }
490 |   }
491 | }
492 | 
493 | void IndexMips::add_mips_neighbors(const unsigned n, const Parameters &parameters, std::vector<IpNeighbor> &pool, SimpleNeighbor *cut_graph_) {
494 |   unsigned range = parameters.Get<unsigned>("R");
495 |   SimpleNeighbor *des_pool = cut_graph_ + (size_t)n * (size_t)range;
496 |   int count = 0;
497 |   // std::cout << pool.size() << std::endl;
498 |   while (des_pool[count].distance != -1 && count < range) count++;
499 |   for (size_t t = 0; t < pool.size(); t++) {
500 |     if (count >= range) break;
501 |     des_pool[count].id = pool[t].id;
502 |     des_pool[count].distance = pool[t].distance;
503 |     count++;
504 |   }
505 |   if (count < range) {
506 |     des_pool[count].distance = -1;
507 |   }
508 | }
509 | 
510 | void IndexMips::Link(const Parameters &parameters, SimpleNeighbor *cut_graph_) {
511 |     unsigned range = parameters.Get<unsigned>("R");
512 |     std::vector<std::mutex> locks(nd_);
513 |     float angle = parameters.Get<float>("A");
514 |     float threshold = std::cos(angle / 180 * kPi);
515 | 
516 |     omp_set_num_threads(48);
517 |     #pragma omp parallel
518 |     {
519 |         std::vector<Neighbor> pool, tmp;
520 |         boost::dynamic_bitset<> flags{nd_, 0};
521 | 
522 |         #pragma omp for schedule(dynamic, 64)  
523 |         for (unsigned n = 0; n < nd_; ++n) {
524 |             flags.reset();
525 |             pool.clear();
526 |             tmp.clear();
527 |             get_neighbors(n, parameters, pool, flags);
528 |             sync_prune(n, pool, parameters, threshold, cut_graph_);
529 |         }
530 |     }
531 |     std::cout << "sync prune done!" << std::endl;
532 | 
533 |     omp_set_num_threads(48);
534 |     #pragma omp parallel
535 |     {
536 |         #pragma omp for schedule(dynamic, 64)  
537 |         for (unsigned n = 0; n < nd_; ++n) {
538 |             InterInsert(n, range, threshold, locks, cut_graph_);
539 |         }
540 |     }
541 |     std::cout << "interinsert done!" << std::endl;
542 | 
543 |     omp_set_num_threads(48);
544 |     #pragma omp parallel
545 |     {
546 |         std::vector<IpNeighbor> ip_pool;
547 | 
548 |         #pragma omp for schedule(dynamic, 64) 
549 |         for (unsigned n = 0; n < nd_; ++n) {
550 |             ip_pool.clear();
551 |             get_mips_neighbors(n, data_ + dimension_ * n, parameters, ip_pool);
552 |             add_mips_neighbors(n, parameters, ip_pool, cut_graph_);
553 |         }
554 |     }
555 |     std::cout << "add mips neighbors done!" << std::endl;
556 | }
557 | 
558 | void IndexMips::Build(size_t n, const float *data,
559 |                      const Parameters &parameters) {
560 |   std::string nn_graph_path = parameters.Get<std::string>("nn_graph_path");
561 |   unsigned range = parameters.Get<unsigned>("R");
562 |   Load_nn_graph(nn_graph_path.c_str());
563 |   std::cout <<"Load nn graph" << std::endl;
564 |   data_ = data;
565 |   init_graph(parameters);
566 |   std::cout << "Init done" << std::endl;
567 |   SimpleNeighbor *cut_graph_ = new SimpleNeighbor[nd_ * (size_t)range];
568 |   Link(parameters, cut_graph_);
569 |   std::cout << "Link done" << std::endl;
570 |   
571 |   final_graph_.resize(nd_);
572 | 
573 |   for (size_t i = 0; i < nd_; i++) {
574 |     SimpleNeighbor *pool = cut_graph_ + i * (size_t)range;
575 |     unsigned pool_size = 0;
576 |     for (unsigned j = 0; j < range; j++) {
577 |       if (pool[j].distance == -1) {
578 |         break;
579 |       }
580 |       pool_size = j;
581 |     }
582 |     ++pool_size;
583 |     final_graph_[i].resize(pool_size);
584 |     for (unsigned j = 0; j < pool_size; j++) {
585 |       final_graph_[i][j] = pool[j].id;
586 |     }
587 |   }
588 | 
589 |   // DFS_expand(parameters);
590 | 
591 | 
592 |   unsigned max, min, avg;
593 |   max = 0;
594 |   min = nd_;
595 |   avg = 0;
596 |   for (size_t i = 0; i < nd_; i++) {
597 |     auto size = final_graph_[i].size();
598 |     max = max < size ? size : max;
599 |     min = min > size ? size : min;
600 |     avg += size;
601 |   }
602 |   avg /= 1.0 * nd_;
603 |   printf("Degree Statistics: Max = %d, Min = %d, Avg = %d\n",
604 |          max, min, avg);
605 | 
606 |   has_built = true;
607 | }
608 | 
609 | void IndexMips::Search(const float *query, const float *x, size_t K,
610 |                       const Parameters &parameters, unsigned *indices) {
611 | }
612 | 
613 | 
614 | int IndexMips::Search_Mips_IP_Cal(const float *query, const float *x, size_t K,
615 |                       const Parameters &parameters, unsigned *indices, std::vector<int> &init_nodes) {
616 |   const unsigned L = parameters.Get<unsigned>("L_search");
617 |   auto dis_cal = 0;
618 |   data_ = x;
619 |   int update_position = 0;
620 |   std::vector<IpNeighbor> retset(L + 1);
621 |   std::vector<unsigned> init_ids(L);
622 |   boost::dynamic_bitset<> flags{nd_, 0};
623 |   std::mt19937 rng(rand());
624 |   GenRandom(rng, init_ids.data(), L, (unsigned)nd_);
625 |   // assert(eps_.size() < L);
626 |   for(unsigned i=0; i<eps_.size(); i++){
627 |     init_ids[i] = eps_[i];
628 |   }
629 |   
630 |   for (unsigned i = 0; i < init_nodes.size() && i < L; i++) {
631 |     init_ids[i] = init_nodes[i];
632 |   }
633 |   DistanceInnerProduct dist_inner;
634 |   for (unsigned i = 0; i < L; i++) {
635 |     unsigned id = init_ids[i];
636 |     float ip = dist_inner.compare(data_ + dimension_ * id, query,
637 |                                     dimension_);
638 |     retset[i] = IpNeighbor(id, ip, true);
639 |     flags[id] = true;
640 |   }
641 | 
642 |   std::sort(retset.begin(), retset.begin() + L);
643 |   
644 |   int k = 0;
645 |   while (k < (int)L) {
646 |     int nk = L;
647 |     if (retset[k].flag) {
648 |       retset[k].flag = false;
649 |       unsigned n = retset[k].id;
650 |       int search_pos = final_graph_[n].size();
651 |       for (unsigned m = 0; m < search_pos; ++m) {
652 |         unsigned id = final_graph_[n][m];
653 |         if (flags[id]) continue;
654 |         flags[id] = 1;        
655 |         float dist = dist_inner.compare(query, data_ + dimension_ * id,
656 |                                         dimension_); 
657 |         dis_cal += 1;
658 |         if (dist <= retset[L - 1].distance) continue;
659 |         IpNeighbor nn(id, dist, true);
660 |         int r = InsertIntoIpPool(retset.data(), L, nn);
661 |         if (r < nk) {
662 |           nk = r;
663 |         }
664 |       }
665 |     }
666 |     if (nk <= k)
667 |       k = nk;
668 |     else
669 |       ++k;
670 |   }
671 |   for (size_t i = 0; i < K; i++) {
672 |     indices[i] = retset[i].id;
673 |   }
674 |   return {dis_cal};
675 | }
676 | 
677 | int IndexMips::Search_Mips_IP_Cal_with_No_SN(const float *query, const float *x, size_t K,
678 |                       const Parameters &parameters, unsigned *indices) {
679 |   const unsigned L = parameters.Get<unsigned>("L_search");
680 |   auto dis_cal = 0;
681 |   data_ = x;
682 |   int update_position = 0;
683 |   std::vector<IpNeighbor> retset(L + 1);
684 |   std::vector<unsigned> init_ids(L);
685 |   boost::dynamic_bitset<> flags{nd_, 0};
686 |   std::mt19937 rng(rand());
687 |   GenRandom(rng, init_ids.data(), L, (unsigned)nd_);
688 |   // assert(eps_.size() < L);
689 |   for(unsigned i=0; i<eps_.size(); i++){
690 |     init_ids[i] = eps_[i];
691 |   }
692 |   
693 |   DistanceInnerProduct dist_inner;
694 |   for (unsigned i = 0; i < L; i++) {
695 |     unsigned id = init_ids[i];
696 |     float ip = dist_inner.compare(data_ + dimension_ * id, query,
697 |                                     dimension_);
698 |     retset[i] = IpNeighbor(id, ip, true);
699 |     flags[id] = true;
700 |   }
701 | 
702 |   std::sort(retset.begin(), retset.begin() + L);
703 |   
704 |   int k = 0;
705 |   while (k < (int)L) {
706 |     int nk = L;
707 |     if (retset[k].flag) {
708 |       retset[k].flag = false;
709 |       unsigned n = retset[k].id;
710 |       int search_pos = final_graph_[n].size();
711 |       for (unsigned m = 0; m < search_pos; ++m) {
712 |         unsigned id = final_graph_[n][m];
713 |         if (flags[id]) continue;
714 |         flags[id] = 1;        
715 |         float dist = dist_inner.compare(query, data_ + dimension_ * id,
716 |                                         dimension_); 
717 |         dis_cal += 1;
718 |         if (dist <= retset[L - 1].distance) continue;
719 |         IpNeighbor nn(id, dist, true);
720 |         int r = InsertIntoIpPool(retset.data(), L, nn);
721 |         if (r < nk) {
722 |           nk = r;
723 |         }
724 |       }
725 |     }
726 |     if (nk <= k)
727 |       k = nk;
728 |     else
729 |       ++k;
730 |   }
731 |   for (size_t i = 0; i < K; i++) {
732 |     indices[i] = retset[i].id;
733 |   }
734 |   return {dis_cal};
735 | }
736 | 
737 | 
738 | 
739 | void IndexMips::DFS(boost::dynamic_bitset<> &flag,
740 |                    std::vector<std::pair<unsigned, unsigned>> &edges,
741 |                    unsigned root, unsigned &cnt) {
742 |   unsigned tmp = root;
743 |   std::stack<unsigned> s;
744 |   s.push(root);
745 |   if (!flag[root]) cnt++;
746 |   flag[root] = true;
747 |   while (!s.empty()) {
748 |     unsigned next = nd_ + 1;
749 |     for (unsigned i = 0; i < final_graph_[tmp].size(); i++) {
750 |       if (flag[final_graph_[tmp][i]] == false) {
751 |         next = final_graph_[tmp][i];
752 |         break;
753 |       }
754 |     }
755 |     if (next == (nd_ + 1)) {
756 |       unsigned head = s.top();
757 |       s.pop();
758 |       if (s.empty()) break;
759 |       tmp = s.top();
760 |       unsigned tail = tmp;
761 |       if (check_edge(head, tail)) {
762 |         edges.push_back(std::make_pair(head, tail));
763 |       }
764 |       continue;
765 |     }
766 |     tmp = next;
767 |     flag[tmp] = true;
768 |     s.push(tmp);
769 |     cnt++;
770 |   }
771 | }
772 | 
773 | 
774 | 
775 | 
776 | int IndexMips::check_connected_component(){
777 |   int cnt = 0;
778 |   boost::dynamic_bitset<> flags{nd_, 0};
779 |   for (unsigned i = 0; i < nd_; i++) {
780 |     if (flags[i] == false) {
781 |       cnt++;
782 |       unsigned root = i;
783 |       std::queue<unsigned> myqueue;
784 |       myqueue.push(root);
785 |       while (!myqueue.empty()) {
786 |         unsigned q_front = myqueue.front();
787 |         myqueue.pop();
788 |         if (flags[q_front]) continue;
789 |         flags[q_front] = true;
790 |         for (unsigned j = 0; j < final_graph_[q_front].size(); j++) {
791 |           unsigned child = final_graph_[q_front][j];
792 |           if (flags[child]) continue;
793 |           myqueue.push(child);
794 |         }
795 |       }
796 | 
797 |     }
798 |   }
799 |   return cnt;
800 | }
801 | 
802 | void IndexMips::findroot(boost::dynamic_bitset<> &flag, unsigned &root,
803 |                         const Parameters &parameter) {
804 |   unsigned id = nd_;
805 |   for (unsigned i = 0; i < nd_; i++) {
806 |     if (flag[i] == false) {
807 |       id = i;
808 |       break;
809 |     }
810 |   }
811 | 
812 |   if (id == nd_) return;  // No Unlinked Node
813 | 
814 |   std::vector<Neighbor> tmp, pool;
815 |   get_neighbors(data_ + dimension_ * id, parameter, tmp, pool);
816 |   std::sort(pool.begin(), pool.end());
817 | 
818 |   bool found = false;
819 |   for (unsigned i = 0; i < pool.size(); i++) {
820 |     if (flag[pool[i].id]) {
821 |       // std::cout << pool[i].id << '\n';
822 |       root = pool[i].id;
823 |       found = true;
824 |       break;
825 |     }
826 |   }
827 |   if (!found) {
828 |     for (int retry = 0; retry < 1000; ++retry) {
829 |       unsigned rid = rand() % nd_;
830 |       if (flag[rid]) {
831 |         root = rid;
832 |         break;
833 |       }
834 |     }
835 |   }
836 |   final_graph_[root].push_back(id);
837 | }
838 | 
839 | bool IndexMips::check_edge(unsigned h, unsigned t) {
840 |   bool flag = true;
841 |   for (unsigned i = 0; i < final_graph_[h].size(); i++) {
842 |     if (t == final_graph_[h][i]) flag = false;
843 |   }
844 |   return flag;
845 | }
846 | 
847 | void IndexMips::strong_connect(const Parameters &parameter) {
848 |   unsigned n_try = parameter.Get<unsigned>("n_try");
849 |   std::vector<std::pair<unsigned, unsigned>> edges_all;
850 |   std::mutex edge_lock;
851 | 
852 | #pragma omp parallel for
853 |   for (unsigned nt = 0; nt < n_try; nt++) {
854 |     unsigned root = rand() % nd_;
855 |     boost::dynamic_bitset<> flags{nd_, 0};
856 |     unsigned unlinked_cnt = 0;
857 |     std::vector<std::pair<unsigned, unsigned>> edges;
858 | 
859 |     while (unlinked_cnt < nd_) {
860 |       DFS(flags, edges, root, unlinked_cnt);
861 |       if (unlinked_cnt >= nd_) break;
862 |       findroot(flags, root, parameter);
863 |     }
864 | 
865 |     LockGuard guard(edge_lock);
866 | 
867 |     for (unsigned i = 0; i < edges.size(); i++) {
868 |       edges_all.push_back(edges[i]);
869 |     }
870 |   }
871 |   unsigned ecnt = 0;
872 |   for (unsigned e = 0; e < edges_all.size(); e++) {
873 |     unsigned start = edges_all[e].first;
874 |     unsigned end = edges_all[e].second;
875 |     unsigned flag = 1;
876 |     for (unsigned j = 0; j < final_graph_[start].size(); j++) {
877 |       if (end == final_graph_[start][j]) {
878 |         flag = 0;
879 |       }
880 |     }
881 |     if (flag) {
882 |       final_graph_[start].push_back(end);
883 |       ecnt++;
884 |     }
885 |   }
886 |   for (size_t i = 0; i < nd_; ++i) {
887 |     if (final_graph_[i].size() > width) {
888 |       width = final_graph_[i].size();
889 |     }
890 |   }
891 | }
892 | 
893 | void IndexMips::DFS_expand(const Parameters &parameter) {
894 |   unsigned n_try = parameter.Get<unsigned>("n_try");
895 |   unsigned range = parameter.Get<unsigned>("R");
896 | 
897 |   std::vector<unsigned> ids(nd_);
898 |   for(unsigned i=0; i<nd_; i++){
899 |     ids[i]=i;
900 |   }
901 |   std::random_shuffle(ids.begin(), ids.end());
902 |   for(unsigned i=0; i<n_try; i++){
903 |     eps_.push_back(ids[i]);
904 |   }
905 | #pragma omp parallel for
906 |   for(unsigned i=0; i<n_try; i++){
907 |     unsigned rootid = eps_[i];
908 |     boost::dynamic_bitset<> flags{nd_, 0};
909 |     std::queue<unsigned> myqueue;
910 |     myqueue.push(rootid);
911 |     flags[rootid]=true;
912 |     std::vector<unsigned> uncheck_set(1);
913 |   
914 |     while(uncheck_set.size() >0){
915 |       while(!myqueue.empty()){
916 |         unsigned q_front=myqueue.front();
917 |         myqueue.pop();
918 | 
919 |         for(unsigned j=0; j<final_graph_[q_front].size(); j++){
920 |           unsigned child = final_graph_[q_front][j];
921 |           if(flags[child])continue;
922 |           flags[child] = true;
923 |           myqueue.push(child);
924 | 
925 |         }
926 |       }
927 | 
928 |       uncheck_set.clear();
929 |       for(unsigned j=0; j<nd_; j++){
930 |         if(flags[j]) continue;
931 |         uncheck_set.push_back(j);
932 |       }
933 |       if(uncheck_set.size()>0){
934 |         for(unsigned j=0; j<nd_; j++){
935 |           if(flags[j] && final_graph_[j].size()< (range + 1)){
936 |             final_graph_[j].push_back(uncheck_set[0]);
937 |             break;
938 |           }
939 |         }
940 |         myqueue.push(uncheck_set[0]);
941 |         flags[uncheck_set[0]]=true;
942 |       }
943 |     }
944 |   }
945 | }
946 | };
947 | 


--------------------------------------------------------------------------------