faiss_benchmark_test.cpp 23.0 KB
Newer Older
Y
yudong.cai 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <gtest/gtest.h>

Y
yudong.cai 已提交
20
#include <cassert>
Y
yudong.cai 已提交
21 22 23 24 25 26 27 28 29 30
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#include <faiss/AutoTune.h>
#include <faiss/Index.h>
#include <faiss/IndexIVF.h>
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexFlat.h>
31
#include <faiss/gpu/GpuIndexIVFSQHybrid.h>
Y
yudong.cai 已提交
32
#include <faiss/gpu/StandardGpuResources.h>
Y
yudong.cai 已提交
33 34 35 36 37 38 39 40 41
#include <faiss/index_io.h>
#include <faiss/utils.h>

#include <hdf5.h>

#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
Y
yudong.cai 已提交
42
#include <vector>
Y
yudong.cai 已提交
43 44 45 46 47 48 49

/*****************************************************
 * To run this test, please download the HDF5 from
 *  https://support.hdfgroup.org/ftp/HDF5/releases/
 * and install it to /usr/local/hdf5 .
 *****************************************************/

Y
yudong.cai 已提交
50 51
double
elapsed() {
Y
yudong.cai 已提交
52 53 54 55 56
    struct timeval tv;
    gettimeofday(&tv, nullptr);
    return tv.tv_sec + tv.tv_usec * 1e-6;
}

Y
yudong.cai 已提交
57 58
void
normalize(float* arr, size_t nq, size_t dim) {
Y
yudong.cai 已提交
59 60 61 62 63 64 65 66
    for (size_t i = 0; i < nq; i++) {
        double vecLen = 0.0;
        for (size_t j = 0; j < dim; j++) {
            double val = arr[i * dim + j];
            vecLen += val * val;
        }
        vecLen = std::sqrt(vecLen);
        for (size_t j = 0; j < dim; j++) {
Y
yudong.cai 已提交
67
            arr[i * dim + j] = (float)(arr[i * dim + j] / vecLen);
Y
yudong.cai 已提交
68 69 70 71
        }
    }
}

Y
yudong.cai 已提交
72 73 74 75 76 77 78 79 80 81 82 83 84 85
void*
hdf5_read(const char* file_name, const char* dataset_name, H5T_class_t dataset_class, size_t& d_out, size_t& n_out) {
    hid_t file, dataset, datatype, dataspace, memspace;
    H5T_class_t t_class;   /* data type class */
    H5T_order_t order;     /* data order */
    size_t size;           /* size of the data element stored in file */
    hsize_t dimsm[3];      /* memory space dimensions */
    hsize_t dims_out[2];   /* dataset dimensions */
    hsize_t count[2];      /* size of the hyperslab in the file */
    hsize_t offset[2];     /* hyperslab offset in the file */
    hsize_t count_out[3];  /* size of the hyperslab in memory */
    hsize_t offset_out[3]; /* hyperslab offset in memory */
    int rank;
    void* data_out; /* output buffer */
Y
yudong.cai 已提交
86 87 88 89 90 91 92 93 94

    /* Open the file and the dataset. */
    file = H5Fopen(file_name, H5F_ACC_RDONLY, H5P_DEFAULT);
    dataset = H5Dopen2(file, dataset_name, H5P_DEFAULT);

    /*
     * Get datatype and dataspace handles and then query
     * dataset class, order, size, rank and dimensions.
     */
Y
yudong.cai 已提交
95
    datatype = H5Dget_type(dataset); /* datatype handle */
Y
yudong.cai 已提交
96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
    t_class = H5Tget_class(datatype);
    assert(t_class == dataset_class || !"Illegal dataset class type");

    order = H5Tget_order(datatype);
    switch (order) {
        case H5T_ORDER_LE:
            printf("Little endian order \n");
            break;
        case H5T_ORDER_BE:
            printf("Big endian order \n");
            break;
        default:
            printf("Illegal endian order \n");
            break;
    }

Y
yudong.cai 已提交
112
    size = H5Tget_size(datatype);
Y
yudong.cai 已提交
113 114
    printf("Data size is %d \n", (int)size);

Y
yudong.cai 已提交
115 116
    dataspace = H5Dget_space(dataset); /* dataspace handle */
    rank = H5Sget_simple_extent_ndims(dataspace);
Y
yudong.cai 已提交
117 118 119 120 121 122 123
    H5Sget_simple_extent_dims(dataspace, dims_out, NULL);
    n_out = dims_out[0];
    d_out = dims_out[1];
    printf("rank %d, dimensions %lu x %lu \n", rank, n_out, d_out);

    /* Define hyperslab in the dataset. */
    offset[0] = offset[1] = 0;
Y
yudong.cai 已提交
124 125
    count[0] = dims_out[0];
    count[1] = dims_out[1];
Y
yudong.cai 已提交
126 127 128 129 130 131 132 133 134 135
    H5Sselect_hyperslab(dataspace, H5S_SELECT_SET, offset, NULL, count, NULL);

    /* Define the memory dataspace. */
    dimsm[0] = dims_out[0];
    dimsm[1] = dims_out[1];
    dimsm[2] = 1;
    memspace = H5Screate_simple(3, dimsm, NULL);

    /* Define memory hyperslab. */
    offset_out[0] = offset_out[1] = offset_out[2] = 0;
Y
yudong.cai 已提交
136 137 138
    count_out[0] = dims_out[0];
    count_out[1] = dims_out[1];
    count_out[2] = 1;
Y
yudong.cai 已提交
139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
    H5Sselect_hyperslab(memspace, H5S_SELECT_SET, offset_out, NULL, count_out, NULL);

    /* Read data from hyperslab in the file into the hyperslab in memory and display. */
    switch (t_class) {
        case H5T_INTEGER:
            data_out = new int[dims_out[0] * dims_out[1]];
            H5Dread(dataset, H5T_NATIVE_INT, memspace, dataspace, H5P_DEFAULT, data_out);
            break;
        case H5T_FLOAT:
            data_out = new float[dims_out[0] * dims_out[1]];
            H5Dread(dataset, H5T_NATIVE_FLOAT, memspace, dataspace, H5P_DEFAULT, data_out);
            break;
        default:
            printf("Illegal dataset class type\n");
            break;
    }

    /* Close/release resources. */
    H5Tclose(datatype);
    H5Dclose(dataset);
    H5Sclose(dataspace);
    H5Sclose(memspace);
    H5Fclose(file);

    return data_out;
}

