diff --git a/faiss/gpu/GpuIcmEncoder.cu b/faiss/gpu/GpuIcmEncoder.cu index 999dd998a0..84202aeb7a 100644 --- a/faiss/gpu/GpuIcmEncoder.cu +++ b/faiss/gpu/GpuIcmEncoder.cu @@ -96,7 +96,7 @@ void GpuIcmEncoder::encode( auto fn = [=](int idx, IcmEncoderImpl* encoder) { size_t i0 = idx * base_shard_size + std::min(size_t(idx), n % nshards); size_t ni = base_shard_size; - if (ni < n % nshards) { + if (idx < n % nshards) { ++ni; } if (ni <= 0) { // only if n < nshards diff --git a/faiss/gpu/test/CMakeLists.txt b/faiss/gpu/test/CMakeLists.txt index c549af3947..6b86695976 100644 --- a/faiss/gpu/test/CMakeLists.txt +++ b/faiss/gpu/test/CMakeLists.txt @@ -43,6 +43,7 @@ faiss_gpu_test(TestGpuIndexFlat.cpp) faiss_gpu_test(TestGpuIndexIVFFlat.cpp) faiss_gpu_test(TestGpuIndexBinaryFlat.cpp) faiss_gpu_test(TestGpuMemoryException.cpp) +faiss_gpu_test(TestGpuIcmEncoder.cpp) faiss_gpu_test(TestGpuIndexIVFPQ.cpp) faiss_gpu_test(TestGpuIndexIVFScalarQuantizer.cpp) faiss_gpu_test(TestGpuResidualQuantizer.cpp) diff --git a/faiss/gpu/test/TestGpuIcmEncoder.cpp b/faiss/gpu/test/TestGpuIcmEncoder.cpp new file mode 100644 index 0000000000..0c793d4f58 --- /dev/null +++ b/faiss/gpu/test/TestGpuIcmEncoder.cpp @@ -0,0 +1,112 @@ +#include +#include +#include +#include + +#include +#include +#include + +using faiss::LocalSearchQuantizer; +using faiss::gpu::GpuIcmEncoder; +using faiss::gpu::GpuResourcesProvider; +using faiss::gpu::StandardGpuResources; + +struct ShardingTestParams { + size_t n; + size_t nshards; +}; + +class GpuIcmEncoderShardingTest + : public ::testing::TestWithParam { + protected: + void SetUp() override { + params = GetParam(); + + lsq.M = 4; + lsq.K = 16; + lsq.d = 32; + + std::uniform_real_distribution dist(-1.0f, 1.0f); + lsq.codebooks.resize(lsq.M * lsq.K * lsq.d); + for (auto& v : lsq.codebooks) { + v = dist(gen); + } + + x.resize(params.n * lsq.d); + codes.resize(params.n * lsq.M); + + for (auto& v : x) { + v = dist(gen); + } + + std::uniform_int_distribution codeDist(0, lsq.K - 1); + for (auto& c : codes) { + c = codeDist(gen); + } + } + + LocalSearchQuantizer lsq; + std::vector x; + std::vector codes; + std::mt19937 gen; + ShardingTestParams params; + static constexpr size_t ils_iters = 4; +}; + +TEST_P(GpuIcmEncoderShardingTest, DataShardingCorrectness) { + std::vector resources(params.nshards); + std::vector provs; + std::vector devices; + + for (size_t i = 0; i < params.nshards; ++i) { + resources[i].noTempMemory(); + provs.push_back(&resources[i]); + devices.push_back(0); // use GPU 0 for testing all shards + } + + GpuIcmEncoder encoder(&lsq, provs, devices); + encoder.set_binary_term(); + + gen.seed(42); + EXPECT_NO_THROW( + encoder.encode(codes.data(), x.data(), gen, params.n, ils_iters)); + + for (auto c : codes) { + EXPECT_GE(c, 0); + EXPECT_LT(c, lsq.K); + } +} + +std::vector GetShardingTestCases() { + return { + {1, 8}, + + {5, 4}, + + {10, 2}, + {10, 3}, + {10, 5}, + {10, 8}, + + {20, 8}, + }; +} + +INSTANTIATE_TEST_SUITE_P( + MultiGpuShardingTests, + GpuIcmEncoderShardingTest, + ::testing::ValuesIn(GetShardingTestCases()), + [](const ::testing::TestParamInfo& info) { + return "n" + std::to_string(info.param.n) + "_shards" + + std::to_string(info.param.nshards); + }); + +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(); +}