mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Alibek Zhakubayev
07979b10e6
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5714 We add the Randomized Ball Carving partition, which splits the dataset into small overlapping leaves, and `gather_rows`, which reads rows of the storage as floats for the kernels. **Motivation** The PiPNN graph build (Rubel et al., "PiPNN: Ultra-Scalable Graph-Based Nearest Neighbor Indexing", KDD '26, arXiv:2602.21247) first partitions the dataset into small overlapping leaves, then builds candidate edges inside each leaf. The partition must read the vectors exactly as search sees them, whatever the storage. It must produce the same leaves for any thread count, so that builds are reproducible, and its memory must stay bounded at billion scale. It must also terminate on degenerate data, such as duplicate points, or inner-product data where one high-norm leader attracts every point. **Implementation** - `faiss/impl/pipnn/kernels.{h,cpp}`: `gather_rows(storage, ids, k, out)` reads rows with `reconstruct()`, except `IndexScalarQuantizer` rows, which it decodes with one `SQuantizer` per call. We do not use `sa_decode` or `reconstruct` for SQ storage: `ScalarQuantizer::decode` opens an OpenMP region, and a nested region inside the partition's `schedule(dynamic)` loops segfaults with the libomp of fbcode builds. - `faiss/impl/pipnn/Partition.{h,cpp}`: Randomized Ball Carving (paper Alg. 5) with multi-level fanout. A root level runs first, then memory-bounded waves of root children. Each wave runs level-synchronous `omp for schedule(dynamic)` item lists that mix stripe items (one single-threaded `stripe_assign` call per 1024-point stripe) and leaf items, which are handed to a `LeafVisitor`. Each level runs four steps that share a `LevelContext`: `plan_level` (leaders and work items), `assign_leaders` (the stripe and leaf loop), `scatter_children` (the counting sort below) and `collect_children` (the merge and degenerate-child rules). - A per-block counting sort scatters the children, so every child lists its points in parent order. Small children are merged in a seeded Fisher-Yates order. An IP-ranked child that holds more than 90% of its parent is re-partitioned by L2, and a child equal to its parent or at the depth limit is cut into id slices. Together these rules guarantee termination. - Leader sampling (Floyd's algorithm), the merge order and every child seed come from `SplitMix64RandomGenerator` and `mix_seed`, and every selection uses a strict total order, so the leaves do not depend on the thread count or on the wave budget. - The stripe and leaf loop captures exceptions per item and rethrows them after the loop. Its master thread polls `InterruptCallback` after its first item and then every 64 items, and the partition checks it between levels. A throwing `InterruptCallback` is captured like any item exception. The classify loop also captures and rethrows exceptions; the counting-sort loops cannot throw. - Memory: the per-level assignment, count and child-id arrays are uninitialized heap arrays, so they are not zero-filled before being overwritten. The per-level node and work-item lists are reserved heap vectors, proportional to the number of nodes in the level. Every partition buffer is released on exit, including on the throwing paths. We reject `c_max > 4096`, because each thread holds a `c_max` x `c_max` float matrix for its leaves. - The `test_pipnn_cpp` test target gains the `openmp` dependency, because the tests include `omp.h` to set the OpenMP thread count. - No caller yet outside the tests. The graph build uses this code in a later diff of the stack. Reviewed By: mnorris11 Differential Revision: D123103150 fbshipit-source-id: df2b82457f9c9840846a4707847f37bf03266ac7