Y
yudong.cai 已提交
166 167
std::string
get_index_file_name(const std::string& ann_test_name, const std::string& index_key, int32_t data_loops) {
Y
yudong.cai 已提交
168 169
    size_t pos = index_key.find_first_of(',', 0);
    std::string file_name = ann_test_name;
Y
yudong.cai 已提交
170
    file_name = file_name + "_" + index_key.substr(0, pos) + "_" + index_key.substr(pos + 1);
Y
yudong.cai 已提交
171 172 173 174
    file_name = file_name + "_" + std::to_string(data_loops) + ".index";
    return file_name;
}

Y
yudong.cai 已提交
175 176
bool
parse_ann_test_name(const std::string& ann_test_name, size_t& dim, faiss::MetricType& metric_type) {
Y
yudong.cai 已提交
177 178
    size_t pos1, pos2;

Y
yudong.cai 已提交
179 180
    if (ann_test_name.empty())
        return false;
Y
yudong.cai 已提交
181 182

    pos1 = ann_test_name.find_first_of('-', 0);
Y
yudong.cai 已提交
183 184
    if (pos1 == std::string::npos)
        return false;
Y
yudong.cai 已提交
185
    pos2 = ann_test_name.find_first_of('-', pos1 + 1);
Y
yudong.cai 已提交
186 187
    if (pos2 == std::string::npos)
        return false;
Y
yudong.cai 已提交
188

Y
yudong.cai 已提交
189 190
    dim = std::stoi(ann_test_name.substr(pos1 + 1, pos2 - pos1 - 1));
    std::string metric_str = ann_test_name.substr(pos2 + 1);
Y
yudong.cai 已提交
191 192 193 194 195 196 197 198 199 200 201
    if (metric_str == "angular") {
        metric_type = faiss::METRIC_INNER_PRODUCT;
    } else if (metric_str == "euclidean") {
        metric_type = faiss::METRIC_L2;
    } else {
        return false;
    }

    return true;
}

