Skip to content

Code-review #104

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gouyt13 merged 1 commit into from
Sep 15, 2026
Merged

Code-review #104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 0 additions & 8 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,14 +135,6 @@ RaBitQ is developed by the
University, Singapore. A GPU implementation is also available in
[cuvs_rabitq](https://github.com/Stardust-SJF/cuvs_rabitq/tree/cuvs_ivf_rabitq).

## Accuracy at a glance

![RaBitQ estimation error benchmark across MSong, YouTube, OpenAI embeddings, Word2Vec, and GIST](docs/docs/assets/img/acc_bench.png)

*Average and maximum relative estimation error across six datasets; lower is
better. Results from the
[SIGMOD camera-ready paper](https://doi.org/10.1145/3725413).*

## RaBitQ across the vector-search ecosystem

The projects below illustrate adoption of RaBitQ techniques across vector
Expand Down
Binary file removed docs/docs/assets/img/acc_bench.png
View file
Open in desktop
Binary file not shown.
8 changes: 0 additions & 8 deletions docs/docs/index.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,14 +33,6 @@ indexes backed by optimized AVX2 and AVX-512 kernels.
</div>
</div>

## Accuracy at a glance

![RaBitQ estimation error benchmark across MSong, YouTube, OpenAI embeddings, Word2Vec, and GIST](assets/img/acc_bench.png)

*Average and maximum relative estimation error across six datasets; lower is
better. Results from the
[SIGMOD camera-ready paper](https://doi.org/10.1145/3725413).*

## Start with Python

Install the latest release from PyPI:
Expand Down
43 changes: 34 additions & 9 deletions include/rabitqlib/index/estimator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,19 @@ inline void split_batch_estdist(
}
}

ConstRowMajorArrayMap<float> f_add_arr(cur_batch.f_add(), 1, fastscan::kBatchSize);
std::array<float, fastscan::kBatchSize> f_add_values;
std::array<float, fastscan::kBatchSize> f_rescale_values;
std::array<float, fastscan::kBatchSize> f_error_values;
cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size());
cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size());
cur_batch.f_error().copy_to(f_error_values.data(), f_error_values.size());
ConstRowMajorArrayMap<float> f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<float> f_rescale_arr(
cur_batch.f_rescale(), 1, fastscan::kBatchSize
f_rescale_values.data(), 1, fastscan::kBatchSize
);
ConstRowMajorArrayMap<float> f_error_arr(
f_error_values.data(), 1, fastscan::kBatchSize
);
ConstRowMajorArrayMap<float> f_error_arr(cur_batch.f_error(), 1, fastscan::kBatchSize);

