C++ AI异构搜索

GitHub - facebookresearch/faiss: A library for efficient similarity search and clustering of dense vectors.

#include <faiss/utils/simdlib.h>

#include <cstddef>
#include <cstdint>
#include <memory>
#include <random>
#include <vector>

#include <omp.h>

#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQFastScan.h>
#include <faiss/impl/AuxIndexStructures.h>

#include <faiss/IndexFlat.h>
//#include <faiss/gpu/GpuIndexFlat.h>
//#include <faiss/gpu/GpuIndexIVFFlat.h>
//#include <faiss/gpu/StandardGpuResources.h>


using namespace faiss;


void testCmpltAndBlendInplace() {
    simd8float32 lowestValues(0, 1, 2, 3, 4, 5, 6, 7);
    simd8uint32 lowestIndices(0, 1, 2, 3, 4, 5, 6, 7);

    simd8float32 candidateValues0(5, 5, 5, 5, 5, 5, 5, 5);
    simd8uint32 candidateIndices0(10, 11, 12, 13, 14, 15, 16, 17);
    cmplt_and_blend_inplace(
        candidateValues0, candidateIndices0, lowestValues, lowestIndices);

    simd8float32 candidateValues1(6, 6, 6, 6, 6, 6, 6, 6);
    simd8uint32 candidateIndices1(20, 21, 22, 23, 24, 25, 26, 27);
    cmplt_and_blend_inplace(
        candidateValues1, candidateIndices1, lowestValues, lowestIndices);

    simd8float32 candidateValues2(0, 1, 2, 3, 4, 5, 5, 5);
    simd8uint32 candidateIndices2(30, 31, 32, 33, 34, 35, 36, 37);
    cmplt_and_blend_inplace(
        candidateValues2, candidateIndices2, lowestValues, lowestIndices);

    simd8float32 expectedValues(0, 1, 2, 3, 4, 5, 5, 5);
    simd8uint32 expectedIndices(0, 1, 2, 3, 4, 5, 16, 17);
    //ASSERT_TRUE(lowestValues.is_same_as(expectedValues));
    //ASSERT_TRUE(lowestIndices.is_same_as(expectedIndices));
}

void testCmpltMinMaxFloat() {
    simd8float32 minValues(0, 0, 0, 0, 0, 0, 0, 0);
    simd8uint32 minIndices(0, 0, 0, 0, 0, 0, 0, 0);
    simd8float32 maxValues(0, 0, 0, 0, 0, 0, 0, 0);
    simd8uint32 maxIndices(0, 0, 0, 0, 0, 0, 0, 0);

    simd8float32 candidateValues0(5, 5, 5, 5, 5, 5, 5, 5);
    simd8uint32 candidateIndices0(10, 11, 12, 13, 14, 15, 16, 17);
    simd8float32 currentValues0(0, 1, 2, 3, 4, 5, 6, 7);
    simd8uint32 currentIndices0(0, 1, 2, 3, 4, 5, 6, 7);

    cmplt_min_max_fast(
        candidateValues0,
        candidateIndices0,
        currentValues0,
        currentIndices0,
        minValues,
        minIndices,
        maxValues,
        maxIndices);

    simd8float32 expectedMinValues(0, 1, 2, 3, 4, 5, 5, 5);
    simd8uint32 expectedMinIndices(0, 1, 2, 3, 4, 5, 16, 17);
    //ASSERT_TRUE(minValues.is_same_as(expectedMinValues));
    //ASSERT_TRUE(minIndices.is_same_as(expectedMinIndices));

    simd8float32 expectedMaxValues(5, 5, 5, 5, 5, 5, 6, 7);
    // the result is not 10,11,12,13,14,5,6,7 because it is _fast version
    simd8uint32 expectedMaxIndices(10, 11, 12, 13, 14, 15, 6, 7);
    //ASSERT_TRUE(maxValues.is_same_as(expectedMaxValues));
    //ASSERT_TRUE(maxIndices.is_same_as(expectedMaxIndices));
}