202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223
int32_t
GetResultHitCount(const faiss::Index::idx_t* ground_index, const faiss::Index::idx_t* index, size_t ground_k, size_t k,
                  size_t nq, int32_t index_add_loops) {
    assert(ground_k <= k);
    int hit = 0;
    for (int i = 0; i < nq; i++) {
        // count the num of results exist in ground truth result set
        // each result replicates INDEX_ADD_LOOPS times
        for (int j_c = 0; j_c < ground_k; j_c++) {
            int r_c = index[i * k + j_c];
            int j_g = 0;
            for (; j_g < ground_k / index_add_loops; j_g++) {
                if (ground_index[i * ground_k + j_g] == r_c) {
                    hit++;
                    continue;
                }
            }
        }
    }
    return hit;
}

Y
yudong.cai 已提交
224 225
void
test_ann_hdf5(const std::string& ann_test_name, const std::string& index_key, int32_t index_add_loops,
226
              const std::vector<size_t>& nprobes, int32_t search_loops) {
Y
yudong.cai 已提交
227 228 229 230 231 232 233 234 235 236 237 238
    double t0 = elapsed();

    const std::string ann_file_name = ann_test_name + ".hdf5";

    faiss::MetricType metric_type;
    size_t dim;

    if (!parse_ann_test_name(ann_test_name, dim, metric_type)) {
        printf("Invalid ann test name: %s\n", ann_test_name.c_str());
        return;
    }

Y
yudong.cai 已提交
239
    faiss::Index* index;
Y
yudong.cai 已提交
240 241 242 243 244 245
    size_t d;

    std::string index_file_name = get_index_file_name(ann_test_name, index_key, index_add_loops);
    try {
        index = faiss::read_index(index_file_name.c_str());
        d = dim;
Y
yudong.cai 已提交
246
    } catch (...) {
Y
yudong.cai 已提交
247 248
        printf("Cannot read index file: %s\n", index_file_name.c_str());

Y
yudong.cai 已提交
249
        printf("[%.3f s] Loading train set\n", elapsed() - t0);
Y
yudong.cai 已提交
250 251

        size_t nb;
Y
yudong.cai 已提交
252
        float* xb = (float*)hdf5_read(ann_file_name.c_str(), "train", H5T_FLOAT, d, nb);
Y
yudong.cai 已提交
253 254
        assert(d == dim || !"dataset does not have correct dimension");

Y
yudong.cai 已提交
255 256 257 258 259
        if (metric_type == faiss::METRIC_INNER_PRODUCT) {
            printf("[%.3f s] Normalizing data set \n", elapsed() - t0);
            normalize(xb, nb, d);
        }

Y
yudong.cai 已提交
260
        printf("[%.3f s] Preparing index \"%s\" d=%ld\n", elapsed() - t0, index_key.c_str(), d);
Y
yudong.cai 已提交
261 262 263

        index = faiss::index_factory(d, index_key.c_str(), metric_type);

Y
yudong.cai 已提交
264
        printf("[%.3f s] Training on %ld vectors\n", elapsed() - t0, nb);
Y
yudong.cai 已提交
265 266 267

        index->train(nb, xb);

Y
yudong.cai 已提交
268
        printf("[%.3f s] Loading database\n", elapsed() - t0);
Y
yudong.cai 已提交
269 270 271

        // add index multiple times to get ~1G data set
        for (int i = 0; i < index_add_loops; i++) {
Y
yudong.cai 已提交
272
            printf("[%.3f s] Indexing database, size %ld*%ld\n", elapsed() - t0, nb, d);
Y
yudong.cai 已提交
273 274 275 276 277
            index->add(nb, xb);
        }

        faiss::write_index(index, index_file_name.c_str());

Y
yudong.cai 已提交
278
        delete[] xb;
Y
yudong.cai 已提交
279 280 281
    }

    size_t nq;
Y
yudong.cai 已提交
282
    float* xq;
Y
yudong.cai 已提交
283
    {
Y
yudong.cai 已提交
284
        printf("[%.3f s] Loading queries\n", elapsed() - t0);
Y
yudong.cai 已提交
285 286 287 288 289 290

        size_t d2;
        xq = (float*)hdf5_read(ann_file_name.c_str(), "test", H5T_FLOAT, d2, nq);
        assert(d == d2 || !"query does not have same dimension as train set");
    }

Y
yudong.cai 已提交
291 292
    size_t k;                 // nb of results per query in the GT
    faiss::Index::idx_t* gt;  // nq * k matrix of ground-truth nearest-neighbors
Y
yudong.cai 已提交
293
    {
Y
yudong.cai 已提交
294
        printf("[%.3f s] Loading ground truth for %ld queries\n", elapsed() - t0, nq);
Y
yudong.cai 已提交
295 296 297

        // load ground-truth and convert int to long
        size_t nq2;
Y
yudong.cai 已提交
298
        int* gt_int = (int*)hdf5_read(ann_file_name.c_str(), "neighbors", H5T_INTEGER, k, nq2);
Y
yudong.cai 已提交
299 300 301
        assert(nq2 == nq || !"incorrect nb of ground truth entries");

        gt = new faiss::Index::idx_t[k * nq];
Y
yudong.cai 已提交
302
        for (int i = 0; i < k * nq; i++) {
Y
yudong.cai 已提交
303 304
            gt[i] = gt_int[i];
        }
Y
yudong.cai 已提交
305
        delete[] gt_int;
Y
yudong.cai 已提交
306 307 308 309 310 311 312 313 314 315 316
    }

    for (auto nprobe : nprobes) {
        faiss::ParameterSpace params;

        std::string nprobe_str = "nprobe=" + std::to_string(nprobe);
        params.set_index_parameters(index, nprobe_str.c_str());

        // output buffers
#if 1
        const size_t NQ = 1000, K = 1000;
Y
yudong.cai 已提交
317 318
        faiss::Index::idx_t* I = new faiss::Index::idx_t[NQ * K];
        float* D = new float[NQ * K];
Y
yudong.cai 已提交
319

Y
yudong.cai 已提交
320
        printf("\n%s | %s | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
321
        printf("======================================================================================\n");
Y
yudong.cai 已提交
322 323
        for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) {   // nq = {10, 100, 1000}
            for (size_t t_k = 100; t_k <= K; t_k *= 10) {  //  k = {100, 1000}
Y
yudong.cai 已提交
324 325 326
                faiss::indexIVF_stats.quantization_time = 0.0;
                faiss::indexIVF_stats.search_time = 0.0;

Y
yudong.cai 已提交
327
                double t_start = elapsed(), t_end;
328 329 330
                for (int i = 0; i < search_loops; i++) {
                    index->search(t_nq, xq, t_k, D, I);
                }
Y
yudong.cai 已提交
331 332 333
                t_end = elapsed();

                // k = 100 for ground truth
334 335 336 337 338 339
                int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);

                printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq, t_k,
                       (t_end - t_start) / search_loops, faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
                       faiss::indexIVF_stats.search_time / 1000 / search_loops,
                       (hit / float(t_nq * k / index_add_loops)));
Y
yudong.cai 已提交
340 341
            }
        }