RowMajorArrayMap<float> est_dist_arr(est_distance, 1, fastscan::kBatchSize);
RowMajorArrayMap<float> ip_x0_qr_arr(ip_x0_qr, 1, fastscan::kBatchSize);
Expand Down Expand Up @@ -145,7 +153,13 @@ inline void split_single_fulldist_direct(
ConstExDataMap<float> cur_ex(ex_data, padded_dim, ex_bits);

// [TODO: optimize this function]
ip_x0_qr = Kernel::mask_ip_x0_q(q_obj.rotated_query(), cur_bin.bin_code(), padded_dim);
// Direct HNSW kernels retain their legacy word-typed ABI. Generic kernels use the
// byte-oriented overload and do not manufacture a typed view of packed storage.
ip_x0_qr = Kernel::mask_ip_x0_q(
q_obj.rotated_query(),
reinterpret_cast<const uint64_t*>(cur_bin.bin_code()),
padded_dim
);

est_dist =
cur_ex.f_add_ex() + g_add +
Expand Down Expand Up @@ -176,9 +190,13 @@ inline void qg_batch_estdist(
);

ConstRowMajorArrayMap<TA> ip_arr(accu_res.data(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<T> f_add_arr(cur_batch.f_add(), 1, fastscan::kBatchSize);
std::array<T, fastscan::kBatchSize> f_add_values;
std::array<T, fastscan::kBatchSize> f_rescale_values;
cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size());
cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size());
ConstRowMajorArrayMap<T> f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<T> f_rescale_arr(
cur_batch.f_rescale(), 1, fastscan::kBatchSize
f_rescale_values.data(), 1, fastscan::kBatchSize
);
RowMajorArrayMap<T> est_dist_arr(est_distance, 1, fastscan::kBatchSize);

Expand Down Expand Up @@ -210,8 +228,14 @@ inline void qg_batch_estdist(
}

ConstRowMajorArrayMap<int32_t> ip_arr(accu_values.data(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<T> f_add_arr(cur_batch.f_add(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<T> f_rescale_arr(cur_batch.f_rescale(), 1, fastscan::kBatchSize);
std::array<T, fastscan::kBatchSize> f_add_values;
std::array<T, fastscan::kBatchSize> f_rescale_values;
cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size());
cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size());
ConstRowMajorArrayMap<T> f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize);
ConstRowMajorArrayMap<T> f_rescale_arr(
f_rescale_values.data(), 1, fastscan::kBatchSize
);

RowMajorArrayMap<T> est_dist_arr(est_distance, 1, fastscan::kBatchSize);

Expand Down Expand Up @@ -294,8 +318,9 @@ inline void split_single_estdist_direct(
) {
ConstBinDataMap<float> cur_bin(bin_data, padded_dim);

// Direct HNSW kernels retain their legacy word-typed ABI; see the note above.
ip_x0_qr = Kernel::warmup_ip_x0_q_512(
cur_bin.bin_code(),
reinterpret_cast<const uint64_t*>(cur_bin.bin_code()),
q_obj.query_bin(),
q_obj.delta(),
q_obj.vl(),
Expand Down
91 changes: 58 additions & 33 deletions include/rabitqlib/index/ivf/initializer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,13 @@
namespace rabitqlib::ivf {
template <class Function>
inline void parallel_for(size_t start, size_t end, size_t numThreads, Function fn) {
if (numThreads <= 0) {
numThreads = std::thread::hardware_concurrency();
if (start >= end) {
return;
}
if (numThreads == 0) {
numThreads = std::max<size_t>(1, std::thread::hardware_concurrency());
}
numThreads = std::min(numThreads, end - start);

if (numThreads == 1) {
for (size_t id = start; id < end; id++) {
Expand All @@ -38,36 +42,49 @@ inline void parallel_for(size_t start, size_t end, size_t numThreads, Function f
std::exception_ptr last_exception = nullptr;
std::mutex last_except_mutex;

threads.reserve(numThreads);
for (size_t thread_id = 0; thread_id < numThreads; ++thread_id) {
threads.push_back(std::thread([&, thread_id] {
while (true) {
size_t id = current.fetch_add(1);

if (id >= end) {
break;
}

try {
fn(id, thread_id);
} catch (...) {
std::unique_lock<std::mutex> last_except_lock(last_except_mutex);
last_exception = std::current_exception();
/*
* This will work even when current is the largest value that
* size_t can fit, because fetch_add returns the previous value
* before the increment (what will result in overflow
* and produce 0 instead of current + 1).
*/
current = end;
break;
}
const auto join_threads = [&threads] {
for (auto& thread : threads) {
if (thread.joinable()) {
thread.join();
}
}));
}
for (auto& thread : threads) {
thread.join();
}
};

try {
threads.reserve(numThreads);
for (size_t thread_id = 0; thread_id < numThreads; ++thread_id) {
threads.emplace_back([&, thread_id] {
while (true) {
size_t id = current.fetch_add(1);

if (id >= end) {
break;
}

try {
fn(id, thread_id);
} catch (...) {
std::unique_lock<std::mutex> last_except_lock(last_except_mutex
);
last_exception = std::current_exception();
/*
* This will work even when current is the largest value that
* size_t can fit, because fetch_add returns the previous value
* before the increment (what will result in overflow
* and produce 0 instead of current + 1).
*/
current = end;
break;
}
}
});
}
} catch (...) {
current = end;
join_threads();
throw;
}
join_threads();
if (last_exception) {
std::rethrow_exception(last_exception);
}
Expand Down Expand Up @@ -97,10 +114,15 @@ inline Initializer::~Initializer() {}
class FlatInitializer : public Initializer {
private:
std::vector<float> centroids_;
MetricType metric_type_;

public:
explicit FlatInitializer(size_t d, size_t k)
: Initializer(d, k), centroids_(num_cluster_ * dim_) {}
explicit FlatInitializer(
size_t d, size_t k, MetricType metric_type = MetricType::METRIC_L2
)
: Initializer(d, k), centroids_(num_cluster_ * dim_), metric_type_(metric_type) {
validate_metric_type(metric_type_);
}

~FlatInitializer() override = default;

Expand All @@ -118,7 +140,10 @@ class FlatInitializer : public Initializer {
std::vector<AnnCandidate<float>> centroid_dist(this->num_cluster_);
for (PID i = 0; i < num_cluster_; ++i) {
centroid_dist[i].id = i;
centroid_dist[i].distance = std::sqrt(euclidean_sqr(query, centroid(i), dim_));
centroid_dist[i].distance =
metric_type_ == METRIC_IP
? dot_product_dis(query, centroid(i), dim_)
: std::sqrt(euclidean_sqr(query, centroid(i), dim_));
}
std::partial_sort(
centroid_dist.begin(),
Expand Down
Loading
Loading