fix algorithm of spreading vectors over shards (#3374)
Summary: simple math: | **input n** | **input nshards** | shard_size | idx | i0 | ni | | -- |-- |-- |-- |-- |-- | | 19 | 6 | 4 | 5 | 20 | **-1** | | 1000 | 37 | 28 | 36 | 1008 | -8 | | 1000 | 64 | 16 | 63 | 1008 | -8 | root cause: integer cause precision loss, `idx * shard_size` overflows, because `(n + nshards - 1) / nshards` is roundup my solution: each shard takes at least `base_shard_size = n / nshards`, then `remain = n % nshards`, we know `0 <= remain < nshards`, next, assign those remain vectors to first `remain` shards, i.e. first `remain` shards take one more vector each. ```c++ auto i0 = idx * base_shard_size; if (i0 < remain) { // if current idx is one of the first `remain` shards i0 += idx; } else { i0 += remain; } ``` simplify above code: `i0 = idx * base_shard_size + std::min(size_t(idx), n % nshards);` Pull Request resolved: https://github.com/facebookresearch/faiss/pull/3374 Reviewed By: fxdawnn Differential Revision: D57867910 Pulled By: junjieqi fbshipit-source-id: 7e72ea5cd197af4f3446fb7a3fd34ad08901dbb2pull/3422/head
parent
db6ff2e095
commit
6e7d9e040f
|
@ -82,7 +82,7 @@ void GpuIcmEncoder::encode(
|
|||
size_t n,
|
||||
size_t ils_iters) const {
|
||||
size_t nshards = shards->size();
|
||||
size_t shard_size = (n + nshards - 1) / nshards;
|
||||
size_t base_shard_size = n / nshards;
|
||||
|
||||
auto codebooks = lsq->codebooks.data();
|
||||
auto M = lsq->M;
|
||||
|
@ -94,8 +94,14 @@ void GpuIcmEncoder::encode(
|
|||
|
||||
// split input data
|
||||
auto fn = [=](int idx, IcmEncoderImpl* encoder) {
|
||||
size_t i0 = idx * shard_size;
|
||||
size_t ni = std::min(shard_size, n - i0);
|
||||
size_t i0 = idx * base_shard_size + std::min(size_t(idx), n % nshards);
|
||||
size_t ni = base_shard_size;
|
||||
if (ni < n % nshards) {
|
||||
++ni;
|
||||
}
|
||||
if (ni <= 0) { // only if n < nshards
|
||||
return;
|
||||
}
|
||||
auto xi = x + i0 * d;
|
||||
auto ci = codes + i0 * M;
|
||||
std::mt19937 geni(idx + seed); // different seed for each shard
|
||||
|
|
Loading…
Reference in New Issue