342
        printf("======================================================================================\n");
Y
yudong.cai 已提交
343
#else
Y
yudong.cai 已提交
344
        printf("[%.3f s] Perform a search on %ld queries\n", elapsed() - t0, nq);
Y
yudong.cai 已提交
345

Y
yudong.cai 已提交
346 347
        faiss::Index::idx_t* I = new faiss::Index::idx_t[nq * k];
        float* D = new float[nq * k];
Y
yudong.cai 已提交
348 349 350

        index->search(nq, xq, k, D, I);

Y
yudong.cai 已提交
351
        printf("[%.3f s] Compute recalls\n", elapsed() - t0);
Y
yudong.cai 已提交
352 353 354

        // evaluate result by hand.
        int n_1 = 0, n_10 = 0, n_100 = 0;
Y
yudong.cai 已提交
355
        for (int i = 0; i < nq; i++) {
Y
yudong.cai 已提交
356
            int gt_nn = gt[i * k];
Y
yudong.cai 已提交
357
            for (int j = 0; j < k; j++) {
Y
yudong.cai 已提交
358
                if (I[i * k + j] == gt_nn) {
Y
yudong.cai 已提交
359 360 361 362 363 364
                    if (j < 1)
                        n_1++;
                    if (j < 10)
                        n_10++;
                    if (j < 100)
                        n_100++;
Y
yudong.cai 已提交
365 366 367 368 369 370 371 372
                }
            }
        }
        printf("R@1 = %.4f\n", n_1 / float(nq));
        printf("R@10 = %.4f\n", n_10 / float(nq));
        printf("R@100 = %.4f\n", n_100 / float(nq));
#endif

Y
yudong.cai 已提交
373
        printf("[%.3f s] Search test done\n\n", elapsed() - t0);
