fix: algorithm of spreading vectors over shards (#4299)

Summary:
#3374 has a typo, my local repo has the correct fixing (find it when sync from recent v1.10), the intended fixing is:

for example, if we have `n = 10`, and `nshards = 8`

- the original buggy code
  ```c++
  size_t shard_size = (n + nshards - 1) / nshards;
  // i.e.  shard_size = (10 + 8 - 1) / 8 = 2
  size_t i0 = idx * shard_size;
  // all i0: [0, 2, 4, 6, 8, 10, 12, 14], wrong!!! 14>n
  size_t ni = std::min(shard_size, n - i0);
  // all ni: [2, 2, 2, 2, 2, 0, -2, -4], wrong!!!
  ```

- should be:
  ```c++
  size_t base_shard_size = n / nshards;
  // i.e.  shard_size = 10 / 8 = 1
  size_t i0 = idx * base_shard_size + std::min(size_t(idx), n % nshards);
  // all i0: [0, 2, 4, 5, 6, 7, 8, 9]
  size_t ni = base_shard_size;
  if (idx < n % nshards) {
      ++ni;
  }
  // all ni: [2, 2, 1, 1, 1, 1, 1, 1]
  ```

Pull Request resolved: #4299

Reviewed By: junjieqi

Differential Revision: D73199869

Pulled By: mnorris11

fbshipit-source-id: 81cbb8284818212781978412262e91d7747f9b97

Author: simshi

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

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)

faiss/gpu/test/TestGpuIcmEncoder.cpp

@@ -0,0 +1,112 @@
+#include <faiss/gpu/GpuIcmEncoder.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/test/TestUtils.h>
+#include <faiss/impl/LocalSearchQuantizer.h>
+
+#include <gtest/gtest.h>
+#include <tuple>
+#include <vector>
+
+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<ShardingTestParams> {
+   protected:
+    void SetUp() override {
+        params = GetParam();
+
+        lsq.M = 4;
+        lsq.K = 16;
+        lsq.d = 32;
+
+        std::uniform_real_distribution<float> 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<int32_t> codeDist(0, lsq.K - 1);
+        for (auto& c : codes) {
+            c = codeDist(gen);
+        }
+    }
+
+    LocalSearchQuantizer lsq;
+    std::vector<float> x;
+    std::vector<int32_t> codes;
+    std::mt19937 gen;
+    ShardingTestParams params;
+    static constexpr size_t ils_iters = 4;
+};
+
+TEST_P(GpuIcmEncoderShardingTest, DataShardingCorrectness) {
+    std::vector<StandardGpuResources> resources(params.nshards);
+    std::vector<GpuResourcesProvider*> provs;
+    std::vector<int> 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<ShardingTestParams> 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<ShardingTestParams>& 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();
+}