void testSearch() {
    // small vectors and database
    int d = 64;
    size_t nb = 4000;

    // ivf centroids
    size_t nlist = 4;

    // more than 2 threads to surface
    // problems related to multi-threading
    omp_set_num_threads(8);

    // random database, also used as queries
    std::vector<float> database(nb * d);
    std::mt19937 rng;
    std::uniform_real_distribution<> distrib;
    for (size_t i = 0; i < nb * d; i++) {
        database[i] = distrib(rng);
    }

    // build index
    faiss::IndexFlatL2 coarse_quantizer(d);
    faiss::IndexIVFPQFastScan index(
        &coarse_quantizer, d, nlist, d / 2, 4, faiss::METRIC_L2, 32);
    index.pq.cp.niter = 10; // speed up train
    index.nprobe = nlist;
    index.train(nb, database.data());
    index.add(nb, database.data());

    std::vector<float> distances(nb);
    std::vector<faiss::idx_t> labels(nb);
    auto t = std::chrono::high_resolution_clock::now();
    int k = 1;
    index.nprobe = k;
    index.search(nb, database.data(), k, distances.data(), labels.data());
    auto knn_time = std::chrono::high_resolution_clock::now() - t;

    int n_ok = 0;
    for (int q = 0; q < nb; q++) {
        for (int i = 0; i < k; i++)
            if (database[q * k + i] == distances[q])
                n_ok++;
    }
    //EXPECT_GT(n_ok, nb * 0.4);
    //faiss::RangeSearchResult rsr(nb);
    //t = std::chrono::high_resolution_clock::now();
    //index.range_search(nb, database.data(), 1.0, &rsr);
    //auto range_time = std::chrono::high_resolution_clock::now() - t;

     we expect the perf of knn and range search
     to be similar, at least within a factor of 4
    //ASSERT_LE(range_time, knn_time * 4);
    //ASSERT_LE(knn_time, range_time * 4);
}

/*
void testGPU() {
    int d = 64;      // dimension
    int nb = 100000; // database size
    int nq = 10000;  // nb of queries

    std::mt19937 rng;
    std::uniform_real_distribution<> distrib;

    float* xb = new float[d * nb];
    float* xq = new float[d * nq];

    for (int i = 0; i < nb; i++) {
        for (int j = 0; j < d; j++)
            xb[d * i + j] = distrib(rng);
        xb[d * i] += i / 1000.;
    }

    for (int i = 0; i < nq; i++) {
        for (int j = 0; j < d; j++)
            xq[d * i + j] = distrib(rng);
        xq[d * i] += i / 1000.;
    }

    faiss::gpu::StandardGpuResources res;

    // Using a flat index

    faiss::gpu::GpuIndexFlatL2 index_flat(&res, d);

    printf("is_trained = %s\n", index_flat.is_trained ? "true" : "false");
    index_flat.add(nb, xb); // add vectors to the index
    printf("ntotal = %ld\n", index_flat.ntotal);

    int k = 4;

    { // search xq
        long* I = new long[k * nq];
        float* D = new float[k * nq];

        index_flat.search(nq, xq, k, D, I);

        // print results
        printf("I (5 first results)=\n");
        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        printf("I (5 last results)=\n");
        for (int i = nq - 5; i < nq; i++) {
            for (int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

    // Using an IVF index

    int nlist = 100;
    faiss::gpu::GpuIndexIVFFlat index_ivf(&res, d, nlist, faiss::METRIC_L2);

    assert(!index_ivf.is_trained);
    index_ivf.train(nb, xb);
    assert(index_ivf.is_trained);
    index_ivf.add(nb, xb); // add vectors to the index

    printf("is_trained = %s\n", index_ivf.is_trained ? "true" : "false");
    printf("ntotal = %ld\n", index_ivf.ntotal);

    { // search xq
        long* I = new long[k * nq];
        float* D = new float[k * nq];

        index_ivf.search(nq, xq, k, D, I);

        // print results
        printf("I (5 first results)=\n");
        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        printf("I (5 last results)=\n");
        for (int i = nq - 5; i < nq; i++) {
            for (int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

    delete[] xb;
    delete[] xq;
}
*/

void test() {
    testCmpltAndBlendInplace();
    testCmpltMinMaxFloat();
    testSearch();
    //testGPU():
}

创作不易，小小的支持一下吧！