Y
yudong.cai 已提交
374

Y
yudong.cai 已提交
375 376
        delete[] I;
        delete[] D;
Y
yudong.cai 已提交
377 378
    }

Y
yudong.cai 已提交
379 380
    delete[] xq;
    delete[] gt;
Y
yudong.cai 已提交
381 382 383 384
    delete index;
}

#ifdef CUSTOMIZATION
Y
yudong.cai 已提交
385
void
386 387
test_ivfsq8h(const std::string& ann_test_name, int32_t index_add_loops, const std::vector<size_t>& nprobes,
             bool pure_gpu_mode, int32_t search_loops) {
Y
yudong.cai 已提交
388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408
    double t0 = elapsed();

    const std::string ann_file_name = ann_test_name + ".hdf5";

    faiss::MetricType metric_type;
    size_t dim;

    if (!parse_ann_test_name(ann_test_name, dim, metric_type)) {
        printf("Invalid ann test name: %s\n", ann_test_name.c_str());
        return;
    }

    faiss::distance_compute_blas_threshold = 800;
    faiss::gpu::StandardGpuResources res;

    const std::string index_key = "IVF16384,SQ8Hybrid";

    faiss::Index* cpu_index = nullptr;
    size_t d;

    std::string index_file_name = get_index_file_name(ann_test_name, index_key, index_add_loops);
Y
yudong.cai 已提交
409
    try {
Y
yudong.cai 已提交
410 411
        cpu_index = faiss::read_index(index_file_name.c_str());
        d = dim;
Y
yudong.cai 已提交
412
    } catch (...) {
Y
yudong.cai 已提交
413 414
        printf("Cannot read index file: %s\n", index_file_name.c_str());

Y
yudong.cai 已提交
415
        printf("[%.3f s] Loading train set\n", elapsed() - t0);
Y
yudong.cai 已提交
416

Y
yudong.cai 已提交
417 418
        size_t nb;
        float* xb = (float*)hdf5_read(ann_file_name.c_str(), "train", H5T_FLOAT, d, nb);
Y
yudong.cai 已提交
419 420
        assert(d == dim || !"dataset does not have correct dimension");

Y
yudong.cai 已提交
421
        printf("[%.3f s] Preparing index \"%s\" d=%ld\n", elapsed() - t0, index_key.c_str(), d);
Y
yudong.cai 已提交
422

Y
yudong.cai 已提交
423
        faiss::Index* ori_index = faiss::index_factory(d, index_key.c_str(), metric_type);
Y
yudong.cai 已提交
424 425 426

        auto device_index = faiss::gpu::index_cpu_to_gpu(&res, 0, ori_index);

Y
yudong.cai 已提交
427
        printf("[%.3f s] Training on %ld vectors\n", elapsed() - t0, nb);
Y
yudong.cai 已提交
428 429 430

        device_index->train(nb, xb);

Y
yudong.cai 已提交
431
        printf("[%.3f s] Loading database\n", elapsed() - t0);
Y
yudong.cai 已提交
432 433

        for (int i = 0; i < index_add_loops; i++) {
Y
yudong.cai 已提交
434
            printf("[%.3f s] Indexing database, size %ld*%ld\n", elapsed() - t0, nb, d);
Y
yudong.cai 已提交
435 436 437 438 439 440
            device_index->add(nb, xb);
        }

        cpu_index = faiss::gpu::index_gpu_to_cpu(device_index);
        faiss::write_index(cpu_index, index_file_name.c_str());

Y
yudong.cai 已提交
441
        delete[] xb;
Y
yudong.cai 已提交
442 443
    }

Y
yudong.cai 已提交
444 445
    faiss::IndexIVF* cpu_ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
    if (cpu_ivf_index != nullptr) {
Y
yudong.cai 已提交
446 447 448 449
        cpu_ivf_index->to_readonly();
    }

    size_t nq;
Y
yudong.cai 已提交
450
    float* xq;
Y
yudong.cai 已提交
451
    {
Y
yudong.cai 已提交
452
        printf("[%.3f s] Loading queries\n", elapsed() - t0);
Y
yudong.cai 已提交
453 454 455 456 457 458 459

        size_t d2;
        xq = (float*)hdf5_read(ann_file_name.c_str(), "test", H5T_FLOAT, d2, nq);
        assert(d == d2 || !"query does not have same dimension as train set");
    }

    size_t k;
Y
yudong.cai 已提交
460
    faiss::Index::idx_t* gt;
Y
yudong.cai 已提交
461
    {
Y
yudong.cai 已提交
462
        printf("[%.3f s] Loading ground truth for %ld queries\n", elapsed() - t0, nq);
Y
yudong.cai 已提交
463 464

        size_t nq2;
Y
yudong.cai 已提交
465
        int* gt_int = (int*)hdf5_read(ann_file_name.c_str(), "neighbors", H5T_INTEGER, k, nq2);
Y
yudong.cai 已提交
466 467 468
        assert(nq2 == nq || !"incorrect nb of ground truth entries");

        gt = new faiss::Index::idx_t[k * nq];
469
        for (uint64_t i = 0; i < k * nq; ++i) {
Y
yudong.cai 已提交
470 471
            gt[i] = gt_int[i];
        }
Y
yudong.cai 已提交
472
        delete[] gt_int;
Y
yudong.cai 已提交
473 474
    }

H
Heisenberg 已提交
475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504
    faiss::gpu::GpuClonerOptions option;
    option.allInGpu = true;

    faiss::IndexComposition index_composition;
    index_composition.index = cpu_index;
    index_composition.quantizer = nullptr;

    faiss::Index* index;
    double copy_time;

    if (!pure_gpu_mode) {
        index_composition.mode = 1;  // 0: all data, 1: copy quantizer, 2: copy data
        index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
        delete index;

        copy_time = elapsed();
        index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
        delete index;
    } else {
        index_composition.mode = 2;
        index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
        delete index;

        copy_time = elapsed();
        index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
    }

    copy_time = elapsed() - copy_time;
    printf("[%.3f s] Copy quantizer completed, cost %f s\n", elapsed() - t0, copy_time);

505 506 507 508 509
    const size_t NQ = 1000, K = 1000;
    if (!pure_gpu_mode) {
        for (auto nprobe : nprobes) {
            auto ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
            ivf_index->nprobe = nprobe;
Y
yudong.cai 已提交
510

511 512 513 514 515
            auto is_gpu_flat_index = dynamic_cast<faiss::gpu::GpuIndexFlat*>(ivf_index->quantizer);
            if (is_gpu_flat_index == nullptr) {
                delete ivf_index->quantizer;
                ivf_index->quantizer = index_composition.quantizer;
            }
Y
yudong.cai 已提交
516

517 518
            int64_t* I = new faiss::Index::idx_t[NQ * K];
            float* D = new float[NQ * K];
Y
yudong.cai 已提交
519

520 521 522 523 524 525
            printf("\n%s | %s-MIX | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
            printf("======================================================================================\n");
            for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) {   // nq = {10, 100, 1000}
                for (size_t t_k = 100; t_k <= K; t_k *= 10) {  //  k = {100, 1000}
                    faiss::indexIVF_stats.quantization_time = 0.0;
                    faiss::indexIVF_stats.search_time = 0.0;
Y
yudong.cai 已提交
526

527 528 529 530 531
                    double t_start = elapsed(), t_end;
                    for (int32_t i = 0; i < search_loops; i++) {
                        cpu_index->search(t_nq, xq, t_k, D, I);
                    }
                    t_end = elapsed();
Y
yudong.cai 已提交
532

533 534
                    // k = 100 for ground truth
                    int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);
Y
yudong.cai 已提交
535

536 537 538 539 540 541 542 543
                    printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq,
                           t_k, (t_end - t_start) / search_loops,
                           faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
                           faiss::indexIVF_stats.search_time / 1000 / search_loops,
                           (hit / float(t_nq * k / index_add_loops)));
                }
            }
            printf("======================================================================================\n");
Y
yudong.cai 已提交
544

545
            printf("[%.3f s] Search test done\n\n", elapsed() - t0);
Y
yudong.cai 已提交
546

547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570
            delete[] I;
            delete[] D;
        }
    } else {
        std::shared_ptr<faiss::Index> gpu_index_ivf_ptr = std::shared_ptr<faiss::Index>(index);

        for (auto nprobe : nprobes) {
            faiss::gpu::GpuIndexIVFSQHybrid* gpu_index_ivf_hybrid =
                dynamic_cast<faiss::gpu::GpuIndexIVFSQHybrid*>(gpu_index_ivf_ptr.get());
            gpu_index_ivf_hybrid->setNumProbes(nprobe);

            int64_t* I = new faiss::Index::idx_t[NQ * K];
            float* D = new float[NQ * K];

            printf("\n%s | %s-GPU | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
            printf("======================================================================================\n");
            for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) {   // nq = {10, 100, 1000}
                for (size_t t_k = 100; t_k <= K; t_k *= 10) {  //  k = {100, 1000}
                    faiss::indexIVF_stats.quantization_time = 0.0;
                    faiss::indexIVF_stats.search_time = 0.0;

                    double t_start = elapsed(), t_end;
                    for (int32_t i = 0; i < search_loops; i++) {
                        gpu_index_ivf_ptr->search(nq, xq, k, D, I);
Y
yudong.cai 已提交
571
                    }
572 573 574 575 576 577 578 579 580 581
                    t_end = elapsed();

                    // k = 100 for ground truth
                    int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);

                    printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq,
                           t_k, (t_end - t_start) / search_loops,
                           faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
                           faiss::indexIVF_stats.search_time / 1000 / search_loops,
                           (hit / float(t_nq * k / index_add_loops)));
Y
yudong.cai 已提交
582 583
                }
            }
