CC-bench-tools/cc_bench.cpp at main · DataSysResearch/CC-bench-tools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#include <getopt.h>

#include <cstring>
#include <iostream>
#include <string>

#include "algorithms/hnsw.hpp"
#include "bench.hpp"
#include "utils.hpp"

int main(int argc, char *argv[]) {
    int retval = PAPI_library_init(PAPI_VER_CURRENT);
    if (retval != PAPI_VER_CURRENT) {
        std::cerr << "PAPI library init error: " << PAPI_strerror(retval)
                  << " (retval=" << retval << ")" << std::endl;
        exit(1);
    }

    std::string data_type, data_path, query_path, batch_res_path, gt_path,
        index_name, dataset_name, stat_path;
    size_t begin_num = 5000, batch_size = 100;
    float write_ratio = 0.5;
    size_t recall_at = 10, R = 16, Ls = 50, Lb = 50,
           num_threads = std::thread::hardware_concurrency();

    struct option long_options[] = {{"dataset_name", required_argument, 0, 0},
                                    {"data_type", required_argument, 0, 0},
                                    {"data_path", required_argument, 0, 0},
                                    {"query_path", required_argument, 0, 0},
                                    {"batch_res_path", required_argument, 0, 0},
                                    {"begin_num", required_argument, 0, 0},
                                    {"max_elements", required_argument, 0, 0},
                                    {"write_ratio", required_argument, 0, 0},
                                    {"batch_size", required_argument, 0, 0},
                                    {"recall_at", required_argument, 0, 0},
                                    {"R", required_argument, 0, 0},
                                    {"Lb", required_argument, 0, 0},
                                    {"Ls", required_argument, 0, 0},
                                    {"dim", required_argument, 0, 0},
                                    {"num_threads", required_argument, 0, 0},
                                    {"gt_path", required_argument, 0, 0},
                                    {"stat_path", required_argument, 0, 0},
                                    {0, 0, 0, 0}};

    int option_index = 0;
    int c;
    while ((c = getopt_long(argc, argv, "", long_options, &option_index)) !=
           -1) {
        if (c == 0) {
            std::string opt_name = long_options[option_index].name;
            if (opt_name == "dataset_name")
                dataset_name = optarg;
            else if (opt_name == "data_type")
                data_type = optarg;
            else if (opt_name == "data_path")
                data_path = optarg;
            else if (opt_name == "query_path")
                query_path = optarg;
            else if (opt_name == "batch_res_path")
                batch_res_path = optarg;
            else if (opt_name == "begin_num")
                begin_num = std::stoul(optarg);
            else if (opt_name == "write_ratio")
                write_ratio = std::stof(optarg);
            else if (opt_name == "batch_size")
                batch_size = std::stoul(optarg);
            else if (opt_name == "recall_at")
                recall_at = std::stoul(optarg);
            else if (opt_name == "R")
                R = std::stoul(optarg);
            else if (opt_name == "Lb")
                Lb = std::stoul(optarg);
            else if (opt_name == "Ls")
                Ls = std::stoul(optarg);
            else if (opt_name == "num_threads")
                num_threads = std::stoul(optarg);
            else if (opt_name == "gt_path")
                gt_path = optarg;
            else if (opt_name == "stat_path")
                stat_path = optarg;
        }
    }

    using TagT = uint32_t;
    using LabelT = uint32_t;
    std::vector<SearchResult<TagT>> search_results;

    size_t data_num, data_dim, aligned_dim;
    get_bin_metadata(data_path, data_num, data_dim);
    search_results.reserve(data_num * (1 / write_ratio - 1));

    Stat stat("HNSW", dataset_name, R, Lb, Ls, write_ratio, num_threads,
              batch_size, batch_res_path);

    if (data_type == "float") {
        using IndexType = HNSW<float, TagT, LabelT>;
        std::unique_ptr<IndexBase<float, TagT, LabelT>> index(
            new IndexType(data_dim, data_num, R, Lb));

        measure_performance(
            [&]() {
                concurrent_bench<float, TagT, LabelT>(
                    data_path, query_path, begin_num, write_ratio, batch_size,
                    recall_at, Ls, num_threads, std::move(index),
                    search_results, stat);
            },
            true);

        overall_recall<float, TagT, LabelT>(query_path, recall_at, Ls,
                                            std::move(index), gt_path);
    } else if (data_type == "int8_t") {
    } else if (data_type == "uint8_t") {
    } else {
        std::cerr << "Unknown data type: " << data_type << "\n";
        return 1;
    }

    save_stat(stat, stat_path);
    write_results(search_results, stat.stagewise_result_path);

    return 0;
}