facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Randomized Ball Carving partition (#5714)

Browse Source 
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
This commit is contained in:
Alibek Zhakubayev
2026-10-09 17:06:32 -07:00
committed by meta-codesync[bot]
parent f46050eb56
commit 07979b10e6
6 changed files with 1157 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/CMakeLists.txt
+2
View File
@@ -172,6 +172,7 @@ set(FAISS_SRC
impl/result_handler/ResultHandler.cpp
impl/NSG.cpp
impl/pipnn/HashPrune.cpp
impl/pipnn/Partition.cpp
impl/pipnn/kernels.cpp
impl/PolysemousTraining.cpp
impl/ProductQuantizer.cpp
@@ -321,6 +322,7 @@ set(FAISS_HEADERS
impl/NNDescent.h
impl/NSG.h
impl/pipnn/HashPrune.h
impl/pipnn/Partition.h
impl/pipnn/kernels.h
impl/Panorama.h
impl/PanoramaStats.h
faiss/impl/pipnn/Partition.cpp
+591
View File
@@ -0,0 +1,591 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/impl/pipnn/Partition.h>
#include <omp.h>
#include <algorithm>
#include <atomic>
#include <cmath>
#include <exception>
#include <iterator>
#include <limits>
#include <memory>
#include <utility>
#include <faiss/Index.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/impl/FaissException.h>
#include <faiss/impl/pipnn/kernels.h>
#include <faiss/utils/AlignedTable.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/random.h>
namespace faiss {
namespace pipnn {
int PartitionParams::fanout_at(int depth) const {
FAISS_ASSERT(depth >= 0);
return static_cast<size_t>(depth) < fanout.size() ? fanout[depth] : 1;
}
uint64_t mix_seed(uint64_t a, uint64_t b) {
return SplitMix64RandomGenerator(int64_t(a ^ (b + (a << 6) + (a >> 2))))
.next();
}
namespace detail {
size_t leader_count(size_t size, const PartitionParams& p, int depth) {
size_t by_fraction = static_cast<size_t>(std::ceil(
static_cast<double>(p.leader_fraction) *
static_cast<double>(size)));
size_t floor4 = 4 * static_cast<size_t>(p.fanout_at(depth));
size_t count = std::max(by_fraction, floor4);
count = std::min(count, static_cast<size_t>(p.leader_cap));
return std::min(count, size);
}
void sample_positions(
size_t N,
size_t k,
uint64_t seed,
std::vector<size_t>& out) {
FAISS_ASSERT(k <= N);
out.clear();
out.reserve(k);
SplitMix64RandomGenerator rng(static_cast<int64_t>(seed));
for (size_t j = N - k; j < N; j++) {
size_t t = static_cast<size_t>(
rng.next() % (static_cast<uint64_t>(j) + 1));
auto it = std::lower_bound(out.begin(), out.end(), t);
if (it != out.end() && *it == t) {
out.push_back(j); // every sampled position is < j
} else {
out.insert(it, t);
}
}
}
Ranking ranking_for(const PartitionParams& p, bool force_l2) {
if (p.metric == METRIC_L2 || force_l2) {
return Ranking::L2;
}
return p.ip_partition_by_angle ? Ranking::Angle : Ranking::IP;
}
} // namespace detail
namespace {
/// Positions per counting-sort block; unlike kStripe, it does not affect
/// leaves.
constexpr size_t kScatterBlock = 65536;
/// Above any leader index, so the merge-order seed never equals a child's.
constexpr uint64_t kMergeShuffleIndex = uint64_t(1) << 33;
constexpr int kMaxCMax = 4096;
constexpr size_t kNoSplit = std::numeric_limits<size_t>::max();
/// `ids` is sorted and owned by a LevelStorage (or the root list).
struct Node {
const int32_t* ids;
size_t size;
uint64_t seed;
int depth;
bool force_l2; // L2 ranking: this or an ancestor was a > 90% IP child
};
/// Owns the id lists of every node of one level.
struct LevelStorage {
std::unique_ptr<int32_t[]> child_ids;
std::vector<std::vector<int32_t>> merge_pools;
};
/// Inputs and stats shared by every level of one partition.
struct LevelContext {
const Index& storage;
const PartitionParams& p;
const LeafVisitor& on_leaf;
PartitionStats& stats;
};
/// A leaf (split == kNoSplit) or kStripe points of a split node from `begin`.
struct WorkItem {
size_t node;
size_t split;
size_t begin;
};
struct SplitInfo {
size_t node = 0;
std::vector<int32_t> leaders; // ascending ids
int fanout = 0;
Ranking ranking = Ranking::L2;
size_t offset = 0; // start in the level assignment / child-id buffers
size_t n_blocks = 0;
size_t count_offset = 0; // start in the level count buffer
std::vector<size_t> child_off; // |L| + 1 offsets into the child-id buffer
};
struct LevelPlan {
std::vector<WorkItem> items;
std::vector<SplitInfo> splits;
size_t n_instances = 0; // (point, chosen leader) pairs of all splits
size_t n_counts = 0; // entries of the count buffer
};
/// Next-level nodes of one split. Merged groups have ids == nullptr until
/// collect_children points them into their pool.
struct SplitOutput {
std::vector<Node> nodes;
std::vector<size_t> pool_offsets;
std::vector<int32_t> pool;
size_t n_slices = 0;
size_t n_l2 = 0;
};
void check_params(const Index& storage, const PartitionParams& p) {
FAISS_THROW_IF_NOT_MSG(
p.c_min >= 1 && p.c_min <= p.c_max,
"PiPNN partition: need 1 <= c_min <= c_max");
FAISS_THROW_IF_NOT_FMT(
p.c_max <= kMaxCMax,
"PiPNN partition: need c_max <= %d (got %d): each thread holds a "
"c_max x c_max float matrix",
kMaxCMax,
p.c_max);
FAISS_THROW_IF_NOT_MSG(
p.leader_cap >= 1 && p.leader_cap <= 65535,
"PiPNN partition: need 1 <= leader_cap <= 65535");
FAISS_THROW_IF_NOT_MSG(
p.leader_fraction >= 0 && p.leader_fraction <= 1,
"PiPNN partition: need 0 <= leader_fraction <= 1");
for (int f : p.fanout) {
FAISS_THROW_IF_NOT_MSG(
f >= 1, "PiPNN partition: every fanout must be >= 1");
}
FAISS_THROW_IF_NOT_MSG(
p.metric == METRIC_L2 || p.metric == METRIC_INNER_PRODUCT,
"PiPNN partition: metric must be L2 or inner product");
FAISS_THROW_IF_NOT_MSG(
storage.ntotal >= 0 &&
storage.ntotal <= std::numeric_limits<int32_t>::max(),
"PiPNN partition: need 0 <= n <= INT32_MAX");
}
void record_leaf(PartitionStats& stats, size_t size, int depth) {
stats.n_leaves++;
stats.n_point_instances += size;
stats.max_leaf = std::max(stats.max_leaf, size);
stats.max_depth_seen = std::max(stats.max_depth_seen, depth);
const size_t bucket = (size - 1) / 64;
if (stats.leaf_size_histogram.size() <= bucket) {
stats.leaf_size_histogram.resize(bucket + 1, 0);
}
stats.leaf_size_histogram[bucket]++;
}
void ensure_size(AlignedTable<float, 64>& table, size_t n) {
if (table.size() < n) {
table.resize(n);
}
}
void push_slices(const Node& child, size_t c_max, SplitOutput& out) {
for (size_t o = 0; o < child.size; o += c_max) {
out.nodes.push_back(
Node{child.ids + o,
std::min(c_max, child.size - o),
child.seed,
child.depth,
child.force_l2});
}
}
LevelPlan plan_level(const LevelContext& ctx, const std::vector<Node>& nodes) {
const PartitionParams& p = ctx.p;
const size_t c_max = static_cast<size_t>(p.c_max);
size_t n_items = 0;
size_t n_splits = 0;
for (const Node& nd : nodes) {
if (nd.size <= c_max) {
n_items++;
} else {
n_items += (nd.size + kStripe - 1) / kStripe;
n_splits++;
}
}
LevelPlan plan;
plan.splits.reserve(n_splits);
plan.items.reserve(n_items);
for (size_t i = 0; i < nodes.size(); i++) {
const Node& nd = nodes[i];
if (nd.size <= c_max) {
record_leaf(ctx.stats, nd.size, nd.depth);
plan.items.push_back(WorkItem{i, kNoSplit, 0});
continue;
}
SplitInfo sp;
sp.node = i;
const size_t n_leaders = detail::leader_count(nd.size, p, nd.depth);
std::vector<size_t> positions;
detail::sample_positions(nd.size, n_leaders, nd.seed, positions);
sp.leaders.resize(n_leaders);
for (size_t l = 0; l < n_leaders; l++) {
sp.leaders[l] = nd.ids[positions[l]];
}
sp.fanout =
std::min(p.fanout_at(nd.depth), static_cast<int>(n_leaders));
sp.ranking = detail::ranking_for(p, nd.force_l2);
sp.offset = plan.n_instances;
plan.n_instances += nd.size * sp.fanout;
sp.n_blocks = (nd.size + kScatterBlock - 1) / kScatterBlock;
sp.count_offset = plan.n_counts;
plan.n_counts += sp.n_blocks * n_leaders;
for (size_t b = 0; b < nd.size; b += kStripe) {
plan.items.push_back(WorkItem{i, plan.splits.size(), b});
}
plan.splits.push_back(std::move(sp));
}
return plan;
}
/// Runs the leaf and stripe items; returns each point's `fanout` nearest
/// leaders.
std::unique_ptr<uint16_t[]> assign_leaders(
const LevelContext& ctx,
const std::vector<Node>& nodes,
const LevelPlan& plan) {
const Index& storage = ctx.storage;
const size_t d = static_cast<size_t>(storage.d);
// new T[n], not std::vector: every entry is written first; zeroing is
// serial.
std::unique_ptr<uint16_t[]> assign(new uint16_t[plan.n_instances]);
std::exception_ptr ex;
std::atomic<bool> interrupt{false};
#pragma omp parallel
{
AlignedTable<float, 64> stripe_vecs, leader_vecs, dist;
std::vector<float> leader_norms;
size_t n_done = 0;
#pragma omp for schedule(dynamic)
for (int64_t i = 0; i < static_cast<int64_t>(plan.items.size()); i++) {
if (interrupt.load(std::memory_order_relaxed)) {
continue;
}
try {
const WorkItem& item = plan.items[i];
const Node& nd = nodes[item.node];
if (item.split == kNoSplit) {
ctx.on_leaf(nd.ids, nd.size);
} else {
const SplitInfo& sp = plan.splits[item.split];
const size_t S = std::min(kStripe, nd.size - item.begin);
const size_t L = sp.leaders.size();
ensure_size(stripe_vecs, S * d);
ensure_size(leader_vecs, L * d);
ensure_size(dist, S * L);
leader_norms.resize(L);
gather_rows(
storage,
nd.ids + item.begin,
S,
stripe_vecs.data());
gather_rows(
storage, sp.leaders.data(), L, leader_vecs.data());
for (size_t l = 0; l < L; l++) {
leader_norms[l] =
fvec_norm_L2sqr(leader_vecs.data() + l * d, d);
}
stripe_assign(
stripe_vecs.data(),
S,
leader_vecs.data(),
leader_norms.data(),
L,
d,
sp.ranking,
sp.fanout,
dist.data(),
assign.get() + sp.offset + item.begin * sp.fanout);
}
if (omp_get_thread_num() == 0 && n_done++ % 64 == 0 &&
InterruptCallback::is_interrupted()) {
FAISS_THROW_MSG("computation interrupted");
}
} catch (...) {
omp_capture_exception(ex, [&] { interrupt = true; });
}
}
}
omp_rethrow_if_exception(ex);
return assign;
}
/// Counting sort per (split, block): count, prefix sums in (leader, block)
/// order, scatter. Every child comes out in parent order. Fills child_off and
/// returns the child ids, owned by next_storage.
const int32_t* scatter_children(
const std::vector<Node>& nodes,
LevelPlan& plan,
const uint16_t* assign,
LevelStorage& next_storage) {
std::vector<SplitInfo>& splits = plan.splits;
std::vector<std::pair<size_t, size_t>> blocks;
for (size_t s = 0; s < splits.size(); s++) {
for (size_t b = 0; b < splits[s].n_blocks; b++) {
blocks.emplace_back(s, b);
}
}
std::unique_ptr<size_t[]> counts(new size_t[plan.n_counts]);
#pragma omp parallel for schedule(dynamic)
for (int64_t i = 0; i < static_cast<int64_t>(blocks.size()); i++) {
const SplitInfo& sp = splits[blocks[i].first];
const size_t b = blocks[i].second;
const Node& nd = nodes[sp.node];
const size_t L = sp.leaders.size();
size_t* cnt = counts.get() + sp.count_offset + b * L;
std::fill(cnt, cnt + L, 0);
const uint16_t* a = assign + sp.offset;
const size_t end = std::min(nd.size, (b + 1) * kScatterBlock);
for (size_t pos = b * kScatterBlock; pos < end; pos++) {
for (int j = 0; j < sp.fanout; j++) {
cnt[a[pos * sp.fanout + j]]++;
}
}
}
for (SplitInfo& sp : splits) {
const size_t L = sp.leaders.size();
sp.child_off.resize(L + 1);
size_t running = sp.offset;
for (size_t l = 0; l < L; l++) {
sp.child_off[l] = running;
for (size_t b = 0; b < sp.n_blocks; b++) {
size_t& c = counts[sp.count_offset + b * L + l];
const size_t block_count = c;
c = running;
running += block_count;
}
}
sp.child_off[L] = running;
}
next_storage.child_ids.reset(new int32_t[plan.n_instances]);
int32_t* child_ids = next_storage.child_ids.get();
#pragma omp parallel for schedule(dynamic)
for (int64_t i = 0; i < static_cast<int64_t>(blocks.size()); i++) {
const SplitInfo& sp = splits[blocks[i].first];
const size_t b = blocks[i].second;
const Node& nd = nodes[sp.node];
size_t* cursor = counts.get() + sp.count_offset + b * sp.leaders.size();
const uint16_t* a = assign + sp.offset;
const size_t end = std::min(nd.size, (b + 1) * kScatterBlock);
for (size_t pos = b * kScatterBlock; pos < end; pos++) {
for (int j = 0; j < sp.fanout; j++) {
child_ids[cursor[a[pos * sp.fanout + j]]++] = nd.ids[pos];
}
}
}
return child_ids;
}
/// Child rules, in order: below c_min it is merged; up to c_max it is a leaf;
/// at max_depth it is sliced; an IP-ranked child above 90% of its parent is
/// re-split by L2; a child equal to its parent is sliced; else it is split.
void classify_children(
const PartitionParams& p,
const Node& nd,
const SplitInfo& sp,
const int32_t* child_ids,
SplitOutput& out) {
const size_t c_min = static_cast<size_t>(p.c_min);
const size_t c_max = static_cast<size_t>(p.c_max);
const int child_depth = nd.depth + 1;
std::vector<uint32_t> small_children;
for (size_t l = 0; l < sp.leaders.size(); l++) {
Node child{
child_ids + sp.child_off[l],
sp.child_off[l + 1] - sp.child_off[l],
mix_seed(nd.seed, l),
child_depth,
nd.force_l2};
if (child.size < c_min) {
small_children.push_back(static_cast<uint32_t>(l));
} else if (child.size <= c_max) {
out.nodes.push_back(child);
} else if (child_depth >= p.max_depth) {
push_slices(child, c_max, out);
out.n_slices++;
} else if (sp.ranking != Ranking::L2 && 10 * child.size > 9 * nd.size) {
child.force_l2 = true;
out.nodes.push_back(child);
out.n_l2++;
} else if (child.size == nd.size) {
push_slices(child, c_max, out);
out.n_slices++;
} else {
out.nodes.push_back(child);
}
}
std::vector<int> merge_order(small_children.size());
rand_perm_splitmix64(
merge_order.data(),
merge_order.size(),
static_cast<int64_t>(mix_seed(nd.seed, kMergeShuffleIndex)));
std::vector<int32_t> group;
auto close_group = [&]() {
std::sort(group.begin(), group.end());
group.erase(std::unique(group.begin(), group.end()), group.end());
out.pool_offsets.push_back(out.pool.size());
out.pool.insert(out.pool.end(), group.begin(), group.end());
// At most c_max points, so it is a leaf and its seed is never used.
out.nodes.push_back(
Node{nullptr, group.size(), 0, child_depth, nd.force_l2});
group.clear();
};
for (int i : merge_order) {
const uint32_t l = small_children[i];
const int32_t* c = child_ids + sp.child_off[l];
const size_t size = sp.child_off[l + 1] - sp.child_off[l];
if (!group.empty() && group.size() + size > c_max) {
close_group();
}
group.insert(group.end(), c, c + size);
}
if (!group.empty()) {
close_group();
}
}
void collect_children(
const LevelContext& ctx,
const std::vector<Node>& nodes,
const std::vector<SplitInfo>& splits,
const int32_t* child_ids,
std::vector<Node>& next,
LevelStorage& next_storage) {
std::vector<SplitOutput> outputs(splits.size());
std::exception_ptr ex;
#pragma omp parallel for schedule(dynamic)
for (int64_t s = 0; s < static_cast<int64_t>(splits.size()); s++) {
try {
classify_children(
ctx.p,
nodes[splits[s].node],
splits[s],
child_ids,
outputs[s]);
} catch (...) {
omp_capture_exception(ex);
}
}
omp_rethrow_if_exception(ex);
size_t n_next = next.size();
for (const SplitOutput& out : outputs) {
n_next += out.nodes.size();
}
next.reserve(n_next);
next_storage.merge_pools.reserve(outputs.size());
for (SplitOutput& out : outputs) {
ctx.stats.n_slice_fallbacks += out.n_slices;
ctx.stats.n_l2_repartitions += out.n_l2;
next_storage.merge_pools.push_back(std::move(out.pool));
const int32_t* pool = next_storage.merge_pools.back().data();
size_t k = 0;
for (Node& child : out.nodes) {
if (child.ids == nullptr) {
child.ids = pool + out.pool_offsets[k++];
}
}
next.insert(
next.end(),
std::make_move_iterator(out.nodes.begin()),
std::make_move_iterator(out.nodes.end()));
std::vector<Node>().swap(out.nodes);
std::vector<size_t>().swap(out.pool_offsets);
}
}
void run_level(
const LevelContext& ctx,
const std::vector<Node>& nodes,
std::vector<Node>& next,
LevelStorage& next_storage) {
LevelPlan plan = plan_level(ctx, nodes);
const std::unique_ptr<uint16_t[]> assign = assign_leaders(ctx, nodes, plan);
if (plan.splits.empty()) {
return;
}
const int32_t* child_ids =
scatter_children(nodes, plan, assign.get(), next_storage);
collect_children(ctx, nodes, plan.splits, child_ids, next, next_storage);
}
} // namespace
void partition_and_visit(
const Index& storage,
const PartitionParams& p,
uint64_t root_seed,
const LeafVisitor& on_leaf,
PartitionStats* stats) {
check_params(storage, p);
PartitionStats st;
const LevelContext ctx{storage, p, on_leaf, st};
const size_t n = static_cast<size_t>(storage.ntotal);
if (n > 0) {
InterruptCallback::check();
std::unique_ptr<int32_t[]> root_ids(new int32_t[n]);
for (size_t i = 0; i < n; i++) {
root_ids[i] = static_cast<int32_t>(i);
}
std::vector<Node> root_children;
LevelStorage root_storage;
run_level(
ctx,
{Node{root_ids.get(), n, root_seed, 0, false}},
root_children,
root_storage);
root_ids.reset();
size_t w0 = 0;
while (w0 < root_children.size()) {
size_t w1 = w0;
size_t load = 0;
while (w1 < root_children.size() &&
(w1 == w0 ||
load + root_children[w1].size <= p.wave_budget)) {
load += root_children[w1].size;
w1++;
}
st.n_waves++;
std::vector<Node> nodes(
root_children.begin() + w0, root_children.begin() + w1);
std::unique_ptr<LevelStorage> level_storage;
while (!nodes.empty()) {
InterruptCallback::check();
std::vector<Node> next;
auto next_storage = std::make_unique<LevelStorage>();
run_level(ctx, nodes, next, *next_storage);
nodes = std::move(next);
level_storage =
std::move(next_storage); // frees the parent lists
}
w0 = w1;
}
}
if (stats) {
*stats = std::move(st);
}
}
} // namespace pipnn
} // namespace faiss
faiss/impl/pipnn/Partition.h
+86
View File
@@ -0,0 +1,86 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#pragma once
#include <faiss/MetricType.h>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <vector>
namespace faiss {
struct Index;
namespace pipnn {
enum class Ranking : uint8_t;
/// Randomized Ball Carving settings (paper Alg. 5).
struct PartitionParams {
int c_max = 1024; // largest leaf
int c_min = 128; // smaller children are merged
int leader_cap = 1000; // most leaders per split
float leader_fraction = 0.005f; // leaders per split, as a fraction of it
std::vector<int> fanout = {10, 3}; // leaders joined per depth, then 1
MetricType metric = METRIC_L2;
bool ip_partition_by_angle = false;
size_t wave_budget = size_t(1) << 30; // point instances per wave
int max_depth = 32; // children at this depth are cut into id slices
int fanout_at(int depth) const;
};
struct PartitionStats {
size_t n_leaves = 0;
size_t n_point_instances = 0; // sum of leaf sizes
size_t max_leaf = 0;
size_t n_slice_fallbacks = 0, n_l2_repartitions = 0;
int max_depth_seen = 0;
/// leaf_size_histogram[i] counts leaves of size in [64 i + 1, 64 (i + 1)].
std::vector<size_t> leaf_size_histogram;
size_t n_waves = 0;
};
/// Called once per leaf with its sorted, unique ids (s <= c_max), concurrently
/// from partition_and_visit's OpenMP region: omp_get_thread_num() identifies
/// the calling thread. `ids` is valid only during the call.
using LeafVisitor = std::function<void(const int32_t* ids, size_t s)>;
/// Randomized Ball Carving with multi-level fanout. Calls `on_leaf` exactly
/// once per leaf and fills `*stats` (if not null). On an exception or an
/// interrupt it rethrows with some leaves unvisited and `*stats` unchanged.
void partition_and_visit(
const Index& storage,
const PartitionParams& p,
uint64_t root_seed,
const LeafVisitor& on_leaf,
PartitionStats* stats);
/// Deterministic hash of (a, b) for derived seeds. For a fixed a, distinct b
/// give distinct results.
uint64_t mix_seed(uint64_t a, uint64_t b);
/// Building blocks of partition_and_visit, exposed for unit tests.
namespace detail {
/// min(max(ceil(leader_fraction * size), 4 * fanout(depth)), leader_cap, size)
size_t leader_count(size_t size, const PartitionParams& p, int depth);
/// k distinct positions of [0, N), ascending (Floyd's algorithm).
void sample_positions(
size_t N,
size_t k,
uint64_t seed,
std::vector<size_t>& out);
Ranking ranking_for(const PartitionParams& p, bool force_l2);
} // namespace detail
} // namespace pipnn
} // namespace faiss
faiss/impl/pipnn/kernels.cpp
+28
View File
@@ -10,8 +10,11 @@
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <typeinfo>
#include <vector>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/impl/FaissAssert.h>
#include <faiss/utils/Heap.h>
@@ -294,5 +297,30 @@ void pairwise_distances(
}
}
void gather_rows(
const Index& storage,
const int32_t* ids,
size_t k,
float* out) {
const size_t d = static_cast<size_t>(storage.d);
if (typeid(storage) == typeid(IndexScalarQuantizer)) {
// Not reconstruct: ScalarQuantizer::decode opens an OpenMP region, and
// a nested region inside the build's dynamic loops crashes libomp.
const auto& sq = static_cast<const IndexScalarQuantizer&>(storage);
std::unique_ptr<ScalarQuantizer::SQuantizer> quant(
sq.sq.select_quantizer());
for (size_t i = 0; i < k; i++) {
quant->decode_vector(
sq.codes.data() +
static_cast<size_t>(ids[i]) * sq.code_size,
out + i * d);
}
return;
}
for (size_t i = 0; i < k; i++) {
storage.reconstruct(ids[i], out + i * d);
}
}
} // namespace pipnn
} // namespace faiss
faiss/impl/pipnn/kernels.h
+11
View File
@@ -23,6 +23,9 @@
*/
namespace faiss {
struct Index;
namespace pipnn {
/// Points per stripe item of the partition kernel. Changing it changes graphs.
@@ -68,5 +71,13 @@ void pairwise_distances(
MetricType metric,
float* out);
/// Rows ids[0..k) of `storage` as floats into out (k x d). SQ rows are decoded
/// directly; other types use reconstruct(), which must not open OpenMP regions.
void gather_rows(
const Index& storage,
const int32_t* ids,
size_t k,
float* out);
} // namespace pipnn
} // namespace faiss
tests/test_pipnn.cpp
+439
View File
@@ -6,6 +6,7 @@
*/
#include <gtest/gtest.h>
#include <omp.h>
#include <algorithm>
#include <cmath>
@@ -13,13 +14,20 @@
#include <cstring>
#include <limits>
#include <map>
#include <memory>
#include <mutex>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include <faiss/IndexFlat.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissException.h>
#include <faiss/impl/pipnn/HashPrune.h>
#include <faiss/impl/pipnn/Partition.h>
#include <faiss/impl/pipnn/kernels.h>
#include <faiss/utils/AlignedTable.h>
#include <faiss/utils/distances.h>
@@ -723,3 +731,434 @@ TEST(PiPNNKernels, RejectBadArguments) {
scratch.get()),
FaissException);
}
/*************************************************************
* PiPNN storage access and Randomized Ball Carving
*************************************************************/
namespace pipnn_partition_test {
std::vector<float> gaussian_data(size_t n, size_t d, int64_t seed) {
std::vector<float> x(n * d);
faiss::float_randn(x.data(), x.size(), seed);
return x;
}
/// x_i = (1 + 10 i / n) * u for one direction u: under max-IP every point ranks
/// leaders by norm, so all points pick the same `fanout` leaders.
std::vector<float> collinear_data(size_t n, size_t d, int64_t seed) {
std::vector<float> u(d);
faiss::float_rand(u.data(), d, seed);
std::vector<float> x(n * d);
for (size_t i = 0; i < n; i++) {
const float scale = 1.0f + 10.0f * float(i) / float(n);
for (size_t j = 0; j < d; j++) {
x[i * d + j] = scale * u[j];
}
}
return x;
}
struct FlatData {
faiss::IndexFlat index;
FlatData(const std::vector<float>& x, size_t d, faiss::MetricType metric)
: index(d, metric) {
index.add(x.size() / d, x.data());
}
};
const uint64_t kSeed = faiss::pipnn::mix_seed(1234, 0);
faiss::pipnn::PartitionParams small_params(faiss::MetricType metric) {
faiss::pipnn::PartitionParams p;
p.c_max = 64;
p.c_min = 16;
p.metric = metric;
return p;
}
/// Sets the OpenMP thread count for its lifetime, then restores it.
struct ScopedOmpThreads {
const int saved = omp_get_max_threads();
explicit ScopedOmpThreads(int n) {
omp_set_num_threads(n);
}
~ScopedOmpThreads() {
omp_set_num_threads(saved);
}
};
using Leaves = std::vector<std::vector<int32_t>>;
/// Leaves are non-empty, strictly increasing, at most c_max long, and cover
/// every point.
void expect_valid_cover(const Leaves& leaves, size_t n, int c_max);
/// Runs partition_and_visit, checks that its leaves are a valid cover, and
/// returns them sorted.
Leaves visit_all(
const faiss::Index& storage,
const faiss::pipnn::PartitionParams& p,
uint64_t seed,
faiss::pipnn::PartitionStats* stats = nullptr) {
std::mutex mu;
Leaves leaves;
faiss::pipnn::partition_and_visit(
storage,
p,
seed,
[&](const int32_t* ids, size_t s) {
std::vector<int32_t> leaf(ids, ids + s);
std::lock_guard<std::mutex> lock(mu);
leaves.push_back(std::move(leaf));
},
stats);
std::sort(leaves.begin(), leaves.end());
expect_valid_cover(leaves, size_t(storage.ntotal), p.c_max);
return leaves;
}
void expect_valid_cover(const Leaves& leaves, size_t n, int c_max) {
std::vector<char> seen(n, 0);
for (const std::vector<int32_t>& leaf : leaves) {
ASSERT_FALSE(leaf.empty());
ASSERT_LE(leaf.size(), size_t(c_max));
for (size_t i = 0; i < leaf.size(); i++) {
ASSERT_GE(leaf[i], 0);
ASSERT_LT(size_t(leaf[i]), n);
if (i > 0) {
ASSERT_LT(leaf[i - 1], leaf[i]);
}
seen[leaf[i]] = 1;
}
}
for (size_t i = 0; i < n; i++) {
ASSERT_TRUE(seen[i]) << "point " << i << " is in no leaf";
}
}
/// Fraction of points that share a leaf with their nearest other point (L2), or
/// with the other point of largest inner product (IP).
double fraction_with_neighbour(
const Leaves& leaves,
const std::vector<float>& x,
size_t d,
faiss::MetricType metric) {
const size_t n = x.size() / d;
faiss::IndexFlat index(d, metric);
index.add(n, x.data());
std::vector<float> D(2 * n);
std::vector<faiss::idx_t> I(2 * n);
index.search(n, x.data(), 2, D.data(), I.data());
std::vector<std::vector<size_t>> leaves_of(n);
for (size_t l = 0; l < leaves.size(); l++) {
for (int32_t id : leaves[l]) {
leaves_of[id].push_back(l);
}
}
size_t hits = 0;
for (size_t i = 0; i < n; i++) {
const faiss::idx_t j =
I[2 * i] == faiss::idx_t(i) ? I[2 * i + 1] : I[2 * i];
const std::vector<size_t>& a = leaves_of[i];
const std::vector<size_t>& b = leaves_of[j];
if (std::find_first_of(a.begin(), a.end(), b.begin(), b.end()) !=
a.end()) {
hits++;
}
}
return double(hits) / double(n);
}
} // namespace pipnn_partition_test
TEST(PiPNNStorage, GatherSQRowsMatchesReconstruct) {
const int d = 12;
const size_t n = 200;
const std::vector<float> x = pipnn_partition_test::gaussian_data(n, d, 7);
faiss::IndexScalarQuantizer sq(d, faiss::ScalarQuantizer::QT_8bit);
sq.train(n, x.data());
sq.add(n, x.data());
const int32_t ids[3] = {5, 0, 199};
faiss::AlignedTable<float, 64> out(3 * d);
faiss::pipnn::gather_rows(sq, ids, 3, out.data());
std::vector<float> expected(d);
for (size_t i = 0; i < 3; i++) {
sq.reconstruct(ids[i], expected.data());
EXPECT_EQ(
0,
std::memcmp(
out.data() + i * d, expected.data(), sizeof(float) * d))
<< "row " << i;
}
}
// Pinned values: changing any of them changes every graph.
TEST(PiPNNPartition, SeededHelpersGolden) {
using faiss::pipnn::mix_seed;
using faiss::pipnn::detail::sample_positions;
EXPECT_EQ(mix_seed(1234, 1), 0x626e95467131d717ULL);
std::vector<size_t> pos;
sample_positions(100, 5, 42, pos);
EXPECT_EQ(pos, (std::vector<size_t>{6, 45, 50, 84, 85}));
sample_positions(10, 10, 7, pos);
EXPECT_EQ(pos, (std::vector<size_t>{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}));
}
TEST(PiPNNPartition, ParameterHelpers) {
using faiss::pipnn::Ranking;
using faiss::pipnn::detail::leader_count;
using faiss::pipnn::detail::ranking_for;
faiss::pipnn::PartitionParams p;
EXPECT_EQ(p.fanout_at(1), 3);
EXPECT_EQ(p.fanout_at(2), 1);
EXPECT_EQ(leader_count(1025, p, 0), 40u); // 4 * fanout floor
EXPECT_EQ(leader_count(100000, p, 0), 500u); // ceil(P_samp * |P|)
EXPECT_EQ(leader_count(10000000, p, 0), 1000u); // leader_cap
EXPECT_EQ(leader_count(300, p, 1), 12u);
EXPECT_EQ(leader_count(3, p, 5), 3u); // never more than |P|
p.leader_cap = 7;
EXPECT_EQ(leader_count(1025, p, 0), 7u);
p.ip_partition_by_angle = true;
EXPECT_EQ(ranking_for(p, false), Ranking::L2);
p.metric = faiss::METRIC_INNER_PRODUCT;
EXPECT_EQ(ranking_for(p, false), Ranking::Angle);
EXPECT_EQ(ranking_for(p, true), Ranking::L2);
p.ip_partition_by_angle = false;
EXPECT_EQ(ranking_for(p, false), Ranking::IP);
}
TEST(PiPNNPartition, SizesAroundCMax) {
namespace ppt = pipnn_partition_test;
const size_t d = 8;
const faiss::pipnn::PartitionParams p = ppt::small_params(faiss::METRIC_L2);
ppt::FlatData empty({}, d, faiss::METRIC_L2);
EXPECT_TRUE(ppt::visit_all(empty.index, p, 1).empty());
const size_t n = size_t(p.c_max);
ppt::FlatData one_leaf(ppt::gaussian_data(n, d, 5), d, faiss::METRIC_L2);
std::vector<int32_t> all(n);
for (size_t i = 0; i < n; i++) {
all[i] = int32_t(i);
}
EXPECT_EQ(ppt::visit_all(one_leaf.index, p, 1), ppt::Leaves{all});
ppt::FlatData above(ppt::gaussian_data(n + 1, d, 11), d, faiss::METRIC_L2);
// One split whose small children are merged: no leaf is below c_min
// (without merging, 25 of its 40 leaves are).
for (const std::vector<int32_t>& leaf : ppt::visit_all(above.index, p, 9)) {
EXPECT_GE(leaf.size(), size_t(p.c_min));
}
}
// The root's 70,000 points span two 65,536-point counting-sort blocks, so the
// per-block counts and cursors must combine into valid children.
TEST(PiPNNPartition, RootSpansCountingSortBlocks) {
namespace ppt = pipnn_partition_test;
const size_t n = 70000, d = 2;
ppt::FlatData data(ppt::gaussian_data(n, d, 31), d, faiss::METRIC_L2);
EXPECT_FALSE(
ppt::visit_all(
data.index, ppt::small_params(faiss::METRIC_L2), ppt::kSeed)
.empty());
}
TEST(PiPNNPartition, LeavesGroupNeighbours) {
namespace ppt = pipnn_partition_test;
const size_t n = 5000, d = 16;
const std::vector<float> x = ppt::gaussian_data(n, d, 2024);
ppt::FlatData l2(x, d, faiss::METRIC_L2);
ppt::FlatData ip(x, d, faiss::METRIC_INNER_PRODUCT);
faiss::IndexScalarQuantizer sq(
d, faiss::ScalarQuantizer::QT_8bit, faiss::METRIC_L2);
sq.train(n, x.data());
sq.add(n, x.data());
struct Case {
const char* name;
const faiss::Index* storage;
faiss::MetricType metric;
int c_min;
bool by_angle;
double min_fraction; // measured: 0.986 for L2, 0.950 for IP
};
const Case cases[] = {
{"L2", &l2.index, faiss::METRIC_L2, 16, false, 0.95},
{"L2, no merging", &l2.index, faiss::METRIC_L2, 1, false, 0.95},
{"IP", &ip.index, faiss::METRIC_INNER_PRODUCT, 16, false, 0.9},
{"IP by angle",
&ip.index,
faiss::METRIC_INNER_PRODUCT,
16,
true,
0.9},
{"SQ8", &sq, faiss::METRIC_L2, 16, false, 0.95}};
for (const Case& c : cases) {
SCOPED_TRACE(c.name);
faiss::pipnn::PartitionParams p = ppt::small_params(c.metric);
p.c_min = c.c_min;
p.ip_partition_by_angle = c.by_angle;
const ppt::Leaves leaves = ppt::visit_all(*c.storage, p, ppt::kSeed);
EXPECT_GE(
ppt::fraction_with_neighbour(leaves, x, d, c.metric),
c.min_fraction);
}
}
TEST(PiPNNPartition, StatsDescribeLeaves) {
namespace ppt = pipnn_partition_test;
const size_t d = 16;
ppt::FlatData data(ppt::gaussian_data(5000, d, 2024), d, faiss::METRIC_L2);
const faiss::pipnn::PartitionParams p = ppt::small_params(faiss::METRIC_L2);
faiss::pipnn::PartitionStats stats;
const ppt::Leaves leaves =
ppt::visit_all(data.index, p, ppt::kSeed, &stats);
size_t instances = 0;
size_t max_leaf = 0;
for (const std::vector<int32_t>& leaf : leaves) {
instances += leaf.size();
max_leaf = std::max(max_leaf, leaf.size());
}
EXPECT_EQ(stats.n_leaves, leaves.size());
EXPECT_EQ(stats.n_point_instances, instances);
EXPECT_EQ(stats.max_leaf, max_leaf);
EXPECT_GE(stats.max_depth_seen, 3);
size_t histogram_total = 0;
for (size_t h : stats.leaf_size_histogram) {
histogram_total += h;
}
EXPECT_EQ(histogram_total, stats.n_leaves);
// Replicas rely on a different seed giving different leaves.
EXPECT_FALSE(
leaves ==
ppt::visit_all(data.index, p, faiss::pipnn::mix_seed(1234, 1)));
}
TEST(PiPNNPartition, IPDegenerateChildUsesL2) {
namespace ppt = pipnn_partition_test;
const size_t d = 16;
ppt::FlatData data(
ppt::collinear_data(5000, d, 31), d, faiss::METRIC_INNER_PRODUCT);
const faiss::pipnn::PartitionParams p =
ppt::small_params(faiss::METRIC_INNER_PRODUCT);
faiss::pipnn::PartitionStats stats;
ppt::visit_all(data.index, p, ppt::kSeed, &stats);
// The children of the 10 largest-norm leaders hold all points and are
// flagged; their descendants rank by L2 and are not flagged again.
EXPECT_EQ(stats.n_l2_repartitions, 10u);
}
TEST(PiPNNPartition, AllZeroDataUsesSliceFallback) {
namespace ppt = pipnn_partition_test;
// All products are 0, so every point picks leaders 0..9, whose children
// equal the root and are cut into id slices.
const size_t n = 3000;
ppt::FlatData data(std::vector<float>(n, 0.0f), 1, faiss::METRIC_L2);
faiss::pipnn::PartitionStats stats;
ppt::visit_all(
data.index,
ppt::small_params(faiss::METRIC_L2),
ppt::kSeed,
&stats);
EXPECT_EQ(stats.n_slice_fallbacks, 10u);
EXPECT_EQ(stats.n_leaves, 10u * ((n + 63) / 64));
}
TEST(PiPNNPartition, DuplicateGroupTerminatesAndCovers) {
namespace ppt = pipnn_partition_test;
const size_t n = 3000, d = 16;
std::vector<float> x = ppt::gaussian_data(n, d, 77);
std::fill(x.begin() + 1000 * d, x.begin() + 1500 * d, 10.0f);
ppt::FlatData data(x, d, faiss::METRIC_L2);
const faiss::pipnn::PartitionParams p = ppt::small_params(faiss::METRIC_L2);
faiss::pipnn::PartitionStats stats;
ppt::visit_all(data.index, p, ppt::kSeed, &stats);
// The group is sliced once a child equals its parent (depth 7 here), well
// before the depth limit.
EXPECT_LT(stats.max_depth_seen, p.max_depth);
}
TEST(PiPNNPartition, DepthLimitPrecedesL2Repartition) {
namespace ppt = pipnn_partition_test;
// 10 root children hold all points; with max_depth = 1 the depth limit
// applies before the L2 re-partition rule, so none is flagged.
const size_t n = 5000, d = 16;
ppt::FlatData data(
ppt::collinear_data(n, d, 31), d, faiss::METRIC_INNER_PRODUCT);
faiss::pipnn::PartitionParams p =
ppt::small_params(faiss::METRIC_INNER_PRODUCT);
p.max_depth = 1;
faiss::pipnn::PartitionStats stats;
ppt::visit_all(data.index, p, ppt::kSeed, &stats);
EXPECT_EQ(stats.n_l2_repartitions, 0u);
EXPECT_EQ(stats.n_slice_fallbacks, 10u);
}
TEST(PiPNNPartition, WavesDoNotChangeLeaves) {
namespace ppt = pipnn_partition_test;
// Waves that carry L2 re-partition flags.
const size_t n = 5000, d = 16;
ppt::FlatData ip(
ppt::collinear_data(n, d, 31), d, faiss::METRIC_INNER_PRODUCT);
faiss::pipnn::PartitionParams p_ip =
ppt::small_params(faiss::METRIC_INNER_PRODUCT);
const ppt::Leaves one_wave = ppt::visit_all(ip.index, p_ip, ppt::kSeed);
p_ip.wave_budget = 3000;
faiss::pipnn::PartitionStats ip_stats;
EXPECT_TRUE(
one_wave == ppt::visit_all(ip.index, p_ip, ppt::kSeed, &ip_stats));
EXPECT_GT(ip_stats.n_waves, 1u);
}
namespace pipnn_partition_test {
struct AlwaysInterrupt : faiss::InterruptCallback {
bool want_interrupt() override {
return true;
}
};
/// Clears the global interrupt callback when the test scope ends.
struct InterruptGuard {
~InterruptGuard() {
faiss::InterruptCallback::clear_instance();
}
};
} // namespace pipnn_partition_test
TEST(PiPNNPartition, InterruptStopsWithin64Items) {
namespace ppt = pipnn_partition_test;
// All-zero data gives one region of 470 slice leaves. With one thread the
// master polls after items 0, 64, ...; the interrupt is set while visiting
// item 9, so items 10..64 still run.
ppt::FlatData data(std::vector<float>(3000, 0.0f), 1, faiss::METRIC_L2);
const faiss::pipnn::PartitionParams p = ppt::small_params(faiss::METRIC_L2);
const size_t install_at = 10;
std::string what;
size_t visited = 0;
{
ppt::InterruptGuard guard;
ppt::ScopedOmpThreads one_thread(1);
try {
faiss::pipnn::partition_and_visit(
data.index,
p,
ppt::kSeed,
[&](const int32_t*, size_t) {
if (++visited == install_at) {
faiss::InterruptCallback::instance.reset(
new ppt::AlwaysInterrupt());
}
},
nullptr);
} catch (const faiss::FaissException& e) {
what = e.what();
}
}
EXPECT_NE(what.find("interrupted"), std::string::npos) << what;
EXPECT_EQ(visited - install_at, 55u);
}