584
            printf("======================================================================================\n");
Y
yudong.cai 已提交
585

586
            printf("[%.3f s] Search test done\n\n", elapsed() - t0);
Y
yudong.cai 已提交
587

588 589 590
            delete[] I;
            delete[] D;
        }
Y
yudong.cai 已提交
591 592
    }

Y
yudong.cai 已提交
593 594
    delete[] xq;
    delete[] gt;
Y
yudong.cai 已提交
595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610
    delete cpu_index;
}
#endif

/************************************************************************************
 * https://github.com/erikbern/ann-benchmarks
 *
 * Dataset 	Dimensions 	Train_size 	Test_size 	Neighbors 	Distance 	Download
 * Fashion-
 *  MNIST   784         60,000      10,000 	    100         Euclidean   HDF5 (217MB)
 * GIST     960         1,000,000   1,000       100         Euclidean   HDF5 (3.6GB)
 * GloVe    100         1,183,514   10,000      100         Angular     HDF5 (463MB)
 * GloVe    200         1,183,514   10,000      100         Angular     HDF5 (918MB)
 * MNIST    784         60,000 	    10,000      100         Euclidean   HDF5 (217MB)
 * NYTimes  256         290,000     10,000      100         Angular     HDF5 (301MB)
 * SIFT     128         1,000,000   10,000      100         Euclidean   HDF5 (501MB)
Y
yudong.cai 已提交
611
 *************************************************************************************/
