All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexBinaryFlat.cpp
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
#include "../../IndexBinaryFlat.h"
#include "../GpuIndexBinaryFlat.h"
#include "../StandardGpuResources.h"
#include "../utils/DeviceUtils.h"
#include "../test/TestUtils.h"
#include "../../utils.h"
#include <gtest/gtest.h>
#include <sstream>
#include <vector>
void compareBinaryDist(const std::vector<int>& cpuDist,
const std::vector<faiss::IndexBinary::idx_t>& cpuLabels,
const std::vector<int>& gpuDist,
const std::vector<faiss::IndexBinary::idx_t>& gpuLabels,
int numQuery,
int k) {
for (int i = 0; i < numQuery; ++i) {
// The index order can be permuted within a group that has the same
// distance, since this is based on the order in which the algorithm
// encounters the values. The last set of equivalent distances seen in the
// min-k might be truncated, so we can't check that set, but all others we
// can check.
std::set<faiss::IndexBinary::idx_t> cpuLabelSet;
std::set<faiss::IndexBinary::idx_t> gpuLabelSet;
int curDist = -1;
for (int j = 0; j < k; ++j) {
int idx = i * k + j;
if (curDist == -1) {
curDist = cpuDist[idx];
}
if (curDist != cpuDist[idx]) {
// Distances must be monotonically increasing
EXPECT_LT(curDist, cpuDist[idx]);
// This is a new set of distances
EXPECT_EQ(cpuLabelSet, gpuLabelSet);
curDist = cpuDist[idx];
cpuLabelSet.clear();
gpuLabelSet.clear();
}
cpuLabelSet.insert(cpuLabels[idx]);
gpuLabelSet.insert(gpuLabels[idx]);
// Because the distances are reproducible, they must be exactly the same
EXPECT_EQ(cpuDist[idx], gpuDist[idx]);
}
}
}
template <int DimMultiple>
void testGpuIndexBinaryFlat() {
faiss::gpu::StandardGpuResources res;
res.noTempMemory();
faiss::gpu::GpuIndexBinaryFlatConfig config;
config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
// multiples of 8 and multiples of 32 use different implementations
int dims = faiss::gpu::randVal(1, 20) * DimMultiple;
faiss::gpu::GpuIndexBinaryFlat gpuIndex(&res, dims, config);
faiss::IndexBinaryFlat cpuIndex(dims);
int numVecs = faiss::gpu::randVal(1, 20000);
int numQuery = faiss::gpu::randVal(1, 1000);
int k = faiss::gpu::randVal(1, 1024);
auto data = faiss::gpu::randBinaryVecs(numVecs, dims);
gpuIndex.add(numVecs, data.data());
cpuIndex.add(numVecs, data.data());
auto query = faiss::gpu::randBinaryVecs(numQuery, dims);
std::vector<int> cpuDist(numQuery * k);
std::vector<faiss::IndexBinary::idx_t> cpuLabels(numQuery * k);
cpuIndex.search(numQuery,
query.data(),
k,
cpuDist.data(),
cpuLabels.data());
std::vector<int> gpuDist(numQuery * k);
std::vector<faiss::IndexBinary::idx_t> gpuLabels(numQuery * k);
gpuIndex.search(numQuery,
query.data(),
k,
gpuDist.data(),
gpuLabels.data());
compareBinaryDist(cpuDist, cpuLabels,
gpuDist, gpuLabels,
numQuery, k);
}
TEST(TestGpuIndexBinaryFlat, Test8) {
for (int tries = 0; tries < 4; ++tries) {
testGpuIndexBinaryFlat<8>();
}
}
TEST(TestGpuIndexBinaryFlat, Test32) {
for (int tries = 0; tries < 4; ++tries) {
testGpuIndexBinaryFlat<32>();
}
}
int main(int argc, char** argv) {
testing::InitGoogleTest(&argc, argv);
// just run with a fixed test seed
faiss::gpu::setTestSeed(100);
return RUN_ALL_TESTS();
}
faiss::IndexBinaryFlat
Definition: IndexBinaryFlat.h:22
faiss::gpu::GpuIndexBinaryFlatConfig
Definition: GpuIndexBinaryFlat.h:19
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:76
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:22
faiss::gpu::GpuIndexBinaryFlat
Definition: GpuIndexBinaryFlat.h:24