From 23aed4e4bcda76de0dee037c23360eed88e43324 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 01/39] chore(cagra): drop local dev artifacts (.clangd, .gitignore entries) --- .gitignore | 7 +++++- cpp/.clangd | 65 ----------------------------------------------------- 2 files changed, 6 insertions(+), 66 deletions(-) delete mode 100644 cpp/.clangd diff --git a/.gitignore b/.gitignore index 3627558ff5..0066d2b89a 100644 --- a/.gitignore +++ b/.gitignore @@ -72,7 +72,9 @@ docs/source/_static/rust # clang tooling compile_commands.json -.clangd/ + + + # serialized ann indexes brute_force_index @@ -86,5 +88,8 @@ ivf_pq_index /datasets/ /*.json +# clangd +*/.clangd + # java .classpath diff --git a/cpp/.clangd b/cpp/.clangd deleted file mode 100644 index 7c4fe036dd..0000000000 --- a/cpp/.clangd +++ /dev/null @@ -1,65 +0,0 @@ -# https://clangd.llvm.org/config - -# Apply a config conditionally to all C files -If: - PathMatch: .*\.(c|h)$ - ---- - -# Apply a config conditionally to all C++ files -If: - PathMatch: .*\.(c|h)pp - ---- - -# Apply a config conditionally to all CUDA files -If: - PathMatch: .*\.cuh? -CompileFlags: - Add: - - "-x" - - "cuda" - # No error on unknown CUDA versions - - "-Wno-unknown-cuda-version" - # Allow variadic CUDA functions - - "-Xclang=-fcuda-allow-variadic-functions" -Diagnostics: - Suppress: - - "variadic_device_fn" - - "attributes_not_allowed" - ---- - -# Tweak the clangd parse settings for all files -CompileFlags: - Add: - # report all errors - - "-ferror-limit=0" - - "-fmacro-backtrace-limit=0" - - "-ftemplate-backtrace-limit=0" - # Skip the CUDA version check - - "--no-cuda-version-check" - Remove: - # remove gcc's -fcoroutines - - -fcoroutines - # remove nvc++ flags unknown to clang - - "-gpu=*" - - "-stdpar*" - # remove nvcc flags unknown to clang - - "-arch*" - - "-gencode*" - - "--generate-code*" - - "-ccbin*" - - "-t=*" - - "--threads*" - - "-Xptxas*" - - "-Xcudafe*" - - "-Xfatbin*" - - "-Xcompiler*" - - "--diag-suppress*" - - "--diag_suppress*" - - "--compiler-options*" - - "--expt-extended-lambda" - - "--expt-relaxed-constexpr" - - "-forward-unknown-to-host-compiler" - - "-Werror=cross-execution-space-call" From a4ebee645418ff566f972ddfd0811147b78905be Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 02/39] feat(cagra): add batched_device_view_from_host utility and unit test --- cpp/src/neighbors/detail/cagra/utils.hpp | 411 +++++++++++++++++- cpp/tests/CMakeLists.txt | 5 +- .../test_batched_device_view_from_host.cu | 205 +++++++++ 3 files changed, 619 insertions(+), 2 deletions(-) create mode 100644 cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu diff --git a/cpp/src/neighbors/detail/cagra/utils.hpp b/cpp/src/neighbors/detail/cagra/utils.hpp index 58bf68bb43..2313631372 100644 --- a/cpp/src/neighbors/detail/cagra/utils.hpp +++ b/cpp/src/neighbors/detail/cagra/utils.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -161,6 +161,22 @@ struct gen_index_msb_1_mask { }; } // namespace utils +template +bool is_ptr_device_accessible(T* ptr) +{ + cudaPointerAttributes attr; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); + return attr.devicePointer != nullptr; +} + +template +bool is_ptr_host_accessible(T* ptr) +{ + cudaPointerAttributes attr; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); + return attr.hostPointer != nullptr; +} + /** * Utility to sync memory from a host_matrix_view to a device_matrix_view * @@ -301,4 +317,397 @@ void copy_with_padding( } } +/** + * Utility to create a batched device view from a host view + * + * This utility will create a batched device view from a host view and will handle the prefetch and + * writeback of the data Each batch can be referenced exactlyonce by calling the next_view() + * function + * + * Usage: + * ``` + * batched_device_view_from_host view(res, host_view, batch_size, host_writeback, + * initialize); while (view.next_view().extent(0) > 0) { auto device_view = view.next_view(); + * // use device_view + * } + * ``` + * + * The call to next_view() will + * * synchronize on all previous operations / increments batch_id_ + * * (optionally) write back the data of the previous batch to the host + * * (optionally) prefetch the data of the next batch + * * return the view of the current batch + * + * @tparam T The type of the data + * @tparam IdxT The type of the index + */ +template +class batched_device_view_from_host { + public: + enum class memory_strategy { + device_only, // data is on device only (no copy needed) + copy_device, // data is explicitly moved to/from device buffers + managed_only, // data is on managed memory (system managed) + }; + + /** + * Create a batched device view from a host view and will handle the prefetch and + * writeback of the data. Each batch can be referenced exactly once by calling the next_view() + * method. + * + * @param res The resources to use + * @param host_view The host view to create the batched device view from + * @param batch_size The batch size + * @param host_writeback Whether to write back the data to the host (only for host memory) + * (default: false) + * @param initialize Whether to initialize the data (only for managed memory) (default: true) + */ + batched_device_view_from_host(raft::resources const& res, + raft::host_matrix_view host_view, + uint64_t batch_size, + bool host_writeback = false, + bool initialize = true) + : res_(res), + host_view_(host_view), + batch_size_(batch_size), + offset_(0), + batch_id_(-2), + num_buffers_(2), + host_writeback_(host_writeback), + initialize_(initialize) + { + if (host_view.extent(0) == 0) { + mem_strategy_ = memory_strategy::device_only; + return; + } + + RAFT_EXPECTS(host_writeback_ || initialize_, + "At least one of host_writeback or initialize must be true"); + + RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr_, host_view.data_handle())); + switch (attr_.type) { + case cudaMemoryTypeUnregistered: + case cudaMemoryTypeHost: + case cudaMemoryTypeManaged: mem_strategy_ = memory_strategy::copy_device; break; + case cudaMemoryTypeDevice: mem_strategy_ = memory_strategy::device_only; break; + } + + RAFT_LOG_DEBUG("Memory strategy: %d for type %d, size %zu", + static_cast(mem_strategy_), + static_cast(attr_.type), + host_view.extent(0) * host_view.extent(1) * sizeof(T)); + + // buffer allocations + if (mem_strategy_ == memory_strategy::copy_device) { + try { + device_mem_[0].emplace(raft::make_device_mdarray( + res, + raft::resource::get_workspace_resource_ref(res), + raft::make_extents(batch_size, host_view.extent(1)))); + device_ptr[0] = device_mem_[0]->data_handle(); + if (batch_size < static_cast(host_view.extent(0))) { + device_mem_[1].emplace(raft::make_device_mdarray( + res, + raft::resource::get_workspace_resource_ref(res), + raft::make_extents(batch_size, host_view.extent(1)))); + device_ptr[1] = device_mem_[1]->data_handle(); + } + if (host_writeback_ && initialize_ && + batch_size * 2 < static_cast(host_view.extent(0))) { + num_buffers_ = 3; + device_mem_[2].emplace(raft::make_device_mdarray( + res, + raft::resource::get_workspace_resource_ref(res), + raft::make_extents(batch_size, host_view.extent(1)))); + device_ptr[2] = device_mem_[2]->data_handle(); + } + } catch (std::bad_alloc& e) { + if (attr_.devicePointer != nullptr) { + RAFT_LOG_DEBUG("Insufficient memory for device buffers, switching to managed memory"); + mem_strategy_ = memory_strategy::managed_only; + } else { + throw std::bad_alloc(); + } + } catch (raft::logic_error& e) { + if (attr_.devicePointer != nullptr) { + RAFT_LOG_DEBUG( + "Insufficient memory for device buffers (logic error), switching to managed memory"); + mem_strategy_ = memory_strategy::managed_only; + } else { + throw raft::logic_error("Insufficient memory for device buffers (logic error)"); + } + } + } + + // setup stream pool if not already present + size_t required_streams = host_writeback_ && initialize_ ? 2 : 1; + if (!res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) || + raft::resource::get_stream_pool_size(res) < required_streams) { + // always create at least 2 streams to account for subsequent iterator calls. + // set_cuda_stream_pool now requires a non-const resource; the referenced resource + // outlives this object, so attaching the pool to it here is safe. + raft::resource::set_cuda_stream_pool(const_cast(res), + std::make_shared(2)); + } + prefetch_stream_ = raft::resource::get_stream_from_stream_pool(res); + writeback_stream_ = raft::resource::get_stream_from_stream_pool(res); + + // if data is managed and not for_write_ we can set the attribute on the device ptr + if (mem_strategy_ == memory_strategy::managed_only) { + location_.type = cudaMemLocationTypeDevice; + location_.id = static_cast(raft::resource::get_device_id(res_)); + if (!host_writeback_) { + advise_read_mostly(host_view_.data_handle(), + host_view_.extent(0) * host_view_.extent(1) * sizeof(T)); + // TODO maybe also reset upon destruction + } + } + + // prefetch next batch (0) + prefetch_next_batch(); + } + + ~batched_device_view_from_host() noexcept + { + raft::resource::sync_stream(res_); + + // if data is on host and for_write --> make sure to copy back last active + // if data is managed and evict --> evict last active + + // make sure to sync on prefetch stream & res + switch (mem_strategy_) { + case memory_strategy::managed_only: + if (!host_writeback_) { + uint32_t discard_pos = batch_id_ % num_buffers_; + size_t discard_size_rows = actual_batch_size_[discard_pos]; + if (batch_id_ > 0) { + discard_pos = (batch_id_ - 1) % num_buffers_; + discard_size_rows += batch_size_; + } + discard_managed_region(device_ptr[discard_pos], + discard_size_rows * host_view_.extent(1) * sizeof(T)); + writeback_stream_.synchronize(); + } + break; + case memory_strategy::copy_device: + if (host_writeback_) { + uint32_t writeback_pos_last = batch_id_ % num_buffers_; + if (batch_id_ > 0) { + uint32_t writeback_pos = (batch_id_ - 1) % num_buffers_; + uint64_t writeback_offset = (batch_id_ - 1) * batch_size_; + writeback_from_device_to_host(device_ptr[writeback_pos], writeback_offset, batch_size_); + } + { + uint64_t writeback_offset_last = batch_id_ * batch_size_; + writeback_from_device_to_host(device_ptr[writeback_pos_last], + writeback_offset_last, + actual_batch_size_[writeback_pos_last]); + } + writeback_stream_.synchronize(); + } + break; + case memory_strategy::device_only: break; + } + } + + /** + * Returns the next view of the batch + * + * This function will ensure the next batch is ready and will trigger the prefetch of the + * subsequent next batch. If writeback is enabled, the last active batch will be written back to + * the host. + * + * @return The next view of the batch + */ + raft::device_matrix_view next_view() + { + bool end_of_data = static_cast((batch_id_ + 1) * batch_size_) >= + static_cast(host_view_.extent(0)); + + // special case for empty host view or last batch surpassed + if (end_of_data) { + return raft::make_device_matrix_view(nullptr, 0, host_view_.extent(1)); + } + + // trigger prefetch of next batch (also increments batch_id_) + prefetch_next_batch(); + + uint32_t current_pos = batch_id_ % num_buffers_; + return raft::make_device_matrix_view( + device_ptr[current_pos], actual_batch_size_[current_pos], host_view_.extent(1)); + } + + private: + /** + * Prefetch the next batch + * + * This function will prefetch the next batch and will handle the writeback of the data. + * + * @return True if the next batch exists, false otherwise + */ + bool prefetch_next_batch() + { + batch_id_++; + + // ensure previous batch at position batch_id_ is ready + if (initialize_) { prefetch_stream_.synchronize(); } + if (host_writeback_) { writeback_stream_.synchronize(); } + + // this step will + // * write back data from batch_id_ - 1 + // * prefetch data for batch_id_ + 1 + + // if data is on host and host_writeback_ is true we will have to copy it back + // if data is on host and initialize_ is true we will have to copy it to the device_ptr + + // if data is managed and !host_writeback_ we can discard the data from device memory + // if data is managed and initialize_ is true we can prefetch it to the device + // if data is managed and !initialize_ we can discard and prefetch the data location + + // if data is on device only this is almost a noop, just prepping the pointers + + RAFT_EXPECTS(static_cast(offset_) <= host_view_.extent(0), "Offset out of bounds"); + + bool next_batch_exists = offset_ < static_cast(host_view_.extent(0)); + + if (next_batch_exists) { + // synchronize to ensure all previous operations are completed + // in particular all work on batch_id_ - 1 + raft::resource::sync_stream(res_); + + int32_t prefetch_pos = (batch_id_ + 1) % num_buffers_; + actual_batch_size_[prefetch_pos] = min(batch_size_, host_view_.extent(0) - offset_); + + switch (mem_strategy_) { + case memory_strategy::managed_only: + if (!host_writeback_ && batch_id_ > 1) { + uint32_t discard_pos = (batch_id_ - 1) % num_buffers_; + size_t discard_size = batch_size_ * host_view_.extent(1) * sizeof(T); + discard_managed_region(device_ptr[discard_pos], discard_size); + } + // prefetch next position + device_ptr[prefetch_pos] = host_view_.data_handle() + offset_ * host_view_.extent(1); + prefetch_managed_region( + device_ptr[prefetch_pos], + actual_batch_size_[prefetch_pos] * host_view_.extent(1) * sizeof(T)); + break; + case memory_strategy::copy_device: + if (host_writeback_ && batch_id_ > 0) { + // copy back last active + uint32_t writeback_pos = (batch_id_ - 1) % num_buffers_; + uint64_t writeback_offset = (batch_id_ - 1) * batch_size_; + writeback_from_device_to_host(device_ptr[writeback_pos], writeback_offset, batch_size_); + } + if (initialize_) { + // prefetch next position + prefetch_from_host_to_device( + device_ptr[prefetch_pos], offset_, actual_batch_size_[prefetch_pos]); + } + + break; + case memory_strategy::device_only: + // just move pointer to next position + device_ptr[prefetch_pos] = host_view_.data_handle() + offset_ * host_view_.extent(1); + break; + } + + offset_ += actual_batch_size_[prefetch_pos]; + } + + return next_batch_exists; + } + + void advise_read_mostly(T* ptr, size_t size) + { +#if CUDA_VERSION >= 13000 + RAFT_CUDA_TRY(cudaMemAdvise(ptr, size, cudaMemAdviseSetReadMostly, location_)); +#else + RAFT_CUDA_TRY(cudaMemAdvise_v2(ptr, size, cudaMemAdviseSetReadMostly, location_)); +#endif + } + + void discard_managed_region(T* dev_ptr, size_t size) + { +#if CUDA_VERSION >= 13000 + void* dptrs[1] = {dev_ptr}; + size_t sizes[1] = {size}; + RAFT_CUDA_TRY(cudaMemDiscardBatchAsync(dptrs, sizes, 1, 0, writeback_stream_)); +#endif + // FIXME: CUDA12 does not support discard + } + + void prefetch_managed_region(T* dev_ptr, size_t size) + { +#if CUDA_VERSION >= 13000 + if (initialize_) { + RAFT_CUDA_TRY(cudaMemPrefetchAsync(dev_ptr, size, location_, 0, prefetch_stream_)); + } else { + void* dptrs[1] = {dev_ptr}; + size_t sizes[1] = {size}; + RAFT_CUDA_TRY( + cudaMemDiscardAndPrefetchBatchAsync(dptrs, sizes, 1, location_, 0, prefetch_stream_)); + } +#else + // FIXME: CUDA12 does not support discard - so we just prefetch + if (initialize_) { + RAFT_CUDA_TRY(cudaMemPrefetchAsync_v2(dev_ptr, size, location_, 0, prefetch_stream_)); + } else { + RAFT_CUDA_TRY(cudaMemPrefetchAsync_v2(dev_ptr, size, location_, 0, prefetch_stream_)); + } +#endif + } + + void prefetch_from_host_to_device(T* dev_ptr, size_t src_row_offset, size_t num_rows) + { + const size_t n_elem = num_rows * host_view_.extent(1); + const size_t n_bytes = n_elem * sizeof(T); + // use memcpy instead of raft::copy to avoid strange behavior with HMM/ATS memory + RAFT_CUDA_TRY(cudaMemcpyAsync(dev_ptr, + host_view_.data_handle() + src_row_offset * host_view_.extent(1), + n_bytes, + cudaMemcpyHostToDevice, + prefetch_stream_)); + } + + void writeback_from_device_to_host(T* dev_ptr, size_t dst_row_offset, size_t num_rows) + { + const size_t n_elem = num_rows * host_view_.extent(1); + const size_t n_bytes = n_elem * sizeof(T); + // use memcpy instead of raft::copy to avoid strange behavior with HMM/ATS memory + RAFT_CUDA_TRY(cudaMemcpyAsync(host_view_.data_handle() + dst_row_offset * host_view_.extent(1), + dev_ptr, + n_bytes, + cudaMemcpyDeviceToHost, + writeback_stream_)); + } + + // stream pool for local streams + std::optional> local_stream_pool_; + rmm::cuda_stream_view prefetch_stream_; + rmm::cuda_stream_view writeback_stream_; + + // configuration + memory_strategy mem_strategy_; + const raft::resources& res_; + bool initialize_; // initialize the data on the device + bool host_writeback_; // write back the data to the host + + // batch position information + uint64_t batch_size_; + int32_t batch_id_; + uint64_t offset_; + + cudaMemLocation location_; + + // input pointer information + raft::host_matrix_view host_view_; + cudaPointerAttributes attr_; + + // internal device buffers + uint64_t num_buffers_; + std::optional> device_mem_[3]; + T* device_ptr[3]; + uint32_t actual_batch_size_[3]; +}; + } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 13a07b10b5..5e6e3b6ea6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -183,6 +183,7 @@ ConfigureTest( neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu neighbors/ann_cagra/bug_iterative_cagra_build.cu neighbors/ann_cagra/bug_issue_93_reproducer.cu + neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu GPUS 1 PERCENT 100 ) @@ -203,7 +204,9 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST - PATH neighbors/ann_cagra/test_optimize_uint32_t.cu neighbors/ann_cagra/test_batch_load_iterator.cu + PATH neighbors/ann_cagra/test_optimize_uint32_t.cu + neighbors/ann_cagra/test_batched_device_view_from_host.cu + neighbors/ann_cagra/test_batch_load_iterator.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu b/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu new file mode 100644 index 0000000000..1e1cc13093 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu @@ -0,0 +1,205 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "../../../src/neighbors/detail/cagra/utils.hpp" + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +using IdxT = uint32_t; + +struct BatchConfig { + bool initialize; + bool host_writeback; +}; + +struct DimsConfig { + int64_t n_rows; + int64_t n_cols; + uint64_t batch_size; +}; + +class BatchedDeviceViewFromHostTest : public ::testing::Test { + protected: + void SetUp() override { raft::resource::sync_stream(res); } + + /** + * Run batched_device_view_from_host over host data, copy device views back, + * and verify against the input. + */ + template + void run_and_verify_batched(InputMatrixView input_view, + uint64_t batch_size, + bool host_writeback, + bool initialize) + { + int64_t n_rows = input_view.extent(0); + int64_t n_cols = input_view.extent(1); + + std::vector readback(n_rows * n_cols); + + int64_t total_processed = 0; + + { + cagra::detail::batched_device_view_from_host batched( + res, + raft::make_host_matrix_view(input_view.data_handle(), n_rows, n_cols), + batch_size, + host_writeback, + initialize); + while (true) { + auto dev_view = batched.next_view(); + if (dev_view.extent(0) == 0) break; + + if (initialize) { + raft::copy(readback.data() + total_processed * n_cols, + dev_view.data_handle(), + dev_view.extent(0) * dev_view.extent(1), + raft::resource::get_cuda_stream(res)); + } + if (host_writeback) { raft::matrix::fill(res, dev_view, IdxT(17)); } + total_processed += dev_view.extent(0); + } + } + raft::resource::sync_stream(res); + + EXPECT_EQ(total_processed, n_rows); + if (initialize) { + for (int64_t i = 0; i < n_rows * n_cols; ++i) { + EXPECT_EQ(readback[i], IdxT(13)) << "Mismatch (initialize) at index " << i; + } + } + if (host_writeback) { + auto readback_view = + raft::make_host_matrix_view(readback.data(), n_rows, n_cols); + raft::copy(res, readback_view, input_view); + raft::resource::sync_stream(res); + for (int64_t i = 0; i < n_rows * n_cols; ++i) { + EXPECT_EQ(readback[i], IdxT(17)) << "Mismatch (host_writeback) at index " << i; + } + } + } + + raft::resources res; +}; + +TEST_F(BatchedDeviceViewFromHostTest, EmptyView) +{ + auto host_empty = raft::make_host_matrix(0, 8); + auto host_view = host_empty.view(); + cagra::detail::batched_device_view_from_host batched( + res, host_view, /*batch_size=*/128, /*host_writeback=*/false, /*initialize=*/true); + + auto view = batched.next_view(); + EXPECT_EQ(view.extent(0), 0); + EXPECT_EQ(view.extent(1), 8); + EXPECT_EQ(view.data_handle(), nullptr); +} + +using BatchDimsParam = std::tuple; + +class BatchedDeviceViewFromHostParameterizedTest + : public BatchedDeviceViewFromHostTest, + public ::testing::WithParamInterface {}; + +TEST_P(BatchedDeviceViewFromHostParameterizedTest, VectorHostData) +{ + auto [batch_config, dims_config] = GetParam(); + auto [initialize, host_writeback] = batch_config; + auto [n_rows, n_cols, batch_size] = dims_config; + + std::vector host_data(n_rows * n_cols); + auto host_view = raft::make_host_matrix_view(host_data.data(), n_rows, n_cols); + + std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); + + run_and_verify_batched(host_view, batch_size, host_writeback, initialize); +} + +TEST_P(BatchedDeviceViewFromHostParameterizedTest, PinnedMemory) +{ + auto [batch_config, dims_config] = GetParam(); + auto [initialize, host_writeback] = batch_config; + auto [n_rows, n_cols, batch_size] = dims_config; + + auto host_matrix = raft::make_pinned_matrix(res, n_rows, n_cols); + auto host_view = host_matrix.view(); + + std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); + + run_and_verify_batched(host_view, batch_size, host_writeback, initialize); +} + +TEST_P(BatchedDeviceViewFromHostParameterizedTest, ManagedMemory) +{ + auto [batch_config, dims_config] = GetParam(); + auto [initialize, host_writeback] = batch_config; + auto [n_rows, n_cols, batch_size] = dims_config; + + auto host_matrix = raft::make_managed_matrix(res, n_rows, n_cols); + auto host_view = host_matrix.view(); + + std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); + + run_and_verify_batched(host_view, batch_size, host_writeback, initialize); +} + +TEST_P(BatchedDeviceViewFromHostParameterizedTest, DeviceMemory) +{ + auto [batch_config, dims_config] = GetParam(); + auto [initialize, host_writeback] = batch_config; + auto [n_rows, n_cols, batch_size] = dims_config; + + auto host_matrix = raft::make_device_matrix(res, n_rows, n_cols); + auto host_view = host_matrix.view(); + + raft::matrix::fill(res, host_view, IdxT(13)); + + run_and_verify_batched(host_view, batch_size, host_writeback, initialize); +} + +static const std::array kBatchConfigs = {{ + {/*initialize=*/true, /*host_writeback=*/false}, + {/*initialize=*/false, /*host_writeback=*/true}, + {/*initialize=*/true, /*host_writeback=*/true}, +}}; + +static const std::array kDimsConfigs = {{ + {/*n_rows=*/64, /*n_cols=*/32, /*batch_size=*/256}, // rows less than batch size, single batch + {/*n_rows=*/64, /*n_cols=*/32, /*batch_size=*/64}, // single batch + {/*n_rows=*/256, /*n_cols=*/32, /*batch_size=*/32}, // multiple batches + {/*n_rows=*/500, + /*n_cols=*/32, + /*batch_size=*/128}, // multiple batches, partial batch in the end +}}; + +INSTANTIATE_TEST_SUITE_P(BatchConfigs, + BatchedDeviceViewFromHostParameterizedTest, + ::testing::Combine(::testing::ValuesIn(kBatchConfigs), + ::testing::ValuesIn(kDimsConfigs))); + +} // namespace cuvs::neighbors::cagra From 90fcf80c2af29a720bc9ff18e46faed36e828169 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 03/39] fix(cagra): bound random seed selection to graph size during build --- cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index e09ef82a39..f1c7305833 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh @@ -511,7 +511,11 @@ struct search hashmap.data(), hash_bitlen, stream, - static_cast(this->dataset_size)); + // Bound random seed selection to the graph size, not the dataset size. + // During iterative / CAGRA-Q build the graph is smaller than the dataset, + // so using dataset_size here selects seeds that index past the graph end + // (out-of-bounds access). See https://github.com/rapidsai/cuvs/pull/1780. + static_cast(graph.extent(0))); std::shared_ptr compute_distance_to_child_nodes_launcher = make_cagra_multi_kernel_jit_launcher Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 04/39] feat(cagra): iterative CAGRA-Q graph build with configurable in-build search - Configurable growth-phase in-build search params (itopk_size, search_width, max_iterations) and internal/smem dtype; itopk auto-forced on the final full-size iteration. - Decouple compression params used during iterative construction from the target index compression. - Add shuffle_dataset option; fix out-of-bounds access from the in-place raft gather by switching to an out-of-place gather. --- cpp/include/cuvs/neighbors/cagra.hpp | 276 ++++++---- cpp/include/cuvs/neighbors/common.hpp | 12 + .../neighbors/detail/cagra/cagra_build.cuh | 480 +++++++++++++++--- .../neighbors/detail/cagra/cagra_search.cuh | 4 + .../detail/cagra/compute_distance.hpp | 2 + 5 files changed, 603 insertions(+), 171 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index d1937cba27..d2dabcd4d7 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -32,10 +32,172 @@ #include #include +namespace CUVS_EXPORT cuvs { +namespace neighbors { +namespace cagra { + +/** + * @defgroup cagra_cpp_search_params CAGRA index search parameters + * @{ + */ + +enum class search_algo { + /** For large batch sizes. */ + SINGLE_CTA = 0, + /** For small batch sizes. */ + MULTI_CTA = 1, + MULTI_KERNEL = 2, + AUTO = 100 +}; + +enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; + +enum class internal_dtype { F16 = 0, E5M2 = 1 }; + +struct search_params : cuvs::neighbors::search_params { + /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ + size_t max_queries = 0; + + /** Number of intermediate search results retained during the search. + * + * This is the main knob to adjust trade off between accuracy and search speed. + * Higher values improve the search accuracy. + */ + size_t itopk_size = 64; + + /** Upper limit of search iterations. Auto select when 0.*/ + size_t max_iterations = 0; + + // In the following we list additional search parameters for fine tuning. + // Reasonable default values are automatically chosen. + + /** Which search implementation to use. */ + search_algo algo = search_algo::AUTO; + + /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ + size_t team_size = 0; + + /** Number of graph nodes to select as the starting point for the search in each iteration. aka + * search width?*/ + size_t search_width = 1; + /** Lower limit of search iterations. */ + size_t min_iterations = 0; + + /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ + size_t thread_block_size = 0; + /** Hashmap type. Auto selection when AUTO. */ + hash_mode hashmap_mode = hash_mode::AUTO; + /** Lower limit of hashmap bit length. More than 8. */ + size_t hashmap_min_bitlen = 0; + /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ + float hashmap_max_fill_rate = 0.5; + + /** Number of iterations of initial random seed node selection. 1 or more. */ + uint32_t num_random_samplings = 1; + /** Bit mask used for initial random seed node selection. */ + uint64_t rand_xor_mask = 0x128394; + + /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ + bool persistent = false; + /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ + float persistent_lifetime = 2; + /** + * Set the fraction of maximum grid size used by persistent kernel. + * Value 1.0 means the kernel grid size is maximum possible for the selected device. + * The value must be greater than 0.0 and not greater than 1.0. + * + * One may need to run other kernels alongside this persistent kernel. This parameter can + * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. + * Note: running any other work on GPU alongside with the persistent kernel makes the setup + * fragile. + * - Running another kernel in another thread usually works, but no progress guaranteed + * - Any CUDA allocations block the context (this issue may be obscured by using pools) + * - Memory copies to not-pinned host memory may block the context + * + * Even when we know there are no other kernels working at the same time, setting + * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. + * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant + * impact on the throughput. + */ + float persistent_device_usage = 1.0; + + /** + * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. + * The value must be equal to or greater than 0.0 and less than 1.0. Default value is + * negative, in which case the filtering rate is automatically calculated when possible. + * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is + * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be + * inferred from the source string. + */ + float filtering_rate = -1.0; + + /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ + * supports FP8. **/ + internal_dtype smem_dtype = internal_dtype::F16; +}; + +/** + * @} + */ + +} // namespace cagra +} // namespace neighbors +} // namespace CUVS_EXPORT cuvs + namespace CUVS_EXPORT cuvs { namespace neighbors { namespace graph_build_params { -using iterative_search_params = cuvs::neighbors::search_params; +/** + * Parameters for the iterative CAGRA graph build algorithm. + * + * Inherits from cagra::search_params so that all search tuning knobs + * (search_width, max_iterations, itopk_size, etc.) are available for + * controlling the search-and-optimize loop during graph construction. + * The defaults are tuned for the build loop (e.g. search_width=1, + * max_iterations=8) and may differ from the regular search defaults. + * + * `build_compression` controls the VPQ parameters applied to the dataset + * *while building the graph*. This is independent of `index_params::compression`, + * which controls the compression of the dataset stored in the final index. + */ +struct iterative_search_params : cuvs::neighbors::cagra::search_params { + /** + * Optional VPQ compression parameters used during iterative graph construction. + * + * When set, the dataset is compressed with these parameters for the + * search-and-optimize loop. When std::nullopt (default), the builder + * falls back to `index_params::compression` (original behaviour). + */ + std::optional build_compression = std::nullopt; + + /** + * Whether to shuffle the dataset before building the graph. + * + * When enabled, the compressed dataset is randomly permuted before graph + * construction begins. This can improve graph quality by breaking any + * spatial locality in the original dataset ordering that might cause + * the iterative builder to get stuck in local optima during early + * iterations. + * + * After graph construction, the node indices in the graph are remapped + * back to the original dataset ordering. + * + * Only applies when compression is enabled (build_compression or + * index_params::compression is set). + */ + bool shuffle_dataset = true; + + iterative_search_params() + { + this->search_width = 1; + this->max_iterations = 8; + // itopk_size controls the search during the *growing* iterations of the build loop. + // 0 (default) means auto-select per iteration (max(graph_degree + 32, 128)); a nonzero + // value overrides it for the growing iterations. The final iteration always uses a fixed + // itopk tied to the output topk, regardless of this value. + this->itopk_size = 0; + } +}; /** Specialized parameters for ACE (Augmented Core Extraction) graph build */ struct ace_params { @@ -192,6 +354,14 @@ struct index_params : cuvs::neighbors::index_params { */ bool guarantee_connectivity = false; + /** + * Whether to skip graph optimization (pruning, reverse edges, MST) during non-final iterations + * of iterative graph building. When true, search results are copied directly into the device + * graph without host round-trips. Only applies to iterative_search_params graph builds; the + * final iteration always runs full optimization. + */ + bool skip_graph_optimization = false; + /** * Whether to add the dataset content to the index, i.e.: * @@ -257,110 +427,6 @@ struct index_params : cuvs::neighbors::index_params { cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded); }; -/** - * @} - */ - -/** - * @defgroup cagra_cpp_search_params CAGRA index search parameters - * @{ - */ - -enum class search_algo { - /** For large batch sizes. */ - SINGLE_CTA = 0, - /** For small batch sizes. */ - MULTI_CTA = 1, - MULTI_KERNEL = 2, - AUTO = 100 -}; - -enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 }; - -enum class internal_dtype { F16 = 0, E5M2 = 1 }; - -struct search_params : cuvs::neighbors::search_params { - /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ - size_t max_queries = 0; - - /** Number of intermediate search results retained during the search. - * - * This is the main knob to adjust trade off between accuracy and search speed. - * Higher values improve the search accuracy. - */ - size_t itopk_size = 64; - - /** Upper limit of search iterations. Auto select when 0.*/ - size_t max_iterations = 0; - - // In the following we list additional search parameters for fine tuning. - // Reasonable default values are automatically chosen. - - /** Which search implementation to use. */ - search_algo algo = search_algo::AUTO; - - /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */ - size_t team_size = 0; - - /** Number of graph nodes to select as the starting point for the search in each iteration. aka - * search width?*/ - size_t search_width = 1; - /** Lower limit of search iterations. */ - size_t min_iterations = 0; - - /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */ - size_t thread_block_size = 0; - /** Hashmap type. Auto selection when AUTO. */ - hash_mode hashmap_mode = hash_mode::AUTO; - /** Lower limit of hashmap bit length. More than 8. */ - size_t hashmap_min_bitlen = 0; - /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/ - float hashmap_max_fill_rate = 0.5; - - /** Number of iterations of initial random seed node selection. 1 or more. */ - uint32_t num_random_samplings = 1; - /** Bit mask used for initial random seed node selection. */ - uint64_t rand_xor_mask = 0x128394; - - /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ - bool persistent = false; - /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ - float persistent_lifetime = 2; - /** - * Set the fraction of maximum grid size used by persistent kernel. - * Value 1.0 means the kernel grid size is maximum possible for the selected device. - * The value must be greater than 0.0 and not greater than 1.0. - * - * One may need to run other kernels alongside this persistent kernel. This parameter can - * be used to reduce the grid size of the persistent kernel to leave a few SMs idle. - * Note: running any other work on GPU alongside with the persistent kernel makes the setup - * fragile. - * - Running another kernel in another thread usually works, but no progress guaranteed - * - Any CUDA allocations block the context (this issue may be obscured by using pools) - * - Memory copies to not-pinned host memory may block the context - * - * Even when we know there are no other kernels working at the same time, setting - * kDeviceUsage to 1.0 surprisingly sometimes hurts performance. Proceed with care. - * If you suspect this is an issue, you can reduce this number to ~0.9 without a significant - * impact on the throughput. - */ - float persistent_device_usage = 1.0; - - /** - * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. - * The value must be equal to or greater than 0.0 and less than 1.0. Default value is - * negative, in which case the filtering rate is automatically calculated when possible. - * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is - * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be - * inferred from the source string. - */ - float filtering_rate = -1.0; - - /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ - * supports FP8. **/ - internal_dtype smem_dtype = internal_dtype::F16; -}; - /** * @} */ diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2fd804f115..1e5ca5a159 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -98,6 +98,18 @@ struct vpq_params { * The max number of data points to use per VQ cluster during training. */ uint32_t max_train_points_per_vq_cluster = 1024; + + friend bool operator==(const vpq_params& a, const vpq_params& b) + { + return a.pq_bits == b.pq_bits && a.pq_dim == b.pq_dim && a.vq_n_centers == b.vq_n_centers && + a.kmeans_n_iters == b.kmeans_n_iters && + a.vq_kmeans_trainset_fraction == b.vq_kmeans_trainset_fraction && + a.pq_kmeans_trainset_fraction == b.pq_kmeans_trainset_fraction && + a.pq_kmeans_type == b.pq_kmeans_type && + a.max_train_points_per_pq_code == b.max_train_points_per_pq_code && + a.max_train_points_per_vq_cluster == b.max_train_points_per_vq_cluster; + } + friend bool operator!=(const vpq_params& a, const vpq_params& b) { return !(a == b); } }; /** @} */ // end group cagra_cpp_index_params diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 774254c84c..f0547beecb 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -19,8 +19,17 @@ #include #include #include +#include +#include +#include +#include #include +#include +#include +#include +#include + #include #include #include @@ -53,6 +62,32 @@ namespace cuvs::neighbors::cagra::detail { constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } +// Functor to remap indices using a permutation lookup table +template +struct remap_indices_op { + const IdxT* perm; + __host__ __device__ IdxT operator()(IdxT idx) const { return perm[idx]; } +}; + +// Functor to compute scattered output index for graph row reordering +template +struct graph_scatter_index_op { + const IdxT* perm; + int64_t degree; + __host__ __device__ int64_t operator()(int64_t idx) const + { + int64_t row = idx / degree; + int64_t col = idx % degree; + return static_cast(perm[row]) * degree + col; + } +}; + +// Functor to convert int64_t to IdxT +template +struct cast_to_idx_op { + __host__ __device__ IdxT operator()(int64_t v) const { return static_cast(v); } +}; + template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { @@ -2016,6 +2051,110 @@ struct mmap_owner { size_t size_; }; +template +__global__ void kern_reconstruct_vpq_queries(const uint8_t* encoded_data, + uint32_t encoded_row_len, + const MathT* vq_codebook, + const MathT* pq_codebook, + uint32_t dim, + uint32_t pq_len, + uint64_t offset, + uint32_t batch_size, + T* output) +{ + const uint64_t batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; + const uint64_t vec_idx = offset + batch_idx; + const uint8_t* vec_data = encoded_data + vec_idx * encoded_row_len; + const uint32_t vq_code = *reinterpret_cast(vec_data); + const uint8_t* pq_codes = vec_data + sizeof(uint32_t); + const MathT* vq_centroid_ptr = vq_codebook + static_cast(vq_code) * dim; + + for (uint32_t d = threadIdx.x; d < dim; d += blockDim.x) { + uint32_t j = d / pq_len; + uint32_t k = d % pq_len; + float val = static_cast(vq_centroid_ptr[d]) + + static_cast(pq_codebook[static_cast(pq_codes[j]) * pq_len + k]); + output[batch_idx * dim + d] = static_cast(val); + } +} + +template +void reconstruct_vpq_queries(raft::resources const& res, + const vpq_dataset& vpq_dset, + uint64_t offset, + uint32_t batch_size, + raft::device_matrix_view output) +{ + const uint32_t dim = vpq_dset.dim(); + const uint32_t pq_len = vpq_dset.pq_len(); + const uint32_t threads = std::min(dim, 256u); + + kern_reconstruct_vpq_queries + <<>>( + vpq_dset.data.data_handle(), + vpq_dset.encoded_row_length(), + vpq_dset.vq_code_book.data_handle(), + vpq_dset.pq_code_book.data_handle(), + dim, + pq_len, + offset, + batch_size, + output.data_handle()); +} + +template +void search_and_optimize(raft::resources const& res, + const cuvs::neighbors::cagra::search_params& search_params, + const index& idx, + raft::device_matrix_view dev_query_view, + raft::device_matrix_view dev_neighbors, + raft::device_matrix_view dev_distances, + raft::device_matrix& dev_output_graph, + size_t curr_query_size, + size_t next_graph_degree, + size_t curr_topk, + uint64_t max_chunk_size) +{ + auto stream = raft::resource::get_cuda_stream(res); + + auto dev_knn_graph = raft::make_device_matrix(res, curr_query_size, curr_topk); + + auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( + res, + dev_query_view.data_handle(), + static_cast(curr_query_size), + static_cast(dev_query_view.extent(1)), + max_chunk_size, + stream, + raft::resource::get_workspace_resource_ref(res)); + for (const auto& batch : query_batch) { + auto batch_dev_query_view = raft::make_device_matrix_view( + batch.data(), batch.size(), dev_query_view.extent(1)); + auto batch_dev_neighbors_view = raft::make_device_matrix_view( + dev_neighbors.data_handle(), batch.size(), curr_topk); + auto batch_dev_distances_view = raft::make_device_matrix_view( + dev_distances.data_handle(), batch.size(), curr_topk); + + cuvs::neighbors::cagra::search(res, + search_params, + idx, + batch_dev_query_view, + batch_dev_neighbors_view, + batch_dev_distances_view); + + raft::copy(dev_knn_graph.data_handle() + batch.offset() * curr_topk, + batch_dev_neighbors_view.data_handle(), + batch.size() * curr_topk, + stream); + } + + dev_output_graph = + raft::make_device_matrix(res, curr_query_size, next_graph_degree); + + graph::optimize(res, dev_knn_graph.view(), dev_output_graph.view(), false); +} + template (params.graph_build_params); + const auto& build_compression = + iter_params.build_compression.has_value() ? iter_params.build_compression : params.compression; + + if (build_compression.has_value()) { + const auto& bc = *build_compression; + RAFT_LOG_INFO( + "Build compression params: pq_bits=%u, pq_dim=%u, vq_n_centers=%u, kmeans_n_iters=%u, " + "vq_kmeans_trainset_fraction=%.4f, pq_kmeans_trainset_fraction=%.4f, " + "max_train_points_per_pq_code=%u, max_train_points_per_vq_cluster=%u%s", + bc.pq_bits, + bc.pq_dim, + bc.vq_n_centers, + bc.kmeans_n_iters, + bc.vq_kmeans_trainset_fraction, + bc.pq_kmeans_trainset_fraction, + bc.max_train_points_per_pq_code, + bc.max_train_points_per_vq_cluster, + iter_params.build_compression.has_value() ? " (from build_compression)" + : " (from compression)"); + } else { + RAFT_LOG_INFO("Build compression: disabled (uncompressed build)"); + } + RAFT_LOG_INFO("Build search params: search_width=%zu, max_iterations=%zu", + iter_params.search_width, + iter_params.max_iterations); + auto cagra_graph = raft::make_host_matrix(0, 0); // Iteratively improve the accuracy of the graph by repeatedly running @@ -2078,6 +2248,17 @@ auto iterative_build_graph( RAFT_LOG_DEBUG("# graph_degree = %lu", (uint64_t)graph_degree); RAFT_LOG_DEBUG("# topk = %lu", (uint64_t)topk); + // A fixed itopk_size (0 = auto) governs the growing iterations, which build graphs of degree + // ~graph_degree/2 and thus request topk ~= graph_degree/2 + 1; the search planner requires + // topk <= itopk_size. (The full-size iterations override itopk internally, so they are not + // constrained by this value.) + RAFT_EXPECTS(iter_params.itopk_size == 0 || + iter_params.itopk_size >= graph_degree / 2 + 1, + "iterative build search itopk_size (%zu) must be 0 (auto) or >= " + "graph_degree / 2 + 1 (%zu)", + (size_t)iter_params.itopk_size, + (size_t)(graph_degree / 2 + 1)); + // Create an initial graph. The initial graph created here is not suitable for // searching, but connectivity is guaranteed. auto offset = raft::make_host_vector(small_graph_degree); @@ -2098,28 +2279,131 @@ auto iterative_build_graph( } } - // Allocate memory for neighbors list using Transparent HugePage - constexpr size_t thp_size = 2 * 1024 * 1024; - size_t byte_size = sizeof(IdxT) * final_graph_size * topk; - if (byte_size % thp_size) { byte_size += thp_size - (byte_size % thp_size); } - mmap_owner neighbors_list(byte_size); - IdxT* neighbors_ptr = (IdxT*)neighbors_list.data(); - memset(neighbors_ptr, 0, byte_size); - bool flag_last = false; auto curr_graph_size = initial_graph_size; + + auto dev_graph = raft::make_device_matrix(res, 0, 0); + bool use_device_graph = false; + + // Generate the compressed index once if compression is enabled + const uint64_t dataset_dim = dev_dataset.extent(1); + std::optional> idx_opt; + + // Optional shuffle permutation for randomizing dataset order during build. + // inverse_perm[shuffled_idx] = original_idx + // perm[shuffled_idx] = original_idx, used to unshuffle the graph after build + auto dev_perm = raft::make_device_vector(res, 0); + bool dataset_shuffled = false; + + // Warn if shuffle is requested but compression is not enabled + if (iter_params.shuffle_dataset && !build_compression.has_value()) { + RAFT_LOG_WARN("shuffle_dataset is only supported with compression enabled; ignoring"); + } + + if (build_compression.has_value()) { + auto start = std::chrono::high_resolution_clock::now(); + RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded, + "VPQ compression is only supported with L2Expanded distance mertric"); + + // Build the VPQ compressed dataset + auto vpq_dset = + cuvs::preprocessing::quantize::pq::vpq_build(res, *build_compression, dev_dataset); + + // Optionally shuffle the compressed dataset to break spatial locality + if (iter_params.shuffle_dataset) { + auto shuffle_start = std::chrono::high_resolution_clock::now(); + RAFT_LOG_INFO("Shuffling compressed dataset to randomize build order..."); + + auto stream = raft::resource::get_cuda_stream(res); + const auto n_rows = vpq_dset.data.extent(0); + const auto row_len = vpq_dset.data.extent(1); + + // Generate random permutation: perm[i] = source index for output row i + // i.e., shuffled_data[i] = original_data[perm[i]] + // So perm maps: shuffled_idx -> original_idx + // Use int64_t for permutation to match vpq_dataset's index type + auto dev_perm_i64 = raft::make_device_vector(res, n_rows); + + // Use legacy permute API to generate permutation indices only (out=nullptr, in=nullptr) + // This just fills dev_perm_i64 with a random permutation of [0, n_rows) + raft::random::permute(dev_perm_i64.data_handle(), + static_cast(nullptr), + static_cast(nullptr), + static_cast(row_len), + static_cast(n_rows), + true, + stream); + + // Apply permutation to VPQ data: shuffled_data[i] = original_data[perm[i]]. + // NOTE: use an out-of-place device gather into a temporary buffer rather than the + // in-place gather overload. The in-place overload uses a host-orchestrated, + // double-buffered, multi-stream path that races here and triggers an asynchronous + // illegal memory access (the crash disappears under CUDA_LAUNCH_BLOCKING=1). + auto shuffled_data = raft::make_device_matrix( + res, vpq_dset.data.extent(0), vpq_dset.data.extent(1)); + raft::matrix::gather(res, + raft::make_const_mdspan(vpq_dset.data.view()), + raft::make_const_mdspan(dev_perm_i64.view()), + shuffled_data.view()); + vpq_dset.data = std::move(shuffled_data); + + // Store perm as IdxT for graph unshuffling later + // perm[shuffled_idx] = original_idx + // This is used for: + // 1. Remapping neighbor values: neighbor j (shuffled) -> perm[j] (original) + // 2. Reordering rows: row i (for shuffled node i) -> position perm[i] (original node) + dev_perm = raft::make_device_vector(res, n_rows); + cast_to_idx_op cast_op; + thrust::transform(raft::resource::get_thrust_policy(res), + dev_perm_i64.data_handle(), + dev_perm_i64.data_handle() + n_rows, + dev_perm.data_handle(), + cast_op); + + dataset_shuffled = true; + + auto shuffle_end = std::chrono::high_resolution_clock::now(); + auto shuffle_ms = + std::chrono::duration_cast(shuffle_end - shuffle_start).count(); + RAFT_LOG_INFO("# Dataset shuffle time: %.3lf sec", (double)shuffle_ms / 1000); + } + + idx_opt.emplace(res, params.metric); + // Use the (optionally shuffled) compressed dataset built above. + idx_opt->update_dataset(res, std::move(vpq_dset)); + auto end = std::chrono::high_resolution_clock::now(); + auto elapsed_ms = std::chrono::duration_cast(end - start).count(); + RAFT_LOG_INFO("# VPQ compression time: %.3lf sec", (double)elapsed_ms / 1000); + + // Free the original dataset -- queries will be reconstructed from VPQ codes. + dev_aligned_dataset.reset(); + RAFT_LOG_INFO( + "# Freed original dataset from device (%.1f MiB); queries will use VPQ reconstruction", + to_mib(final_graph_size * dataset_dim * sizeof(T))); + } while (true) { auto start = std::chrono::high_resolution_clock::now(); auto curr_query_size = std::min(2 * curr_graph_size, final_graph_size); auto next_graph_degree = small_graph_degree; if (curr_graph_size == final_graph_size) { next_graph_degree = graph_degree; } + RAFT_LOG_INFO("Current graph size %lu: # current graph degree = %lu", + (uint64_t)curr_graph_size, + (uint64_t)next_graph_degree); // The search count (topk) is set to the next graph degree + 1, because // pruning is not used except in the last iteration. // (*) The appropriate setting for itopk_size requires careful consideration. - auto curr_topk = next_graph_degree + 1; - auto curr_itopk_size = next_graph_degree + 32; + auto curr_topk = next_graph_degree + 1; + // The configurable itopk (iter_params.itopk_size, 0 = auto) applies only to the true growing + // iterations, where the degree being built is small_graph_degree. When the graph reaches its + // full size the search builds a graph_degree-degree graph (topk = graph_degree + 1); that + // iteration needs a larger itopk, so it overrides the configured value with the auto formula. + // The final iteration (flag_last) uses a fixed itopk tied to the output topk. + auto curr_itopk_size = + (iter_params.itopk_size > 0 && next_graph_degree == small_graph_degree) + ? (uint64_t)iter_params.itopk_size + : std::max(next_graph_degree + 32, (uint64_t)128); if (flag_last) { curr_topk = topk; curr_itopk_size = curr_topk + 32; @@ -2134,71 +2418,135 @@ auto iterative_build_graph( (uint64_t)curr_itopk_size, (uint64_t)curr_topk); - cuvs::neighbors::cagra::search_params search_params; - search_params.algo = cuvs::neighbors::cagra::search_algo::AUTO; - search_params.max_queries = max_chunk_size; - search_params.itopk_size = curr_itopk_size; - - // Create an index (idx), a query view (dev_query_view), and a mdarray for - // search results (neighbors). - auto dev_dataset_view = raft::make_device_matrix_view( - dev_dataset.data_handle(), (int64_t)curr_graph_size, dev_dataset.extent(1)); - - auto idx = index( - res, params.metric, dev_dataset_view, raft::make_const_mdspan(cagra_graph.view())); - - auto dev_query_view = raft::make_device_matrix_view( - dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1)); - - auto neighbors_view = - raft::make_host_matrix_view(neighbors_ptr, curr_query_size, curr_topk); - - // Search. - // Since there are many queries, divide them into batches and search them. - auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - res, - dev_query_view.data_handle(), - static_cast(curr_query_size), - static_cast(dev_query_view.extent(1)), - max_chunk_size, - raft::resource::get_cuda_stream(res), - raft::resource::get_workspace_resource_ref(res)); - for (const auto& batch : query_batch) { - auto batch_dev_query_view = raft::make_device_matrix_view( - batch.data(), batch.size(), dev_query_view.extent(1)); - auto batch_dev_neighbors_view = raft::make_device_matrix_view( - dev_neighbors.data_handle(), batch.size(), curr_topk); - auto batch_dev_distances_view = raft::make_device_matrix_view( - dev_distances.data_handle(), batch.size(), curr_topk); - - cuvs::neighbors::cagra::search(res, - search_params, - idx, - batch_dev_query_view, - batch_dev_neighbors_view, - batch_dev_distances_view); - - auto batch_neighbors_view = raft::make_host_matrix_view( - neighbors_view.data_handle() + batch.offset() * curr_topk, batch.size(), curr_topk); - raft::copy(res, batch_neighbors_view, batch_dev_neighbors_view); + cuvs::neighbors::cagra::search_params search_params = iter_params; + search_params.max_queries = max_chunk_size; + search_params.itopk_size = curr_itopk_size; + + // Create index and query views. + if (!build_compression.has_value()) { + auto dev_dataset_view = raft::make_device_matrix_view( + dev_dataset.data_handle(), (int64_t)curr_graph_size, dev_dataset.extent(1)); + if (use_device_graph) { + idx_opt.emplace( + res, params.metric, dev_dataset_view, raft::make_const_mdspan(dev_graph.view())); + } else { + idx_opt.emplace( + res, params.metric, dev_dataset_view, raft::make_const_mdspan(cagra_graph.view())); + } + } else { + if (use_device_graph) { + idx_opt->update_graph(res, raft::make_const_mdspan(dev_graph.view())); + } else { + idx_opt->update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + } } - - // Optimize graph - auto next_graph_size = curr_query_size; - cagra_graph = raft::make_host_matrix(0, 0); // delete existing grahp - cagra_graph = raft::make_host_matrix(next_graph_size, next_graph_degree); - optimize( - res, neighbors_view, cagra_graph.view(), flag_last ? params.guarantee_connectivity : 0); + const auto& idx = *idx_opt; + + // When compression is enabled, reconstruct queries from VPQ codes instead of + // reading from the (freed) original dataset. + auto dev_reconstructed_queries = + build_compression.has_value() + ? raft::make_device_matrix(res, curr_query_size, dataset_dim) + : raft::make_device_matrix(res, 0, 0); + if (build_compression.has_value()) { + auto* vpq_dset = dynamic_cast*>(&idx.data()); + RAFT_EXPECTS(vpq_dset != nullptr, "Expected VPQ dataset in compressed index"); + reconstruct_vpq_queries( + res, *vpq_dset, 0, curr_query_size, dev_reconstructed_queries.view()); + } + auto dev_query_view = + build_compression.has_value() + ? raft::make_device_matrix_view( + dev_reconstructed_queries.data_handle(), (int64_t)curr_query_size, dataset_dim) + : raft::make_device_matrix_view( + dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1)); + + auto dev_optimized_graph = raft::make_device_matrix(res, 0, 0); + + search_and_optimize(res, + search_params, + idx, + dev_query_view, + dev_neighbors.view(), + dev_distances.view(), + dev_optimized_graph, + curr_query_size, + next_graph_degree, + curr_topk, + max_chunk_size); + + dev_graph = std::move(dev_optimized_graph); + use_device_graph = true; auto end = std::chrono::high_resolution_clock::now(); auto elapsed_ms = std::chrono::duration_cast(end - start).count(); RAFT_LOG_DEBUG("# elapsed time: %.3lf sec", (double)elapsed_ms / 1000); if (flag_last) { break; } - flag_last = (curr_graph_size == final_graph_size); - curr_graph_size = next_graph_size; + flag_last = (curr_graph_size == final_graph_size); + auto next_graph_size = curr_query_size; + curr_graph_size = next_graph_size; } + // TODO: when build_compression matches params.compression, the dataset is compressed twice + // (once for the build loop and once in build()'s shared tail). We could avoid this by returning + // the index directly (with its VPQ dataset and device-side graph) instead of just the host graph. + auto stream = raft::resource::get_cuda_stream(res); + + // If the dataset was shuffled, we need to unshuffle the graph: + // Recall: perm[shuffled_idx] = original_idx (stored in dev_perm) + // 1. Remap neighbor indices from shuffled space to original space + // 2. Reorder rows from shuffled order to original order + if (dataset_shuffled) { + auto unshuffle_start = std::chrono::high_resolution_clock::now(); + RAFT_LOG_INFO("Unshuffling graph to restore original dataset ordering..."); + + const auto n_rows = dev_graph.extent(0); + const auto degree = dev_graph.extent(1); + + // Step 1: Remap all neighbor indices using perm + // graph[i][j] contains shuffled index j; we need original index = perm[j] + remap_indices_op remap_op{dev_perm.data_handle()}; + thrust::transform(raft::resource::get_thrust_policy(res), + dev_graph.data_handle(), + dev_graph.data_handle() + n_rows * degree, + dev_graph.data_handle(), + remap_op); + + // Step 2: Reorder rows back to original order + // Row i in dev_graph is for shuffled node i, which is original node perm[i]. + // We want this row to be at position perm[i] in the final graph. + // scatter: output[map[i]] = input[i], so map[i] = perm[i] + auto dev_unshuffled_graph = raft::make_device_matrix(res, n_rows, degree); + + // Use thrust::scatter to reorder: for each row i, place it at position perm[i] + // We scatter row-by-row conceptually, but do it element-wise with computed output indices + graph_scatter_index_op scatter_idx_op{dev_perm.data_handle(), degree}; + auto output_indices = + thrust::make_transform_iterator(thrust::make_counting_iterator(0), scatter_idx_op); + + thrust::scatter(raft::resource::get_thrust_policy(res), + dev_graph.data_handle(), + dev_graph.data_handle() + n_rows * degree, + output_indices, + dev_unshuffled_graph.data_handle()); + + dev_graph = std::move(dev_unshuffled_graph); + + auto unshuffle_end = std::chrono::high_resolution_clock::now(); + auto unshuffle_ms = + std::chrono::duration_cast(unshuffle_end - unshuffle_start) + .count(); + RAFT_LOG_INFO("# Graph unshuffle time: %.3lf sec", (double)unshuffle_ms / 1000); + } + + cagra_graph = raft::make_host_matrix(dev_graph.extent(0), dev_graph.extent(1)); + raft::copy(cagra_graph.data_handle(), + dev_graph.data_handle(), + dev_graph.extent(0) * dev_graph.extent(1), + stream); + raft::resource::sync_stream(res); + return cagra_graph; } diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 4d09e3683b..265df89c59 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -73,11 +73,13 @@ void search_main_core( topk, queries.extent(1)); + RAFT_LOG_DEBUG("search_main_core: creating plan with max_node_id=%u", params.max_node_id); using CagraSampleFilterT_s = typename CagraSampleFilterT_Selector::type; std::unique_ptr< search_plan_impl> plan = factory::create( res, params, dataset_desc, queries.extent(1), graph.extent(0), graph.extent(1), topk); + RAFT_LOG_DEBUG("search_main_core: plan created, plan->max_node_id=%u", plan->max_node_id); plan->check(topk); @@ -158,6 +160,7 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } // Search using a plain (strided) row-major dataset + RAFT_LOG_DEBUG("Searching with strided dataset"); RAFT_EXPECTS(index.metric() != cuvs::distance::DistanceType::CosineExpanded || index.dataset_norms().has_value(), "Dataset norms must be provided for CosineExpanded metric"); @@ -184,6 +187,7 @@ void search_main(raft::resources const& res, RAFT_FAIL("FP32 VPQ dataset support is coming soon"); } else if (auto* vpq_dset = dynamic_cast*>(&index.data()); vpq_dset != nullptr) { + RAFT_LOG_DEBUG("Searching with VPQ dataset"); if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && raft::getComputeCapability().first < 9) { RAFT_LOG_WARN( diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index a99ec64bc0..28cc6b6eba 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -229,6 +229,8 @@ struct dataset_descriptor_host { ~state() noexcept { if (std::holds_alternative(value)) { + // RAFT_LOG_INFO("trying to free descriptor state %p", + // reinterpret_cast(this)); auto& [ptr, stream] = std::get(value); RAFT_CUDA_TRY_NO_THROW(cudaFreeAsync(ptr, stream)); } From 90d121963e6d3348ce66280d5958ef2331639a41 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 05/39] feat(bench): expose iterative CAGRA-Q build/search params in cuvs_bench --- .../src/cuvs/cuvs_ann_bench_param_parser.h | 98 ++++++++++++++++++- python/cuvs_bench/cuvs_bench/run/__main__.py | 8 ++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 57b47d97db..82db80d2e7 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -367,10 +367,12 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } // Parse build-algo-specific parameters and use them to decide on the algo type - nlohmann::json ivf_pq_build_conf = collect_conf_with_prefix(conf, "ivf_pq_build_"); - nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); - nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); - nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); + nlohmann::json ivf_pq_build_conf = collect_conf_with_prefix(conf, "ivf_pq_build_"); + nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); + nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); + nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); + nlohmann::json build_compression_conf = collect_conf_with_prefix(conf, "build_compression_"); + nlohmann::json build_search_conf = collect_conf_with_prefix(conf, "build_search_"); // When graph_build_algo is not specified, leave graph_build_params as monostate so the // CAGRA build uses AUTO selection (NN_DESCENT or IVF_PQ based on dataset/heuristics). @@ -401,6 +403,94 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } else if constexpr (std::is_same_v) { parse_build_param(nn_descent_conf, arg); + } else if constexpr (std::is_same_v< + U, + cuvs::neighbors::graph_build_params::iterative_search_params>) { + if (!build_compression_conf.empty()) { + auto vpq_pams = arg.build_compression.value_or(cuvs::neighbors::vpq_params{}); + parse_build_param(build_compression_conf, vpq_pams); + arg.build_compression.emplace(vpq_pams); + } + if (build_search_conf.contains("width")) { + arg.search_width = build_search_conf.at("width"); + } + if (build_search_conf.contains("max_iterations")) { + arg.max_iterations = build_search_conf.at("max_iterations"); + } + if (build_search_conf.contains("min_iterations")) { + arg.min_iterations = build_search_conf.at("min_iterations"); + } + if (build_search_conf.contains("itopk")) { arg.itopk_size = build_search_conf.at("itopk"); } + if (build_search_conf.contains("max_queries")) { + arg.max_queries = build_search_conf.at("max_queries"); + } + if (build_search_conf.contains("team_size")) { + arg.team_size = build_search_conf.at("team_size"); + } + if (build_search_conf.contains("thread_block_size")) { + arg.thread_block_size = build_search_conf.at("thread_block_size"); + } + if (build_search_conf.contains("hashmap_min_bitlen")) { + arg.hashmap_min_bitlen = build_search_conf.at("hashmap_min_bitlen"); + } + if (build_search_conf.contains("hashmap_max_fill_rate")) { + arg.hashmap_max_fill_rate = build_search_conf.at("hashmap_max_fill_rate"); + } + if (build_search_conf.contains("num_random_samplings")) { + arg.num_random_samplings = build_search_conf.at("num_random_samplings"); + } + if (build_search_conf.contains("persistent")) { + arg.persistent = build_search_conf.at("persistent"); + } + if (build_search_conf.contains("persistent_lifetime")) { + arg.persistent_lifetime = build_search_conf.at("persistent_lifetime"); + } + if (build_search_conf.contains("persistent_device_usage")) { + arg.persistent_device_usage = build_search_conf.at("persistent_device_usage"); + } + if (build_search_conf.contains("algo")) { + std::string algo = build_search_conf.at("algo"); + if (algo == "single_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA; + } else if (algo == "multi_cta") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + } else if (algo == "multi_kernel") { + arg.algo = cuvs::neighbors::cagra::search_algo::MULTI_KERNEL; + } else if (algo == "auto") { + arg.algo = cuvs::neighbors::cagra::search_algo::AUTO; + } + } + if (build_search_conf.contains("hashmap_mode")) { + std::string mode = build_search_conf.at("hashmap_mode"); + if (mode == "hash") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::HASH; + } else if (mode == "small") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::SMALL; + } else if (mode == "auto") { + arg.hashmap_mode = cuvs::neighbors::cagra::hash_mode::AUTO; + } + } + // Whether to shuffle the (compressed) dataset before the iterative build loop. + if (build_search_conf.contains("shuffle_dataset")) { + arg.shuffle_dataset = build_search_conf.at("shuffle_dataset").get(); + } + // Precision of the codebook/query in shared memory for the VPQ search used during + // the iterative build. Accepts an integer code (0=F16, 1=E5M2) or a string. + if (build_search_conf.contains("smem_dtype")) { + const auto& sd = build_search_conf.at("smem_dtype"); + if (sd.is_number_integer()) { + arg.smem_dtype = static_cast(sd.get()); + } else { + std::string s = sd.get(); + if (s == "f16" || s == "F16" || s == "fp16" || s == "half") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + } else if (s == "e5m2" || s == "E5M2" || s == "fp8") { + arg.smem_dtype = cuvs::neighbors::cagra::internal_dtype::E5M2; + } else { + throw std::runtime_error("invalid value for build_search smem_dtype: " + s); + } + } + } } }, params.graph_build_params); diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 6950ff7202..58d1b604bd 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -5,6 +5,7 @@ import json import os +import warnings from pathlib import Path from typing import Optional @@ -257,6 +258,13 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ + warnings.warn( + "The 'cuvs_bench.run' CLI is deprecated and will be removed in a future release. " + "Use BenchmarkOrchestrator from cuvs_bench.orchestrator instead.", + FutureWarning, + stacklevel=2, + ) + if not data_export: # Determine backend type and extra kwargs from --backend-config backend_type = "cpp_gbench" From a9389588737ec7cf15688dad83b23a05db823c4e Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 00:53:16 -0700 Subject: [PATCH 06/39] test(cagra): VPQ/iterative build test updates --- cpp/tests/neighbors/ann_cagra.cuh | 104 +++++++++--------- .../bug_graph_smaller_than_dataset.cu | 20 ++-- cpp/tests/neighbors/ann_utils.cuh | 22 ++-- 3 files changed, 73 insertions(+), 73 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 7b86cc70ad..1f7b5c6977 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1547,38 +1547,38 @@ inline std::vector generate_inputs() {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // Corner cases for small datasets - inputs2 = raft::util::itertools::product( - {2}, - {3, 6, 31, 32, 64, 101}, - {1, 10}, - {2}, // k - {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, - {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL}, - {0}, // query size - {0}, - {256}, - {1}, - {cuvs::distance::DistanceType::L2Expanded}, - {false}, - {true}, - {true}, - {0.995}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); - inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + // // Corner cases for small datasets + // inputs2 = raft::util::itertools::product( + // {2}, + // {3, 6, 31, 32, 64, 101}, + // {1, 10}, + // {2}, // k + // {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, + // {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL}, + // {0}, // query size + // {0}, + // {256}, + // {1}, + // {cuvs::distance::DistanceType::L2Expanded}, + // {false}, + // {true}, + // {true}, + // {0.995}, + // {std::optional{std::nullopt}}, + // {std::optional{std::nullopt}}, + // {std::optional{std::nullopt}}, + // {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + // cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); + // inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying dim and build algo. inputs2 = raft::util::itertools::product( {100}, - {1000}, - {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 1024}, // dim - {16}, // k - {graph_build_algo::IVF_PQ, - graph_build_algo::NN_DESCENT, + {1000000}, + {768}, // dim + {16}, // k + { // graph_build_algo::IVF_PQ, + // graph_build_algo::NN_DESCENT, graph_build_algo::ITERATIVE_CAGRA_SEARCH}, {search_algo::AUTO}, {10}, @@ -1592,7 +1592,7 @@ inline std::vector generate_inputs() {false}, {true}, {false}, - {0.995}, + {0.01}, {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, @@ -1657,29 +1657,29 @@ inline std::vector generate_inputs() {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // Varying n_rows, host_dataset - inputs2 = raft::util::itertools::product( - {100}, - {10000}, - {32}, - {10}, - {graph_build_algo::AUTO}, - {search_algo::AUTO}, - {10}, - {0}, // team_size - {64}, - {1}, - {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, - {false, true}, - {false}, - {true}, - {0.985}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {std::optional{std::nullopt}}, - {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); - inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + // // Varying n_rows, host_dataset + // inputs2 = raft::util::itertools::product( + // {100}, + // {10000}, + // {32}, + // {10}, + // {graph_build_algo::AUTO}, + // {search_algo::AUTO}, + // {10}, + // {0}, // team_size + // {64}, + // {1}, + // {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, + // {false, true}, + // {false}, + // {true}, + // {0.985}, + // {std::optional{std::nullopt}}, + // {std::optional{std::nullopt}}, + // {std::optional{std::nullopt}}, + // {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + // cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); + // inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // A few PQ configurations. // Varying dim, vq_n_centers diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index adeb774a8b..b06c1cba92 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -38,8 +38,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { protected: void run() { - // Create a dataset with 1000 points - constexpr int64_t n_dataset = 1000; + // Create a dataset with 10000 points + constexpr int64_t n_dataset = 10000; constexpr int64_t n_dim = 128; constexpr int64_t n_queries = 100; constexpr int64_t k = 10; @@ -63,9 +63,9 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Recreate the bug scenario: LARGE dataset, SMALL graph // (like iterative_build_graph does in intermediate iterations) - constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph + constexpr int64_t n_graph = n_dataset / 2; // Only 5000 nodes in graph - // Step 1: Build index on SMALL subset (500 points) + // Step 1: Build index on SMALL subset (5000 points) auto small_dataset_view = raft::make_device_matrix_view( dataset.data_handle(), n_graph, n_dim); @@ -74,13 +74,13 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { auto small_index = cagra::build(res, small_index_params, small_dataset_view); raft::resource::sync_stream(res); - // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) - // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 + // Step 2: Update to FULL dataset (10000 points) but keep small graph (5000 nodes) + // This creates the exact bug scenario: dataset.size=10000, graph.extent(0)=5000 small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); // Verify the mismatch - THIS IS THE BUG SCENARIO! - ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes - ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! // Create queries @@ -100,8 +100,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { search_params.algo = cagra::search_algo::SINGLE_CTA; // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS - // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! - // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! + // Before fix: random seeds use dataset.size (10000) -> tries to access graph[7000] -> CRASH! + // After fix: random seeds use graph.extent(0) (5000) -> only accesses graph[0-4999] -> SAFE! cagra::search(res, search_params, small_index, diff --git a/cpp/tests/neighbors/ann_utils.cuh b/cpp/tests/neighbors/ann_utils.cuh index 7240363ee4..e3dcbea6c6 100644 --- a/cpp/tests/neighbors/ann_utils.cuh +++ b/cpp/tests/neighbors/ann_utils.cuh @@ -127,10 +127,10 @@ struct idx_dist_pair { /** Calculate recall value using only neighbor indices */ template -auto calc_recall(const std::vector& expected_idx, - const std::vector& actual_idx, - size_t rows, - size_t cols) +std::tuple calc_recall(const std::vector& expected_idx, + const std::vector& actual_idx, + size_t rows, + size_t cols) { size_t match_count = 0; size_t total_count = static_cast(rows) * static_cast(cols); @@ -219,13 +219,13 @@ auto eval_recall(const std::vector& expected_idx, /** Overload of calc_recall to account for distances */ template -auto calc_recall(const std::vector& expected_idx, - const std::vector& actual_idx, - const std::vector& expected_dist, - const std::vector& actual_dist, - size_t rows, - size_t cols, - double eps) +std::tuple calc_recall(const std::vector& expected_idx, + const std::vector& actual_idx, + const std::vector& expected_dist, + const std::vector& actual_dist, + size_t rows, + size_t cols, + double eps) { size_t match_count = 0; size_t index_match_count = 0; From 4f4068ab15be7391680161c1b1d72b47b7b8e416 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 14 Jul 2026 03:18:17 -0700 Subject: [PATCH 07/39] fix(cagra): use in-place gather for dataset shuffle (remove raft workaround) The shuffle_dataset path used an out-of-place gather into a temporary buffer to work around an illegal memory access in raft's in-place gather overload when n_rows * row_len exceeded 2^31 (32-bit index overflow). That bug is now fixed upstream in raft (rapidsai/raft#3059, closes #3055), which the cuvs raft pin now includes. Revert to the in-place gather to drop the extra full-size temporary allocation and copy. --- .../neighbors/detail/cagra/cagra_build.cuh | 29 +++++++------------ 1 file changed, 11 insertions(+), 18 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index f0547beecb..7021132a5d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2252,8 +2252,7 @@ auto iterative_build_graph( // ~graph_degree/2 and thus request topk ~= graph_degree/2 + 1; the search planner requires // topk <= itopk_size. (The full-size iterations override itopk internally, so they are not // constrained by this value.) - RAFT_EXPECTS(iter_params.itopk_size == 0 || - iter_params.itopk_size >= graph_degree / 2 + 1, + RAFT_EXPECTS(iter_params.itopk_size == 0 || iter_params.itopk_size >= graph_degree / 2 + 1, "iterative build search itopk_size (%zu) must be 0 (auto) or >= " "graph_degree / 2 + 1 (%zu)", (size_t)iter_params.itopk_size, @@ -2334,18 +2333,13 @@ auto iterative_build_graph( true, stream); - // Apply permutation to VPQ data: shuffled_data[i] = original_data[perm[i]]. - // NOTE: use an out-of-place device gather into a temporary buffer rather than the - // in-place gather overload. The in-place overload uses a host-orchestrated, - // double-buffered, multi-stream path that races here and triggers an asynchronous - // illegal memory access (the crash disappears under CUDA_LAUNCH_BLOCKING=1). - auto shuffled_data = raft::make_device_matrix( - res, vpq_dset.data.extent(0), vpq_dset.data.extent(1)); - raft::matrix::gather(res, - raft::make_const_mdspan(vpq_dset.data.view()), - raft::make_const_mdspan(dev_perm_i64.view()), - shuffled_data.view()); - vpq_dset.data = std::move(shuffled_data); + // Apply permutation to VPQ data in place: data[i] = original_data[perm[i]]. + // Previously this used an out-of-place gather into a temporary buffer to work around + // an illegal memory access in the in-place gather overload when n_rows * row_len + // exceeded 2^31 (32-bit index overflow). That bug is fixed upstream in raft + // (rapidsai/raft#3059, issue #3055), which cuvs now pins, so the in-place gather is + // safe again and avoids the extra full-size temporary allocation and copy. + raft::matrix::gather(res, vpq_dset.data.view(), raft::make_const_mdspan(dev_perm_i64.view())); // Store perm as IdxT for graph unshuffling later // perm[shuffled_idx] = original_idx @@ -2400,10 +2394,9 @@ auto iterative_build_graph( // full size the search builds a graph_degree-degree graph (topk = graph_degree + 1); that // iteration needs a larger itopk, so it overrides the configured value with the auto formula. // The final iteration (flag_last) uses a fixed itopk tied to the output topk. - auto curr_itopk_size = - (iter_params.itopk_size > 0 && next_graph_degree == small_graph_degree) - ? (uint64_t)iter_params.itopk_size - : std::max(next_graph_degree + 32, (uint64_t)128); + auto curr_itopk_size = (iter_params.itopk_size > 0 && next_graph_degree == small_graph_degree) + ? (uint64_t)iter_params.itopk_size + : std::max(next_graph_degree + 32, (uint64_t)128); if (flag_last) { curr_topk = topk; curr_itopk_size = curr_topk + 32; From cc6529185191cad3ecf1bcf81bb43e89c706b002 Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:07:33 +0000 Subject: [PATCH 08/39] fix style --- cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h | 2 +- cpp/include/cuvs/neighbors/common.hpp | 2 +- cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh | 2 +- cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu | 2 +- .../neighbors/ann_cagra/test_batched_device_view_from_host.cu | 2 +- python/cuvs_bench/cuvs_bench/run/__main__.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 82db80d2e7..cd830cea68 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 1e5ca5a159..f7b935a86d 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index f1c7305833..ee834a0b24 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index b06c1cba92..844471dd2d 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu b/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu index 1e1cc13093..eb72dbec92 100644 --- a/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu +++ b/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 58d1b604bd..7f52b9c49b 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # From 25e6d7dedd6f5dd53835d4f144cf625b431f0aa4 Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:19:06 +0000 Subject: [PATCH 09/39] restore clangd and gitignore changes --- .gitignore | 7 +----- cpp/.clangd | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 6 deletions(-) create mode 100644 cpp/.clangd diff --git a/.gitignore b/.gitignore index 0066d2b89a..3627558ff5 100644 --- a/.gitignore +++ b/.gitignore @@ -72,9 +72,7 @@ docs/source/_static/rust # clang tooling compile_commands.json - - - +.clangd/ # serialized ann indexes brute_force_index @@ -88,8 +86,5 @@ ivf_pq_index /datasets/ /*.json -# clangd -*/.clangd - # java .classpath diff --git a/cpp/.clangd b/cpp/.clangd new file mode 100644 index 0000000000..7c4fe036dd --- /dev/null +++ b/cpp/.clangd @@ -0,0 +1,65 @@ +# https://clangd.llvm.org/config + +# Apply a config conditionally to all C files +If: + PathMatch: .*\.(c|h)$ + +--- + +# Apply a config conditionally to all C++ files +If: + PathMatch: .*\.(c|h)pp + +--- + +# Apply a config conditionally to all CUDA files +If: + PathMatch: .*\.cuh? +CompileFlags: + Add: + - "-x" + - "cuda" + # No error on unknown CUDA versions + - "-Wno-unknown-cuda-version" + # Allow variadic CUDA functions + - "-Xclang=-fcuda-allow-variadic-functions" +Diagnostics: + Suppress: + - "variadic_device_fn" + - "attributes_not_allowed" + +--- + +# Tweak the clangd parse settings for all files +CompileFlags: + Add: + # report all errors + - "-ferror-limit=0" + - "-fmacro-backtrace-limit=0" + - "-ftemplate-backtrace-limit=0" + # Skip the CUDA version check + - "--no-cuda-version-check" + Remove: + # remove gcc's -fcoroutines + - -fcoroutines + # remove nvc++ flags unknown to clang + - "-gpu=*" + - "-stdpar*" + # remove nvcc flags unknown to clang + - "-arch*" + - "-gencode*" + - "--generate-code*" + - "-ccbin*" + - "-t=*" + - "--threads*" + - "-Xptxas*" + - "-Xcudafe*" + - "-Xfatbin*" + - "-Xcompiler*" + - "--diag-suppress*" + - "--diag_suppress*" + - "--compiler-options*" + - "--expt-extended-lambda" + - "--expt-relaxed-constexpr" + - "-forward-unknown-to-host-compiler" + - "-Werror=cross-execution-space-call" From 6bca27512894cfcf22a4cd5af5d3d6e658145c40 Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:34:03 +0000 Subject: [PATCH 10/39] revert cuvs_bench warning --- python/cuvs_bench/cuvs_bench/run/__main__.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 7f52b9c49b..29dcbd1f41 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -1,11 +1,10 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # import json import os -import warnings from pathlib import Path from typing import Optional @@ -258,12 +257,6 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ - warnings.warn( - "The 'cuvs_bench.run' CLI is deprecated and will be removed in a future release. " - "Use BenchmarkOrchestrator from cuvs_bench.orchestrator instead.", - FutureWarning, - stacklevel=2, - ) if not data_export: # Determine backend type and extra kwargs from --backend-config From 6ed4ab5967737a09b9a7d2f4b6ae9e25a320a7d8 Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:35:08 +0000 Subject: [PATCH 11/39] remove whitespace --- python/cuvs_bench/cuvs_bench/run/__main__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cuvs_bench/cuvs_bench/run/__main__.py b/python/cuvs_bench/cuvs_bench/run/__main__.py index 29dcbd1f41..6950ff7202 100644 --- a/python/cuvs_bench/cuvs_bench/run/__main__.py +++ b/python/cuvs_bench/cuvs_bench/run/__main__.py @@ -257,7 +257,6 @@ def main( and any backend-specific connection parameters (host, port, etc.). """ - if not data_export: # Determine backend type and extra kwargs from --backend-config backend_type = "cpp_gbench" From 9d44b2374fc3417099e135021956374acf353640 Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:50:44 +0000 Subject: [PATCH 12/39] remove duplicate file in cmakelists.txt --- cpp/tests/CMakeLists.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 46c53c4f49..e50470ecdc 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -187,7 +187,6 @@ ConfigureTest( neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu neighbors/ann_cagra/bug_iterative_cagra_build.cu neighbors/ann_cagra/bug_issue_93_reproducer.cu - neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu GPUS 1 PERCENT 100 ) From 73c9bbcdbe23cdd9993e730b91b080b9c6f34d0e Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:52:38 +0000 Subject: [PATCH 13/39] revert to auto for type deduction --- cpp/tests/neighbors/ann_utils.cuh | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/cpp/tests/neighbors/ann_utils.cuh b/cpp/tests/neighbors/ann_utils.cuh index e3dcbea6c6..7240363ee4 100644 --- a/cpp/tests/neighbors/ann_utils.cuh +++ b/cpp/tests/neighbors/ann_utils.cuh @@ -127,10 +127,10 @@ struct idx_dist_pair { /** Calculate recall value using only neighbor indices */ template -std::tuple calc_recall(const std::vector& expected_idx, - const std::vector& actual_idx, - size_t rows, - size_t cols) +auto calc_recall(const std::vector& expected_idx, + const std::vector& actual_idx, + size_t rows, + size_t cols) { size_t match_count = 0; size_t total_count = static_cast(rows) * static_cast(cols); @@ -219,13 +219,13 @@ auto eval_recall(const std::vector& expected_idx, /** Overload of calc_recall to account for distances */ template -std::tuple calc_recall(const std::vector& expected_idx, - const std::vector& actual_idx, - const std::vector& expected_dist, - const std::vector& actual_dist, - size_t rows, - size_t cols, - double eps) +auto calc_recall(const std::vector& expected_idx, + const std::vector& actual_idx, + const std::vector& expected_dist, + const std::vector& actual_dist, + size_t rows, + size_t cols, + double eps) { size_t match_count = 0; size_t index_match_count = 0; From 23cf81575eca343de0fec701cf8d83ab062a638e Mon Sep 17 00:00:00 2001 From: aamijar Date: Wed, 22 Jul 2026 22:59:22 +0000 Subject: [PATCH 14/39] remove commented out code --- cpp/src/neighbors/detail/cagra/compute_distance.hpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index 28cc6b6eba..a99ec64bc0 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -229,8 +229,6 @@ struct dataset_descriptor_host { ~state() noexcept { if (std::holds_alternative(value)) { - // RAFT_LOG_INFO("trying to free descriptor state %p", - // reinterpret_cast(this)); auto& [ptr, stream] = std::get(value); RAFT_CUDA_TRY_NO_THROW(cudaFreeAsync(ptr, stream)); } From c2f9b6a6ae7311e5923746e2a4f061154c9ad069 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Thu, 23 Jul 2026 06:49:18 -0700 Subject: [PATCH 15/39] fix(cagra): pass graph_size to persistent single-CTA kernel to bound build-time random seeds --- .../neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp | 1 + .../detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh | 6 ++++-- .../cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in | 2 ++ .../detail/cagra/search_single_cta_kernel_launcher_jit.cuh | 3 ++- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp index c9702d3b72..a686f29921 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/kernel_def.hpp @@ -76,6 +76,7 @@ using search_single_cta_p_kernel_func_t = const std::uint32_t, const std::uint32_t, const dataset_descriptor_base_t*, + const IndexT, cagra_sample_filter); } // namespace single_cta_search diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh index d481c44946..8cc39de507 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -551,6 +551,7 @@ __device__ void search_single_cta_p_impl( const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id_offset, // Offset to add to query_id when calling filter const dataset_descriptor_base_t* dataset_desc, + const IndexT graph_size, cagra_sample_filter filter_payload) { using job_desc_type = job_desc_t>; @@ -625,7 +626,8 @@ __device__ void search_single_cta_p_impl( query_id, query_id_offset, dataset_desc, - filter_payload); + filter_payload, + graph_size); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in index 9986f7abc1..b003220497 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_single_cta_p_kernel.cu.in @@ -51,6 +51,7 @@ extern "C" __global__ __launch_bounds__(1024, 1) void search_single_cta_p( const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id_offset, const dataset_desc_base* dataset_desc, + const index_t graph_size, cagra_sample_filter_t filter_payload) { search_single_cta_p_impl(graph.extent(0)), filter_payload); last_touch.store(std::chrono::system_clock::now(), std::memory_order_relaxed); From cad2e8f5608c1c8f37e3b7f9414ec6c4cf405329 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Thu, 6 Aug 2026 00:28:41 -0700 Subject: [PATCH 16/39] Brought back the tests --- cpp/tests/neighbors/ann_cagra.cuh | 96 +++++++++++++++---------------- 1 file changed, 48 insertions(+), 48 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 0419f30166..f3b45ee5b8 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1598,30 +1598,30 @@ inline std::vector generate_inputs() {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // // Corner cases for small datasets - // inputs2 = raft::util::itertools::product( - // {2}, - // {3, 6, 31, 32, 64, 101}, - // {1, 10}, - // {2}, // k - // {32}, // degree - // {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, - // {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL}, - // {0}, // query size - // {0}, - // {256}, - // {1}, - // {cuvs::distance::DistanceType::L2Expanded}, - // {false}, - // {true}, - // {true}, - // {0.995}, - // {std::optional{std::nullopt}}, - // {std::optional{std::nullopt}}, - // {std::optional{std::nullopt}}, - // {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - // cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); - // inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + // Corner cases for small datasets + inputs2 = raft::util::itertools::product( + {2}, + {3, 6, 31, 32, 64, 101}, + {1, 10}, + {2}, // k + {32}, // degree + {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, + {search_algo::SINGLE_CTA, search_algo::MULTI_CTA, search_algo::MULTI_KERNEL}, + {0}, // query size + {0}, + {256}, + {1}, + {cuvs::distance::DistanceType::L2Expanded}, + {false}, + {true}, + {true}, + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); + inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying dim and build algo. inputs2 = raft::util::itertools::product( @@ -1712,30 +1712,30 @@ inline std::vector generate_inputs() {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); - // // Varying n_rows, host_dataset - // inputs2 = raft::util::itertools::product( - // {100}, - // {10000}, - // {32}, - // {10}, - // {32}, // degree - // {graph_build_algo::AUTO}, - // {search_algo::AUTO}, - // {10}, - // {0}, // team_size - // {64}, - // {1}, - // {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, - // {false, true}, - // {false}, - // {true}, - // {0.985}, - // {std::optional{std::nullopt}}, - // {std::optional{std::nullopt}}, - // {std::optional{std::nullopt}}, - // {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - // cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); - // inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + // Varying n_rows, host_dataset + inputs2 = raft::util::itertools::product( + {100}, + {10000}, + {32}, + {10}, + {32}, // degree + {graph_build_algo::AUTO}, + {search_algo::AUTO}, + {10}, + {0}, // team_size + {64}, + {1}, + {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, + {false, true}, + {false}, + {true}, + {0.985}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); + inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // A few PQ configurations. // Varying dim, vq_n_centers From 0e98b34efea8f771e7b0b6f1c3fddcba422d7ec2 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Thu, 6 Aug 2026 07:33:21 -0700 Subject: [PATCH 17/39] Fixed the test --- cpp/tests/neighbors/ann_cagra.cuh | 12 +++++------ .../bug_graph_smaller_than_dataset.cu | 20 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index f3b45ee5b8..f85e36a6ef 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1626,13 +1626,13 @@ inline std::vector generate_inputs() // Varying dim and build algo. inputs2 = raft::util::itertools::product( {100}, - {1000000}, - {768}, // dim + {1000}, + {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 768, 1024}, // dim {16}, // k {32}, // degree - { // graph_build_algo::IVF_PQ, - // graph_build_algo::NN_DESCENT, - graph_build_algo::ITERATIVE_CAGRA_SEARCH}, + {graph_build_algo::IVF_PQ, + graph_build_algo::NN_DESCENT, + graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build {search_algo::AUTO}, {10}, {0}, @@ -1645,7 +1645,7 @@ inline std::vector generate_inputs() {false}, {true}, {false}, - {0.01}, + {0.995}, {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, {std::optional{std::nullopt}}, diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index 844471dd2d..16fa93f47b 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -38,8 +38,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { protected: void run() { - // Create a dataset with 10000 points - constexpr int64_t n_dataset = 10000; + // Create a dataset with 1000 points + constexpr int64_t n_dataset = 1000; constexpr int64_t n_dim = 128; constexpr int64_t n_queries = 100; constexpr int64_t k = 10; @@ -63,9 +63,9 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Recreate the bug scenario: LARGE dataset, SMALL graph // (like iterative_build_graph does in intermediate iterations) - constexpr int64_t n_graph = n_dataset / 2; // Only 5000 nodes in graph + constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph - // Step 1: Build index on SMALL subset (5000 points) + // Step 1: Build index on SMALL subset (500 points) auto small_dataset_view = raft::make_device_matrix_view( dataset.data_handle(), n_graph, n_dim); @@ -74,13 +74,13 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { auto small_index = cagra::build(res, small_index_params, small_dataset_view); raft::resource::sync_stream(res); - // Step 2: Update to FULL dataset (10000 points) but keep small graph (5000 nodes) - // This creates the exact bug scenario: dataset.size=10000, graph.extent(0)=5000 + // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) + // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); // Verify the mismatch - THIS IS THE BUG SCENARIO! - ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes - ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! // Create queries @@ -100,8 +100,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { search_params.algo = cagra::search_algo::SINGLE_CTA; // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS - // Before fix: random seeds use dataset.size (10000) -> tries to access graph[7000] -> CRASH! - // After fix: random seeds use graph.extent(0) (5000) -> only accesses graph[0-4999] -> SAFE! + // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! + // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! cagra::search(res, search_params, small_index, From 8a378cb53319ccac0f6b755aa05f2e4cd78a5f63 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 10 Aug 2026 05:33:22 -0700 Subject: [PATCH 18/39] Removed dataset shuffle --- .../neighbors/detail/cagra/cagra_build.cuh | 146 +----------------- 1 file changed, 7 insertions(+), 139 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 820cbef016..1446bb81a5 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -62,32 +62,6 @@ namespace cuvs::neighbors::cagra::detail { constexpr double to_mib(size_t bytes) { return static_cast(bytes) / (1 << 20); } constexpr double to_gib(size_t bytes) { return static_cast(bytes) / (1 << 30); } -// Functor to remap indices using a permutation lookup table -template -struct remap_indices_op { - const IdxT* perm; - __host__ __device__ IdxT operator()(IdxT idx) const { return perm[idx]; } -}; - -// Functor to compute scattered output index for graph row reordering -template -struct graph_scatter_index_op { - const IdxT* perm; - int64_t degree; - __host__ __device__ int64_t operator()(int64_t idx) const - { - int64_t row = idx / degree; - int64_t col = idx % degree; - return static_cast(perm[row]) * degree + col; - } -}; - -// Functor to convert int64_t to IdxT -template -struct cast_to_idx_op { - __host__ __device__ IdxT operator()(int64_t v) const { return static_cast(v); } -}; - template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { @@ -2157,6 +2131,12 @@ void search_and_optimize(raft::resources const& res, graph::optimize(res, dev_knn_graph.view(), dev_output_graph.view(), false); } +// Builds a CAGRA graph iteratively, growing the graph while repeatedly running CAGRA's search() +// and optimize(). When compression is enabled the compressed dataset is consumed as-is. +// +// NOTE: This function EXPECTS the dataset to already be shuffled by the caller when a randomized +// build order is desired. It no longer performs any dataset shuffling (or graph unshuffling) +// internally, so the returned graph's node ordering matches the order of the input dataset. template > idx_opt; - // Optional shuffle permutation for randomizing dataset order during build. - // inverse_perm[shuffled_idx] = original_idx - // perm[shuffled_idx] = original_idx, used to unshuffle the graph after build - auto dev_perm = raft::make_device_vector(res, 0); - bool dataset_shuffled = false; - - // Warn if shuffle is requested but compression is not enabled - if (iter_params.shuffle_dataset && !build_compression.has_value()) { - RAFT_LOG_WARN("shuffle_dataset is only supported with compression enabled; ignoring"); - } - if (build_compression.has_value()) { auto start = std::chrono::high_resolution_clock::now(); RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded, @@ -2310,62 +2279,8 @@ auto iterative_build_graph( auto vpq_dset = cuvs::preprocessing::quantize::pq::vpq_build(res, *build_compression, dev_dataset); - // Optionally shuffle the compressed dataset to break spatial locality - if (iter_params.shuffle_dataset) { - auto shuffle_start = std::chrono::high_resolution_clock::now(); - RAFT_LOG_INFO("Shuffling compressed dataset to randomize build order..."); - - auto stream = raft::resource::get_cuda_stream(res); - const auto n_rows = vpq_dset.data.extent(0); - const auto row_len = vpq_dset.data.extent(1); - - // Generate random permutation: perm[i] = source index for output row i - // i.e., shuffled_data[i] = original_data[perm[i]] - // So perm maps: shuffled_idx -> original_idx - // Use int64_t for permutation to match vpq_dataset's index type - auto dev_perm_i64 = raft::make_device_vector(res, n_rows); - - // Use legacy permute API to generate permutation indices only (out=nullptr, in=nullptr) - // This just fills dev_perm_i64 with a random permutation of [0, n_rows) - raft::random::permute(dev_perm_i64.data_handle(), - static_cast(nullptr), - static_cast(nullptr), - static_cast(row_len), - static_cast(n_rows), - true, - stream); - - // Apply permutation to VPQ data in place: data[i] = original_data[perm[i]]. - // Previously this used an out-of-place gather into a temporary buffer to work around - // an illegal memory access in the in-place gather overload when n_rows * row_len - // exceeded 2^31 (32-bit index overflow). That bug is fixed upstream in raft - // (rapidsai/raft#3059, issue #3055), which cuvs now pins, so the in-place gather is - // safe again and avoids the extra full-size temporary allocation and copy. - raft::matrix::gather(res, vpq_dset.data.view(), raft::make_const_mdspan(dev_perm_i64.view())); - - // Store perm as IdxT for graph unshuffling later - // perm[shuffled_idx] = original_idx - // This is used for: - // 1. Remapping neighbor values: neighbor j (shuffled) -> perm[j] (original) - // 2. Reordering rows: row i (for shuffled node i) -> position perm[i] (original node) - dev_perm = raft::make_device_vector(res, n_rows); - cast_to_idx_op cast_op; - thrust::transform(raft::resource::get_thrust_policy(res), - dev_perm_i64.data_handle(), - dev_perm_i64.data_handle() + n_rows, - dev_perm.data_handle(), - cast_op); - - dataset_shuffled = true; - - auto shuffle_end = std::chrono::high_resolution_clock::now(); - auto shuffle_ms = - std::chrono::duration_cast(shuffle_end - shuffle_start).count(); - RAFT_LOG_INFO("# Dataset shuffle time: %.3lf sec", (double)shuffle_ms / 1000); - } - idx_opt.emplace(res, params.metric); - // Use the (optionally shuffled) compressed dataset built above. + // Use the compressed dataset built above (expected to be pre-shuffled by the caller). idx_opt->update_dataset(res, std::move(vpq_dset)); auto end = std::chrono::high_resolution_clock::now(); auto elapsed_ms = std::chrono::duration_cast(end - start).count(); @@ -2488,53 +2403,6 @@ auto iterative_build_graph( // the index directly (with its VPQ dataset and device-side graph) instead of just the host graph. auto stream = raft::resource::get_cuda_stream(res); - // If the dataset was shuffled, we need to unshuffle the graph: - // Recall: perm[shuffled_idx] = original_idx (stored in dev_perm) - // 1. Remap neighbor indices from shuffled space to original space - // 2. Reorder rows from shuffled order to original order - if (dataset_shuffled) { - auto unshuffle_start = std::chrono::high_resolution_clock::now(); - RAFT_LOG_INFO("Unshuffling graph to restore original dataset ordering..."); - - const auto n_rows = dev_graph.extent(0); - const auto degree = dev_graph.extent(1); - - // Step 1: Remap all neighbor indices using perm - // graph[i][j] contains shuffled index j; we need original index = perm[j] - remap_indices_op remap_op{dev_perm.data_handle()}; - thrust::transform(raft::resource::get_thrust_policy(res), - dev_graph.data_handle(), - dev_graph.data_handle() + n_rows * degree, - dev_graph.data_handle(), - remap_op); - - // Step 2: Reorder rows back to original order - // Row i in dev_graph is for shuffled node i, which is original node perm[i]. - // We want this row to be at position perm[i] in the final graph. - // scatter: output[map[i]] = input[i], so map[i] = perm[i] - auto dev_unshuffled_graph = raft::make_device_matrix(res, n_rows, degree); - - // Use thrust::scatter to reorder: for each row i, place it at position perm[i] - // We scatter row-by-row conceptually, but do it element-wise with computed output indices - graph_scatter_index_op scatter_idx_op{dev_perm.data_handle(), degree}; - auto output_indices = - thrust::make_transform_iterator(thrust::make_counting_iterator(0), scatter_idx_op); - - thrust::scatter(raft::resource::get_thrust_policy(res), - dev_graph.data_handle(), - dev_graph.data_handle() + n_rows * degree, - output_indices, - dev_unshuffled_graph.data_handle()); - - dev_graph = std::move(dev_unshuffled_graph); - - auto unshuffle_end = std::chrono::high_resolution_clock::now(); - auto unshuffle_ms = - std::chrono::duration_cast(unshuffle_end - unshuffle_start) - .count(); - RAFT_LOG_INFO("# Graph unshuffle time: %.3lf sec", (double)unshuffle_ms / 1000); - } - cagra_graph = raft::make_host_matrix(dev_graph.extent(0), dev_graph.extent(1)); raft::copy(cagra_graph.data_handle(), dev_graph.data_handle(), From aa2227b8fb08d1c2718d21c8b0ff0af171258ed2 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 10 Aug 2026 05:52:09 -0700 Subject: [PATCH 19/39] Removed unused pointer residency helper; use memory_type_from_pointer instead --- cpp/src/neighbors/detail/cagra/utils.hpp | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/utils.hpp b/cpp/src/neighbors/detail/cagra/utils.hpp index 2313631372..7b31fbdee3 100644 --- a/cpp/src/neighbors/detail/cagra/utils.hpp +++ b/cpp/src/neighbors/detail/cagra/utils.hpp @@ -161,22 +161,6 @@ struct gen_index_msb_1_mask { }; } // namespace utils -template -bool is_ptr_device_accessible(T* ptr) -{ - cudaPointerAttributes attr; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); - return attr.devicePointer != nullptr; -} - -template -bool is_ptr_host_accessible(T* ptr) -{ - cudaPointerAttributes attr; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); - return attr.hostPointer != nullptr; -} - /** * Utility to sync memory from a host_matrix_view to a device_matrix_view * From d9c6bfd10941470b85c22735834749b363cc5ebb Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 10 Aug 2026 06:01:04 -0700 Subject: [PATCH 20/39] Pre-commit changes --- c/src/neighbors/brute_force.cpp | 2 +- c/src/neighbors/cagra.cpp | 2 +- c/src/neighbors/ivf_flat.cpp | 2 +- c/src/neighbors/mg_cagra.cpp | 2 +- c/src/neighbors/mg_ivf_flat.cpp | 2 +- c/src/neighbors/mg_ivf_pq.cpp | 2 +- ci/build_cpp.sh | 2 +- ci/build_go.sh | 2 +- ci/build_java.sh | 2 +- ci/build_python.sh | 2 +- ci/build_rust.sh | 2 +- ci/build_wheel_cuvs.sh | 2 +- ci/build_wheel_libcuvs.sh | 2 +- ci/test_cpp.sh | 2 +- ci/test_wheel_cuvs.sh | 2 +- cpp/bench/ann/CMakeLists.txt | 2 +- cpp/bench/ann/src/cuvs/cuvs_benchmark.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h | 2 +- cpp/cmake/modules/compute_matrix_product.cmake | 2 +- cpp/include/cuvs/detail/jit_lto/common_fragments.hpp | 2 +- .../cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp | 2 +- cpp/include/cuvs/neighbors/ivf_pq.hpp | 2 +- cpp/include/cuvs/neighbors/ivf_rabitq.hpp | 2 +- cpp/include/cuvs/neighbors/nn_descent.hpp | 2 +- cpp/include/cuvs/util/file_io.hpp | 2 +- cpp/src/neighbors/brute_force_serialize.cu | 2 +- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/cagra_helpers.hpp | 2 +- cpp/src/neighbors/detail/cagra/cagra_serialize.cuh | 2 +- cpp/src/neighbors/detail/cagra/graph_core.cuh | 2 +- .../detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh | 2 +- .../detail/cagra/jit_lto_kernels/search_multi_jit.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in | 2 +- .../detail/cagra/search_multi_cta_kernel_launcher_jit.cuh | 2 +- .../detail/cagra/search_multi_kernel_launcher_jit.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in | 2 +- cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp | 2 +- cpp/src/neighbors/detail/hnsw.hpp | 2 +- cpp/src/neighbors/ivf_pq_index.cu | 2 +- cpp/src/neighbors/ivf_rabitq.cu | 2 +- cpp/src/neighbors/ivf_rabitq/defines.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh | 2 +- .../neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh | 2 +- .../ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu | 2 +- .../ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu | 2 +- .../bitwise_block_sort_emit_topk_kernel.cu.in | 2 +- .../jit_lto_kernels/bitwise_emit_distances_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh | 2 +- .../compute_bitwise_quantized_ip_for_vec_kernel.cu.in | 2 +- .../compute_inner_products_with_bitwise_block_sort_impl.cuh | 2 +- ...pute_inner_products_with_bitwise_block_sort_kernel.cu.in | 2 +- ...mpute_inner_products_with_bitwise_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_bitwise_impl.cuh | 2 +- .../compute_inner_products_with_bitwise_kernel.cu.in | 2 +- .../compute_inner_products_with_bitwise_planner.hpp | 2 +- ...ompute_inner_products_with_lut16_opt_block_sort_impl.cuh | 2 +- ...te_inner_products_with_lut16_opt_block_sort_kernel.cu.in | 2 +- ...ute_inner_products_with_lut16_opt_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_lut16_opt_impl.cuh | 2 +- .../compute_inner_products_with_lut16_opt_kernel.cu.in | 2 +- .../compute_inner_products_with_lut16_opt_planner.hpp | 2 +- .../compute_inner_products_with_lut_block_sort_impl.cuh | 2 +- .../compute_inner_products_with_lut_block_sort_kernel.cu.in | 2 +- .../compute_inner_products_with_lut_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_lut_impl.cuh | 2 +- .../compute_inner_products_with_lut_kernel.cu.in | 2 +- .../compute_inner_products_with_lut_planner.hpp | 2 +- .../jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in | 2 +- .../ivf_rabitq/jit_lto_kernels/device_functions.cuh | 2 +- .../ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp | 2 +- .../ivf_rabitq/jit_lto_kernels/launcher_factory.hpp | 2 +- .../jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in | 2 +- .../jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in | 2 +- .../jit_lto_kernels/lut_emit_distances_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/utils/IO.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/memory.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu | 2 +- cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/space.hpp | 2 +- cpp/src/neighbors/mg/snmg.cuh | 2 +- cpp/src/neighbors/nn_descent.cu | 2 +- cpp/tests/neighbors/ann_cagra.cuh | 6 +++--- cpp/tests/neighbors/ann_cagra/test_filter_udf.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace.cuh | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_ivf_rabitq.cuh | 2 +- cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu | 2 +- examples/build.sh | 2 +- examples/cpp/src/cagra_filter_udf_example.cu | 2 +- examples/cpp/src/cagra_hnsw_ace_build.cu | 2 +- examples/cpp/src/hnsw_openai_example.cu | 2 +- python/cuvs/cuvs/tests/test_cagra_ace.py | 2 +- python/cuvs/cuvs/tests/test_hnsw_ace.py | 2 +- 113 files changed, 115 insertions(+), 115 deletions(-) diff --git a/c/src/neighbors/brute_force.cpp b/c/src/neighbors/brute_force.cpp index 081f433a3e..d926a56cd4 100644 --- a/c/src/neighbors/brute_force.cpp +++ b/c/src/neighbors/brute_force.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 004b810c78..a01f7b051f 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/ivf_flat.cpp b/c/src/neighbors/ivf_flat.cpp index 729c810aa0..7ae9073f9e 100644 --- a/c/src/neighbors/ivf_flat.cpp +++ b/c/src/neighbors/ivf_flat.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_cagra.cpp b/c/src/neighbors/mg_cagra.cpp index 495eff8a34..30be2f76a8 100644 --- a/c/src/neighbors/mg_cagra.cpp +++ b/c/src/neighbors/mg_cagra.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_ivf_flat.cpp b/c/src/neighbors/mg_ivf_flat.cpp index 4e1b2883ea..cdf261a049 100644 --- a/c/src/neighbors/mg_ivf_flat.cpp +++ b/c/src/neighbors/mg_ivf_flat.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_ivf_pq.cpp b/c/src/neighbors/mg_ivf_pq.cpp index 41aa323138..8570f07c8b 100644 --- a/c/src/neighbors/mg_ivf_pq.cpp +++ b/c/src/neighbors/mg_ivf_pq.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/ci/build_cpp.sh b/ci/build_cpp.sh index bb497692bd..cda3d10bd9 100755 --- a/ci/build_cpp.sh +++ b/ci/build_cpp.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_go.sh b/ci/build_go.sh index 056f40b345..e04f6f1d8a 100755 --- a/ci/build_go.sh +++ b/ci/build_go.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_java.sh b/ci/build_java.sh index 2e363bb452..167c642dd2 100755 --- a/ci/build_java.sh +++ b/ci/build_java.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_python.sh b/ci/build_python.sh index 6823cbbec5..34be953de5 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_rust.sh b/ci/build_rust.sh index e9218a8adc..e84e454e05 100755 --- a/ci/build_rust.sh +++ b/ci/build_rust.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_wheel_cuvs.sh b/ci/build_wheel_cuvs.sh index 2acbab8a2d..2e174bc87e 100755 --- a/ci/build_wheel_cuvs.sh +++ b/ci/build_wheel_cuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_wheel_libcuvs.sh b/ci/build_wheel_libcuvs.sh index e012935749..9bc1ac0a37 100755 --- a/ci/build_wheel_libcuvs.sh +++ b/ci/build_wheel_libcuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index 0905cd64f0..be62cfc754 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/test_wheel_cuvs.sh b/ci/test_wheel_cuvs.sh index 36fefdf852..227857dd8c 100755 --- a/ci/test_wheel_cuvs.sh +++ b/ci/test_wheel_cuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 90d23d9aef..80f116f586 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu index 3056ddc365..1a334924ec 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu index c903b39fcc..8c5854051d 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index db618f6559..c7733c293e 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu index 1c6772d994..e3c38025c7 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "cuvs_ivf_rabitq_wrapper.h" diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index ca8f77b808..542f0bc6dd 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp index 1a99e56028..f40ad67a63 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp +++ b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h index 282d57dc2e..cf2bb7608e 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 82a34f9242..6d81821b88 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index ef2a8e6002..a55052c4a6 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp index 30885947b5..2b2f3db5a7 100644 --- a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index 57f8a258fb..686c3ff108 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp index d26bc04022..c33d466613 100644 --- a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 4c031049e2..929a099cf1 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index a7d67ec2c0..b0afeed732 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/brute_force_serialize.cu b/cpp/src/neighbors/brute_force_serialize.cu index 1b7595ee11..e3a4a2c041 100644 --- a/cpp/src/neighbors/brute_force_serialize.cu +++ b/cpp/src/neighbors/brute_force_serialize.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..34b2e72f90 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp index ee78930970..5c6a63f9fc 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f106b82500..e80c6b6932 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 52b4542798..1762bdbfb0 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh index 4c4f2e4f62..0e9e981e69 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh index a714ced5c2..02cec6ee5e 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in index 7c642fe406..4e94922566 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh index 8a673405b7..ff1064f24c 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh index bc341b9082..549474d045 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in index 4616a9652b..7869e4f05e 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp index e5157ffa6a..e448bebeb4 100644 --- a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp +++ b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 88580de929..649886f924 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_pq_index.cu b/cpp/src/neighbors/ivf_pq_index.cu index 28b985eec8..114b98bf0c 100644 --- a/cpp/src/neighbors/ivf_pq_index.cu +++ b/cpp/src/neighbors/ivf_pq_index.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu index d5572d4039..14a9678bad 100644 --- a/cpp/src/neighbors/ivf_rabitq.cu +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/defines.hpp b/cpp/src/neighbors/ivf_rabitq/defines.hpp index f91c06a1ff..aac296cb12 100644 --- a/cpp/src/neighbors/ivf_rabitq/defines.hpp +++ b/cpp/src/neighbors/ivf_rabitq/defines.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu index 6755bfc3de..6e1e2c7d18 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh index 0640148814..135117d4a9 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index 49070acbcf..c071fb9a0a 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index 75ddaee865..d275c84eb3 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index af0876acca..ea40f02931 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh index 90ef71ed12..21aaf414ac 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu index 59a3e575d9..c9c0a1d275 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh index 8db4bb464c..b7a485b7be 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu index 22bb74c682..2c81c40e85 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh index 73557f57ea..68ca342370 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh index 1e004d9164..125a759ca0 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index 92906fb0e6..fde565e0e9 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu index 8357af5fbb..bee53b54a2 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in index d8dcff324d..6a9b0d9576 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in index 20c5aec6a1..fb6a130512 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh index a16058c2af..fb78d0d560 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in index e4e2aa6d79..d5a63ee313 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh index a577ec44ce..e3b2da9733 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in index 7bbdc24139..785aa22b26 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp index dbe6ec5fd4..b126e723cc 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh index 4ff53b067c..cf87f17314 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in index 1f7355c7fe..53897c47db 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index 3ec4809395..ce7197e7e1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh index 45c8793b59..cc393f82db 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in index 8e45f99c0f..6aa138b917 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp index acf4e27e3a..e71903764b 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh index fc88d79fc4..575766d054 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in index 4160aa7a33..713a15b249 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp index 95e0694c9c..8bf18dc94a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh index 615824621f..23cf3fe49e 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in index 014cd12ac2..88a35a8b65 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp index 9db3839880..6ea2783d2c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh index 1a44453887..12c050d91b 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in index 7761879ce0..bb06616f9b 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp index 3559a9bee1..a58f4d8835 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in index 0828b0581a..807899b36e 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh index 7939b8ec13..8f338ab6ea 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in index c2b7c21726..3faddca95f 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp index 2314d20b62..db49dc7c72 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp index d9c94821f1..333243bb45 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in index 24aee45e03..fccc80ed86 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in index 5828027b6d..f177e0cc57 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in index 28d508e977..0d03cbb15b 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp index 0a6b00ba14..8ac69845b0 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp index 8b2934b771..07fb7f1285 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp index 0050ddee72..d012caa84b 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh index 51d3651b45..287aecc656 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu index 2bb9fb2174..6b82473128 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp index 5126cfbdff..1e616ab107 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/space.hpp b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp index df756de7e0..bee35cce2e 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/space.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index 43e4aa4471..288a03ebcf 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/nn_descent.cu b/cpp/src/neighbors/nn_descent.cu index eb2541b553..9405d4e608 100644 --- a/cpp/src/neighbors/nn_descent.cu +++ b/cpp/src/neighbors/nn_descent.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index f85e36a6ef..84c94eaee7 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1628,11 +1628,11 @@ inline std::vector generate_inputs() {100}, {1000}, {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 768, 1024}, // dim - {16}, // k - {32}, // degree + {16}, // k + {32}, // degree {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT, - graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build + graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build {search_algo::AUTO}, {10}, {0}, diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index 093727d318..e5dd1f77fc 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index c75b3555f6..30ac24c852 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu index 4cde210d62..da6ba5c969 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu index d8664d4e14..af167fb4e2 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu index 4c95192d8a..76f5b8cb71 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu index 3e4b91e759..433366f05b 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 938f41f846..3c825c9333 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 2b412f3401..5725856d9a 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/build.sh b/examples/build.sh index 1be41c01e4..0dc7e2760f 100755 --- a/examples/build.sh +++ b/examples/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cuvs empty project template build script diff --git a/examples/cpp/src/cagra_filter_udf_example.cu b/examples/cpp/src/cagra_filter_udf_example.cu index 0ab42dd580..5da0c10b9e 100644 --- a/examples/cpp/src/cagra_filter_udf_example.cu +++ b/examples/cpp/src/cagra_filter_udf_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/cpp/src/cagra_hnsw_ace_build.cu b/examples/cpp/src/cagra_hnsw_ace_build.cu index d23c08e22d..1602b98513 100644 --- a/examples/cpp/src/cagra_hnsw_ace_build.cu +++ b/examples/cpp/src/cagra_hnsw_ace_build.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/cpp/src/hnsw_openai_example.cu b/examples/cpp/src/hnsw_openai_example.cu index abb8346218..3e71f9f1e5 100644 --- a/examples/cpp/src/hnsw_openai_example.cu +++ b/examples/cpp/src/hnsw_openai_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py index c1633e3cad..5ea45781ce 100644 --- a/python/cuvs/cuvs/tests/test_cagra_ace.py +++ b/python/cuvs/cuvs/tests/test_cagra_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py index 183d530e7c..663640e50d 100644 --- a/python/cuvs/cuvs/tests/test_hnsw_ace.py +++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # From ae89f9f69cc2eb4e9053d89a6b4264dd49be0224 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 10 Aug 2026 08:22:04 -0700 Subject: [PATCH 21/39] perf(cagra): per-batch VPQ query reconstruction + device-pool graph temporaries in iterative build (~2.7x faster build) --- .../neighbors/detail/cagra/cagra_build.cuh | 168 ++++++++++++------ 1 file changed, 111 insertions(+), 57 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 1446bb81a5..df4c8ecb87 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2079,56 +2079,115 @@ void reconstruct_vpq_queries(raft::resources const& res, output.data_handle()); } +// Runs CAGRA search for `curr_query_size` queries against `idx` in chunks of `max_chunk_size`, +// stacks the results into a kNN graph, and optimizes it into the next graph (returned). +// +// Query source: +// - `vpq_queries == nullptr`: queries are read directly from `dev_query_view` (uncompressed +// build; the view is a slice of the resident device dataset). +// - `vpq_queries != nullptr`: `dev_query_view` is ignored and each chunk of queries is +// reconstructed on the fly from the VPQ codes into a small reusable scratch buffer, so we +// never materialize the whole (up to N x dim) reconstructed dataset. template -void search_and_optimize(raft::resources const& res, - const cuvs::neighbors::cagra::search_params& search_params, - const index& idx, - raft::device_matrix_view dev_query_view, - raft::device_matrix_view dev_neighbors, - raft::device_matrix_view dev_distances, - raft::device_matrix& dev_output_graph, - size_t curr_query_size, - size_t next_graph_degree, - size_t curr_topk, - uint64_t max_chunk_size) +raft::device_matrix search_and_optimize( + raft::resources const& res, + const cuvs::neighbors::cagra::search_params& search_params, + const index& idx, + raft::device_matrix_view dev_query_view, + raft::device_matrix_view dev_neighbors, + raft::device_matrix_view dev_distances, + raft::device_matrix prev_graph, + const vpq_dataset* vpq_queries, + size_t curr_query_size, + size_t next_graph_degree, + size_t curr_topk, + uint64_t max_chunk_size) { auto stream = raft::resource::get_cuda_stream(res); - auto dev_knn_graph = raft::make_device_matrix(res, curr_query_size, curr_topk); - - auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - res, - dev_query_view.data_handle(), - static_cast(curr_query_size), - static_cast(dev_query_view.extent(1)), - max_chunk_size, - stream, - raft::resource::get_workspace_resource_ref(res)); - for (const auto& batch : query_batch) { - auto batch_dev_query_view = raft::make_device_matrix_view( - batch.data(), batch.size(), dev_query_view.extent(1)); + // These buffers scale with N (e.g. N * (intermediate_degree+1) for the kNN graph). Allocate them + // from the default device resource (a pool over device memory): allocating from the large + // workspace resource here would use an unpooled managed_memory_resource, paying a synchronous + // cudaMallocManaged/cudaFree every iteration for multi-GB buffers. + auto dev_knn_graph = + raft::make_device_matrix(res, curr_query_size, curr_topk); + + // Query row length: the VPQ dim when reconstructing, otherwise the dataset view's stride. + const int64_t query_dim = + vpq_queries != nullptr ? static_cast(vpq_queries->dim()) : dev_query_view.extent(1); + + // Scratch for one reconstructed chunk (only when compressing). Reused across chunks; safe because + // all reconstruct/search/copy work is serialized on `stream`. + auto batch_queries = + vpq_queries != nullptr + ? raft::make_device_matrix(res, static_cast(max_chunk_size), query_dim) + : raft::make_device_matrix(res, 0, 0); + + auto run_batch = [&](int64_t offset, + int64_t batch_size, + raft::device_matrix_view batch_query_view) { auto batch_dev_neighbors_view = raft::make_device_matrix_view( - dev_neighbors.data_handle(), batch.size(), curr_topk); + dev_neighbors.data_handle(), batch_size, curr_topk); auto batch_dev_distances_view = raft::make_device_matrix_view( - dev_distances.data_handle(), batch.size(), curr_topk); + dev_distances.data_handle(), batch_size, curr_topk); cuvs::neighbors::cagra::search(res, search_params, idx, - batch_dev_query_view, + batch_query_view, batch_dev_neighbors_view, batch_dev_distances_view); - raft::copy(dev_knn_graph.data_handle() + batch.offset() * curr_topk, + raft::copy(dev_knn_graph.data_handle() + offset * curr_topk, batch_dev_neighbors_view.data_handle(), - batch.size() * curr_topk, + batch_size * curr_topk, stream); + }; + + if (vpq_queries != nullptr) { + // Reconstruct-and-search one chunk at a time: reconstruct source rows [offset, offset+bs) into + // the scratch, then search that chunk. + for (int64_t offset = 0; offset < static_cast(curr_query_size); + offset += static_cast(max_chunk_size)) { + const int64_t batch_size = + std::min(static_cast(max_chunk_size), + static_cast(curr_query_size) - offset); + reconstruct_vpq_queries(res, + *vpq_queries, + static_cast(offset), + static_cast(batch_size), + batch_queries.view()); + auto batch_query_view = raft::make_device_matrix_view( + batch_queries.data_handle(), batch_size, query_dim); + run_batch(offset, batch_size, batch_query_view); + } + } else { + auto query_batch = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( + res, + dev_query_view.data_handle(), + static_cast(curr_query_size), + query_dim, + max_chunk_size, + stream, + raft::resource::get_workspace_resource_ref(res)); + for (const auto& batch : query_batch) { + auto batch_query_view = raft::make_device_matrix_view( + batch.data(), static_cast(batch.size()), query_dim); + run_batch( + static_cast(batch.offset()), static_cast(batch.size()), batch_query_view); + } } - dev_output_graph = + // The previous-iteration graph (which `idx` was built on) is no longer needed now that the + // search has produced `dev_knn_graph`. Release it before allocating the full-size output graph + // so we never hold two large graph buffers at once. + prev_graph = raft::make_device_matrix(res, 0, 0); + + auto dev_output_graph = raft::make_device_matrix(res, curr_query_size, next_graph_degree); graph::optimize(res, dev_knn_graph.view(), dev_output_graph.view(), false); + return dev_output_graph; } // Builds a CAGRA graph iteratively, growing the graph while repeatedly running CAGRA's search() @@ -2292,6 +2351,7 @@ auto iterative_build_graph( "# Freed original dataset from device (%.1f MiB); queries will use VPQ reconstruction", to_mib(final_graph_size * dataset_dim * sizeof(T))); } + while (true) { auto start = std::chrono::high_resolution_clock::now(); auto curr_query_size = std::min(2 * curr_graph_size, final_graph_size); @@ -2352,40 +2412,34 @@ auto iterative_build_graph( } const auto& idx = *idx_opt; - // When compression is enabled, reconstruct queries from VPQ codes instead of - // reading from the (freed) original dataset. - auto dev_reconstructed_queries = - build_compression.has_value() - ? raft::make_device_matrix(res, curr_query_size, dataset_dim) - : raft::make_device_matrix(res, 0, 0); + // With compression, search_and_optimize reconstructs queries from the VPQ codes per batch, so + // pass the VPQ dataset and leave the query view empty. Without compression, queries are slices + // of the resident device dataset. + const vpq_dataset* vpq_queries = nullptr; if (build_compression.has_value()) { - auto* vpq_dset = dynamic_cast*>(&idx.data()); - RAFT_EXPECTS(vpq_dset != nullptr, "Expected VPQ dataset in compressed index"); - reconstruct_vpq_queries( - res, *vpq_dset, 0, curr_query_size, dev_reconstructed_queries.view()); + vpq_queries = dynamic_cast*>(&idx.data()); + RAFT_EXPECTS(vpq_queries != nullptr, "Expected VPQ dataset in compressed index"); } auto dev_query_view = build_compression.has_value() - ? raft::make_device_matrix_view( - dev_reconstructed_queries.data_handle(), (int64_t)curr_query_size, dataset_dim) + ? raft::make_device_matrix_view(static_cast(nullptr), 0, 0) : raft::make_device_matrix_view( dev_dataset.data_handle(), (int64_t)curr_query_size, dev_dataset.extent(1)); - auto dev_optimized_graph = raft::make_device_matrix(res, 0, 0); - - search_and_optimize(res, - search_params, - idx, - dev_query_view, - dev_neighbors.view(), - dev_distances.view(), - dev_optimized_graph, - curr_query_size, - next_graph_degree, - curr_topk, - max_chunk_size); - - dev_graph = std::move(dev_optimized_graph); + // Hand the current graph to search_and_optimize so it can release it as soon as the search + // consumes it (before the new output graph is allocated), then take back the new graph. + dev_graph = search_and_optimize(res, + search_params, + idx, + dev_query_view, + dev_neighbors.view(), + dev_distances.view(), + std::move(dev_graph), + vpq_queries, + curr_query_size, + next_graph_degree, + curr_topk, + max_chunk_size); use_device_graph = true; auto end = std::chrono::high_resolution_clock::now(); From 355d240731b803cf3056d0abbaf9876279281452 Mon Sep 17 00:00:00 2001 From: aamijar Date: Mon, 10 Aug 2026 21:59:18 +0000 Subject: [PATCH 22/39] Revert "Pre-commit changes" This reverts commit d9c6bfd10941470b85c22735834749b363cc5ebb. --- c/src/neighbors/brute_force.cpp | 2 +- c/src/neighbors/cagra.cpp | 2 +- c/src/neighbors/ivf_flat.cpp | 2 +- c/src/neighbors/mg_cagra.cpp | 2 +- c/src/neighbors/mg_ivf_flat.cpp | 2 +- c/src/neighbors/mg_ivf_pq.cpp | 2 +- ci/build_cpp.sh | 2 +- ci/build_go.sh | 2 +- ci/build_java.sh | 2 +- ci/build_python.sh | 2 +- ci/build_rust.sh | 2 +- ci/build_wheel_cuvs.sh | 2 +- ci/build_wheel_libcuvs.sh | 2 +- ci/test_cpp.sh | 2 +- ci/test_wheel_cuvs.sh | 2 +- cpp/bench/ann/CMakeLists.txt | 2 +- cpp/bench/ann/src/cuvs/cuvs_benchmark.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu | 2 +- cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp | 2 +- cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h | 2 +- cpp/cmake/modules/compute_matrix_product.cmake | 2 +- cpp/include/cuvs/detail/jit_lto/common_fragments.hpp | 2 +- .../cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp | 2 +- cpp/include/cuvs/neighbors/ivf_pq.hpp | 2 +- cpp/include/cuvs/neighbors/ivf_rabitq.hpp | 2 +- cpp/include/cuvs/neighbors/nn_descent.hpp | 2 +- cpp/include/cuvs/util/file_io.hpp | 2 +- cpp/src/neighbors/brute_force_serialize.cu | 2 +- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/cagra_helpers.hpp | 2 +- cpp/src/neighbors/detail/cagra/cagra_serialize.cuh | 2 +- cpp/src/neighbors/detail/cagra/graph_core.cuh | 2 +- .../detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh | 2 +- .../detail/cagra/jit_lto_kernels/search_multi_jit.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in | 2 +- .../detail/cagra/search_multi_cta_kernel_launcher_jit.cuh | 2 +- .../detail/cagra/search_multi_kernel_launcher_jit.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in | 2 +- cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp | 2 +- cpp/src/neighbors/detail/hnsw.hpp | 2 +- cpp/src/neighbors/ivf_pq_index.cu | 2 +- cpp/src/neighbors/ivf_rabitq.cu | 2 +- cpp/src/neighbors/ivf_rabitq/defines.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu | 2 +- cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh | 2 +- .../neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh | 2 +- .../ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu | 2 +- .../ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu | 2 +- .../bitwise_block_sort_emit_topk_kernel.cu.in | 2 +- .../jit_lto_kernels/bitwise_emit_distances_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh | 2 +- .../compute_bitwise_quantized_ip_for_vec_kernel.cu.in | 2 +- .../compute_inner_products_with_bitwise_block_sort_impl.cuh | 2 +- ...pute_inner_products_with_bitwise_block_sort_kernel.cu.in | 2 +- ...mpute_inner_products_with_bitwise_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_bitwise_impl.cuh | 2 +- .../compute_inner_products_with_bitwise_kernel.cu.in | 2 +- .../compute_inner_products_with_bitwise_planner.hpp | 2 +- ...ompute_inner_products_with_lut16_opt_block_sort_impl.cuh | 2 +- ...te_inner_products_with_lut16_opt_block_sort_kernel.cu.in | 2 +- ...ute_inner_products_with_lut16_opt_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_lut16_opt_impl.cuh | 2 +- .../compute_inner_products_with_lut16_opt_kernel.cu.in | 2 +- .../compute_inner_products_with_lut16_opt_planner.hpp | 2 +- .../compute_inner_products_with_lut_block_sort_impl.cuh | 2 +- .../compute_inner_products_with_lut_block_sort_kernel.cu.in | 2 +- .../compute_inner_products_with_lut_block_sort_planner.hpp | 2 +- .../compute_inner_products_with_lut_impl.cuh | 2 +- .../compute_inner_products_with_lut_kernel.cu.in | 2 +- .../compute_inner_products_with_lut_planner.hpp | 2 +- .../jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in | 2 +- .../ivf_rabitq/jit_lto_kernels/device_functions.cuh | 2 +- .../ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp | 2 +- .../ivf_rabitq/jit_lto_kernels/launcher_factory.hpp | 2 +- .../jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in | 2 +- .../jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in | 2 +- .../jit_lto_kernels/lut_emit_distances_kernel.cu.in | 2 +- cpp/src/neighbors/ivf_rabitq/utils/IO.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/memory.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh | 2 +- cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu | 2 +- cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp | 2 +- cpp/src/neighbors/ivf_rabitq/utils/space.hpp | 2 +- cpp/src/neighbors/mg/snmg.cuh | 2 +- cpp/src/neighbors/nn_descent.cu | 2 +- cpp/tests/neighbors/ann_cagra.cuh | 6 +++--- cpp/tests/neighbors/ann_cagra/test_filter_udf.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace.cuh | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu | 2 +- cpp/tests/neighbors/ann_ivf_rabitq.cuh | 2 +- cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu | 2 +- examples/build.sh | 2 +- examples/cpp/src/cagra_filter_udf_example.cu | 2 +- examples/cpp/src/cagra_hnsw_ace_build.cu | 2 +- examples/cpp/src/hnsw_openai_example.cu | 2 +- python/cuvs/cuvs/tests/test_cagra_ace.py | 2 +- python/cuvs/cuvs/tests/test_hnsw_ace.py | 2 +- 113 files changed, 115 insertions(+), 115 deletions(-) diff --git a/c/src/neighbors/brute_force.cpp b/c/src/neighbors/brute_force.cpp index d926a56cd4..081f433a3e 100644 --- a/c/src/neighbors/brute_force.cpp +++ b/c/src/neighbors/brute_force.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index a01f7b051f..004b810c78 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/ivf_flat.cpp b/c/src/neighbors/ivf_flat.cpp index 7ae9073f9e..729c810aa0 100644 --- a/c/src/neighbors/ivf_flat.cpp +++ b/c/src/neighbors/ivf_flat.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_cagra.cpp b/c/src/neighbors/mg_cagra.cpp index 30be2f76a8..495eff8a34 100644 --- a/c/src/neighbors/mg_cagra.cpp +++ b/c/src/neighbors/mg_cagra.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_ivf_flat.cpp b/c/src/neighbors/mg_ivf_flat.cpp index cdf261a049..4e1b2883ea 100644 --- a/c/src/neighbors/mg_ivf_flat.cpp +++ b/c/src/neighbors/mg_ivf_flat.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/c/src/neighbors/mg_ivf_pq.cpp b/c/src/neighbors/mg_ivf_pq.cpp index 8570f07c8b..41aa323138 100644 --- a/c/src/neighbors/mg_ivf_pq.cpp +++ b/c/src/neighbors/mg_ivf_pq.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/ci/build_cpp.sh b/ci/build_cpp.sh index cda3d10bd9..bb497692bd 100755 --- a/ci/build_cpp.sh +++ b/ci/build_cpp.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_go.sh b/ci/build_go.sh index e04f6f1d8a..056f40b345 100755 --- a/ci/build_go.sh +++ b/ci/build_go.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_java.sh b/ci/build_java.sh index 167c642dd2..2e363bb452 100755 --- a/ci/build_java.sh +++ b/ci/build_java.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_python.sh b/ci/build_python.sh index 34be953de5..6823cbbec5 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_rust.sh b/ci/build_rust.sh index e84e454e05..e9218a8adc 100755 --- a/ci/build_rust.sh +++ b/ci/build_rust.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_wheel_cuvs.sh b/ci/build_wheel_cuvs.sh index 2e174bc87e..2acbab8a2d 100755 --- a/ci/build_wheel_cuvs.sh +++ b/ci/build_wheel_cuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/build_wheel_libcuvs.sh b/ci/build_wheel_libcuvs.sh index 9bc1ac0a37..e012935749 100755 --- a/ci/build_wheel_libcuvs.sh +++ b/ci/build_wheel_libcuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index be62cfc754..0905cd64f0 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/ci/test_wheel_cuvs.sh b/ci/test_wheel_cuvs.sh index 227857dd8c..36fefdf852 100755 --- a/ci/test_wheel_cuvs.sh +++ b/ci/test_wheel_cuvs.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 set -euo pipefail diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 80f116f586..90d23d9aef 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu index 1a334924ec..3056ddc365 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_benchmark.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu index 8c5854051d..c903b39fcc 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h index c7733c293e..db618f6559 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_hnswlib_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu index e3c38025c7..1c6772d994 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #include "cuvs_ivf_rabitq_wrapper.h" diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index 542f0bc6dd..ca8f77b808 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp index f40ad67a63..1a99e56028 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp +++ b/cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h index cf2bb7608e..282d57dc2e 100644 --- a/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/cmake/modules/compute_matrix_product.cmake b/cpp/cmake/modules/compute_matrix_product.cmake index 6d81821b88..82a34f9242 100644 --- a/cpp/cmake/modules/compute_matrix_product.cmake +++ b/cpp/cmake/modules/compute_matrix_product.cmake @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index a55052c4a6..ef2a8e6002 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp index 2b2f3db5a7..30885947b5 100644 --- a/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index 686c3ff108..57f8a258fb 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp index c33d466613..d26bc04022 100644 --- a/cpp/include/cuvs/neighbors/ivf_rabitq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_rabitq.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 929a099cf1..4c031049e2 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/util/file_io.hpp b/cpp/include/cuvs/util/file_io.hpp index b0afeed732..a7d67ec2c0 100644 --- a/cpp/include/cuvs/util/file_io.hpp +++ b/cpp/include/cuvs/util/file_io.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/brute_force_serialize.cu b/cpp/src/neighbors/brute_force_serialize.cu index e3a4a2c041..1b7595ee11 100644 --- a/cpp/src/neighbors/brute_force_serialize.cu +++ b/cpp/src/neighbors/brute_force_serialize.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 34b2e72f90..ee87c2c0ab 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp index 5c6a63f9fc..ee78930970 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index e80c6b6932..f106b82500 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 1762bdbfb0..52b4542798 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh index 0e9e981e69..4c4f2e4f62 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_cta_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh index 02cec6ee5e..a714ced5c2 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/search_multi_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in index 4e94922566..7c642fe406 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh index ff1064f24c..8a673405b7 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel_launcher_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh index 549474d045..bc341b9082 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel_launcher_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in index 7869e4f05e..4616a9652b 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp index e448bebeb4..e5157ffa6a 100644 --- a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp +++ b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 649886f924..88580de929 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_pq_index.cu b/cpp/src/neighbors/ivf_pq_index.cu index 114b98bf0c..28b985eec8 100644 --- a/cpp/src/neighbors/ivf_pq_index.cu +++ b/cpp/src/neighbors/ivf_pq_index.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu index 14a9678bad..d5572d4039 100644 --- a/cpp/src/neighbors/ivf_rabitq.cu +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/defines.hpp b/cpp/src/neighbors/ivf_rabitq/defines.hpp index aac296cb12..f91c06a1ff 100644 --- a/cpp/src/neighbors/ivf_rabitq/defines.hpp +++ b/cpp/src/neighbors/ivf_rabitq/defines.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu index 6e1e2c7d18..6755bfc3de 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh index 135117d4a9..0640148814 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index c071fb9a0a..49070acbcf 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index d275c84eb3..75ddaee865 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index ea40f02931..af0876acca 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh index 21aaf414ac..90ef71ed12 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu index c9c0a1d275..59a3e575d9 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh index b7a485b7be..8db4bb464c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu index 2c81c40e85..22bb74c682 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh index 68ca342370..73557f57ea 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh index 125a759ca0..1e004d9164 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_common.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index fde565e0e9..92906fb0e6 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu index bee53b54a2..8357af5fbb 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in index 6a9b0d9576..d8dcff324d 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in index fb6a130512..20c5aec6a1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/bitwise_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh index fb78d0d560..a16058c2af 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/block_sort.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in index d5a63ee313..e4e2aa6d79 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_bitwise_quantized_ip_for_vec_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh index e3b2da9733..a577ec44ce 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in index 785aa22b26..7bbdc24139 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp index b126e723cc..dbe6ec5fd4 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh index cf87f17314..4ff53b067c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in index 53897c47db..1f7355c7fe 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp index ce7197e7e1..3ec4809395 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_bitwise_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh index cc393f82db..45c8793b59 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in index 6aa138b917..8e45f99c0f 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp index e71903764b..acf4e27e3a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh index 575766d054..fc88d79fc4 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in index 713a15b249..4160aa7a33 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp index 8bf18dc94a..95e0694c9c 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut16_opt_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh index 23cf3fe49e..615824621f 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in index 88a35a8b65..014cd12ac2 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp index 6ea2783d2c..9db3839880 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_block_sort_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh index 12c050d91b..1a44453887 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_impl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in index bb06616f9b..7761879ce0 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp index a58f4d8835..3559a9bee1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_inner_products_with_lut_planner.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in index 807899b36e..0828b0581a 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/compute_lut_ip_for_vec_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh index 8f338ab6ea..7939b8ec13 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/device_functions.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in index 3faddca95f..c2b7c21726 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/extract_code_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp index db49dc7c72..2314d20b62 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/kernel_def.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp index 333243bb45..d9c94821f1 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/launcher_factory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in index fccc80ed86..24aee45e03 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut16_opt_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in index f177e0cc57..5828027b6d 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_block_sort_emit_topk_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in index 0d03cbb15b..28d508e977 100644 --- a/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in +++ b/cpp/src/neighbors/ivf_rabitq/jit_lto_kernels/lut_emit_distances_kernel.cu.in @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp index 8ac69845b0..0a6b00ba14 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/IO.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp index 07fb7f1285..8b2934b771 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/StopW.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp index d012caa84b..0050ddee72 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/memory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh index 287aecc656..51d3651b45 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu index 6b82473128..2bb9fb2174 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp index 1e616ab107..5126cfbdff 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/ivf_rabitq/utils/space.hpp b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp index bee35cce2e..df756de7e0 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/space.hpp +++ b/cpp/src/neighbors/ivf_rabitq/utils/space.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index 288a03ebcf..43e4aa4471 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/nn_descent.cu b/cpp/src/neighbors/nn_descent.cu index 9405d4e608..eb2541b553 100644 --- a/cpp/src/neighbors/nn_descent.cu +++ b/cpp/src/neighbors/nn_descent.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 84c94eaee7..f85e36a6ef 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1628,11 +1628,11 @@ inline std::vector generate_inputs() {100}, {1000}, {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 768, 1024}, // dim - {16}, // k - {32}, // degree + {16}, // k + {32}, // degree {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT, - graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build + graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build {search_algo::AUTO}, {10}, {0}, diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index e5dd1f77fc..093727d318 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 30ac24c852..c75b3555f6 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu index da6ba5c969..4cde210d62 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu index af167fb4e2..d8664d4e14 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu index 76f5b8cb71..4c95192d8a 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu index 433366f05b..3e4b91e759 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 3c825c9333..938f41f846 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 5725856d9a..2b412f3401 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/build.sh b/examples/build.sh index 0dc7e2760f..1be41c01e4 100755 --- a/examples/build.sh +++ b/examples/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # cuvs empty project template build script diff --git a/examples/cpp/src/cagra_filter_udf_example.cu b/examples/cpp/src/cagra_filter_udf_example.cu index 5da0c10b9e..0ab42dd580 100644 --- a/examples/cpp/src/cagra_filter_udf_example.cu +++ b/examples/cpp/src/cagra_filter_udf_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/cpp/src/cagra_hnsw_ace_build.cu b/examples/cpp/src/cagra_hnsw_ace_build.cu index 1602b98513..d23c08e22d 100644 --- a/examples/cpp/src/cagra_hnsw_ace_build.cu +++ b/examples/cpp/src/cagra_hnsw_ace_build.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/cpp/src/hnsw_openai_example.cu b/examples/cpp/src/hnsw_openai_example.cu index 3e71f9f1e5..abb8346218 100644 --- a/examples/cpp/src/hnsw_openai_example.cu +++ b/examples/cpp/src/hnsw_openai_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py index 5ea45781ce..c1633e3cad 100644 --- a/python/cuvs/cuvs/tests/test_cagra_ace.py +++ b/python/cuvs/cuvs/tests/test_cagra_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py index 663640e50d..183d530e7c 100644 --- a/python/cuvs/cuvs/tests/test_hnsw_ace.py +++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # From 8cd41911c4450a912ecb53ed31d901594962ea54 Mon Sep 17 00:00:00 2001 From: aamijar Date: Mon, 10 Aug 2026 22:04:22 +0000 Subject: [PATCH 23/39] revert another spdx change --- cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index 16fa93f47b..adeb774a8b 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ From 6ad9234655a846385ea8b839e5f569335fe95f6d Mon Sep 17 00:00:00 2001 From: aamijar Date: Mon, 10 Aug 2026 22:12:59 +0000 Subject: [PATCH 24/39] revert test to minimize diff --- cpp/tests/neighbors/ann_cagra.cuh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index f85e36a6ef..1e969ec5fe 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -1627,12 +1627,12 @@ inline std::vector generate_inputs() inputs2 = raft::util::itertools::product( {100}, {1000}, - {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 768, 1024}, // dim - {16}, // k - {32}, // degree + {1, 3, 5, 7, 8, 17, 64, 128, 137, 192, 256, 512, 1024}, // dim + {16}, // k + {32}, // degree {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT, - graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // Iterative cagra q build + graph_build_algo::ITERATIVE_CAGRA_SEARCH}, {search_algo::AUTO}, {10}, {0}, From 8a9fda5c02c1fae5b6ddf9d99476f42980ac00f4 Mon Sep 17 00:00:00 2001 From: aamijar Date: Thu, 13 Aug 2026 00:30:53 +0000 Subject: [PATCH 25/39] remove unrelated utility and test --- cpp/src/neighbors/detail/cagra/utils.hpp | 395 +----------------- cpp/tests/CMakeLists.txt | 4 +- .../test_batched_device_view_from_host.cu | 205 --------- 3 files changed, 2 insertions(+), 602 deletions(-) delete mode 100644 cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu diff --git a/cpp/src/neighbors/detail/cagra/utils.hpp b/cpp/src/neighbors/detail/cagra/utils.hpp index 7b31fbdee3..58bf68bb43 100644 --- a/cpp/src/neighbors/detail/cagra/utils.hpp +++ b/cpp/src/neighbors/detail/cagra/utils.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -301,397 +301,4 @@ void copy_with_padding( } } -/** - * Utility to create a batched device view from a host view - * - * This utility will create a batched device view from a host view and will handle the prefetch and - * writeback of the data Each batch can be referenced exactlyonce by calling the next_view() - * function - * - * Usage: - * ``` - * batched_device_view_from_host view(res, host_view, batch_size, host_writeback, - * initialize); while (view.next_view().extent(0) > 0) { auto device_view = view.next_view(); - * // use device_view - * } - * ``` - * - * The call to next_view() will - * * synchronize on all previous operations / increments batch_id_ - * * (optionally) write back the data of the previous batch to the host - * * (optionally) prefetch the data of the next batch - * * return the view of the current batch - * - * @tparam T The type of the data - * @tparam IdxT The type of the index - */ -template -class batched_device_view_from_host { - public: - enum class memory_strategy { - device_only, // data is on device only (no copy needed) - copy_device, // data is explicitly moved to/from device buffers - managed_only, // data is on managed memory (system managed) - }; - - /** - * Create a batched device view from a host view and will handle the prefetch and - * writeback of the data. Each batch can be referenced exactly once by calling the next_view() - * method. - * - * @param res The resources to use - * @param host_view The host view to create the batched device view from - * @param batch_size The batch size - * @param host_writeback Whether to write back the data to the host (only for host memory) - * (default: false) - * @param initialize Whether to initialize the data (only for managed memory) (default: true) - */ - batched_device_view_from_host(raft::resources const& res, - raft::host_matrix_view host_view, - uint64_t batch_size, - bool host_writeback = false, - bool initialize = true) - : res_(res), - host_view_(host_view), - batch_size_(batch_size), - offset_(0), - batch_id_(-2), - num_buffers_(2), - host_writeback_(host_writeback), - initialize_(initialize) - { - if (host_view.extent(0) == 0) { - mem_strategy_ = memory_strategy::device_only; - return; - } - - RAFT_EXPECTS(host_writeback_ || initialize_, - "At least one of host_writeback or initialize must be true"); - - RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr_, host_view.data_handle())); - switch (attr_.type) { - case cudaMemoryTypeUnregistered: - case cudaMemoryTypeHost: - case cudaMemoryTypeManaged: mem_strategy_ = memory_strategy::copy_device; break; - case cudaMemoryTypeDevice: mem_strategy_ = memory_strategy::device_only; break; - } - - RAFT_LOG_DEBUG("Memory strategy: %d for type %d, size %zu", - static_cast(mem_strategy_), - static_cast(attr_.type), - host_view.extent(0) * host_view.extent(1) * sizeof(T)); - - // buffer allocations - if (mem_strategy_ == memory_strategy::copy_device) { - try { - device_mem_[0].emplace(raft::make_device_mdarray( - res, - raft::resource::get_workspace_resource_ref(res), - raft::make_extents(batch_size, host_view.extent(1)))); - device_ptr[0] = device_mem_[0]->data_handle(); - if (batch_size < static_cast(host_view.extent(0))) { - device_mem_[1].emplace(raft::make_device_mdarray( - res, - raft::resource::get_workspace_resource_ref(res), - raft::make_extents(batch_size, host_view.extent(1)))); - device_ptr[1] = device_mem_[1]->data_handle(); - } - if (host_writeback_ && initialize_ && - batch_size * 2 < static_cast(host_view.extent(0))) { - num_buffers_ = 3; - device_mem_[2].emplace(raft::make_device_mdarray( - res, - raft::resource::get_workspace_resource_ref(res), - raft::make_extents(batch_size, host_view.extent(1)))); - device_ptr[2] = device_mem_[2]->data_handle(); - } - } catch (std::bad_alloc& e) { - if (attr_.devicePointer != nullptr) { - RAFT_LOG_DEBUG("Insufficient memory for device buffers, switching to managed memory"); - mem_strategy_ = memory_strategy::managed_only; - } else { - throw std::bad_alloc(); - } - } catch (raft::logic_error& e) { - if (attr_.devicePointer != nullptr) { - RAFT_LOG_DEBUG( - "Insufficient memory for device buffers (logic error), switching to managed memory"); - mem_strategy_ = memory_strategy::managed_only; - } else { - throw raft::logic_error("Insufficient memory for device buffers (logic error)"); - } - } - } - - // setup stream pool if not already present - size_t required_streams = host_writeback_ && initialize_ ? 2 : 1; - if (!res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) || - raft::resource::get_stream_pool_size(res) < required_streams) { - // always create at least 2 streams to account for subsequent iterator calls. - // set_cuda_stream_pool now requires a non-const resource; the referenced resource - // outlives this object, so attaching the pool to it here is safe. - raft::resource::set_cuda_stream_pool(const_cast(res), - std::make_shared(2)); - } - prefetch_stream_ = raft::resource::get_stream_from_stream_pool(res); - writeback_stream_ = raft::resource::get_stream_from_stream_pool(res); - - // if data is managed and not for_write_ we can set the attribute on the device ptr - if (mem_strategy_ == memory_strategy::managed_only) { - location_.type = cudaMemLocationTypeDevice; - location_.id = static_cast(raft::resource::get_device_id(res_)); - if (!host_writeback_) { - advise_read_mostly(host_view_.data_handle(), - host_view_.extent(0) * host_view_.extent(1) * sizeof(T)); - // TODO maybe also reset upon destruction - } - } - - // prefetch next batch (0) - prefetch_next_batch(); - } - - ~batched_device_view_from_host() noexcept - { - raft::resource::sync_stream(res_); - - // if data is on host and for_write --> make sure to copy back last active - // if data is managed and evict --> evict last active - - // make sure to sync on prefetch stream & res - switch (mem_strategy_) { - case memory_strategy::managed_only: - if (!host_writeback_) { - uint32_t discard_pos = batch_id_ % num_buffers_; - size_t discard_size_rows = actual_batch_size_[discard_pos]; - if (batch_id_ > 0) { - discard_pos = (batch_id_ - 1) % num_buffers_; - discard_size_rows += batch_size_; - } - discard_managed_region(device_ptr[discard_pos], - discard_size_rows * host_view_.extent(1) * sizeof(T)); - writeback_stream_.synchronize(); - } - break; - case memory_strategy::copy_device: - if (host_writeback_) { - uint32_t writeback_pos_last = batch_id_ % num_buffers_; - if (batch_id_ > 0) { - uint32_t writeback_pos = (batch_id_ - 1) % num_buffers_; - uint64_t writeback_offset = (batch_id_ - 1) * batch_size_; - writeback_from_device_to_host(device_ptr[writeback_pos], writeback_offset, batch_size_); - } - { - uint64_t writeback_offset_last = batch_id_ * batch_size_; - writeback_from_device_to_host(device_ptr[writeback_pos_last], - writeback_offset_last, - actual_batch_size_[writeback_pos_last]); - } - writeback_stream_.synchronize(); - } - break; - case memory_strategy::device_only: break; - } - } - - /** - * Returns the next view of the batch - * - * This function will ensure the next batch is ready and will trigger the prefetch of the - * subsequent next batch. If writeback is enabled, the last active batch will be written back to - * the host. - * - * @return The next view of the batch - */ - raft::device_matrix_view next_view() - { - bool end_of_data = static_cast((batch_id_ + 1) * batch_size_) >= - static_cast(host_view_.extent(0)); - - // special case for empty host view or last batch surpassed - if (end_of_data) { - return raft::make_device_matrix_view(nullptr, 0, host_view_.extent(1)); - } - - // trigger prefetch of next batch (also increments batch_id_) - prefetch_next_batch(); - - uint32_t current_pos = batch_id_ % num_buffers_; - return raft::make_device_matrix_view( - device_ptr[current_pos], actual_batch_size_[current_pos], host_view_.extent(1)); - } - - private: - /** - * Prefetch the next batch - * - * This function will prefetch the next batch and will handle the writeback of the data. - * - * @return True if the next batch exists, false otherwise - */ - bool prefetch_next_batch() - { - batch_id_++; - - // ensure previous batch at position batch_id_ is ready - if (initialize_) { prefetch_stream_.synchronize(); } - if (host_writeback_) { writeback_stream_.synchronize(); } - - // this step will - // * write back data from batch_id_ - 1 - // * prefetch data for batch_id_ + 1 - - // if data is on host and host_writeback_ is true we will have to copy it back - // if data is on host and initialize_ is true we will have to copy it to the device_ptr - - // if data is managed and !host_writeback_ we can discard the data from device memory - // if data is managed and initialize_ is true we can prefetch it to the device - // if data is managed and !initialize_ we can discard and prefetch the data location - - // if data is on device only this is almost a noop, just prepping the pointers - - RAFT_EXPECTS(static_cast(offset_) <= host_view_.extent(0), "Offset out of bounds"); - - bool next_batch_exists = offset_ < static_cast(host_view_.extent(0)); - - if (next_batch_exists) { - // synchronize to ensure all previous operations are completed - // in particular all work on batch_id_ - 1 - raft::resource::sync_stream(res_); - - int32_t prefetch_pos = (batch_id_ + 1) % num_buffers_; - actual_batch_size_[prefetch_pos] = min(batch_size_, host_view_.extent(0) - offset_); - - switch (mem_strategy_) { - case memory_strategy::managed_only: - if (!host_writeback_ && batch_id_ > 1) { - uint32_t discard_pos = (batch_id_ - 1) % num_buffers_; - size_t discard_size = batch_size_ * host_view_.extent(1) * sizeof(T); - discard_managed_region(device_ptr[discard_pos], discard_size); - } - // prefetch next position - device_ptr[prefetch_pos] = host_view_.data_handle() + offset_ * host_view_.extent(1); - prefetch_managed_region( - device_ptr[prefetch_pos], - actual_batch_size_[prefetch_pos] * host_view_.extent(1) * sizeof(T)); - break; - case memory_strategy::copy_device: - if (host_writeback_ && batch_id_ > 0) { - // copy back last active - uint32_t writeback_pos = (batch_id_ - 1) % num_buffers_; - uint64_t writeback_offset = (batch_id_ - 1) * batch_size_; - writeback_from_device_to_host(device_ptr[writeback_pos], writeback_offset, batch_size_); - } - if (initialize_) { - // prefetch next position - prefetch_from_host_to_device( - device_ptr[prefetch_pos], offset_, actual_batch_size_[prefetch_pos]); - } - - break; - case memory_strategy::device_only: - // just move pointer to next position - device_ptr[prefetch_pos] = host_view_.data_handle() + offset_ * host_view_.extent(1); - break; - } - - offset_ += actual_batch_size_[prefetch_pos]; - } - - return next_batch_exists; - } - - void advise_read_mostly(T* ptr, size_t size) - { -#if CUDA_VERSION >= 13000 - RAFT_CUDA_TRY(cudaMemAdvise(ptr, size, cudaMemAdviseSetReadMostly, location_)); -#else - RAFT_CUDA_TRY(cudaMemAdvise_v2(ptr, size, cudaMemAdviseSetReadMostly, location_)); -#endif - } - - void discard_managed_region(T* dev_ptr, size_t size) - { -#if CUDA_VERSION >= 13000 - void* dptrs[1] = {dev_ptr}; - size_t sizes[1] = {size}; - RAFT_CUDA_TRY(cudaMemDiscardBatchAsync(dptrs, sizes, 1, 0, writeback_stream_)); -#endif - // FIXME: CUDA12 does not support discard - } - - void prefetch_managed_region(T* dev_ptr, size_t size) - { -#if CUDA_VERSION >= 13000 - if (initialize_) { - RAFT_CUDA_TRY(cudaMemPrefetchAsync(dev_ptr, size, location_, 0, prefetch_stream_)); - } else { - void* dptrs[1] = {dev_ptr}; - size_t sizes[1] = {size}; - RAFT_CUDA_TRY( - cudaMemDiscardAndPrefetchBatchAsync(dptrs, sizes, 1, location_, 0, prefetch_stream_)); - } -#else - // FIXME: CUDA12 does not support discard - so we just prefetch - if (initialize_) { - RAFT_CUDA_TRY(cudaMemPrefetchAsync_v2(dev_ptr, size, location_, 0, prefetch_stream_)); - } else { - RAFT_CUDA_TRY(cudaMemPrefetchAsync_v2(dev_ptr, size, location_, 0, prefetch_stream_)); - } -#endif - } - - void prefetch_from_host_to_device(T* dev_ptr, size_t src_row_offset, size_t num_rows) - { - const size_t n_elem = num_rows * host_view_.extent(1); - const size_t n_bytes = n_elem * sizeof(T); - // use memcpy instead of raft::copy to avoid strange behavior with HMM/ATS memory - RAFT_CUDA_TRY(cudaMemcpyAsync(dev_ptr, - host_view_.data_handle() + src_row_offset * host_view_.extent(1), - n_bytes, - cudaMemcpyHostToDevice, - prefetch_stream_)); - } - - void writeback_from_device_to_host(T* dev_ptr, size_t dst_row_offset, size_t num_rows) - { - const size_t n_elem = num_rows * host_view_.extent(1); - const size_t n_bytes = n_elem * sizeof(T); - // use memcpy instead of raft::copy to avoid strange behavior with HMM/ATS memory - RAFT_CUDA_TRY(cudaMemcpyAsync(host_view_.data_handle() + dst_row_offset * host_view_.extent(1), - dev_ptr, - n_bytes, - cudaMemcpyDeviceToHost, - writeback_stream_)); - } - - // stream pool for local streams - std::optional> local_stream_pool_; - rmm::cuda_stream_view prefetch_stream_; - rmm::cuda_stream_view writeback_stream_; - - // configuration - memory_strategy mem_strategy_; - const raft::resources& res_; - bool initialize_; // initialize the data on the device - bool host_writeback_; // write back the data to the host - - // batch position information - uint64_t batch_size_; - int32_t batch_id_; - uint64_t offset_; - - cudaMemLocation location_; - - // input pointer information - raft::host_matrix_view host_view_; - cudaPointerAttributes attr_; - - // internal device buffers - uint64_t num_buffers_; - std::optional> device_mem_[3]; - T* device_ptr[3]; - uint32_t actual_batch_size_[3]; -}; - } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index f08ce55bf4..8421c3219e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -213,9 +213,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST - PATH neighbors/ann_cagra/test_optimize_uint32_t.cu - neighbors/ann_cagra/test_batched_device_view_from_host.cu - neighbors/ann_cagra/test_batch_load_iterator.cu + PATH neighbors/ann_cagra/test_optimize_uint32_t.cu neighbors/ann_cagra/test_batch_load_iterator.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu b/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu deleted file mode 100644 index eb72dbec92..0000000000 --- a/cpp/tests/neighbors/ann_cagra/test_batched_device_view_from_host.cu +++ /dev/null @@ -1,205 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -#include "../../../src/neighbors/detail/cagra/utils.hpp" - -#include -#include -#include -#include - -namespace cuvs::neighbors::cagra { - -using IdxT = uint32_t; - -struct BatchConfig { - bool initialize; - bool host_writeback; -}; - -struct DimsConfig { - int64_t n_rows; - int64_t n_cols; - uint64_t batch_size; -}; - -class BatchedDeviceViewFromHostTest : public ::testing::Test { - protected: - void SetUp() override { raft::resource::sync_stream(res); } - - /** - * Run batched_device_view_from_host over host data, copy device views back, - * and verify against the input. - */ - template - void run_and_verify_batched(InputMatrixView input_view, - uint64_t batch_size, - bool host_writeback, - bool initialize) - { - int64_t n_rows = input_view.extent(0); - int64_t n_cols = input_view.extent(1); - - std::vector readback(n_rows * n_cols); - - int64_t total_processed = 0; - - { - cagra::detail::batched_device_view_from_host batched( - res, - raft::make_host_matrix_view(input_view.data_handle(), n_rows, n_cols), - batch_size, - host_writeback, - initialize); - while (true) { - auto dev_view = batched.next_view(); - if (dev_view.extent(0) == 0) break; - - if (initialize) { - raft::copy(readback.data() + total_processed * n_cols, - dev_view.data_handle(), - dev_view.extent(0) * dev_view.extent(1), - raft::resource::get_cuda_stream(res)); - } - if (host_writeback) { raft::matrix::fill(res, dev_view, IdxT(17)); } - total_processed += dev_view.extent(0); - } - } - raft::resource::sync_stream(res); - - EXPECT_EQ(total_processed, n_rows); - if (initialize) { - for (int64_t i = 0; i < n_rows * n_cols; ++i) { - EXPECT_EQ(readback[i], IdxT(13)) << "Mismatch (initialize) at index " << i; - } - } - if (host_writeback) { - auto readback_view = - raft::make_host_matrix_view(readback.data(), n_rows, n_cols); - raft::copy(res, readback_view, input_view); - raft::resource::sync_stream(res); - for (int64_t i = 0; i < n_rows * n_cols; ++i) { - EXPECT_EQ(readback[i], IdxT(17)) << "Mismatch (host_writeback) at index " << i; - } - } - } - - raft::resources res; -}; - -TEST_F(BatchedDeviceViewFromHostTest, EmptyView) -{ - auto host_empty = raft::make_host_matrix(0, 8); - auto host_view = host_empty.view(); - cagra::detail::batched_device_view_from_host batched( - res, host_view, /*batch_size=*/128, /*host_writeback=*/false, /*initialize=*/true); - - auto view = batched.next_view(); - EXPECT_EQ(view.extent(0), 0); - EXPECT_EQ(view.extent(1), 8); - EXPECT_EQ(view.data_handle(), nullptr); -} - -using BatchDimsParam = std::tuple; - -class BatchedDeviceViewFromHostParameterizedTest - : public BatchedDeviceViewFromHostTest, - public ::testing::WithParamInterface {}; - -TEST_P(BatchedDeviceViewFromHostParameterizedTest, VectorHostData) -{ - auto [batch_config, dims_config] = GetParam(); - auto [initialize, host_writeback] = batch_config; - auto [n_rows, n_cols, batch_size] = dims_config; - - std::vector host_data(n_rows * n_cols); - auto host_view = raft::make_host_matrix_view(host_data.data(), n_rows, n_cols); - - std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); - - run_and_verify_batched(host_view, batch_size, host_writeback, initialize); -} - -TEST_P(BatchedDeviceViewFromHostParameterizedTest, PinnedMemory) -{ - auto [batch_config, dims_config] = GetParam(); - auto [initialize, host_writeback] = batch_config; - auto [n_rows, n_cols, batch_size] = dims_config; - - auto host_matrix = raft::make_pinned_matrix(res, n_rows, n_cols); - auto host_view = host_matrix.view(); - - std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); - - run_and_verify_batched(host_view, batch_size, host_writeback, initialize); -} - -TEST_P(BatchedDeviceViewFromHostParameterizedTest, ManagedMemory) -{ - auto [batch_config, dims_config] = GetParam(); - auto [initialize, host_writeback] = batch_config; - auto [n_rows, n_cols, batch_size] = dims_config; - - auto host_matrix = raft::make_managed_matrix(res, n_rows, n_cols); - auto host_view = host_matrix.view(); - - std::fill(host_view.data_handle(), host_view.data_handle() + n_rows * n_cols, IdxT(13)); - - run_and_verify_batched(host_view, batch_size, host_writeback, initialize); -} - -TEST_P(BatchedDeviceViewFromHostParameterizedTest, DeviceMemory) -{ - auto [batch_config, dims_config] = GetParam(); - auto [initialize, host_writeback] = batch_config; - auto [n_rows, n_cols, batch_size] = dims_config; - - auto host_matrix = raft::make_device_matrix(res, n_rows, n_cols); - auto host_view = host_matrix.view(); - - raft::matrix::fill(res, host_view, IdxT(13)); - - run_and_verify_batched(host_view, batch_size, host_writeback, initialize); -} - -static const std::array kBatchConfigs = {{ - {/*initialize=*/true, /*host_writeback=*/false}, - {/*initialize=*/false, /*host_writeback=*/true}, - {/*initialize=*/true, /*host_writeback=*/true}, -}}; - -static const std::array kDimsConfigs = {{ - {/*n_rows=*/64, /*n_cols=*/32, /*batch_size=*/256}, // rows less than batch size, single batch - {/*n_rows=*/64, /*n_cols=*/32, /*batch_size=*/64}, // single batch - {/*n_rows=*/256, /*n_cols=*/32, /*batch_size=*/32}, // multiple batches - {/*n_rows=*/500, - /*n_cols=*/32, - /*batch_size=*/128}, // multiple batches, partial batch in the end -}}; - -INSTANTIATE_TEST_SUITE_P(BatchConfigs, - BatchedDeviceViewFromHostParameterizedTest, - ::testing::Combine(::testing::ValuesIn(kBatchConfigs), - ::testing::ValuesIn(kDimsConfigs))); - -} // namespace cuvs::neighbors::cagra From befe09c2fb7c934fa89e0ea875edaf2db47beea4 Mon Sep 17 00:00:00 2001 From: aamijar Date: Fri, 14 Aug 2026 01:38:31 +0000 Subject: [PATCH 26/39] PQ dataset API --- .../src/cuvs/cuvs_ann_bench_param_parser.h | 16 ++-- cpp/include/cuvs/neighbors/cagra.hpp | 53 ++++++++---- cpp/src/neighbors/cagra.cuh | 36 ++++++++- cpp/src/neighbors/cagra_build_inst.cu.in | 8 ++ cpp/src/neighbors/cagra_build_matrix.json | 12 ++- .../neighbors/detail/cagra/cagra_build.cuh | 80 +++++-------------- 6 files changed, 111 insertions(+), 94 deletions(-) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 17589d6675..220ffd86ab 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -360,12 +360,11 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } // Parse build-algo-specific parameters and use them to decide on the algo type - nlohmann::json ivf_pq_build_conf = collect_conf_with_prefix(conf, "ivf_pq_build_"); - nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); - nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); - nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); - nlohmann::json build_compression_conf = collect_conf_with_prefix(conf, "build_compression_"); - nlohmann::json build_search_conf = collect_conf_with_prefix(conf, "build_search_"); + nlohmann::json ivf_pq_build_conf = collect_conf_with_prefix(conf, "ivf_pq_build_"); + nlohmann::json ivf_pq_search_conf = collect_conf_with_prefix(conf, "ivf_pq_search_"); + nlohmann::json nn_descent_conf = collect_conf_with_prefix(conf, "nn_descent_"); + nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_"); + nlohmann::json build_search_conf = collect_conf_with_prefix(conf, "build_search_"); // When graph_build_algo is not specified, leave graph_build_params as monostate so the // CAGRA build uses AUTO selection (NN_DESCENT or IVF_PQ based on dataset/heuristics). @@ -399,11 +398,6 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index } else if constexpr (std::is_same_v< U, cuvs::neighbors::graph_build_params::iterative_search_params>) { - if (!build_compression_conf.empty()) { - auto vpq_pams = arg.build_compression.value_or(cuvs::neighbors::vpq_params{}); - parse_build_param(build_compression_conf, vpq_pams); - arg.build_compression.emplace(vpq_pams); - } if (build_search_conf.contains("width")) { arg.search_width = build_search_conf.at("width"); } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index fe391e0a44..e6e9e340a5 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include @@ -161,19 +162,8 @@ namespace graph_build_params { * The defaults are tuned for the build loop (e.g. search_width=1, * max_iterations=8) and may differ from the regular search defaults. * - * `build_compression` controls the VPQ parameters applied to the dataset - * *while building the graph*. It does not change the dataset view attached to the final index. */ struct iterative_search_params : cuvs::neighbors::cagra::search_params { - /** - * Optional VPQ compression parameters used during iterative graph construction. - * - * When set, the dataset is compressed with these parameters for the - * search-and-optimize loop. When std::nullopt (default), the builder - * uses the dense input dataset. - */ - std::optional build_compression = std::nullopt; - iterative_search_params() { this->search_width = 1; @@ -944,6 +934,10 @@ using device_standard_index = template using host_standard_index = index>; +/** CAGRA index with float queries and a device-resident VPQ dataset view. */ +template +using vpq_index = index>; + /** CAGRA index with a device-resident VPQ dataset (f16 codebook vectors). */ template using vpq_f16_index = index>; @@ -954,9 +948,12 @@ using vpq_f32_index = index -using cagra_index_t = index, - uint32_t, - cuvs::neighbors::dataset_view_type_t>; +using cagra_index_t = + std::conditional_t, + vpq_index, + index, + uint32_t, + cuvs::neighbors::dataset_view_type_t>>; /** * @} @@ -968,10 +965,11 @@ using cagra_index_t = index` + */ +auto build(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::neighbors::device_vpq_dataset_view const& dataset) + -> cuvs::neighbors::cagra::vpq_index; + /** * @brief Build from a device padded dataset view (`float`). * @param[in] res raft resources diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 80e2f2a07e..fc1e29c033 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -296,13 +296,43 @@ template auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::cagra_index_t { - using T = cuvs::neighbors::cagra_view_element_type_t; - using IdxT = uint32_t; + using index_type = cuvs::neighbors::cagra::cagra_index_t; + using T = typename index_type::value_type; + using IdxT = uint32_t; // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until attach_dataset(...) supplies a device-padded dataset. if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { - RAFT_FAIL("cagra::build: VPQ-compressed dataset cannot be used for dense graph construction."); + auto effective_params = params; + if (std::holds_alternative(effective_params.graph_build_params)) { + effective_params.graph_build_params = graph_build_params::iterative_search_params{}; + } + + RAFT_EXPECTS(std::holds_alternative( + effective_params.graph_build_params), + "cagra::build: a VPQ dataset requires iterative_search_params graph construction"); + RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, + "cagra::build: a VPQ dataset supports only L2Expanded distance"); + RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); + RAFT_EXPECTS(dataset.dset().pq_bits() == 8, + "cagra::build: VPQ dataset requires pq_bits == 8, got %u", + dataset.dset().pq_bits()); + auto const pq_len = dataset.dset().pq_len(); + RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, + "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", + pq_len); + + detail::check_graph_degree(effective_params.intermediate_graph_degree, + effective_params.graph_degree, + static_cast(dataset.n_rows())); + auto cagra_graph = detail::iterative_build_graph(res, effective_params, dataset); + + index_type idx(res, effective_params.metric); + idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); + if (effective_params.attach_dataset_on_build) { + idx.update_device_dataset_same_layout(res, dataset); + } + return idx; } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( res, params, dataset); diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index acaaa942c1..178486d107 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -18,6 +18,9 @@ using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view< using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; +#if @emit_vpq_build@ +using inst_device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +#endif } // namespace namespace cuvs::neighbors::cagra { @@ -55,6 +58,11 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_padded_view_t, CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_standard_view_t, cuvs::neighbors::cagra::host_standard_index); +#if @emit_vpq_build@ +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_device_vpq_view_t, + cuvs::neighbors::cagra::vpq_index); +#endif + #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_build_matrix.json b/cpp/src/neighbors/cagra_build_matrix.json index a7995005c4..9fae2b33f9 100644 --- a/cpp/src/neighbors/cagra_build_matrix.json +++ b/cpp/src/neighbors/cagra_build_matrix.json @@ -2,19 +2,23 @@ "_data": [ { "data_type": "float", - "data_abbrev": "f" + "data_abbrev": "f", + "emit_vpq_build": 1 }, { "data_type": "half", - "data_abbrev": "h" + "data_abbrev": "h", + "emit_vpq_build": 0 }, { "data_type": "int8_t", - "data_abbrev": "i8" + "data_abbrev": "i8", + "emit_vpq_build": 0 }, { "data_type": "uint8_t", - "data_abbrev": "u8" + "data_abbrev": "u8", + "emit_vpq_build": 0 } ], "_index": [ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 2403dc0d75..1674ac1768 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -7,7 +7,6 @@ #include "../../../core/nvtx.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "graph_core.cuh" -#include #include #include @@ -2158,7 +2157,8 @@ auto ensure_device_padded_for_iterative_search( } template - requires cuvs::neighbors::is_dense_row_major_dataset_view_v + requires(cuvs::neighbors::is_dense_row_major_dataset_view_v || + cuvs::neighbors::is_device_vpq_f16_dataset_view_v) auto iterative_build_graph(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> raft::host_matrix @@ -2168,45 +2168,34 @@ auto iterative_build_graph(raft::resources const& res, const auto& iter_params = std::get(params.graph_build_params); - const auto& build_compression = iter_params.build_compression; - - if (build_compression.has_value()) { - const auto& bc = *build_compression; - RAFT_LOG_INFO( - "Build compression params: pq_bits=%u, pq_dim=%u, vq_n_centers=%u, kmeans_n_iters=%u, " - "vq_kmeans_trainset_fraction=%.4f, pq_kmeans_trainset_fraction=%.4f, " - "max_train_points_per_pq_code=%u, max_train_points_per_vq_cluster=%u", - bc.pq_bits, - bc.pq_dim, - bc.vq_n_centers, - bc.kmeans_n_iters, - bc.vq_kmeans_trainset_fraction, - bc.pq_kmeans_trainset_fraction, - bc.max_train_points_per_pq_code, - bc.max_train_points_per_vq_cluster); - } else { - RAFT_LOG_INFO("Build compression: disabled (uncompressed build)"); - } RAFT_LOG_INFO("Build search params: search_width=%zu, max_iterations=%zu", iter_params.search_width, iter_params.max_iterations); auto cagra_graph = raft::make_host_matrix(0, 0); - // Iteratively improve the accuracy of the graph by repeatedly running - // CAGRA's search() and optimize(). Host or non-CAGRA-aligned device inputs are uploaded - // and padded here only for the internal search loop — same role as main's - // make_aligned_dataset() inside iterative_build_graph. IVF-PQ / NN-descent never take this path. + // Iteratively improve the graph by repeatedly running CAGRA search and optimize. Dense inputs are + // padded on device; VPQ inputs are searched directly and reconstructed per query batch. RAFT_LOG_INFO("Iteratively creating/improving graph index using CAGRA's search() and optimize()"); std::unique_ptr> padded_own; - auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); - - auto dev_dataset = search_dataset.view(); - uint32_t logical_dim = search_dataset.dim(); + auto dev_dataset = + raft::make_device_matrix_view(static_cast(nullptr), 0, 0); + uint32_t logical_dim = dataset.dim(); + uint64_t final_graph_size; + const cuvs::neighbors::device_vpq_dataset* vpq_dataset = nullptr; + + if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + final_graph_size = static_cast(dataset.n_rows()); + vpq_dataset = &dataset.dset(); + } else { + auto search_dataset = ensure_device_padded_for_iterative_search(res, dataset, padded_own); + dev_dataset = search_dataset.view(); + logical_dim = search_dataset.dim(); + final_graph_size = static_cast(search_dataset.n_rows()); + } // Determine initial graph size. - uint64_t final_graph_size = (uint64_t)search_dataset.n_rows(); uint64_t initial_graph_size = (final_graph_size + 1) / 2; while (initial_graph_size > graph_degree * 64) { initial_graph_size = (initial_graph_size + 1) / 2; @@ -2264,33 +2253,6 @@ auto iterative_build_graph(raft::resources const& res, auto dev_graph = raft::make_device_matrix(res, 0, 0); bool use_device_graph = false; - // Generate the compressed dataset once if compression is enabled. The owner remains alive for - // the complete iterative loop while each temporary index stores only its non-owning view. - const uint64_t dataset_row_width = dev_dataset.extent(1); - std::optional> vpq_dataset; - - if (build_compression.has_value()) { - auto start = std::chrono::high_resolution_clock::now(); - RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded, - "VPQ compression is only supported with L2Expanded distance metric"); - - vpq_dataset.emplace( - cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, *build_compression, search_dataset)); - auto end = std::chrono::high_resolution_clock::now(); - [[maybe_unused]] auto elapsed_ms = - std::chrono::duration_cast(end - start).count(); - RAFT_LOG_INFO("# VPQ compression time: %.3lf sec", (double)elapsed_ms / 1000); - - // Release only storage created internally. Caller-owned device storage remains untouched. - if (padded_own) { - padded_own.reset(); - RAFT_LOG_INFO( - "# Freed internally padded dataset from device (%.1f MiB); queries will use VPQ " - "reconstruction", - to_mib(final_graph_size * dataset_row_width * sizeof(T))); - } - } - while (true) { auto start = std::chrono::high_resolution_clock::now(); auto curr_query_size = std::min(2 * curr_graph_size, final_graph_size); @@ -2333,7 +2295,7 @@ auto iterative_build_graph(raft::resources const& res, // Each index holds non-owning dataset and graph views. The local dataset owner and the graph // passed to search_and_optimize keep those views alive for the duration of the search. - if (vpq_dataset.has_value()) { + if (vpq_dataset != nullptr) { auto idx = cuvs::neighbors::cagra::vpq_f16_index(res, params.metric); idx.update_device_dataset_same_layout(res, vpq_dataset->as_dataset_view()); if (use_device_graph) { @@ -2351,7 +2313,7 @@ auto iterative_build_graph(raft::resources const& res, dev_neighbors.view(), dev_distances.view(), std::move(dev_graph), - &*vpq_dataset, + vpq_dataset, curr_query_size, next_graph_degree, curr_topk, From 492d765556de9a03ccc9a67fb85ad78ed8482d0e Mon Sep 17 00:00:00 2001 From: aamijar Date: Fri, 14 Aug 2026 01:53:17 +0000 Subject: [PATCH 27/39] remove alias --- cpp/include/cuvs/neighbors/cagra.hpp | 20 ++++++++------------ cpp/src/neighbors/cagra_build_inst.cu.in | 6 ++++-- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index e6e9e340a5..dc18a6b792 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -934,10 +934,6 @@ using device_standard_index = template using host_standard_index = index>; -/** CAGRA index with float queries and a device-resident VPQ dataset view. */ -template -using vpq_index = index>; - /** CAGRA index with a device-resident VPQ dataset (f16 codebook vectors). */ template using vpq_f16_index = index>; @@ -948,12 +944,12 @@ using vpq_f32_index = index -using cagra_index_t = - std::conditional_t, - vpq_index, - index, - uint32_t, - cuvs::neighbors::dataset_view_type_t>>; +using cagra_index_t = std::conditional_t< + cuvs::neighbors::is_device_vpq_f16_dataset_view_v, + index>, + index, + uint32_t, + cuvs::neighbors::dataset_view_type_t>>; /** * @} @@ -1001,12 +997,12 @@ using cagra_index_t = * @param[in] res raft resources * @param[in] params CAGRA index build parameters * @param[in] dataset device VPQ dataset view - * @return built `vpq_index` + * @return built `index>` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, cuvs::neighbors::device_vpq_dataset_view const& dataset) - -> cuvs::neighbors::cagra::vpq_index; + -> index>; /** * @brief Build from a device padded dataset view (`float`). diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index 178486d107..90d63c3ca9 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -59,8 +59,10 @@ CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_host_standard_view_t, cuvs::neighbors::cagra::host_standard_index); #if @emit_vpq_build@ -CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(inst_device_vpq_view_t, - cuvs::neighbors::cagra::vpq_index); +CUVS_DEFINE_CAGRA_BUILD_OVERLOAD( + inst_device_vpq_view_t, + cuvs::neighbors::cagra:: + index>); #endif #undef CUVS_DEFINE_CAGRA_BUILD_OVERLOAD From dc2ccfeaf5e8130ac9e3dc7e853027ae6150371c Mon Sep 17 00:00:00 2001 From: aamijar Date: Sun, 16 Aug 2026 01:17:57 +0000 Subject: [PATCH 28/39] remove unused operators --- cpp/include/cuvs/neighbors/common.hpp | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index eb938cfd26..935938c9b0 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -109,18 +109,6 @@ struct vpq_params { * The max number of data points to use per VQ cluster during training. */ uint32_t max_train_points_per_vq_cluster = 1024; - - friend bool operator==(const vpq_params& a, const vpq_params& b) - { - return a.pq_bits == b.pq_bits && a.pq_dim == b.pq_dim && a.vq_n_centers == b.vq_n_centers && - a.kmeans_n_iters == b.kmeans_n_iters && - a.vq_kmeans_trainset_fraction == b.vq_kmeans_trainset_fraction && - a.pq_kmeans_trainset_fraction == b.pq_kmeans_trainset_fraction && - a.pq_kmeans_type == b.pq_kmeans_type && - a.max_train_points_per_pq_code == b.max_train_points_per_pq_code && - a.max_train_points_per_vq_cluster == b.max_train_points_per_vq_cluster; - } - friend bool operator!=(const vpq_params& a, const vpq_params& b) { return !(a == b); } }; /** @} */ // end group cagra_cpp_index_params From 79bd7bb67787362e318c6878f4469263849f900c Mon Sep 17 00:00:00 2001 From: aamijar Date: Sun, 16 Aug 2026 01:26:49 +0000 Subject: [PATCH 29/39] remove some small diff --- cpp/src/neighbors/detail/cagra/cagra_search.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 34a1762e85..165e478337 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -85,6 +85,7 @@ void search_main_core( search_plan_impl> plan = factory::create( res, params, dataset_desc, queries.extent(1), graph.extent(0), graph.extent(1), topk); + plan->check(topk); RAFT_LOG_DEBUG("Cagra search"); @@ -207,7 +208,6 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } // Search using a plain (strided) row-major dataset - RAFT_LOG_DEBUG("Searching with strided dataset"); RAFT_EXPECTS(index.metric() != cuvs::distance::DistanceType::CosineExpanded || index.dataset_norms().has_value(), "Dataset norms must be provided for CosineExpanded metric"); @@ -238,7 +238,6 @@ void search_main(raft::resources const& res, RAFT_FAIL("FP32 VPQ dataset support is coming soon"); } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { auto const& vv = index.dataset(); - RAFT_LOG_DEBUG("Searching with VPQ dataset"); if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && raft::getComputeCapability().first < 9) { RAFT_LOG_WARN( From bb1fcde35add677f6ae78a44aedf62bdf78b6b2a Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 17 Aug 2026 07:08:51 -0700 Subject: [PATCH 30/39] VPQ dataset serialization so CAGRA-Q can build from a compressed dataset on disk. Also added tests --- c/tests/CMakeLists.txt | 4 +- .../cuvs/preprocessing/quantize/pq.hpp | 79 +++++ .../neighbors/detail/dataset_serialize.hpp | 69 +++++ cpp/src/preprocessing/quantize/pq.cu | 56 ++++ cpp/tests/CMakeLists.txt | 3 +- .../ann_cagra/bug_iterative_cagra_build.cu | 43 ++- .../ann_cagra/test_iterative_cagra_q.cu | 280 ++++++++++++++++++ cpp/tests/preprocessing/vpq_serialization.cu | 260 ++++++++++++++++ 8 files changed, 788 insertions(+), 6 deletions(-) create mode 100644 cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu create mode 100644 cpp/tests/preprocessing/vpq_serialization.cu diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index 7d6c588bd9..ff8f807a6a 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -89,7 +89,9 @@ ConfigureTest(NAME IVF_FLAT_C_TEST PATH neighbors/run_ivf_flat_c.c neighbors/ann ConfigureTest(NAME IVF_PQ_C_TEST PATH neighbors/run_ivf_pq_c.c neighbors/ann_ivf_pq_c.cu) ConfigureTest(NAME IVF_SQ_C_TEST PATH neighbors/run_ivf_sq_c.c neighbors/ann_ivf_sq_c.cu) ConfigureTest(NAME CAGRA_C_TEST PATH neighbors/ann_cagra_c.cu) -ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu) +if(BUILD_MG_ALGOS) + ConfigureTest(NAME MG_C_TEST PATH neighbors/run_mg_c.c neighbors/ann_mg_c.cu) +endif() ConfigureTest( NAME ALL_NEIGHBORS_C_TEST PATH neighbors/run_all_neighbors_c.c neighbors/all_neighbors_c.cu ) diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index 112341f2ad..f6456624d6 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -14,6 +14,9 @@ #include #include +#include +#include +#include #include #include @@ -331,6 +334,82 @@ template } } +/** Current VPQ dataset serialization format version. */ +inline constexpr int vpq_serialization_version = 1; + +/** + * @brief Write a VPQ dataset (both codebooks plus the encoded rows) to a stream. + * + * Lets compression be done once, offline, and reused: the encoded rows are what CAGRA-Q builds and + * searches over, so a stored VPQ dataset removes the need to keep the dense vectors around or + * re-quantize them on every run. + * + * The file opens with the same preamble as `cagra::serialize` — a 4-byte NumPy dtype prefix then + * `vpq_serialization_version` — followed by a dataset kind tag and the codebook element type. A file + * of the wrong kind, or one written by an older format, is rejected rather than misread. Bump the + * version whenever the encoded row layout changes, since that layout is a library convention and is + * not otherwise described by the file. + * + * @code{.cpp} + * #include + * #include + * + * // Offline, once. + * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, rows); + * cuvs::preprocessing::quantize::pq::serialize(res, "base.vpq", vpq); + * + * // Later, per run: load the compressed rows and build a CAGRA-Q graph over them. + * std::unique_ptr> loaded; + * cuvs::preprocessing::quantize::pq::deserialize(res, "base.vpq", &loaded); + * auto index = cuvs::neighbors::cagra::build(res, index_params, loaded->as_dataset_view()); + * // `loaded` must outlive `index`, which only holds a view of it. + * @endcode + * + * @param[in] res raft resource + * @param[in] os output stream, opened in binary mode + * @param[in] dataset the VPQ dataset to write + */ +void serialize(raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::device_vpq_dataset& dataset); + +/** + * @copydoc serialize + * + * @param[in] res raft resource + * @param[in] filename path to write, truncated if it exists + * @param[in] dataset the VPQ dataset to write + */ +void serialize(raft::resources const& res, + const std::string& filename, + const cuvs::neighbors::device_vpq_dataset& dataset); + +/** + * @brief Read a VPQ dataset written by `serialize`. + * + * Returned through an out-parameter because the dataset owns device allocations and has no default + * constructor, matching how `cagra::deserialize` hands back its dataset. Throws if the blob was not + * written by `serialize` or holds codebooks of a different element type. + * + * @param[in] res raft resource + * @param[in] is input stream, opened in binary mode + * @param[out] out_dataset receives the loaded dataset; must not be null + */ +void deserialize(raft::resources const& res, + std::istream& is, + std::unique_ptr>* out_dataset); + +/** + * @copydoc deserialize + * + * @param[in] res raft resource + * @param[in] filename path to read + * @param[out] out_dataset receives the loaded dataset; must not be null + */ +void deserialize(raft::resources const& res, + const std::string& filename, + std::unique_ptr>* out_dataset); + /** @} */ // end of group product } // namespace pq diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 05c71e0213..6e73f36d10 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -279,6 +279,40 @@ auto deserialize_host_dense(raft::resources const& res, std::istream& is) return std::make_unique(std::move(storage), metadata.dim); } +/** VPQ codebooks are floating point; the encoded rows are always uint8 and carry no dtype. */ +template +constexpr auto vpq_wire_dtype() -> cudaDataType_t +{ + static_assert(std::is_same_v || std::is_same_v, + "serialize_vpq: codebook element type must be float or half"); + return std::is_same_v ? CUDA_R_16F : CUDA_R_32F; +} + +/** + * Write the payload of a VPQ dataset: six scalars followed by the two codebooks and the encoded + * rows. + * + * Stays on `raft::serialize_mdspan` rather than the `write_dense_bytes` scheme used by the dense + * path above, because `deserialize_vpq` reads with `raft::deserialize_mdspan`, which expects the + * NumPy header that helper embeds per matrix. The scalar types must also match the reader exactly: + * `n_rows` is `IdxT` and the remaining five are `uint32_t`. + */ +template +void serialize_vpq(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + raft::serialize_scalar(res, os, dataset.n_rows()); + raft::serialize_scalar(res, os, dataset.dim()); + raft::serialize_scalar(res, os, dataset.vq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_n_centers()); + raft::serialize_scalar(res, os, dataset.pq_len()); + raft::serialize_scalar(res, os, dataset.encoded_row_length()); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.vq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.pq_code_book.view())); + raft::serialize_mdspan(res, os, raft::make_const_mdspan(dataset.data.view())); +} + template auto deserialize_vpq(raft::resources const& res, std::istream& is) -> std::unique_ptr> @@ -305,6 +339,41 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) std::move(vq_code_book), std::move(pq_code_book), std::move(data)); } +/** + * Write a self-describing VPQ dataset blob: tag + codebook dtype + payload. + * + * The tag and dtype are deliberately written here rather than inside `serialize_vpq`, mirroring how + * `serialize_cagra_dense_dataset` wraps the dense payload, so that a reader can identify the blob + * before committing to a `DataT`. + */ +template +void serialize_vpq_dataset(raft::resources const& res, + std::ostream& os, + device_vpq_dataset const& dataset) +{ + raft::serialize_scalar(res, os, kSerializeVPQDataset); + raft::serialize_scalar(res, os, vpq_wire_dtype()); + serialize_vpq(res, os, dataset); +} + +/** Read a blob written by `serialize_vpq_dataset`, validating the tag and codebook dtype. */ +template +auto deserialize_vpq_dataset(raft::resources const& res, std::istream& is) + -> std::unique_ptr> +{ + const auto tag = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(tag == kSerializeVPQDataset, + "deserialize_vpq_dataset: expected VPQ tag (%u), got %u", + static_cast(kSerializeVPQDataset), + static_cast(tag)); + const auto dtype = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(dtype == vpq_wire_dtype(), + "deserialize_vpq_dataset: codebook dtype (%d) does not match expected (%d)", + static_cast(dtype), + static_cast(vpq_wire_dtype())); + return deserialize_vpq(res, is); +} + template auto deserialize_dense_dataset(raft::resources const& res, std::istream& is) -> std::unique_ptr diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..673e49759b 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -3,13 +3,20 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../neighbors/detail/dataset_serialize.hpp" +#include "../../util/serialize_validation.hpp" #include "./detail/pq.cuh" #include +#include #include #include +#include +#include +#include + namespace cuvs::preprocessing::quantize::pq { #define CUVS_INST_QUANTIZATION(T, QuantI) \ @@ -76,6 +83,55 @@ CUVS_INST_VPQ_BUILD(uint8_t); #undef CUVS_INST_VPQ_BUILD +void serialize(raft::resources const& res, + std::ostream& os, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + // Same file preamble as cagra::serialize. The nested blob carries only a kind tag and dtype, + // matching serialize_cagra_dense_dataset, because a nested blob relies on its enclosing file for + // the version; a standalone .vpq has no enclosing file, so the version is written here. + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, vpq_serialization_version); + ::cuvs::neighbors::detail::serialize_vpq_dataset(res, os, dataset); +} + +void serialize(raft::resources const& res, + const std::string& filename, + const cuvs::neighbors::device_vpq_dataset& dataset) +{ + std::ofstream os(filename, std::ios::out | std::ios::binary | std::ios::trunc); + RAFT_EXPECTS(os.good(), "pq::serialize: cannot open %s for writing", filename.c_str()); + serialize(res, os, dataset); +} + +void deserialize(raft::resources const& res, + std::istream& is, + std::unique_ptr>* out_dataset) +{ + RAFT_EXPECTS(out_dataset != nullptr, "pq::deserialize: out_dataset must not be null"); + char dtype_string[4]; + RAFT_EXPECTS(is.read(dtype_string, 4), "pq::deserialize: failed to read the dtype prefix"); + RAFT_EXPECTS(cuvs::util::validate_serialized_dtype(dtype_string, sizeof(dtype_string)), + "pq::deserialize: dtype prefix does not match a VPQ dataset with half codebooks"); + auto const version = raft::deserialize_scalar(res, is); + RAFT_EXPECTS(version == vpq_serialization_version, + "pq::deserialize: serialization version mismatch, expected %d, got %d", + vpq_serialization_version, + version); + *out_dataset = ::cuvs::neighbors::detail::deserialize_vpq_dataset(res, is); +} + +void deserialize(raft::resources const& res, + const std::string& filename, + std::unique_ptr>* out_dataset) +{ + std::ifstream is(filename, std::ios::in | std::ios::binary); + RAFT_EXPECTS(is.good(), "pq::deserialize: cannot open %s for reading", filename.c_str()); + deserialize(res, is, out_dataset); +} + namespace detail { template diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 8421c3219e..eb7bd8024e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -199,7 +199,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST - PATH neighbors/ann_cagra/test_float_uint32_t.cu + PATH neighbors/ann_cagra/test_float_uint32_t.cu neighbors/ann_cagra/test_iterative_cagra_q.cu GPUS 1 PERCENT 100 ) @@ -415,6 +415,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu + preprocessing/vpq_serialization.cu preprocessing/pca.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu index 3d6b5c98a2..de9575058c 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu @@ -7,6 +7,7 @@ #include "../cagra_padded_build_helpers.cuh" #include +#include #include #include @@ -23,17 +24,22 @@ class CagraIterativeBuildBugTest : public ::testing::Test { using data_type = DataT; protected: - void run() + // The bug manifests when graph_degree is equal to intermediate_graph_degree + // see issue https://github.com/rapidsai/cuvs/issues/1818 + static auto bug_index_params() -> cagra::index_params { - // Set up iterative CAGRA graph building cagra::index_params index_params; - // The bug manifests when graph_degree is equal to intermediate_graph_degree - // see issue https://github.com/rapidsai/cuvs/issues/1818 index_params.graph_degree = 16; index_params.intermediate_graph_degree = 16; // Use iterative CAGRA search for graph building index_params.graph_build_params = graph_build_params::iterative_search_params(); + return index_params; + } + + void run() + { + auto index_params = bug_index_params(); cuvs::neighbors::test::padded_device_matrix_for_cagra padded( res, raft::make_const_mdspan(dataset->view())); @@ -46,6 +52,30 @@ class CagraIterativeBuildBugTest : public ::testing::Test { ASSERT_EQ(cagra_index.dim(), n_dim); } + // Same bug, reached through iterative CAGRA-Q: the graph is built from a PQ-compressed dataset, + // so the searches driving the build run on compressed rows instead of dense ones. + void run_compressed() + { + cuvs::neighbors::vpq_params vpq_params; + // pq_len = n_dim / pq_dim must be 2, 4 or 8 for CAGRA-Q. Codebook quality is irrelevant here, + // since only graph construction is under test, so training stays short. + vpq_params.pq_dim = static_cast(n_dim / 4); + vpq_params.vq_n_centers = 64; + vpq_params.kmeans_n_iters = 5; + + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, vpq_params, raft::make_const_mdspan(dataset->view())); + + // No padding and no attach step: a compressed dataset is read through its own view, which the + // index holds on to, so `compressed` has to outlive the index. + auto cagra_index = cagra::build(res, bug_index_params(), compressed.as_dataset_view()); + raft::resource::sync_stream(res); + + ASSERT_GT(cagra_index.size(), 0); + ASSERT_EQ(cagra_index.dim(), n_dim); + ASSERT_EQ(cagra_index.graph_degree(), 16u); + } + void SetUp() override { dataset.emplace(raft::make_device_matrix(res, n_samples, n_dim)); @@ -85,4 +115,9 @@ TYPED_TEST_SUITE(CagraIterativeBuildBugTest, TestTypes); TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildTest) { this->run(); } +TYPED_TEST(CagraIterativeBuildBugTest, IterativeBuildFromCompressedDatasetTest) +{ + this->run_compressed(); +} + } // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu new file mode 100644 index 0000000000..bad504b589 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -0,0 +1,280 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * Iterative CAGRA-Q: building a CAGRA graph directly from a PQ-compressed dataset. + * + * This path takes a `device_vpq_dataset_view` instead of dense rows, so the caller owns + * compression and the inner searches of the iterative build run against the compressed data. + * It only accepts `L2Expanded`, `pq_bits == 8` and `pq_len` in {2, 4, 8}, so it does not fit + * the dtype-templated suites in ann_cagra.cuh and lives in its own file. + * + * What is checked here is that a compressed dataset, freshly compressed or loaded from disk, + * builds a usable graph, and that the constraints above are rejected rather than accepted and + * quietly ignored. Serialization fidelity itself is covered by preprocessing/vpq_serialization.cu. + */ + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +namespace { + +auto compress(const raft::resources& res, + raft::device_matrix_view dataset, + uint32_t pq_dim, + uint32_t pq_bits = 8) -> vpq_dataset_t +{ + cuvs::neighbors::vpq_params params; + params.pq_dim = pq_dim; + params.pq_bits = pq_bits; + params.vq_n_centers = 32; + params.kmeans_n_iters = 5; // Codebooks need to be well defined here, not optimal. + return cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, params, dataset); +} + +auto iterative_params(uint32_t graph_degree = 32) -> index_params +{ + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = graph_degree; + params.intermediate_graph_degree = graph_degree * 2; + params.graph_build_params = graph_build_params::iterative_search_params(); + return params; +} + +/** + * Fraction of queries that retrieve their own row, where the queries are dataset rows. + * + * A sanity signal rather than a quality metric: search quality is a benchmark's job, and an + * iterative build varies by a few points run to run over identical input, so callers assert a + * loose floor that only a broken dataset would miss. + */ +template +auto self_recall_at_1(const raft::resources& res, + const IndexT& idx, + raft::device_matrix_view queries) -> double +{ + constexpr int64_t k = 10; + const auto n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + search_params params; + params.itopk_size = 64; + search(res, params, idx, queries, neighbors.view(), distances.view()); + + std::vector ids(static_cast(n_queries * k)); + raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + int64_t hits = 0; + for (int64_t q = 0; q < n_queries; q++) { + if (ids[q * k] == static_cast(q)) { hits++; } + } + return static_cast(hits) / static_cast(n_queries); +} + +} // namespace + +struct CagraQInputs { + int64_t n_rows; + int64_t dim; + uint32_t pq_dim; // pq_len = dim / pq_dim, which must land in {2, 4, 8} +}; + +std::ostream& operator<<(std::ostream& os, const CagraQInputs& in) +{ + return os << "n_rows:" << in.n_rows << " dim:" << in.dim << " pq_dim:" << in.pq_dim + << " pq_len:" << (in.dim / in.pq_dim); +} + +/** Shared clustered dataset; each test compresses it itself so nothing leaks between cases. */ +class CagraQCompressedTestBase : public ::testing::Test { + protected: + void make_dataset(int64_t n_rows, int64_t dim) + { + dataset_.emplace(raft::make_device_matrix(res_, n_rows, dim)); + auto labels = raft::make_device_vector(res_, n_rows); + raft::random::make_blobs(res_, + dataset_->view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + 1234ULL); + raft::resource::sync_stream(res_); + } + + auto dataset() -> raft::device_matrix_view + { + return raft::make_const_mdspan(dataset_->view()); + } + + /** The first rows of the dataset, reused as queries. */ + auto queries(int64_t n_queries) -> raft::device_matrix_view + { + return raft::make_device_matrix_view( + dataset_->data_handle(), std::min(n_queries, dataset_->extent(0)), dataset_->extent(1)); + } + + void TearDown() override + { + dataset_.reset(); + raft::resource::sync_stream(res_); + } + + raft::resources res_; + std::optional> dataset_ = std::nullopt; +}; + +class CagraQBuildTest : public CagraQCompressedTestBase, + public ::testing::WithParamInterface { + protected: + void SetUp() override + { + params_ = GetParam(); + make_dataset(params_.n_rows, params_.dim); + } + + CagraQInputs params_{}; +}; + +TEST_P(CagraQBuildTest, BuildsAndSearchesAFreshlyCompressedDataset) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + ASSERT_EQ(compressed.pq_len(), static_cast(params_.dim / params_.pq_dim)); + + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + ASSERT_EQ(idx.size(), params_.n_rows); + ASSERT_EQ(idx.dim(), params_.dim); + ASSERT_EQ(idx.graph_degree(), 32u); + + // Searchable straight after build: the index keeps the compressed view it was built from, so + // unlike a dense standard-layout dataset there is nothing to attach first. + EXPECT_GT(self_recall_at_1(res_, idx, queries(1000)), 0.5); +} + +TEST_P(CagraQBuildTest, BuildsFromADeserializedDataset) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + + // The path the benchmarks take: compress offline, store, then build from the file. + std::stringstream stored; + cuvs::preprocessing::quantize::pq::serialize(res_, stored, compressed); + std::unique_ptr loaded; + cuvs::preprocessing::quantize::pq::deserialize(res_, stored, &loaded); + ASSERT_NE(loaded, nullptr); + + auto idx = cagra::build(res_, iterative_params(), loaded->as_dataset_view()); + ASSERT_EQ(idx.size(), params_.n_rows); + ASSERT_EQ(idx.dim(), params_.dim); + EXPECT_GT(self_recall_at_1(res_, idx, queries(1000)), 0.5); +} + +TEST_P(CagraQBuildTest, PromotesUnsetGraphBuildParamsToIterative) +{ + auto compressed = compress(res_, dataset(), params_.pq_dim); + + // Iterative search is the only construction a compressed dataset supports, so leaving + // graph_build_params at its default must select it, not fall back to IVF-PQ or NN-descent. + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 32; + ASSERT_TRUE(std::holds_alternative(params.graph_build_params)); + + auto idx = cagra::build(res_, params, compressed.as_dataset_view()); + EXPECT_EQ(idx.size(), params_.n_rows); +} + +INSTANTIATE_TEST_CASE_P(CagraQBuildTests, + CagraQBuildTest, + ::testing::ValuesIn(std::vector{ + {2000, 64, 32}, // pq_len 2 + {2000, 128, 32}, // pq_len 4 + {2000, 256, 32}, // pq_len 8 + })); + +/** The constraints the VPQ build overload documents, each of which must be rejected loudly. */ +class CagraQContractTest : public CagraQCompressedTestBase { + protected: + void SetUp() override { make_dataset(n_rows, dim); } + + static constexpr int64_t n_rows = 1000; + static constexpr int64_t dim = 64; +}; + +TEST_F(CagraQContractTest, RejectsNonIterativeGraphBuilder) +{ + auto compressed = compress(res_, dataset(), 32); + auto params = iterative_params(); + params.graph_build_params = + graph_build_params::nn_descent_params(params.intermediate_graph_degree); + EXPECT_THROW(cagra::build(res_, params, compressed.as_dataset_view()), raft::exception); +} + +TEST_F(CagraQContractTest, RejectsMetricOtherThanL2Expanded) +{ + auto compressed = compress(res_, dataset(), 32); + auto params = iterative_params(); + params.metric = cuvs::distance::DistanceType::InnerProduct; + EXPECT_THROW(cagra::build(res_, params, compressed.as_dataset_view()), raft::exception); +} + +TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) +{ + auto compressed = compress(res_, dataset(), 32, /* pq_bits */ 6); + ASSERT_EQ(compressed.pq_bits(), 6u); + EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), + raft::exception); +} + +TEST_F(CagraQContractTest, RejectsPqLenOutsideSupportedSet) +{ + auto compressed = compress(res_, dataset(), /* pq_dim */ 4); // pq_len = 64 / 4 = 16 + ASSERT_EQ(compressed.pq_len(), 16u); + EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), + raft::exception); +} + +TEST_F(CagraQContractTest, RejectsEmptyDataset) +{ + // Hand-built rather than compressed, since make_vpq_dataset rejects an empty input of its own + // accord. Every other constraint is satisfied so that only the emptiness can trip. + const auto width = static_cast(dim); + auto vq_code_book = raft::make_device_matrix(res_, 1, width); + auto pq_code_book = raft::make_device_matrix(res_, 256, 2); + auto codes = raft::make_device_matrix(res_, 0, 4 + dim / 2); + vpq_dataset_t empty{std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + ASSERT_EQ(empty.n_rows(), 0); + + EXPECT_THROW(cagra::build(res_, iterative_params(), empty.as_dataset_view()), raft::exception); +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/preprocessing/vpq_serialization.cu b/cpp/tests/preprocessing/vpq_serialization.cu new file mode 100644 index 0000000000..6e41cb48c7 --- /dev/null +++ b/cpp/tests/preprocessing/vpq_serialization.cu @@ -0,0 +1,260 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../neighbors/vpq_utils.cuh" +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::pq { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +struct VpqSerializationInputs { + int64_t n_rows; + int64_t dim; + uint32_t pq_bits; + uint32_t pq_dim; + uint32_t vq_n_centers; // 0 lets the heuristic choose + uint64_t seed; +}; + +std::ostream& operator<<(std::ostream& os, const VpqSerializationInputs& in) +{ + return os << "n_rows:" << in.n_rows << " dim:" << in.dim << " pq_bits:" << in.pq_bits + << " pq_dim:" << in.pq_dim << " vq_n_centers:" << in.vq_n_centers + << " seed:" << in.seed; +} + +template +auto to_host(const raft::resources& res, raft::device_matrix_view m) -> std::vector +{ + std::vector host(static_cast(m.extent(0)) * static_cast(m.extent(1))); + raft::copy(host.data(), m.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return host; +} + +/** Bitwise, not approximate: serialization is expected not to perturb a single bit. */ +template +void expect_same_bits(const raft::resources& res, + raft::device_matrix_view expected, + raft::device_matrix_view actual, + const char* what) +{ + ASSERT_EQ(expected.extent(0), actual.extent(0)) << what; + ASSERT_EQ(expected.extent(1), actual.extent(1)) << what; + const auto lhs = to_host(res, expected); + const auto rhs = to_host(res, actual); + EXPECT_EQ(0, std::memcmp(lhs.data(), rhs.data(), lhs.size() * sizeof(T))) << what; +} + +class VpqSerializationTest : public ::testing::TestWithParam { + public: + VpqSerializationTest() + : params_(::testing::TestWithParam::GetParam()), + dataset_(raft::make_device_matrix(res_, params_.n_rows, params_.dim)) + { + } + + protected: + void SetUp() override + { + auto labels = raft::make_device_vector(res_, params_.n_rows); + raft::random::make_blobs(res_, + dataset_.view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + params_.seed); + raft::resource::sync_stream(res_); + } + + auto compress() -> vpq_dataset_t + { + cuvs::neighbors::vpq_params vpq; + vpq.pq_bits = params_.pq_bits; + vpq.pq_dim = params_.pq_dim; + vpq.vq_n_centers = params_.vq_n_centers; + // The codebooks only have to be well defined here, not good, so keep training short. + vpq.kmeans_n_iters = 5; + return make_vpq_dataset(res_, vpq, raft::make_const_mdspan(dataset_.view())); + } + + void expect_equivalent(const vpq_dataset_t& expected, const vpq_dataset_t& actual) + { + ASSERT_EQ(expected.n_rows(), actual.n_rows()); + ASSERT_EQ(expected.dim(), actual.dim()); + ASSERT_EQ(expected.vq_n_centers(), actual.vq_n_centers()); + ASSERT_EQ(expected.pq_n_centers(), actual.pq_n_centers()); + ASSERT_EQ(expected.pq_len(), actual.pq_len()); + ASSERT_EQ(expected.encoded_row_length(), actual.encoded_row_length()); + ASSERT_EQ(expected.pq_bits(), actual.pq_bits()); + ASSERT_EQ(expected.pq_dim(), actual.pq_dim()); + + expect_same_bits(res_, + raft::make_const_mdspan(expected.vq_code_book.view()), + raft::make_const_mdspan(actual.vq_code_book.view()), + "vq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.pq_code_book.view()), + raft::make_const_mdspan(actual.pq_code_book.view()), + "pq_code_book"); + expect_same_bits(res_, + raft::make_const_mdspan(expected.data.view()), + raft::make_const_mdspan(actual.data.view()), + "encoded rows"); + } + + /** + * Decodes both datasets and compares the reconstructions, which checks that a kernel can consume + * the deserialized extents and strides rather than only that the numbers match. + */ + void expect_same_decoded(const vpq_dataset_t& expected, const vpq_dataset_t& actual) + { + if (expected.pq_bits() != 8) { return; } // decode_vpq_dataset implements pq_bits == 8 only + auto stream = raft::resource::get_cuda_stream(res_); + auto lhs = raft::make_device_matrix(res_, expected.n_rows(), expected.dim()); + auto rhs = raft::make_device_matrix(res_, actual.n_rows(), actual.dim()); + cuvs::neighbors::decode_vpq_dataset(lhs.view(), expected, stream); + cuvs::neighbors::decode_vpq_dataset(rhs.view(), actual, stream); + raft::resource::sync_stream(res_); + expect_same_bits(res_, + raft::make_const_mdspan(lhs.view()), + raft::make_const_mdspan(rhs.view()), + "decoded rows"); + } + + raft::resources res_; + VpqSerializationInputs params_; + raft::device_matrix dataset_; +}; + +TEST_P(VpqSerializationTest, RoundTrip) +{ + auto original = compress(); + + { + SCOPED_TRACE("through a stream"); + std::stringstream stream; + serialize(res_, stream, original); + std::unique_ptr restored; + deserialize(res_, stream, &restored); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + expect_same_decoded(original, *restored); + } + + { + SCOPED_TRACE("through a file"); + const std::string path = "cuvs_vpq_serialization_test.bin"; + serialize(res_, path, original); + std::unique_ptr restored; + deserialize(res_, path, &restored); + std::remove(path.c_str()); + ASSERT_NE(restored, nullptr); + expect_equivalent(original, *restored); + } +} + +// Named for this suite rather than `inputs`: product_quantization.cu declares a variable template of +// that name in this same namespace, which would collide under a unity build. +const std::vector vpq_serialization_inputs = { + // pq_len = dim / pq_dim of 2, 4 and 8: the three values CAGRA-Q accepts. + {1000, 64, 8, 32, 0, 42ULL}, + {1000, 128, 8, 32, 0, 42ULL}, + {1000, 256, 8, 32, 0, 42ULL}, + // An explicit VQ codebook size rather than the heuristic. + {2000, 128, 8, 64, 64, 42ULL}, + // pq_bits below 8 packs several codes per byte, so encoded_row_length stops being pq_dim. + {500, 96, 6, 24, 0, 42ULL}, + {500, 32, 4, 16, 0, 42ULL}, +}; + +INSTANTIATE_TEST_CASE_P(VpqSerializationTests, + VpqSerializationTest, + ::testing::ValuesIn(vpq_serialization_inputs)); + +/** Writes the preamble that `serialize` emits, so only the field under test differs. */ +static void write_preamble(const raft::resources& res, std::ostream& os, int version) +{ + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + os << dtype_string; + raft::serialize_scalar(res, os, version); +} + +TEST(VpqSerialization, RejectsEmptyStream) +{ + raft::resources res; + std::stringstream stream; + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsForeignDtypePrefix) +{ + raft::resources res; + std::stringstream stream; + std::string dtype_string = raft::numpy_serializer::get_numpy_dtype().to_string(); + dtype_string.resize(4); + stream << dtype_string; + raft::serialize_scalar(res, stream, vpq_serialization_version); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsFutureVersion) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version + 1); + + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsTruncatedPayload) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version); + // A correct preamble followed by nothing: the payload reader must fail rather than return a + // dataset built from whatever the scalars happened to deserialize to. + std::unique_ptr restored; + EXPECT_THROW(deserialize(res, stream, &restored), raft::exception); +} + +TEST(VpqSerialization, RejectsNullOutParameter) +{ + raft::resources res; + std::stringstream stream; + write_preamble(res, stream, vpq_serialization_version); + EXPECT_THROW(deserialize(res, stream, nullptr), raft::exception); +} + +} // namespace cuvs::preprocessing::quantize::pq From 31c6322ef34bcf55b4af3b859b0b537686c9ef67 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 17 Aug 2026 08:13:48 -0700 Subject: [PATCH 31/39] Recognise a compressed .vpq base_file in the ann benchmark and hand it to the algorithm as a path --- cpp/bench/ann/src/common/ann_types.hpp | 31 ++++++ cpp/bench/ann/src/common/benchmark.hpp | 142 ++++++++++++++++--------- cpp/bench/ann/src/common/conf.hpp | 34 +++++- cpp/bench/ann/src/common/dataset.hpp | 49 +++++++++ 4 files changed, 206 insertions(+), 50 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index bd669dff78..52105a2dc1 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -159,6 +159,37 @@ class algo : public algo_base { // and should not release dataset before searching is finished. virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) {}; + /* ### Base sets the benchmark cannot read ### + + Some algorithms build from a base set that has been compressed for them offline, which is + neither dense nor made of `T` values and so cannot be passed as `build`'s `const T*`. Such a + base set is handed over as a file path and the algorithm owns whatever it decodes. + + A path rather than a library type on purpose: this header is shared with the faiss, hnswlib and + diskann wrappers, and must not acquire their unrelated dependencies. + + Loading is separate from building because the benchmark times only `build_from_base_set_file`. + Deserializing a compressed base set is benchmark setup, the same as reading a dense one, and + folding it into the measured build would inflate build times by however long the file takes to + read. `set_base_set_file` is also called in search mode, before `load`, for algorithms whose + index file holds only part of the picture and needs the base set reattached. + */ + + /** + * Hand over a compressed base set as a file path. Returns the number of rows in it, which the + * benchmark has no way of reading for itself. Called outside the timed sections. + */ + virtual auto set_base_set_file(const std::string& /*file*/) -> size_t + { + throw std::runtime_error{"This algorithm cannot read a compressed base set from a file."}; + } + + /** Build the index from the base set handed over by `set_base_set_file`. */ + virtual void build_from_base_set_file() + { + throw std::runtime_error{"This algorithm cannot build from a compressed base set."}; + } + /** * Make a shallow copy of the algo wrapper that shares the resources and ensures thread-safe * access to them. */ diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index a588b1e2a6..9e76192b95 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -134,8 +134,22 @@ void bench_build(::benchmark::State& state, const auto algo_property = parse_algo_property(algo->get_preference(), index.build_param); - const T* base_set = dataset->base_set(algo_property.dataset_memory_type); - std::size_t index_size = dataset->base_set_size(); + // Loading the base set is setup, not part of the build, so a compressed one is read here rather + // than inside the timed loop below; its row count comes back from the algorithm, since the + // benchmark cannot read the file. + const bool base_compressed = dataset->base_is_compressed(); + const T* base_set = nullptr; + std::size_t index_size = 0; + try { + if (base_compressed) { + index_size = algo->set_base_set_file(dataset->base_file()); + } else { + base_set = dataset->base_set(algo_property.dataset_memory_type); + index_size = dataset->base_set_size(); + } + } catch (const std::exception& e) { + return state.SkipWithError("Failed to load the base set: " + std::string(e.what())); + } cuda_timer gpu_timer{algo}; { @@ -158,7 +172,11 @@ void bench_build(::benchmark::State& state, [[maybe_unused]] auto ntx_lap = nvtx.lap(); [[maybe_unused]] auto gpu_lap = gpu_timer.lap(!no_lap_sync); try { - algo->build(base_set, index_size); + if (base_compressed) { + algo->build_from_base_set_file(); + } else { + algo->build(base_set, index_size); + } } catch (const std::exception& e) { state.SkipWithError(std::string(e.what())); } @@ -238,6 +256,11 @@ void bench_search(::benchmark::State& state, auto ualgo = create_algo(index.algo, dataset->distance(), dataset->dim(), index.build_param); a = ualgo.get(); + // An index built from a compressed base set stores only its graph, so `load` alone would + // leave it with nothing to search over. Handing the file over first lets `load` attach the + // same rows the graph was built from and return a complete index. The row count it returns + // is of no use here; only the build reports that. + if (dataset->base_is_compressed()) { a->set_base_set_file(dataset->base_file()); } a->load(index_file); current_algo = std::move(ualgo); } @@ -251,6 +274,12 @@ void bench_search(::benchmark::State& state, std::make_unique(std::move(parse_algo_property(a->get_preference(), sp_json))); if (search_param->needs_dataset()) { + if (dataset->base_is_compressed()) { + state.SkipWithError("The search parameters of '" + index.name + + "' require the dense base set, which a compressed base_file does not " + "provide."); + return; + } try { a->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); @@ -540,6 +569,7 @@ void dispatch_benchmark(std::string cmdline, auto dataset = std::make_shared>(dataset_conf.name, base_file, + dataset_conf.base_compressed, dataset_conf.subset_first_row, dataset_conf.subset_size, query_file, @@ -552,7 +582,13 @@ void dispatch_benchmark(std::string cmdline, if (build_mode) { if (file_exists(base_file)) { log_info("Using the dataset file '%s'", base_file.c_str()); - ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); + if (dataset_conf.base_compressed) { + // The row count sits inside the compressed file, so it is reported per benchmark as + // `index_size` once the algorithm has read it, rather than up front here. + ::benchmark::AddCustomContext("base_format", "vpq"); + } else { + ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); + } ::benchmark::AddCustomContext("dim", std::to_string(dataset->dim())); } else { log_warn("dataset file '%s' does not exist; benchmarking index building is impossible.", @@ -718,51 +754,59 @@ inline auto run_main(int argc, char** argv) -> int log_warn("cudart library is not found, GPU-based indices won't work."); } - auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); - std::string dtype = conf.get_dataset_conf().dtype; - - if (dtype == "float") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "half") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "uint8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "int8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else { - log_error("datatype '%s' is not supported", dtype.c_str()); + // A rejected configuration reaches us as an exception, from the json parser or from the dataset + // itself. Reporting it here keeps that a legible error and a non-zero exit code, rather than an + // abort from an uncaught exception. + try { + auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); + std::string dtype = conf.get_dataset_conf().dtype; + + if (dtype == "float") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "half") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "uint8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "int8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else { + log_error("datatype '%s' is not supported", dtype.c_str()); + return -1; + } + } catch (const std::exception& e) { + log_error("%s", e.what()); return -1; } diff --git a/cpp/bench/ann/src/common/conf.hpp b/cpp/bench/ann/src/common/conf.hpp index afc7bc0a1f..0ed63d0ba6 100644 --- a/cpp/bench/ann/src/common/conf.hpp +++ b/cpp/bench/ann/src/common/conf.hpp @@ -8,12 +8,19 @@ #include #include +#include #include #include #include namespace cuvs::bench { +inline auto has_suffix(const std::string& str, const std::string& suffix) -> bool +{ + return str.size() >= suffix.size() && + str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; +} + class configuration { public: struct index { @@ -40,6 +47,12 @@ class configuration { std::string distance; std::optional groundtruth_neighbors_file{std::nullopt}; + // The base_file holds rows already compressed for the algorithm (a .vpq written by the offline + // VPQ compression tool) rather than dense vectors. The benchmark cannot read such a file: its + // rows are not `dtype` values, so they cannot travel through `algo::build`. The path is + // handed to the algorithm instead. Queries stay dense, and `dtype` keeps describing them. + bool base_compressed{false}; + // data type of input dataset, possible values ["float", "int8", "uint8"] std::string dtype; @@ -99,11 +112,30 @@ class configuration { } if (conf.contains("subset_size")) { dataset_conf_.subset_size = conf.at("subset_size"); } + // Decided separately from the dtype inference below, so that an explicit "dtype" does not stop + // us noticing that the base set is compressed. + if (conf.contains("base_format")) { + const auto base_format = conf.at("base_format").get(); + if (base_format == "vpq") { + dataset_conf_.base_compressed = true; + } else if (base_format != "dense") { + throw std::runtime_error("Unknown base_format '" + base_format + + "', expected \"vpq\" or \"dense\""); + } + } else { + dataset_conf_.base_compressed = has_suffix(dataset_conf_.base_file, ".vpq"); + } + if (conf.contains("dtype")) { dataset_conf_.dtype = conf.at("dtype"); } else { auto filename = dataset_conf_.base_file; - if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { + if (dataset_conf_.base_compressed) { + // A VPQ dataset stores its codebooks as half, but it is searched with float queries and + // yields a float index, so float is the type the benchmark instantiates. Keyed off the flag + // rather than the suffix, so that an explicit base_format also gets a dtype. + dataset_conf_.dtype = "float"; + } else if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { dataset_conf_.dtype = "half"; } else if (filename.size() > 9 && filename.compare(filename.size() - 9, 9, "fp16.fbin") == 0) { diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 4dc43c343c..93a193cad4 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -152,6 +153,10 @@ struct dataset { private: std::string name_; std::string distance_; + std::string base_file_; + // A compressed base set is opaque to the benchmark: `base_set_` stays lazy and untouched, and the + // path is handed to the algorithm, which is the only thing able to decode it. + bool base_compressed_; blob base_set_; blob query_set_; std::optional> filter_bitset_; @@ -171,9 +176,22 @@ struct dataset { } } + // Reading a compressed base set as a .bin would not fail, it would succeed on garbage: the first + // eight bytes of a .vpq are a dtype prefix and a numpy magic, which parse as an absurd shape. + // Hence an explicit error, and callers that can proceed must ask `base_is_compressed()` first. + inline void throw_if_base_compressed(const char* what) const + { + if (base_compressed_) { + throw std::runtime_error{std::string{"dataset::"} + what + + "() is not available for the compressed base_file '" + base_file_ + + "'; only the algorithm can read that file."}; + } + } + public: dataset(std::string name, std::string base_file, + bool base_compressed, uint32_t subset_first_row, uint32_t subset_size, std::string query_file, @@ -182,9 +200,27 @@ struct dataset { std::optional filtering_rate = std::nullopt) : name_{std::move(name)}, distance_{std::move(distance)}, + base_file_{base_file}, + base_compressed_{base_compressed}, base_set_{base_file, subset_first_row, subset_size}, query_set_{query_file} { + if (base_compressed_) { + // Which rows went into the file was decided when it was written, so subsetting here is not + // merely unsupported: there is nothing left to select from. + if (subset_first_row != 0 || subset_size != 0) { + throw std::runtime_error{ + "A compressed base_file cannot be subset by the benchmark; choose the rows when " + "compressing it (the tool's --subset_first_row / --subset_size) and drop these keys."}; + } + // The bitset is sized from the base set row count, which is inside the compressed file. + if (filtering_rate.has_value()) { + throw std::runtime_error{ + "filtering_rate is not supported with a compressed base_file: generating the filter " + "bitset needs the base set size, which only the algorithm can read."}; + } + } + if (filtering_rate.has_value()) { // Generate a random bitset for filtering auto n_rows = static_cast(subset_size) + static_cast(subset_first_row); @@ -210,6 +246,8 @@ struct dataset { [[nodiscard]] auto name() const -> std::string { return name_; } [[nodiscard]] auto distance() const -> std::string { return distance_; } + [[nodiscard]] auto base_file() const -> std::string { return base_file_; } + [[nodiscard]] auto base_is_compressed() const -> bool { return base_compressed_; } [[nodiscard]] auto dim() const -> int { auto d = dim_.load(std::memory_order_relaxed); @@ -221,6 +259,14 @@ struct dataset { } catch (const std::runtime_error& e) { // Any exception raised above will re-raise next time we try to access the query set. query_set_.reset_lazy_state(); + // A compressed base set has no dense header to fall back on, and reading it as one would + // yield a nonsense dimension rather than an error. + if (base_compressed_) { + throw std::runtime_error{ + "Cannot determine the dataset dimension: the query set is not readable and the base set " + "is compressed. " + + std::string{e.what()}}; + } // If the query set is not accessible, use the base set. // Don't catch the exception here, because we have nothing else to do anyway. d = static_cast(base_set_.n_cols()); @@ -235,6 +281,7 @@ struct dataset { } [[nodiscard]] auto base_set_size() const -> size_t { + throw_if_base_compressed("base_set_size"); std::lock_guard lock(mutex_); auto r = base_set_.n_rows(); cache_dim(base_set_); @@ -272,6 +319,7 @@ struct dataset { [[nodiscard]] auto base_set() const -> const DataT* { + throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(); cache_dim(base_set_); @@ -281,6 +329,7 @@ struct dataset { HugePages request_hugepages_2mb = HugePages::kDisable) const -> const DataT* { + throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(memory_type, request_hugepages_2mb); cache_dim(base_set_); From 843425b3fbbd69d21f003a04879a4303897e8056 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 17 Aug 2026 08:31:32 -0700 Subject: [PATCH 32/39] Don't refuse a compressed base set in search mode: cuvs_cagra's needs_dataset() is always true --- cpp/bench/ann/src/common/ann_types.hpp | 4 ++++ cpp/bench/ann/src/common/benchmark.hpp | 14 +++++++------- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index 52105a2dc1..6baee5b52f 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -157,6 +157,10 @@ class algo : public algo_base { // and set_search_dataset() should save the passed-in pointer somewhere. // The client code should call set_search_dataset() before searching, // and should not release dataset before searching is finished. + // + // A compressed base set is never handed over this way, as it has no dense rows to pass, so + // needs_dataset() says nothing about one. An algorithm that cannot search the base set it was + // given with the parameters it was given has to reject them itself, from set_search_param(). virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) {}; /* ### Base sets the benchmark cannot read ### diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 9e76192b95..483ed524d9 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -273,13 +273,13 @@ void bench_search(::benchmark::State& state, current_algo_props = std::make_unique(std::move(parse_algo_property(a->get_preference(), sp_json))); - if (search_param->needs_dataset()) { - if (dataset->base_is_compressed()) { - state.SkipWithError("The search parameters of '" + index.name + - "' require the dense base set, which a compressed base_file does not " - "provide."); - return; - } + // Not a reliable signal for a compressed base set: cuvs_cagra answers true unconditionally, + // because its index file carries no dataset and the dense rows are re-attached here instead. + // There are no dense rows to attach for a compressed base, and the algorithm already has the + // file from `set_base_set_file` above. An algorithm that truly cannot search without the dense + // rows, such as CAGRA with refine_ratio > 1, has to reject that combination itself: only it + // knows which of its search parameters read them. + if (search_param->needs_dataset() && !dataset->base_is_compressed()) { try { a->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); From 689a882099d99f46e066b50b4aca646508c6a89a Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 17 Aug 2026 08:58:43 -0700 Subject: [PATCH 33/39] Serialize the VPQ dataset with a CAGRA-Q index so a loaded index searches without the dense rows --- c/src/neighbors/cagra.cpp | 8 +- cpp/include/cuvs/neighbors/cagra.hpp | 115 +++++++++++++- cpp/src/neighbors/cagra_serialize.cuh | 37 +++++ cpp/src/neighbors/cagra_serialize_inst.cu.in | 3 + .../detail/cagra/cagra_serialize.cuh | 26 +++- .../ann_cagra/test_iterative_cagra_q.cu | 141 ++++++++++++++++-- 6 files changed, 303 insertions(+), 27 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 99e456e23c..42200258b6 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -1011,7 +1011,7 @@ static auto read_serialized_header(cuvsResources_t res, const char *filename) "serialization version mismatch, expected %d, got %d", cuvs::neighbors::cagra::cagra_serialization_version, version); using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - RAFT_EXPECTS(dataset_kind_raw <= static_cast(kind::host_standard), + RAFT_EXPECTS(dataset_kind_raw <= static_cast(kind::device_vpq_f16), "Invalid serialized dataset kind %u in file %s", dataset_kind_raw, filename); return {output_dtype, static_cast(dataset_kind_raw)}; @@ -1058,6 +1058,12 @@ void dispatch_serialized_dataset_kind( fn.template operator()< cuvs::neighbors::device_padded_dataset_view>(); break; + case serialized_kind::device_vpq_f16: + // A recognised file the C API has no index layout for, as opposed to an unreadable one. + // cuvsDatasetLayout_t covers standard and padded only, and every C entry point dispatches + // on that layout, so there is nothing here to hand a VPQ index to yet. + RAFT_FAIL("File holds a VPQ-compressed (CAGRA-Q) dataset, which the C API has no dataset " + "layout for; load it through the C++ API"); } } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index dc18a6b792..0b4ae3f090 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2307,7 +2307,7 @@ void search( * @{ */ -/** Dense dataset storage kind recorded in a serialized CAGRA index. */ +/** Dataset storage kind recorded in a serialized CAGRA index. */ enum class serialized_dataset_kind : std::uint32_t { /** The serialized index does not contain a dataset payload. */ none = 0, @@ -2319,16 +2319,19 @@ enum class serialized_dataset_kind : std::uint32_t { host_padded = 3, /** Host-resident dataset using its standard row layout. */ host_standard = 4, + /** Device-resident VPQ-compressed dataset with f16 codebooks (CAGRA-Q). */ + device_vpq_f16 = 5, }; /** Current experimental CAGRA serialization format version. */ inline constexpr int cagra_serialization_version = 6; -// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes. -// They use the same strided dataset payload; the serialized dataset kind selects the matching -// owning dataset type during deserialization. To support a new dataset kind (e.g. vpq_f16_index), -// add matching overloads here and a corresponding deserialize_ in -// detail/dataset_serialize.hpp (dense views use serialize_cagra_dense_dataset). +// Serialize and deserialize are overloaded for device/host and padded/standard dense indexes, +// which share the same strided dataset payload, and for vpq_f16_index, which writes a VPQ payload +// instead. The serialized dataset kind selects the matching owning dataset type during +// deserialization. To support a further kind, add matching overloads here and a corresponding +// serialize_/deserialize_ in detail/dataset_serialize.hpp (dense views use +// serialize_cagra_dense_dataset, VPQ ones serialize_vpq_dataset). /** * Save the index to file. @@ -2879,6 +2882,106 @@ void deserialize(raft::resources const& handle, std::unique_ptr>* out_dataset = nullptr); +/* vpq_f16_index overloads (CAGRA-Q). + * + * The compressed rows travel with the index, so that a deserialized index can be searched without + * the dense dataset it was compressed from and without retraining the codebooks. As everywhere + * else, the index holds a view: `deserialize` returns the owning dataset through `out_dataset`, + * which the caller has to keep alive for as long as the index is used. + * + * Unlike the dense overloads, `out_dataset` is required. Nothing can be searched in a VPQ index + * whose rows were dropped, so there is no use for a graph-only load, and asking for one is an + * error rather than a silently unusable index. For the same reason `include_dataset = false` + * produces an index that only `update_dataset` can make searchable again. + */ +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + const std::string& filename, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + +void serialize(raft::resources const& handle, + std::ostream& os, + const cuvs::neighbors::cagra::vpq_f16_index& index, + bool include_dataset = true); + +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::vpq_f16_index* index, + std::unique_ptr>* out_dataset); + /** @copydoc serialize */ void serialize(raft::resources const& handle, const std::string& filename, diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index 83d047560b..9d7614e498 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -155,6 +155,43 @@ namespace cuvs::neighbors::cagra { cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ \ + void serialize(raft::resources const& handle, \ + const std::string& filename, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, filename, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize( \ + handle, filename, index, out_dataset); \ + } \ + \ + void serialize(raft::resources const& handle, \ + std::ostream& os, \ + const cuvs::neighbors::cagra::vpq_f16_index& index, \ + bool include_dataset) \ + { \ + cuvs::neighbors::cagra::detail::serialize( \ + handle, os, index, include_dataset); \ + } \ + \ + void deserialize( \ + raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::vpq_f16_index* index, \ + std::unique_ptr>* out_dataset) \ + { \ + cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ + } \ + \ void serialize_to_hnswlib( \ raft::resources const& handle, \ std::ostream& os, \ diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 3d34adb36f..58e555d17e 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/cpp/src/neighbors/cagra_serialize_inst.cu.in @@ -12,6 +12,7 @@ namespace { using data_t = @data_type@; using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; } // namespace @@ -21,6 +22,8 @@ extern template void index::compute raft::resources const&); extern template void index::compute_dataset_norms_( raft::resources const&); +extern template void index::compute_dataset_norms_( + raft::resources const&); CUVS_INST_CAGRA_SERIALIZE(data_t); diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index f2e0c4f07b..add4bb2532 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -60,6 +60,8 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser return kind::host_padded; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { return kind::host_standard; + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + return kind::device_vpq_f16; } else { static_assert(sizeof(DatasetViewT) == 0, "serialized_dataset_kind_for_view: unsupported dataset view type"); @@ -69,7 +71,7 @@ constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::ser constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) { using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - return raw <= static_cast(kind::host_standard); + return raw <= static_cast(kind::device_vpq_f16); } /** @@ -123,9 +125,14 @@ void serialize(raft::resources const& res, RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + // The payload describes its own codebook type, which is `half` here regardless of T: the + // dtype prefix written above is the type of the queries this index answers, not of its rows. + // `dset()` is safe to call because a view over no rows left include_dataset false above. + neighbors::detail::serialize_vpq_dataset(res, os, index_.dataset().dset()); } else { - // Future dataset types (e.g. VPQ) require a new branch here and a corresponding - // deserialize overload. Use static_assert to catch unsupported types at compile time. + // A further dataset type requires a new branch here and a corresponding deserialize branch. + // Use static_assert to catch unsupported types at compile time. static_assert( sizeof(DatasetViewT) == 0, "serialize: dataset serialization is not yet implemented for this DatasetViewT"); @@ -401,7 +408,16 @@ void deserialize( std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { - cuvs::neighbors::detail::skip_dense_dataset(res, is); + // Dropping the rows leaves a searchable index for a dense view, whose dataset can be + // reattached from the caller's own copy, but not for a VPQ one: the compressed rows exist + // nowhere else. Refuse rather than hand back an index that cannot answer a query, and skip + // the dense payload only when it is in fact dense. + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL( + "cagra::deserialize: a VPQ index cannot be loaded without its dataset; pass out_dataset"); + } else { + cuvs::neighbors::detail::skip_dense_dataset(res, is); + } } else { auto const expected_kind = serialized_dataset_kind_for_view(); RAFT_EXPECTS( @@ -419,6 +435,8 @@ void deserialize( } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_host_standard_dataset(res, is); + } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + dataset_owner = cuvs::neighbors::detail::deserialize_vpq_dataset(res, is); } else { static_assert(sizeof(DatasetViewT) == 0, "deserialize: dataset deserialization is not implemented for this view"); diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu index bad504b589..5ac32b72b4 100644 --- a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -12,8 +12,9 @@ * the dtype-templated suites in ann_cagra.cuh and lives in its own file. * * What is checked here is that a compressed dataset, freshly compressed or loaded from disk, - * builds a usable graph, and that the constraints above are rejected rather than accepted and - * quietly ignored. Serialization fidelity itself is covered by preprocessing/vpq_serialization.cu. + * builds a usable graph, that such an index survives a trip through a file with its rows, and that + * the constraints above are rejected rather than accepted and quietly ignored. Fidelity of the + * dataset payload itself is covered by preprocessing/vpq_serialization.cu. */ #include @@ -65,6 +66,28 @@ auto iterative_params(uint32_t graph_degree = 32) -> index_params return params; } +constexpr int64_t kSearchK = 10; + +/** Neighbour ids for `queries`, row-major [n_queries, kSearchK]. */ +template +auto neighbor_ids(const raft::resources& res, + const IndexT& idx, + raft::device_matrix_view queries) -> std::vector +{ + const auto n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(res, n_queries, kSearchK); + auto distances = raft::make_device_matrix(res, n_queries, kSearchK); + + search_params params; + params.itopk_size = 64; + search(res, params, idx, queries, neighbors.view(), distances.view()); + + std::vector ids(static_cast(n_queries * kSearchK)); + raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return ids; +} + /** * Fraction of queries that retrieve their own row, where the queries are dataset rows. * @@ -77,22 +100,11 @@ auto self_recall_at_1(const raft::resources& res, const IndexT& idx, raft::device_matrix_view queries) -> double { - constexpr int64_t k = 10; + auto ids = neighbor_ids(res, idx, queries); const auto n_queries = queries.extent(0); - auto neighbors = raft::make_device_matrix(res, n_queries, k); - auto distances = raft::make_device_matrix(res, n_queries, k); - - search_params params; - params.itopk_size = 64; - search(res, params, idx, queries, neighbors.view(), distances.view()); - - std::vector ids(static_cast(n_queries * k)); - raft::copy(ids.data(), neighbors.data_handle(), ids.size(), raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - - int64_t hits = 0; + int64_t hits = 0; for (int64_t q = 0; q < n_queries; q++) { - if (ids[q * k] == static_cast(q)) { hits++; } + if (ids[q * kSearchK] == static_cast(q)) { hits++; } } return static_cast(hits) / static_cast(n_queries); } @@ -221,6 +233,103 @@ INSTANTIATE_TEST_CASE_P(CagraQBuildTests, {2000, 256, 32}, // pq_len 8 })); +/** + * An index over compressed rows is serialized with those rows, so that a loaded index can be + * searched without the dense dataset it came from and without retraining the codebooks. The + * ownership split is the usual one: the file yields an owning dataset, the index only views it. + */ +class CagraQSerializeTest : public CagraQCompressedTestBase { + protected: + void SetUp() override { make_dataset(n_rows, dim); } + + static constexpr int64_t n_rows = 2000; + static constexpr int64_t dim = 128; + static constexpr uint32_t pq_dim = 32; // pq_len 4 +}; + +TEST_F(CagraQSerializeTest, RoundTripsThroughAFileWithItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + auto before = neighbor_ids(res_, idx, queries(500)); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + ASSERT_NE(owner, nullptr); + EXPECT_EQ(owner->n_rows(), compressed.n_rows()); + EXPECT_EQ(owner->dim(), compressed.dim()); + EXPECT_EQ(owner->pq_len(), compressed.pq_len()); + EXPECT_EQ(owner->pq_bits(), compressed.pq_bits()); + EXPECT_EQ(owner->vq_n_centers(), compressed.vq_n_centers()); + EXPECT_EQ(owner->encoded_row_length(), compressed.encoded_row_length()); + + ASSERT_EQ(restored.size(), idx.size()); + ASSERT_EQ(restored.dim(), idx.dim()); + ASSERT_EQ(restored.graph_degree(), idx.graph_degree()); + EXPECT_EQ(restored.metric(), idx.metric()); + + // Same graph over the same rows, so the results are identical rather than merely comparable. + auto after = neighbor_ids(res_, restored, queries(500)); + ASSERT_EQ(after.size(), before.size()); + size_t mismatches = 0; + for (size_t i = 0; i < before.size(); i++) { + mismatches += static_cast(after[i] != before[i]); + } + EXPECT_EQ(mismatches, 0u) << mismatches << " of " << before.size() << " neighbour ids changed"; +} + +TEST_F(CagraQSerializeTest, RefusesToLoadWithoutItsDataset) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // Dropping the rows on load is fine for a dense index, whose caller can attach its own copy, but + // it would leave a VPQ index unsearchable with no way back: the rows exist nowhere else. + vpq_f16_index restored{res_}; + EXPECT_THROW(cagra::deserialize(res_, stored, &restored, nullptr), raft::exception); +} + +TEST_F(CagraQSerializeTest, SerializesTheGraphAloneWhenAsked) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx, /* include_dataset */ false); + + vpq_f16_index restored{res_}; + std::unique_ptr owner; + cagra::deserialize(res_, stored, &restored, &owner); + + // Nothing to own, and a graph that only update_dataset() can make searchable again. + EXPECT_EQ(owner, nullptr); + EXPECT_EQ(restored.size(), idx.size()); + EXPECT_EQ(restored.graph_degree(), idx.graph_degree()); +} + +TEST_F(CagraQSerializeTest, RejectsLoadingACompressedIndexAsDense) +{ + auto compressed = compress(res_, dataset(), pq_dim); + auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); + + std::stringstream stored; + cagra::serialize(res_, stored, idx); + + // The dtype prefix says float either way, so it is the recorded dataset kind that has to stop + // the dense reader from interpreting VPQ codes as rows of floats. + device_padded_index dense{res_}; + std::unique_ptr> dense_owner; + EXPECT_THROW(cagra::deserialize(res_, stored, &dense, &dense_owner), raft::exception); +} + /** The constraints the VPQ build overload documents, each of which must be rejected loudly. */ class CagraQContractTest : public CagraQCompressedTestBase { protected: From 8863d1d4fe01b0ae0067ba3ac1d6fd1ce86d6055 Mon Sep 17 00:00:00 2001 From: aamijar Date: Tue, 18 Aug 2026 00:36:38 +0000 Subject: [PATCH 34/39] C and Python APIs --- c/include/cuvs/core/dataset.h | 14 +- c/include/cuvs/neighbors/cagra.h | 6 +- c/src/neighbors/cagra.cpp | 190 ++++++++++++++++-- python/cuvs/cuvs/common/dataset.pxd | 11 + python/cuvs/cuvs/common/dataset.pyx | 2 + python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 14 ++ python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 8 +- .../preprocessing/quantize/pq/__init__.py | 11 +- .../cuvs/preprocessing/quantize/pq/pq.pyx | 64 +++++- 9 files changed, 292 insertions(+), 28 deletions(-) diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h index 78d3547495..0d28a9977c 100644 --- a/c/include/cuvs/core/dataset.h +++ b/c/include/cuvs/core/dataset.h @@ -20,7 +20,8 @@ extern "C" { */ typedef enum { CUVS_DATASET_LAYOUT_STANDARD = 0, - CUVS_DATASET_LAYOUT_PADDED = 1 + CUVS_DATASET_LAYOUT_PADDED = 1, + CUVS_DATASET_LAYOUT_VPQ = 2 } cuvsDatasetLayout_t; /** @@ -48,6 +49,9 @@ typedef struct { } cuvsDataset; typedef cuvsDataset* cuvsDataset_t; +struct cuvsCagraCompressionParams; +typedef struct cuvsCagraCompressionParams* cuvsCagraCompressionParams_t; + /** * @brief Create an empty owning dataset handle. * @@ -72,6 +76,14 @@ CUVS_EXPORT cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res, cuvsDatasetMemType_t target_mem_type, cuvsDataset_t* padded_dataset); +/** + * @brief Compress a dense dataset into a device VPQ dataset. + */ +CUVS_EXPORT cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res, + cuvsCagraCompressionParams_t params, + cuvsDataset_t dataset, + cuvsDataset_t* vpq_dataset); + /** * @brief Create a non-owning padded dataset view from a host- or device-resident tensor. * diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 350711d069..66e222d747 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -118,8 +118,6 @@ struct cuvsCagraCompressionParams { double pq_kmeans_trainset_fraction; }; -typedef struct cuvsCagraCompressionParams* cuvsCagraCompressionParams_t; - struct cuvsIvfPqParams { cuvsIvfPqIndexParams_t ivf_pq_build_params; cuvsIvfPqSearchParams_t ivf_pq_search_params; @@ -655,6 +653,10 @@ CUVS_EXPORT cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * + * A `CUVS_DATASET_LAYOUT_VPQ` dataset created by `cuvsDatasetMakeVpq` builds an iterative CAGRA-Q + * index. VPQ input requires `L2Expanded` and `ITERATIVE_CAGRA_SEARCH` (or `AUTO_SELECT`), and the + * VPQ dataset must outlive the index because the index stores a non-owning view. + * * @param[in] res cuvsResources_t opaque C handle * @param[in] params cuvsCagraIndexParams_t used to build CAGRA index * @param[in] dataset cuvsDataset_t training dataset or dataset view diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 99e456e23c..4eb83cb7bb 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include "../core/exceptions.hpp" #include "../core/interop.hpp" @@ -52,7 +53,13 @@ struct cuvs_cagra_c_api_index_lifetime_holder { /** Owns how to delete co-located index storage; `cuvsCagraIndex::addr` points here. */ struct sg_cagra_c_api_index_box { void* index_ptr; - enum class dataset_layout : uint8_t { device_padded, device_standard, host_padded, host_standard } layout; + enum class dataset_layout : uint8_t { + device_padded, + device_standard, + device_vpq, + host_padded, + host_standard + } layout; cuvs::neighbors::c_api::detail::owner_record owner_rec; }; @@ -63,6 +70,8 @@ constexpr auto sg_cagra_index_layout_from_view() return sg_cagra_c_api_index_box::dataset_layout::device_standard; } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_padded; + } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + return sg_cagra_c_api_index_box::dataset_layout::device_vpq; } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::host_standard; } else { @@ -100,6 +109,13 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, fn(*idx); break; } + case sg_cagra_c_api_index_box::dataset_layout::device_vpq: { + using index_t = cuvs::neighbors::cagra:: + index>; + auto* idx = reinterpret_cast(box->index_ptr); + fn(*idx); + break; + } case sg_cagra_c_api_index_box::dataset_layout::host_standard: { if constexpr (AllowHost) { auto* idx = @@ -369,6 +385,77 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) } } +using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset; +using device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; + +static void bind_vpq_owner_to_dataset(std::unique_ptr owner, + cuvsDataset_t* output) +{ + RAFT_EXPECTS(output != nullptr, "VPQ output dataset pointer must not be null"); + auto* out = new cuvsDataset{}; + out->addr = reinterpret_cast(owner.release()); + out->destroy_addr = &destroy_typed_addr; + out->dtype.code = kDLFloat; + out->dtype.bits = 32; + out->dtype.lanes = 1; + out->mem_type = CUVS_DATASET_MEM_TYPE_DEVICE; + out->layout = CUVS_DATASET_LAYOUT_VPQ; + out->is_owning = true; + *output = out; +} + +static auto make_cpp_vpq_params(cuvsCagraCompressionParams const& params) + -> cuvs::neighbors::vpq_params +{ + auto out = cuvs::neighbors::vpq_params{}; + out.pq_bits = params.pq_bits; + out.pq_dim = params.pq_dim; + out.vq_n_centers = params.vq_n_centers; + out.kmeans_n_iters = params.kmeans_n_iters; + out.vq_kmeans_trainset_fraction = params.vq_kmeans_trainset_fraction; + out.pq_kmeans_trainset_fraction = params.pq_kmeans_trainset_fraction; + return out; +} + +template +static auto make_vpq_from_dense_dataset(raft::resources* res_ptr, + cuvsCagraCompressionParams const& params, + cuvsDataset_t dataset) + -> std::unique_ptr +{ + RAFT_EXPECTS(dataset->layout == CUVS_DATASET_LAYOUT_STANDARD || + dataset->layout == CUVS_DATASET_LAYOUT_PADDED, + "cuvsDatasetMakeVpq: source dataset must have STANDARD or PADDED layout"); + auto cpp_params = make_cpp_vpq_params(params); + std::unique_ptr owner; + auto make = [&](auto const& view) { + owner = std::make_unique( + cuvs::preprocessing::quantize::pq::make_vpq_dataset(*res_ptr, cpp_params, view)); + }; + + const bool padded = dataset->layout == CUVS_DATASET_LAYOUT_PADDED; + if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) { + if (padded) { + with_dataset_view, + cuvs::neighbors::device_padded_dataset_view>(dataset, make); + } else { + with_dataset_view, + cuvs::neighbors::device_standard_dataset_view>(dataset, make); + } + } else if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST) { + if (padded) { + with_dataset_view, + cuvs::neighbors::host_padded_dataset_view>(dataset, make); + } else { + with_dataset_view, + cuvs::neighbors::host_standard_dataset_view>(dataset, make); + } + } else { + RAFT_FAIL("cuvsDatasetMakeVpq: invalid source dataset memory type"); + } + return owner; +} + template static void make_device_padded_dataset(raft::resources* res_ptr, DLManagedTensor* dataset_tensor, @@ -547,12 +634,17 @@ static void attach_dataset(raft::resources* res_ptr, "cuvsCagraUpdateDataset: null index handle", "cuvsCagraUpdateDataset: host index layout is allowed for this operation", [&](auto& idx) { - auto padded_idx = cuvs::neighbors::cagra::attach_dataset(*res_ptr, idx, padded_view); - auto* holder = - new cuvs_cagra_c_api_index_lifetime_holder{std::move(padded_idx)}; - destroy_sg_cagra_c_api_box(index->addr); - index->addr = 0; - bind_index_lifetime_holder_to_C_index(index, index->dtype, holder); + using index_dataset_view_t = std::remove_cvref_t; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraUpdateDataset: replacing a VPQ dataset is not supported"); + } else { + auto padded_idx = cuvs::neighbors::cagra::attach_dataset(*res_ptr, idx, padded_view); + auto* holder = + new cuvs_cagra_c_api_index_lifetime_holder{std::move(padded_idx)}; + destroy_sg_cagra_c_api_box(index->addr); + index->addr = 0; + bind_index_lifetime_holder_to_C_index(index, index->dtype, holder); + } }); }); } @@ -959,13 +1051,19 @@ void _serialize(cuvsResources_t res, const char *filename, : "cuvsCagraSerializeGraph: null index handle"; with_index_by_layout(box, null_handle_err, "", [&](auto &idx) { - if (include_dataset) { - RAFT_EXPECTS( - idx.dataset().n_rows() > 0, - "cuvsCagraSerializeGraphAndDataset: index has no attached dataset"); + using index_dataset_view_t = std::remove_cvref_t; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL( + "CAGRA index serialization is not supported for VPQ indices"); + } else { + if (include_dataset) { + RAFT_EXPECTS( + idx.dataset().n_rows() > 0, + "cuvsCagraSerializeGraphAndDataset: index has no attached dataset"); + } + cuvs::neighbors::cagra::serialize( + *res_ptr, std::string(filename), idx, include_dataset); } - cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx, - include_dataset); }); } @@ -1122,8 +1220,13 @@ void _serialize_to_hnswlib(cuvsResources_t res, const char *filename, box, "cuvsCagraSerializeToHnswlib: null index handle", "cuvsCagraSerializeToHnswlib: host indices are allowed", [&](auto &idx) { - cuvs::neighbors::cagra::serialize_to_hnswlib( - *res_ptr, std::string(filename), idx); + using index_dataset_view_t = std::remove_cvref_t; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for VPQ indices"); + } else { + cuvs::neighbors::cagra::serialize_to_hnswlib( + *res_ptr, std::string(filename), idx); + } }); } template @@ -1197,7 +1300,14 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) box, "cuvsCagraIndexGetDataset: null index handle", "cuvsCagraIndexGetDataset: host indices are allowed", - [&](auto& idx) { cuvs::core::to_dlpack(idx.dataset().view(), dataset); }); + [&](auto& idx) { + using index_dataset_view_t = std::remove_cvref_t; + if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + RAFT_FAIL("cuvsCagraIndexGetDataset does not expose VPQ datasets as dense DLPack tensors"); + } else { + cuvs::core::to_dlpack(idx.dataset().view(), dataset); + } + }); } template @@ -1578,6 +1688,36 @@ extern "C" cuvsError_t cuvsDatasetMakeStandardView(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsDatasetMakeVpq(cuvsResources_t res, + cuvsCagraCompressionParams_t params, + cuvsDataset_t dataset, + cuvsDataset_t* vpq_dataset) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(params != nullptr, "cuvsDatasetMakeVpq: null compression params"); + RAFT_EXPECTS(dataset != nullptr && dataset->addr != 0, + "cuvsDatasetMakeVpq: null source dataset"); + RAFT_EXPECTS(vpq_dataset != nullptr, "cuvsDatasetMakeVpq: null output dataset"); + *vpq_dataset = nullptr; + auto* res_ptr = reinterpret_cast(res); + std::unique_ptr owner; + if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 32) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLFloat && dataset->dtype.bits == 16) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLInt && dataset->dtype.bits == 8) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else if (dataset->dtype.code == kDLUInt && dataset->dtype.bits == 8) { + owner = make_vpq_from_dense_dataset(res_ptr, *params, dataset); + } else { + RAFT_FAIL("cuvsDatasetMakeVpq: unsupported source dtype: code=%d, bits=%d", + dataset->dtype.code, + dataset->dtype.bits); + } + bind_vpq_owner_to_dataset(std::move(owner), vpq_dataset); + }); +} + static cuvsError_t dispatch_attach_dataset(cuvsResources_t res, cuvsDataset_t device_padded_dataset, cuvsCagraIndex_t index) @@ -1741,7 +1881,15 @@ extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res, index->addr = 0; index->dtype = dtype; - if (dtype.code == kDLFloat && dtype.bits == 32) { + if (dataset->layout == CUVS_DATASET_LAYOUT_VPQ) { + RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE, + "cuvsCagraBuild: VPQ dataset must be device-resident"); + RAFT_EXPECTS(dtype.code == kDLFloat && dtype.bits == 32, + "cuvsCagraBuild: VPQ dataset query dtype must be float32"); + with_dataset_view(dataset, [&](auto const& view) { + build_index_from_dataset_view(res_ptr, params, view, index); + }); + } else if (dtype.code == kDLFloat && dtype.bits == 32) { build_dispatch_on_mem_type_and_layout(res_ptr, params, dataset, index); } else if (dtype.code == kDLFloat && dtype.bits == 16) { build_dispatch_on_mem_type_and_layout(res_ptr, params, dataset, index); @@ -1839,10 +1987,12 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, auto index = *index_c_ptr; auto* box = reinterpret_cast(index.addr); RAFT_EXPECTS(box != nullptr, "cuvsCagraSearch: null index handle"); - RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded, - "cuvsCagraSearch: index must be device-padded. For standard indices, call " + RAFT_EXPECTS(box->layout == sg_cagra_c_api_index_box::dataset_layout::device_padded || + box->layout == sg_cagra_c_api_index_box::dataset_layout::device_vpq, + "cuvsCagraSearch: index must be device-padded or VPQ. For standard indices, call " "cuvsCagraUpdateDataset first."); - RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); + RAFT_EXPECTS(queries.dtype.code == index.dtype.code && queries.dtype.bits == index.dtype.bits, + "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _search( diff --git a/python/cuvs/cuvs/common/dataset.pxd b/python/cuvs/cuvs/common/dataset.pxd index ac2d76ec18..742f783766 100644 --- a/python/cuvs/cuvs/common/dataset.pxd +++ b/python/cuvs/cuvs/common/dataset.pxd @@ -14,6 +14,7 @@ cdef extern from "cuvs/core/dataset.h" nogil: ctypedef enum cuvsDatasetLayout_t: CUVS_DATASET_LAYOUT_STANDARD CUVS_DATASET_LAYOUT_PADDED + CUVS_DATASET_LAYOUT_VPQ ctypedef enum cuvsDatasetMemType_t: CUVS_DATASET_MEM_TYPE_HOST @@ -23,6 +24,10 @@ cdef extern from "cuvs/core/dataset.h" nogil: pass ctypedef cuvsDataset* cuvsDataset_t + cdef struct cuvsCagraCompressionParams: + pass + ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t + cuvsError_t cuvsDatasetCreate(cuvsDataset_t* dataset) cuvsError_t cuvsDatasetMakePadded(cuvsResources_t res, @@ -38,6 +43,12 @@ cdef extern from "cuvs/core/dataset.h" nogil: DLManagedTensor* dataset, cuvsDataset_t* standard_dataset) + cuvsError_t cuvsDatasetMakeVpq( + cuvsResources_t res, + cuvsCagraCompressionParams_t params, + cuvsDataset_t dataset, + cuvsDataset_t* vpq_dataset) + cuvsError_t cuvsDatasetDestroy(cuvsDataset_t dataset) cuvsError_t cuvsDatasetGetMemType(cuvsDataset_t dataset, diff --git a/python/cuvs/cuvs/common/dataset.pyx b/python/cuvs/cuvs/common/dataset.pyx index 0c83633d13..b27f061b56 100644 --- a/python/cuvs/cuvs/common/dataset.pyx +++ b/python/cuvs/cuvs/common/dataset.pyx @@ -46,6 +46,8 @@ cdef class Dataset: if self.dataset == NULL: return None check_cuvs(cuvsDatasetGetLayout(self.dataset, &layout)) + if layout == CUVS_DATASET_LAYOUT_VPQ: + return "vpq" if layout == CUVS_DATASET_LAYOUT_PADDED: return "padded" return "standard" diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 9e4dbdb6f3..bd838c8754 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -43,6 +43,20 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: ITERATIVE_CAGRA_SEARCH ACE + ctypedef struct cuvsCagraCompressionParams: + uint32_t pq_bits + uint32_t pq_dim + uint32_t vq_n_centers + uint32_t kmeans_n_iters + double vq_kmeans_trainset_fraction + double pq_kmeans_trainset_fraction + + ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t + + cuvsError_t cuvsCagraCompressionParamsCreate( + cuvsCagraCompressionParams_t* params) + cuvsError_t cuvsCagraCompressionParamsDestroy( + cuvsCagraCompressionParams_t params) ctypedef struct cuvsIvfPqParams: cuvsIvfPqIndexParams_t ivf_pq_build_params cuvsIvfPqSearchParams_t ivf_pq_search_params diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index dd481df259..beeff0b93a 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -474,6 +474,9 @@ def build(IndexParams index_params, dataset, resources=None): Supported dtype [float, half, int8, uint8] **Note:** For ACE build algorithm, the dataset MUST be in host memory. Use NumPy arrays or call .get() on CuPy arrays before passing. + A ``Dataset`` with ``layout == "vpq"`` builds an iterative CAGRA-Q + index and requires ``metric="sqeuclidean"`` plus + ``build_algo="iterative_cagra_search"``. {resources_docstring} Returns @@ -527,7 +530,10 @@ def build(IndexParams index_params, dataset, resources=None): dl_data_type_to_numpy(idx.index.dtype)).name idx._dataset_source = dataset_obj - if not is_ace_build: + if dataset_obj.layout == "vpq": + _keep_dataset_alive(idx, dataset_obj) + idx._dataset_source = None + elif not is_ace_build: if (dataset_obj.layout == "padded" and dataset_obj.memory_type == "device" and dataset_obj.is_owning): diff --git a/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py b/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py index 7db0c383fd..833caedead 100644 --- a/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py +++ b/python/cuvs/cuvs/preprocessing/quantize/pq/__init__.py @@ -1,12 +1,17 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from .pq import Quantizer, QuantizerParams, build, transform, inverse_transform +from .pq import ( + Quantizer, QuantizerParams, VpqParams, build, inverse_transform, + make_vpq_dataset, transform, +) __all__ = [ "Quantizer", "QuantizerParams", + "VpqParams", "build", "transform", - "inverse_transform" + "inverse_transform", + "make_vpq_dataset", ] diff --git a/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx b/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx index fd1a5326e0..67749fae41 100644 --- a/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx +++ b/python/cuvs/cuvs/preprocessing/quantize/pq/pq.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -7,6 +7,16 @@ import numpy as np from cuvs.common cimport cydlpack +from cuvs.common.dataset cimport ( + Dataset, + cuvsDatasetMakeStandardView, + cuvsDatasetMakeVpq, +) +from cuvs.neighbors.cagra.cagra cimport ( + cuvsCagraCompressionParams, + cuvsCagraCompressionParamsCreate, + cuvsCagraCompressionParamsDestroy, +) from pylibraft.common import auto_convert_output, device_ndarray from pylibraft.common.cai_wrapper import wrap_array @@ -22,6 +32,32 @@ PQ_KMEANS_TYPES = { PQ_KMEANS_NAMES = {v: k for k, v in PQ_KMEANS_TYPES.items()} + +cdef class VpqParams: + """Parameters for creating a CAGRA-Q VPQ dataset.""" + + cdef cuvsCagraCompressionParams* params + + def __cinit__(self): + self.params = NULL + check_cuvs(cuvsCagraCompressionParamsCreate(&self.params)) + + def __dealloc__(self): + if self.params != NULL: + cuvsCagraCompressionParamsDestroy(self.params) + + def __init__(self, *, pq_bits=8, pq_dim=0, vq_n_centers=0, + kmeans_n_iters=25, vq_kmeans_trainset_fraction=0.0, + pq_kmeans_trainset_fraction=0.0): + self.params.pq_bits = pq_bits + self.params.pq_dim = pq_dim + self.params.vq_n_centers = vq_n_centers + self.params.kmeans_n_iters = kmeans_n_iters + self.params.vq_kmeans_trainset_fraction = \ + vq_kmeans_trainset_fraction + self.params.pq_kmeans_trainset_fraction = \ + pq_kmeans_trainset_fraction + cdef class QuantizerParams: """ Parameters for product quantization @@ -377,3 +413,29 @@ def inverse_transform(Quantizer quantizer, codes, output=None, vq_labels=None, r vq_labels_dlpack)) return output + + +@auto_sync_resources +def make_vpq_dataset(VpqParams params, dataset, resources=None): + """Create an owning device VPQ dataset for iterative CAGRA-Q.""" + cdef Dataset dense + cdef Dataset vpq = Dataset() + cdef cuvsResources_t res = resources.get_c_obj() + cdef cydlpack.DLManagedTensor* dataset_dlpack = NULL + + if isinstance(dataset, Dataset): + dense = dataset + else: + dataset_ai = wrap_array(dataset) + _check_input_array( + dataset_ai, + [np.dtype("float32"), np.dtype("float16"), + np.dtype("int8"), np.dtype("uint8")]) + dataset_dlpack = cydlpack.dlpack_c(dataset_ai) + dense = Dataset() + check_cuvs(cuvsDatasetMakeStandardView( + res, dataset_dlpack, &dense.dataset)) + + check_cuvs(cuvsDatasetMakeVpq( + res, params.params, dense.dataset, &vpq.dataset)) + return vpq From 28ca46b979dc237e8c220be0fbdf2d191591cea2 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 05:44:35 -0700 Subject: [PATCH 35/39] Reverse cuvsbench changes --- cpp/bench/ann/src/common/ann_types.hpp | 35 ------ cpp/bench/ann/src/common/benchmark.hpp | 144 +++++++++---------------- cpp/bench/ann/src/common/conf.hpp | 34 +----- cpp/bench/ann/src/common/dataset.hpp | 49 --------- 4 files changed, 51 insertions(+), 211 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index 6baee5b52f..bd669dff78 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -157,43 +157,8 @@ class algo : public algo_base { // and set_search_dataset() should save the passed-in pointer somewhere. // The client code should call set_search_dataset() before searching, // and should not release dataset before searching is finished. - // - // A compressed base set is never handed over this way, as it has no dense rows to pass, so - // needs_dataset() says nothing about one. An algorithm that cannot search the base set it was - // given with the parameters it was given has to reject them itself, from set_search_param(). virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) {}; - /* ### Base sets the benchmark cannot read ### - - Some algorithms build from a base set that has been compressed for them offline, which is - neither dense nor made of `T` values and so cannot be passed as `build`'s `const T*`. Such a - base set is handed over as a file path and the algorithm owns whatever it decodes. - - A path rather than a library type on purpose: this header is shared with the faiss, hnswlib and - diskann wrappers, and must not acquire their unrelated dependencies. - - Loading is separate from building because the benchmark times only `build_from_base_set_file`. - Deserializing a compressed base set is benchmark setup, the same as reading a dense one, and - folding it into the measured build would inflate build times by however long the file takes to - read. `set_base_set_file` is also called in search mode, before `load`, for algorithms whose - index file holds only part of the picture and needs the base set reattached. - */ - - /** - * Hand over a compressed base set as a file path. Returns the number of rows in it, which the - * benchmark has no way of reading for itself. Called outside the timed sections. - */ - virtual auto set_base_set_file(const std::string& /*file*/) -> size_t - { - throw std::runtime_error{"This algorithm cannot read a compressed base set from a file."}; - } - - /** Build the index from the base set handed over by `set_base_set_file`. */ - virtual void build_from_base_set_file() - { - throw std::runtime_error{"This algorithm cannot build from a compressed base set."}; - } - /** * Make a shallow copy of the algo wrapper that shares the resources and ensures thread-safe * access to them. */ diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 483ed524d9..a588b1e2a6 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -134,22 +134,8 @@ void bench_build(::benchmark::State& state, const auto algo_property = parse_algo_property(algo->get_preference(), index.build_param); - // Loading the base set is setup, not part of the build, so a compressed one is read here rather - // than inside the timed loop below; its row count comes back from the algorithm, since the - // benchmark cannot read the file. - const bool base_compressed = dataset->base_is_compressed(); - const T* base_set = nullptr; - std::size_t index_size = 0; - try { - if (base_compressed) { - index_size = algo->set_base_set_file(dataset->base_file()); - } else { - base_set = dataset->base_set(algo_property.dataset_memory_type); - index_size = dataset->base_set_size(); - } - } catch (const std::exception& e) { - return state.SkipWithError("Failed to load the base set: " + std::string(e.what())); - } + const T* base_set = dataset->base_set(algo_property.dataset_memory_type); + std::size_t index_size = dataset->base_set_size(); cuda_timer gpu_timer{algo}; { @@ -172,11 +158,7 @@ void bench_build(::benchmark::State& state, [[maybe_unused]] auto ntx_lap = nvtx.lap(); [[maybe_unused]] auto gpu_lap = gpu_timer.lap(!no_lap_sync); try { - if (base_compressed) { - algo->build_from_base_set_file(); - } else { - algo->build(base_set, index_size); - } + algo->build(base_set, index_size); } catch (const std::exception& e) { state.SkipWithError(std::string(e.what())); } @@ -256,11 +238,6 @@ void bench_search(::benchmark::State& state, auto ualgo = create_algo(index.algo, dataset->distance(), dataset->dim(), index.build_param); a = ualgo.get(); - // An index built from a compressed base set stores only its graph, so `load` alone would - // leave it with nothing to search over. Handing the file over first lets `load` attach the - // same rows the graph was built from and return a complete index. The row count it returns - // is of no use here; only the build reports that. - if (dataset->base_is_compressed()) { a->set_base_set_file(dataset->base_file()); } a->load(index_file); current_algo = std::move(ualgo); } @@ -273,13 +250,7 @@ void bench_search(::benchmark::State& state, current_algo_props = std::make_unique(std::move(parse_algo_property(a->get_preference(), sp_json))); - // Not a reliable signal for a compressed base set: cuvs_cagra answers true unconditionally, - // because its index file carries no dataset and the dense rows are re-attached here instead. - // There are no dense rows to attach for a compressed base, and the algorithm already has the - // file from `set_base_set_file` above. An algorithm that truly cannot search without the dense - // rows, such as CAGRA with refine_ratio > 1, has to reject that combination itself: only it - // knows which of its search parameters read them. - if (search_param->needs_dataset() && !dataset->base_is_compressed()) { + if (search_param->needs_dataset()) { try { a->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); @@ -569,7 +540,6 @@ void dispatch_benchmark(std::string cmdline, auto dataset = std::make_shared>(dataset_conf.name, base_file, - dataset_conf.base_compressed, dataset_conf.subset_first_row, dataset_conf.subset_size, query_file, @@ -582,13 +552,7 @@ void dispatch_benchmark(std::string cmdline, if (build_mode) { if (file_exists(base_file)) { log_info("Using the dataset file '%s'", base_file.c_str()); - if (dataset_conf.base_compressed) { - // The row count sits inside the compressed file, so it is reported per benchmark as - // `index_size` once the algorithm has read it, rather than up front here. - ::benchmark::AddCustomContext("base_format", "vpq"); - } else { - ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); - } + ::benchmark::AddCustomContext("n_records", std::to_string(dataset->base_set_size())); ::benchmark::AddCustomContext("dim", std::to_string(dataset->dim())); } else { log_warn("dataset file '%s' does not exist; benchmarking index building is impossible.", @@ -754,59 +718,51 @@ inline auto run_main(int argc, char** argv) -> int log_warn("cudart library is not found, GPU-based indices won't work."); } - // A rejected configuration reaches us as an exception, from the json parser or from the dataset - // itself. Reporting it here keeps that a legible error and a non-zero exit code, rather than an - // abort from an uncaught exception. - try { - auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); - std::string dtype = conf.get_dataset_conf().dtype; - - if (dtype == "float") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "half") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "uint8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else if (dtype == "int8") { - dispatch_benchmark(cmdline, - conf, - force_overwrite, - build_mode, - search_mode, - override_kv, - metric_objective, - threads, - no_lap_sync); - } else { - log_error("datatype '%s' is not supported", dtype.c_str()); - return -1; - } - } catch (const std::exception& e) { - log_error("%s", e.what()); + auto& conf = bench::configuration::initialize(conf_stream, data_prefix, index_prefix); + std::string dtype = conf.get_dataset_conf().dtype; + + if (dtype == "float") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "half") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "uint8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else if (dtype == "int8") { + dispatch_benchmark(cmdline, + conf, + force_overwrite, + build_mode, + search_mode, + override_kv, + metric_objective, + threads, + no_lap_sync); + } else { + log_error("datatype '%s' is not supported", dtype.c_str()); return -1; } diff --git a/cpp/bench/ann/src/common/conf.hpp b/cpp/bench/ann/src/common/conf.hpp index 0ed63d0ba6..afc7bc0a1f 100644 --- a/cpp/bench/ann/src/common/conf.hpp +++ b/cpp/bench/ann/src/common/conf.hpp @@ -8,19 +8,12 @@ #include #include -#include #include #include #include namespace cuvs::bench { -inline auto has_suffix(const std::string& str, const std::string& suffix) -> bool -{ - return str.size() >= suffix.size() && - str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; -} - class configuration { public: struct index { @@ -47,12 +40,6 @@ class configuration { std::string distance; std::optional groundtruth_neighbors_file{std::nullopt}; - // The base_file holds rows already compressed for the algorithm (a .vpq written by the offline - // VPQ compression tool) rather than dense vectors. The benchmark cannot read such a file: its - // rows are not `dtype` values, so they cannot travel through `algo::build`. The path is - // handed to the algorithm instead. Queries stay dense, and `dtype` keeps describing them. - bool base_compressed{false}; - // data type of input dataset, possible values ["float", "int8", "uint8"] std::string dtype; @@ -112,30 +99,11 @@ class configuration { } if (conf.contains("subset_size")) { dataset_conf_.subset_size = conf.at("subset_size"); } - // Decided separately from the dtype inference below, so that an explicit "dtype" does not stop - // us noticing that the base set is compressed. - if (conf.contains("base_format")) { - const auto base_format = conf.at("base_format").get(); - if (base_format == "vpq") { - dataset_conf_.base_compressed = true; - } else if (base_format != "dense") { - throw std::runtime_error("Unknown base_format '" + base_format + - "', expected \"vpq\" or \"dense\""); - } - } else { - dataset_conf_.base_compressed = has_suffix(dataset_conf_.base_file, ".vpq"); - } - if (conf.contains("dtype")) { dataset_conf_.dtype = conf.at("dtype"); } else { auto filename = dataset_conf_.base_file; - if (dataset_conf_.base_compressed) { - // A VPQ dataset stores its codebooks as half, but it is searched with float queries and - // yields a float index, so float is the type the benchmark instantiates. Keyed off the flag - // rather than the suffix, so that an explicit base_format also gets a dtype. - dataset_conf_.dtype = "float"; - } else if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { + if (filename.size() > 6 && filename.compare(filename.size() - 6, 6, "f16bin") == 0) { dataset_conf_.dtype = "half"; } else if (filename.size() > 9 && filename.compare(filename.size() - 9, 9, "fp16.fbin") == 0) { diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 93a193cad4..4dc43c343c 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -13,7 +13,6 @@ #include #include #include -#include #include #include @@ -153,10 +152,6 @@ struct dataset { private: std::string name_; std::string distance_; - std::string base_file_; - // A compressed base set is opaque to the benchmark: `base_set_` stays lazy and untouched, and the - // path is handed to the algorithm, which is the only thing able to decode it. - bool base_compressed_; blob base_set_; blob query_set_; std::optional> filter_bitset_; @@ -176,22 +171,9 @@ struct dataset { } } - // Reading a compressed base set as a .bin would not fail, it would succeed on garbage: the first - // eight bytes of a .vpq are a dtype prefix and a numpy magic, which parse as an absurd shape. - // Hence an explicit error, and callers that can proceed must ask `base_is_compressed()` first. - inline void throw_if_base_compressed(const char* what) const - { - if (base_compressed_) { - throw std::runtime_error{std::string{"dataset::"} + what + - "() is not available for the compressed base_file '" + base_file_ + - "'; only the algorithm can read that file."}; - } - } - public: dataset(std::string name, std::string base_file, - bool base_compressed, uint32_t subset_first_row, uint32_t subset_size, std::string query_file, @@ -200,27 +182,9 @@ struct dataset { std::optional filtering_rate = std::nullopt) : name_{std::move(name)}, distance_{std::move(distance)}, - base_file_{base_file}, - base_compressed_{base_compressed}, base_set_{base_file, subset_first_row, subset_size}, query_set_{query_file} { - if (base_compressed_) { - // Which rows went into the file was decided when it was written, so subsetting here is not - // merely unsupported: there is nothing left to select from. - if (subset_first_row != 0 || subset_size != 0) { - throw std::runtime_error{ - "A compressed base_file cannot be subset by the benchmark; choose the rows when " - "compressing it (the tool's --subset_first_row / --subset_size) and drop these keys."}; - } - // The bitset is sized from the base set row count, which is inside the compressed file. - if (filtering_rate.has_value()) { - throw std::runtime_error{ - "filtering_rate is not supported with a compressed base_file: generating the filter " - "bitset needs the base set size, which only the algorithm can read."}; - } - } - if (filtering_rate.has_value()) { // Generate a random bitset for filtering auto n_rows = static_cast(subset_size) + static_cast(subset_first_row); @@ -246,8 +210,6 @@ struct dataset { [[nodiscard]] auto name() const -> std::string { return name_; } [[nodiscard]] auto distance() const -> std::string { return distance_; } - [[nodiscard]] auto base_file() const -> std::string { return base_file_; } - [[nodiscard]] auto base_is_compressed() const -> bool { return base_compressed_; } [[nodiscard]] auto dim() const -> int { auto d = dim_.load(std::memory_order_relaxed); @@ -259,14 +221,6 @@ struct dataset { } catch (const std::runtime_error& e) { // Any exception raised above will re-raise next time we try to access the query set. query_set_.reset_lazy_state(); - // A compressed base set has no dense header to fall back on, and reading it as one would - // yield a nonsense dimension rather than an error. - if (base_compressed_) { - throw std::runtime_error{ - "Cannot determine the dataset dimension: the query set is not readable and the base set " - "is compressed. " + - std::string{e.what()}}; - } // If the query set is not accessible, use the base set. // Don't catch the exception here, because we have nothing else to do anyway. d = static_cast(base_set_.n_cols()); @@ -281,7 +235,6 @@ struct dataset { } [[nodiscard]] auto base_set_size() const -> size_t { - throw_if_base_compressed("base_set_size"); std::lock_guard lock(mutex_); auto r = base_set_.n_rows(); cache_dim(base_set_); @@ -319,7 +272,6 @@ struct dataset { [[nodiscard]] auto base_set() const -> const DataT* { - throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(); cache_dim(base_set_); @@ -329,7 +281,6 @@ struct dataset { HugePages request_hugepages_2mb = HugePages::kDisable) const -> const DataT* { - throw_if_base_compressed("base_set"); std::lock_guard lock(mutex_); auto* r = base_set_.data(memory_type, request_hugepages_2mb); cache_dim(base_set_); From 682ee360b5835c91cc61858ecf946a0627c39147 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 07:14:03 -0700 Subject: [PATCH 36/39] Share the iterative CAGRA build query chunk size as a constant --- cpp/src/neighbors/detail/cagra/cagra_build.cuh | 2 +- cpp/src/neighbors/detail/cagra/cagra_helpers.hpp | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 32786b4e1b..95fca2ccc0 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -2230,7 +2230,7 @@ auto iterative_build_graph(raft::resources const& res, RAFT_LOG_DEBUG("# initial graph size = %lu", (uint64_t)initial_graph_size); // Allocate memory for search results. - constexpr uint64_t max_chunk_size = 8192; + constexpr uint64_t max_chunk_size = helpers::kIterativeBuildChunkSize; // +1 because the search may return the query node itself as a neighbor; // this is consistent with the per-iteration curr_topk = next_graph_degree + 1 auto topk = intermediate_degree + 1; diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp index ee78930970..1e12fec983 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.hpp @@ -11,4 +11,7 @@ namespace cuvs::neighbors::cagra::helpers { /** The batch size for the CAGRA optimize stage. */ constexpr static size_t kOptimizeBatchSize = 256 * 1024; +/** The query chunk size (search batch size) of the iterative CAGRA graph build. */ +constexpr static size_t kIterativeBuildChunkSize = 8192; + } // namespace cuvs::neighbors::cagra::helpers From b3c551f761a4c84c49669b66e3aa0b00c85ade52 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 07:28:13 -0700 Subject: [PATCH 37/39] Don't count graph staging buffers in optimize_workspace_size when the graphs are already on the device --- cpp/include/cuvs/neighbors/cagra.hpp | 16 +++++++---- .../neighbors/detail/cagra/cagra_helpers.cpp | 28 +++++++++++++------ 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 0b4ae3f090..2a6b162d69 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -4771,14 +4771,20 @@ namespace helpers { * @param[in] intermediate_graph_degree degree of the input graph for the optimization process * @param[in] index_size * @param[in] mst_optimize whether to use MST optimization + * @param[in] device_resident_graphs whether the input and output graphs are already device + * resident. In that case `optimize` reads and writes them in place instead of staging + * them through device buffers, so those staging allocations are left out of the + * estimate. * @return tuple of [host_size, device_size, host_fixed_size, device_fixed_size] memory sizes in * bytes */ -std::tuple optimize_workspace_size(size_t n_rows, - size_t graph_degree, - size_t intermediate_degree, - size_t index_size, - bool mst_optimize = false); +std::tuple optimize_workspace_size( + size_t n_rows, + size_t graph_degree, + size_t intermediate_degree, + size_t index_size, + bool mst_optimize = false, + bool device_resident_graphs = false); /** * Calculate memory usage of CAGRA build. diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp index fa31bccbc3..c9b6b38103 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp @@ -35,7 +35,8 @@ std::tuple optimize_workspace_size(size_t n_rows size_t graph_degree, size_t intermediate_degree, size_t index_size, - bool mst_optimize) + bool mst_optimize, + bool device_resident_graphs) { RAFT_EXPECTS(graph_degree > 0, "graph_degree must be greater than 0"); RAFT_EXPECTS(intermediate_degree >= graph_degree, @@ -59,17 +60,25 @@ std::tuple optimize_workspace_size(size_t n_rows // Prune stage memory // We neglect 8 bytes (both on host and device) for stats - size_t prune_dev_fixed = batch_size * intermediate_degree; // detour count (uint8_t) - prune_dev_fixed += batch_size * sizeof(uint32_t); // d_num_detour_edges - prune_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) - - size_t prune_dev = n_rows * intermediate_degree * index_size; // d_input_graph + size_t prune_dev_fixed = batch_size * intermediate_degree; // detour count (uint8_t) + prune_dev_fixed += batch_size * sizeof(uint32_t); // d_num_detour_edges + + // Buffers that only exist to stage host-resident graphs to the device. When the caller already + // owns device graphs, batch_load_iterator passes them through and the kernels read and write + // them in place. + size_t prune_dev = 0; + if (!device_resident_graphs) { + prune_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + prune_dev += n_rows * intermediate_degree * index_size; // d_input_graph + } prune_dev += prune_dev_fixed; // Reverse graph stage memory size_t rev_dev = n_rows * graph_degree * index_size; // d_rev_graph rev_dev += n_rows * sizeof(uint32_t); // d_rev_graph_count - rev_dev += n_rows * index_size; // d_dest_nodes + if (!device_resident_graphs) { + rev_dev += n_rows * index_size; // d_dest_nodes + } // Memory for merging graphs (host only optional) size_t combine_host_fixed = graph_degree * sizeof(uint32_t); // histogram @@ -77,7 +86,10 @@ std::tuple optimize_workspace_size(size_t n_rows combine_host += combine_host_fixed; // additional memory for combine stage on device (3 batches) - size_t combine_dev_fixed = 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + size_t combine_dev_fixed = 0; + if (!device_resident_graphs) { + combine_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_output_graph(2*batch) + } if (mst_optimize) { combine_dev_fixed += 2 * batch_size * graph_degree * index_size; // d_mst_graph(2*batch) combine_dev_fixed += 2 * batch_size * sizeof(uint32_t); // d_mst_graph_num_edges(2*batch) From c030c573615f9e775d312b992624e5c11de2b695 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 07:33:16 -0700 Subject: [PATCH 38/39] Model the iterative CAGRA build's peak host and device memory, including the CAGRA-Q path --- cpp/include/cuvs/neighbors/cagra.hpp | 29 ++- .../neighbors/detail/cagra/cagra_helpers.cpp | 235 +++++++++++++++++- 2 files changed, 254 insertions(+), 10 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 2a6b162d69..06de5a6a83 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -4786,6 +4786,23 @@ std::tuple optimize_workspace_size( bool mst_optimize = false, bool device_resident_graphs = false); +/** + * Calculate the device memory footprint of a VPQ-compressed (CAGRA-Q) dataset. + * + * The footprint is the sum of the VQ codebook, the PQ codebook and the encoded rows. Parameters + * left at 0 are resolved with the same heuristics that `vpq_build` applies. + * + * @param[in] dataset shape of the uncompressed dataset + * @param[in] params VPQ compression parameters + * @param[in] codebook_element_size size in bytes of a codebook element (2 for the f16 codebooks + * used by CAGRA-Q) + * + * @return compressed dataset size in bytes + */ +size_t vpq_dataset_size(raft::matrix_extent dataset, + cuvs::neighbors::vpq_params params, + size_t codebook_element_size = 2); + /** * Calculate memory usage of CAGRA build. * @@ -4794,13 +4811,17 @@ std::tuple optimize_workspace_size( * @param[in] dtype element type of the dataset * (e.g. `CUDA_R_32F`, `CUDA_R_16F`, `CUDA_R_8I`, `CUDA_R_8U`) * @param[in] cparams CAGRA index building parameters + * @param[in] compression when set, the build consumes a VPQ-compressed (CAGRA-Q) dataset with + * these parameters rather than the dense dataset described by `dataset` and `dtype` * * @return pair of [host_size, device_size] memory sizes in bytes */ -std::pair cagra_build_mem_usage(raft::resources const& res, - raft::matrix_extent dataset, - cudaDataType_t dtype, - cuvs::neighbors::cagra::index_params cparams); +std::pair cagra_build_mem_usage( + raft::resources const& res, + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::cagra::index_params cparams, + std::optional compression = std::nullopt); /** * @brief Optimize a KNN graph into a CAGRA graph. diff --git a/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp index c9b6b38103..6c391e5a0b 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp +++ b/cpp/src/neighbors/detail/cagra/cagra_helpers.cpp @@ -6,10 +6,12 @@ #include "cagra_helpers.hpp" #include +#include #include #include #include #include +#include #include namespace cuvs::neighbors::cagra::helpers { @@ -113,6 +115,128 @@ std::tuple optimize_workspace_size(size_t n_rows return std::make_tuple(total_host, total_dev, total_host_fixed, total_dev_fixed); } +size_t vpq_dataset_size(raft::matrix_extent dataset, + cuvs::neighbors::vpq_params params, + size_t codebook_element_size) +{ + const size_t n_rows = dataset.extent(0); + const size_t dim = dataset.extent(1); + + // Mirror detail::fill_missing_params_heuristics for the fields that affect the footprint. + const size_t pq_bits = params.pq_bits == 0 ? 8 : params.pq_bits; + const size_t pq_dim = + params.pq_dim == 0 ? raft::div_rounding_up_safe(dim, size_t{4}) : params.pq_dim; + const size_t vq_n_centers = + params.vq_n_centers == 0 + ? raft::round_up_safe(static_cast(std::sqrt(static_cast(n_rows))), 8) + : params.vq_n_centers; + + const size_t pq_len = raft::div_rounding_up_safe(dim, pq_dim); + const size_t pq_n_centers = size_t{1} << pq_bits; + + // Every encoded row starts with its inlined VQ label (vpq_build always requests inline labels) + // and is followed by the bit-packed PQ codes. + using label_type = uint32_t; + constexpr size_t kLabelBits = 8 * sizeof(label_type); + const size_t encoded_row_length = + sizeof(label_type) * (1 + raft::div_rounding_up_safe(pq_dim * pq_bits, kLabelBits)); + + return vq_n_centers * dim * codebook_element_size // vq_code_book + + pq_n_centers * pq_len * codebook_element_size // pq_code_book + + n_rows * encoded_row_length; // encoded rows +} + +namespace { + +/** + * Device memory held by a single CAGRA search plan, in bytes. + * + * Mirrors `search_plan_impl_base` (algorithm selection), `search_plan_impl::adjust_search_params` + * and `search_plan_impl::calc_hashmap_params` together with the per-algorithm buffers of + * detail/cagra/search_single_cta.cuh and detail/cagra/search_multi_cta.cuh. + */ +size_t search_plan_mem_usage(cuvs::neighbors::cagra::search_params params, + size_t max_queries, + size_t itopk_size, + size_t graph_degree, + size_t dataset_size, + size_t index_size) +{ + // The iterative build always searches with max_queries = kIterativeBuildChunkSize, which is far + // above the occupancy threshold of the AUTO heuristic, so the algorithm is decided by itopk + // alone. + const bool multi_cta = + params.algo == cuvs::neighbors::cagra::search_algo::MULTI_CTA || + (params.algo == cuvs::neighbors::cagra::search_algo::AUTO && itopk_size > 512); + + constexpr size_t kMultiCtaItopkSize = 32; + + const size_t search_width = std::max(1, params.search_width); + + size_t max_iterations = params.max_iterations; + if (params.max_iterations == 0) { + max_iterations = multi_cta ? kMultiCtaItopkSize : itopk_size / search_width; + size_t reachable_nodes = 1; + while (reachable_nodes < dataset_size) { + reachable_nodes *= std::max(2, graph_degree / 2); + max_iterations += 1; + } + } + if (params.max_iterations < params.min_iterations) { max_iterations = params.min_iterations; } + max_iterations = std::max(max_iterations, params.max_iterations); + + // The internal topk is rounded up to a multiple of 32. + if (itopk_size % 32 != 0) { itopk_size += 32 - (itopk_size % 32); } + + // Smallest hash table that keeps the expected number of entries under the maximum fill rate. + auto hash_bitlen = [&](size_t min_bitlen, size_t expected_nodes) { + size_t bitlen = std::max(min_bitlen, params.hashmap_min_bitlen); + while (static_cast(expected_nodes) > + static_cast(size_t{1} << bitlen) * params.hashmap_max_fill_rate) { + bitlen += 1; + } + return bitlen; + }; + + // num_executed_iterations, allocated for every non-persistent plan. + size_t dev = max_queries * sizeof(uint32_t); + + if (!multi_cta) { + // AUTO and SMALL hash modes keep the visited-node table in shared memory, so nothing is + // allocated globally. AUTO falls back to a global table once the small one would exceed 8K + // entries. + if (params.hashmap_mode != cuvs::neighbors::cagra::hash_mode::HASH) { + const size_t small_bitlen = hash_bitlen(8, itopk_size + search_width * graph_degree); + if (small_bitlen <= 13) { return dev; } + } + const size_t bitlen = + hash_bitlen(11, itopk_size + search_width * graph_degree * max_iterations); + return dev + index_size * max_queries * (size_t{1} << bitlen); // hashmap + } + + // Multi-CTA keeps the per-CTA visited table in shared memory but shares the traversed-node + // table across the CTAs of a query, so that one is always global. + const size_t num_cta_per_query = + std::max(search_width, raft::div_rounding_up_safe(itopk_size, kMultiCtaItopkSize)); + const size_t num_intermediate = num_cta_per_query * kMultiCtaItopkSize; + const size_t bitlen = + hash_bitlen(11, num_cta_per_query * std::max(kMultiCtaItopkSize, max_iterations)); + + dev += index_size * max_queries * (size_t{1} << bitlen); // hashmap + // intermediate_indices / intermediate_distances + dev += num_intermediate * max_queries * (index_size + sizeof(float)); + // topk_workspace: one state byte per 8 candidates per thread of a 1024-thread block + // (_cuann_find_topk_bufferSize). + constexpr size_t kTopkThreads = 1024; + constexpr size_t kTopkStateBits = 8; + dev += raft::div_rounding_up_safe( + raft::div_rounding_up_safe(num_intermediate, kTopkThreads), kTopkStateBits) * + kTopkThreads * max_queries; + return dev; +} + +} // namespace + // All sizes are in bytes inline std::pair ivf_pq_build_mem_usage( raft::resources const& res, @@ -220,10 +344,97 @@ inline std::pair nn_descent_build_mem_usage(raft::resources cons return std::make_pair(total_host, total_dev); } -std::pair cagra_build_mem_usage(raft::resources const& res, - raft::matrix_extent dataset, - cudaDataType_t dtype, - cuvs::neighbors::cagra::index_params cparams) +// All sizes are in bytes +inline std::pair iterative_build_mem_usage( + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::graph_build_params::iterative_search_params params, + size_t graph_degree, + size_t intermediate_graph_degree, + bool guarantee_connectivity, + std::optional compression) +{ + // Mirrors detail::iterative_build_graph and detail::search_and_optimize. The build grows the + // graph by repeatedly searching the graph it has so far and optimizing the result; every + // allocation peaks on the final iteration, where the query set, the kNN graph and the output + // graph all span the whole dataset. + constexpr size_t kIndexSize = sizeof(uint32_t); // IdxT + + const size_t n_rows = dataset.extent(0); + const size_t dim = dataset.extent(1); + const size_t chunk = kIterativeBuildChunkSize; + const size_t dtype_size = cuda_data_type_size(dtype); + // The search may return the query node itself, hence the extra column. + const size_t topk = intermediate_graph_degree + 1; + + // The dataset stays resident on the device for the whole build: either VPQ-compressed, or padded + // to CAGRA's row alignment. + size_t dataset_dev; + size_t query_scratch; + if (compression.has_value()) { + dataset_dev = vpq_dataset_size(dataset, compression.value()); + // Queries are reconstructed from the codes one chunk at a time rather than materialized for + // the whole dataset. + query_scratch = chunk * dim * dtype_size; + } else { + const size_t stride = + cuvs::neighbors::cagra_required_row_width(static_cast(dim), dtype_size); + dataset_dev = n_rows * stride * dtype_size; + // Padded rows are depadded into a per-chunk scratch buffer before being used as queries. + query_scratch = stride == dim ? 0 : chunk * dim * dtype_size; + } + + // Search results for one chunk, live for the whole loop. + size_t results_dev = chunk * topk * kIndexSize; // dev_neighbors + results_dev += chunk * topk * sizeof(float); // dev_distances + + // The graph produced by the previous iteration is still alive while the current search fills the + // kNN graph, and the kNN graph is still alive while optimize writes the output graph. So one + // full-size graph and one full-size kNN graph always coexist. search_and_optimize releases the + // previous graph before allocating the output graph, so a third full-size buffer never appears. + const size_t graph_dev = n_rows * graph_degree * kIndexSize; // dev_graph / dev_output_graph + const size_t knn_dev = n_rows * topk * kIndexSize; // dev_knn_graph + + // The final iteration searches a graph_degree graph, requests topk neighbors and derives its + // internal topk from that. + cuvs::neighbors::cagra::search_params search_params = params; + search_params.max_queries = chunk; + const size_t search_dev = + search_plan_mem_usage(search_params, chunk, topk + 32, graph_degree, n_rows, kIndexSize); + + auto [host_workspace_size, gpu_workspace_size, host_ws_fixed, gpu_ws_fixed] = + optimize_workspace_size(n_rows, + graph_degree, + std::max(topk, graph_degree), + kIndexSize, + guarantee_connectivity, + /* device_resident_graphs = */ true); + + // Searching and optimizing run sequentially within an iteration, so the search plan and the + // optimize workspace never coexist and are combined with max() rather than summed. The query + // scratch is not part of that: search_and_optimize holds it at function scope, so it is still + // alive while optimize runs. + // + // Two transients are left out. The dataset copy made by make_device_padded_dataset briefly + // coexists with its source, and cagra::search re-pads a query chunk when its rows are not + // CAGRA-aligned; both are caller-owned or chunk-sized, and counting them would inflate the + // estimate enough to push callers to an out-of-core build unnecessarily. + size_t total_dev = dataset_dev + results_dev + graph_dev + knn_dev + query_scratch + + std::max(search_dev, gpu_workspace_size); + + // On the host the optimize workspace of the last iteration and the returned graph are also + // sequential: the workspace is released before the device graph is copied back. + size_t total_host = std::max(host_workspace_size, n_rows * graph_degree * kIndexSize); + + return std::make_pair(total_host, total_dev); +} + +std::pair cagra_build_mem_usage( + raft::resources const& res, + raft::matrix_extent dataset, + cudaDataType_t dtype, + cuvs::neighbors::cagra::index_params cparams, + std::optional compression) { using namespace cuvs::neighbors; @@ -249,9 +460,21 @@ std::pair cagra_build_mem_usage(raft::resources const& res, cparams.graph_degree, cparams.intermediate_graph_degree, cparams.guarantee_connectivity); + } else if (std::holds_alternative( + cparams.graph_build_params)) { + RAFT_LOG_INFO("Considering CAGRA in memory build with iterative CAGRA search"); + std::tie(total_host, total_dev) = iterative_build_mem_usage( + dataset, + dtype, + std::get(cparams.graph_build_params), + cparams.graph_degree, + cparams.intermediate_graph_degree, + cparams.guarantee_connectivity, + compression); } else { - // iterative build - // TODO(tfeher): proper estimate + // No graph build algorithm selected yet (std::monostate) or an out-of-core (ACE) build, whose + // requirements are modelled by check_ace_memory_requirements instead. Fall back to the size of + // the dataset plus the graphs. total_host = dataset.extent(0) * dataset.extent(1) * cuda_data_type_size(dtype) + dataset.extent(0) * (cparams.graph_degree + cparams.intermediate_graph_degree) * sizeof(uint32_t); From f4790ca2dba07d377d2a9618309c01e284038c9d Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 18 Aug 2026 08:06:16 -0700 Subject: [PATCH 39/39] added tests for memory estimation --- cpp/tests/CMakeLists.txt | 1 + .../neighbors/ann_cagra/test_mem_usage.cu | 272 ++++++++++++++++++ 2 files changed, 273 insertions(+) create mode 100644 cpp/tests/neighbors/ann_cagra/test_mem_usage.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c8e64a7ab6..377494946f 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -214,6 +214,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HELPERS_TEST PATH neighbors/ann_cagra/test_optimize_uint32_t.cu neighbors/ann_cagra/test_batch_load_iterator.cu + neighbors/ann_cagra/test_mem_usage.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/test_mem_usage.cu b/cpp/tests/neighbors/ann_cagra/test_mem_usage.cu new file mode 100644 index 0000000000..0d7223d553 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_mem_usage.cu @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/* + * Memory-usage estimators for the CAGRA build. + * + * These predict how much host and device memory a build needs without running it, so that callers + * such as hnsw::build and the ACE partition heuristic can choose between an in-memory and an + * out-of-core build. An estimate cannot be checked against itself, so each test here compares one + * against something derived independently: the VPQ footprint against a dataset that was actually + * compressed, and the iterative estimate against the buffers that provably coexist during the build + * and against the peak that raft::memory_stats_resources measures while a real build runs. + */ + +#include + +#include "../cagra_padded_build_helpers.cuh" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +namespace { + +using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; + +/** Size of the graph index type the CAGRA build uses internally. */ +constexpr size_t kIndexSize = sizeof(uint32_t); + +auto extents_of(int64_t n_rows, int64_t dim) { return raft::make_extents(n_rows, dim); } + +/** Device bytes actually owned by a compressed dataset: both codebooks plus the encoded rows. */ +auto compressed_bytes(const vpq_dataset_t& v) -> size_t +{ + return static_cast(v.vq_code_book.extent(0)) * v.vq_code_book.extent(1) * sizeof(half) + + static_cast(v.pq_code_book.extent(0)) * v.pq_code_book.extent(1) * sizeof(half) + + static_cast(v.data.extent(0)) * v.data.extent(1); +} + +auto make_clustered(const raft::resources& res, int64_t n_rows, int64_t dim) + -> raft::device_matrix +{ + auto dataset = raft::make_device_matrix(res, n_rows, dim); + auto labels = raft::make_device_vector(res, n_rows); + raft::random::make_blobs(res, + dataset.view(), + labels.view(), + 5, // clusters + std::nullopt, // random centers + std::nullopt, // scalar std + 1.0F, // cluster std + true, // shuffle + -10.0F, // center box min + 10.0F, // center box max + 1234ULL); + raft::resource::sync_stream(res); + return dataset; +} + +auto iterative_params(size_t graph_degree, size_t intermediate_graph_degree) -> index_params +{ + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = graph_degree; + params.intermediate_graph_degree = intermediate_graph_degree; + params.graph_build_params = graph_build_params::iterative_search_params(); + return params; +} + +auto device_estimate(const raft::resources& res, + int64_t n_rows, + int64_t dim, + const index_params& params, + std::optional compression = std::nullopt) -> size_t +{ + return helpers::cagra_build_mem_usage( + res, extents_of(n_rows, dim), CUDA_R_32F, params, compression) + .second; +} + +} // namespace + +// --------------------------------------------------------------------------- +// vpq_dataset_size: predicted footprint vs a dataset that was really compressed +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, VpqDatasetSizeMatchesCompressedDataset) +{ + raft::resources res; + constexpr int64_t n_rows = 512; + constexpr int64_t dim = 32; + + cuvs::neighbors::vpq_params params; + params.pq_bits = 8; + params.pq_dim = 8; // dim must be divisible by pq_dim + params.vq_n_centers = 32; + params.kmeans_n_iters = 2; + + auto dataset = make_clustered(res, n_rows, dim); + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + EXPECT_EQ(helpers::vpq_dataset_size(extents_of(n_rows, dim), params), + compressed_bytes(compressed)); +} + +TEST(CagraMemUsage, VpqDatasetSizeResolvesUnsetParams) +{ + raft::resources res; + constexpr int64_t n_rows = 1024; + constexpr int64_t dim = 32; + + // pq_dim and vq_n_centers left at 0, so both the estimator and the build have to derive them. + cuvs::neighbors::vpq_params params; + params.pq_bits = 8; + params.kmeans_n_iters = 2; + + auto dataset = make_clustered(res, n_rows, dim); + auto compressed = cuvs::preprocessing::quantize::pq::make_vpq_dataset( + res, params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + EXPECT_EQ(helpers::vpq_dataset_size(extents_of(n_rows, dim), params), + compressed_bytes(compressed)); +} + +// --------------------------------------------------------------------------- +// optimize_workspace_size: device-resident graphs need no staging buffers +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, OptimizeWorkspaceSkipsStagingForDeviceResidentGraphs) +{ + constexpr size_t n_rows = 1000000; + constexpr size_t graph_degree = 32; + constexpr size_t intermediate_degree = 256; + + auto [host_staged, dev_staged, host_fixed_staged, dev_fixed_staged] = + helpers::optimize_workspace_size( + n_rows, graph_degree, intermediate_degree, kIndexSize, false, false); + auto [host_resident, dev_resident, host_fixed_resident, dev_fixed_resident] = + helpers::optimize_workspace_size( + n_rows, graph_degree, intermediate_degree, kIndexSize, false, true); + + // The flag only removes device-side staging buffers, so the host estimate must not move. + EXPECT_EQ(host_resident, host_staged); + EXPECT_EQ(host_fixed_resident, host_fixed_staged); + + // At this shape the prune stage dominates the device total, and the d_input_graph staging copy + // it no longer needs is the largest single term the flag drops. + EXPECT_LT(dev_resident, dev_staged); + EXPECT_LE(dev_fixed_resident, dev_fixed_staged); +} + +TEST(CagraMemUsage, OptimizeWorkspaceDefaultsToStagedGraphs) +{ + constexpr size_t n_rows = 4096; + EXPECT_EQ(helpers::optimize_workspace_size(n_rows, 32, 64, kIndexSize), + helpers::optimize_workspace_size(n_rows, 32, 64, kIndexSize, false, false)); +} + +// --------------------------------------------------------------------------- +// cagra_build_mem_usage: the iterative graph build +// --------------------------------------------------------------------------- + +TEST(CagraMemUsage, IterativeEstimateCoversCoexistingBuffers) +{ + raft::resources res; + constexpr int64_t n_rows = 1000000; + constexpr int64_t dim = 128; // 16-byte aligned for float, so never padded + constexpr size_t graph_degree = 64; + constexpr size_t intermediate_degree = 128; + + const size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); + const size_t graph_bytes = static_cast(n_rows) * graph_degree * kIndexSize; + const size_t knn_bytes = static_cast(n_rows) * (intermediate_degree + 1) * kIndexSize; + + const auto estimated = + device_estimate(res, n_rows, dim, iterative_params(graph_degree, intermediate_degree)); + + // The resident dataset, the graph carried over from the previous iteration and the kNN graph the + // final search fills are all live at the same moment, so the estimate cannot be below their sum. + EXPECT_GE(estimated, dataset_bytes + graph_bytes + knn_bytes); +} + +TEST(CagraMemUsage, IterativeEstimateGrowsWithDatasetAndDegree) +{ + raft::resources res; + constexpr int64_t dim = 128; + const auto params = iterative_params(64, 128); + + EXPECT_GT(device_estimate(res, 2000000, dim, params), + device_estimate(res, 1000000, dim, params)); + EXPECT_GT(device_estimate(res, 1000000, 256, params), + device_estimate(res, 1000000, dim, params)); + EXPECT_GT(device_estimate(res, 1000000, dim, iterative_params(128, 256)), + device_estimate(res, 1000000, dim, params)); +} + +TEST(CagraMemUsage, IterativeEstimateIsSmallerForCompressedDataset) +{ + raft::resources res; + constexpr int64_t n_rows = 1000000; + constexpr int64_t dim = 128; + const auto params = iterative_params(64, 128); + + cuvs::neighbors::vpq_params compression; + compression.pq_bits = 8; + compression.pq_dim = 32; + compression.vq_n_centers = 1024; + + // CAGRA-Q keeps codes instead of dense rows, so the resident dataset term shrinks by an order of + // magnitude even though the per-chunk reconstruction scratch is new. + EXPECT_LT(device_estimate(res, n_rows, dim, params, compression), + device_estimate(res, n_rows, dim, params)); +} + +TEST(CagraMemUsage, IterativeEstimateAcceptsEqualDegrees) +{ + raft::resources res; + // graph_degree == intermediate_graph_degree once broke the iterative build (issue #1818). The + // estimator has to accept it rather than trip the degree check inside optimize_workspace_size. + EXPECT_GT(device_estimate(res, 100000, 64, iterative_params(16, 16)), 0U); +} + +TEST(CagraMemUsage, IterativeEstimateBoundsMeasuredPeak) +{ + raft::resources res; + constexpr int64_t n_rows = 20000; + constexpr int64_t dim = 64; + const auto params = iterative_params(32, 64); + + const auto estimated = device_estimate(res, n_rows, dim, params); + + size_t measured = 0; + { + // The tracking handle replaces the global device resource, so it has to outlive everything + // allocated below; declaration order here gives it exactly that. + raft::memory_stats_resources tracked{res}; + auto dataset = make_clustered(tracked, n_rows, dim); + cuvs::neighbors::test::padded_device_matrix_for_cagra padded( + tracked, raft::make_const_mdspan(dataset.view())); + auto built = cagra::build(tracked, params, padded.view); + raft::resource::sync_stream(tracked); + ASSERT_GT(built.size(), 0); + + const auto peak = tracked.get_bytes_peak(); + measured = peak.device_global + peak.device_workspace + peak.device_large_workspace + + peak.device_managed; + } + + // One-sided on purpose: hnsw::build uses this estimate to decide whether an in-memory build + // fits, so under-predicting is the failure that matters. No upper bound is asserted because the + // estimate also carries the whole workspace allowance of the resources handle. + EXPECT_GE(estimated, measured); +} + +} // namespace cuvs::neighbors::cagra