Y
yudong.cai 已提交
612

Y
yudong.cai 已提交
613
TEST(FAISSTEST, BENCHMARK) {
614 615 616 617 618 619 620
    std::vector<size_t> param_nprobes = {8, 128};
    const int32_t SEARCH_LOOPS = 5;
    const int32_t SIFT_INSERT_LOOPS = 2;  // insert twice to get ~1G data set
    const int32_t GLOVE_INSERT_LOOPS = 1;

    test_ann_hdf5("sift-128-euclidean", "IVF4096,Flat", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
    test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
Y
yudong.cai 已提交
621
#ifdef CUSTOMIZATION
622 623 624
    test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8Hybrid", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
    test_ivfsq8h("sift-128-euclidean", SIFT_INSERT_LOOPS, param_nprobes, false, SEARCH_LOOPS);
    test_ivfsq8h("sift-128-euclidean", SIFT_INSERT_LOOPS, param_nprobes, true, SEARCH_LOOPS);
Y
yudong.cai 已提交
625 626
#endif

627 628
    test_ann_hdf5("glove-200-angular", "IVF4096,Flat", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
    test_ann_hdf5("glove-200-angular", "IVF16384,SQ8", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
Y
yudong.cai 已提交
629
#ifdef CUSTOMIZATION
630 631 632
    test_ann_hdf5("glove-200-angular", "IVF16384,SQ8Hybrid", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
    test_ivfsq8h("glove-200-angular", GLOVE_INSERT_LOOPS, param_nprobes, false, SEARCH_LOOPS);
    test_ivfsq8h("glove-200-angular", GLOVE_INSERT_LOOPS, param_nprobes, true, SEARCH_LOOPS);
Y
yudong.cai 已提交
633 634
#endif
}