diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a6438c1a8d..9e49dce763 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -36,6 +36,7 @@ jobs: - wheel-tests-cuvs - devcontainer - telemetry-setup + - stress-rerun-anchor uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@main permissions: contents: read @@ -53,6 +54,14 @@ jobs: - name: Telemetry setup if: ${{ vars.TELEMETRY_ENABLED == 'true' }} uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main + stress-rerun-anchor: + # Stress testing only: every test job needs this job, so re-running it re-runs all + # test jobs against the existing build artifacts without re-running any builds. + runs-on: ubuntu-latest + permissions: {} + steps: + - name: Anchor for test job re-runs + run: echo "Re-running the test jobs" build-details: uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main check-nightly-ci: @@ -469,7 +478,7 @@ jobs: node_type: cpu16 script: ci/build_cpp.sh conda-cpp-tests: - needs: [conda-cpp-build, changed-files] + needs: [conda-cpp-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -478,7 +487,6 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp # Split the C++ tests across several jobs for each CUDA/GPU configuration strategy: fail-fast: false @@ -488,6 +496,7 @@ jobs: with: build_type: pull-request script: "ci/test_cpp.sh ${{ matrix.shard }} ${{ matrix.num_shards }}" + matrix_type: "pull-request,nightly" cache-paths: .cache/cuda-jit cache-key-prefix: libcuvs-cudajit-v1 # Caches are immutable. Re-key when the JIT-relevant sources change so a completed run can @@ -548,7 +557,7 @@ jobs: # Build a conda package for each CUDA x ARCH x minimum supported Python version matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) conda-python-tests: - needs: [conda-python-build, changed-files] + needs: [conda-python-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -557,10 +566,10 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda with: build_type: pull-request script: ci/test_python.sh + matrix_type: "pull-request,nightly" rocky8-clib-standalone-build: needs: checks permissions: @@ -573,7 +582,7 @@ jobs: with: build_type: pull-request rocky8-clib-tests: - needs: [rocky8-clib-standalone-build, changed-files] + needs: [rocky8-clib-standalone-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -581,7 +590,6 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_standalone_c with: build_type: pull-request container-image-repository: "rapidsai/citestwheel" @@ -595,7 +603,6 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. with: @@ -606,7 +613,7 @@ jobs: artifact-name-prefix: "cuvs-java" file_to_upload: "java/cuvs-java/target/" conda-java-tests: - needs: conda-java-build + needs: [conda-java-build, stress-rerun-anchor] permissions: actions: read contents: read @@ -637,7 +644,7 @@ jobs: artifact-name-prefix: "cuvs-lucene" file_to_upload: "java/cuvs-lucene/target/" conda-lucene-tests: - needs: conda-lucene-build + needs: [conda-lucene-build, stress-rerun-anchor] permissions: actions: read contents: read @@ -746,7 +753,7 @@ jobs: # Build a wheel for each CUDA x ARCH x minimum supported Python version matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) wheel-tests-cuvs: - needs: [wheel-build-cuvs, changed-files] + needs: [wheel-build-cuvs, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -755,10 +762,10 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels with: build_type: pull-request script: ci/test_wheel_cuvs.sh + matrix_type: "pull-request,nightly" devcontainer: permissions: actions: read diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 08009b731d..d6e67dbb4d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1472,6 +1472,7 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/detail/cagra/topk_for_cagra/topk.cu ${cuvs_cagra_search_cuda_inst_files} src/neighbors/detail/cagra/cagra_helpers.cpp + src/neighbors/detail/pointer_attributes.cpp src/neighbors/dynamic_batching.cu src/neighbors/composite/index.cu $<$:src/neighbors/cagra.cpp> diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index da5e5baa9d..2e202b5c89 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -66,13 +66,15 @@ inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search) threads_rec); } -/** Whether the GPU can dereference this pointer: device, managed or mapped host memory. */ +/** + * Whether the GPU can dereference this pointer: device, managed or mapped host memory, or pageable + * memory on systems where the device reads it through the host page tables (ATS). + */ inline auto is_device_accessible(const void* ptr) -> bool { cudaPointerAttributes attrs{}; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attrs, ptr)); - return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged || - attrs.devicePointer != nullptr; + return cuvs::neighbors::detail::is_device_accessible(attrs); } /** diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 0578dd0589..df82d51193 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -34,14 +34,23 @@ function(_cutile_detect) return() endif() + # bin2c is installed next to the real nvcc. In conda environments, CUDAToolkit_BIN_DIR is + # targets//bin, which only contains an nvcc symlink, so also search the symlink target's + # directory. + set(bin2c_dirs ${CUDAToolkit_BIN_DIR}) + if(CUDAToolkit_NVCC_EXECUTABLE) + get_filename_component(nvcc_real_path "${CUDAToolkit_NVCC_EXECUTABLE}" REALPATH) + get_filename_component(nvcc_real_dir "${nvcc_real_path}" DIRECTORY) + list(APPEND bin2c_dirs "${nvcc_real_dir}") + endif() find_program( CUVS_CUTILE_BIN2C NAMES bin2c - PATHS ${CUDAToolkit_BIN_DIR} + PATHS ${bin2c_dirs} NO_DEFAULT_PATH ) if(NOT CUVS_CUTILE_BIN2C) - message(STATUS "cuTile disabled: bin2c not found in ${CUDAToolkit_BIN_DIR}.") + message(STATUS "cuTile disabled: bin2c not found in ${bin2c_dirs}.") return() endif() diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index baa5b44da2..2b6b2aad14 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1094,6 +1094,18 @@ template namespace detail { +/** + * Whether memory described by `attr` may be dereferenced in device code on the current device. + * + * Device and managed memory are device-accessible, and registered (pinned) host memory is when it + * is mapped (`devicePointer` is set). Pageable (unregistered) host memory is device-accessible only + * when the device reads it through the host page tables (ATS). On systems with HMM, + * `cudaPointerGetAttributes` can also report a non-null `devicePointer` for pageable memory (for + * example after the first GPU access to pageable memory in the process), but device accesses to it + * are then serviced through GPU page faults, so it is treated as host-only there. + */ +[[nodiscard]] auto is_device_accessible(const cudaPointerAttributes& attr) -> bool; + template [[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) { @@ -1115,7 +1127,10 @@ template if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { return const_cast(src.data_handle()); } - auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + // Host memory needs a device mapping; pageable memory only counts when the device reads it + // through the host page tables (ATS), not through HMM page faults. + auto* device_ptr = + is_device_accessible(ptr_attrs) ? reinterpret_cast(ptr_attrs.devicePointer) : nullptr; RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); return device_ptr; } diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 332fed16b5..4478cec2a6 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -1082,7 +1082,8 @@ void kmeans_fit( compute_centroid_shift(handle, raft::make_const_mdspan(centroids_const), raft::make_const_mdspan(new_centroids_view), - sqrdNormError.view()); + sqrdNormError.view(), + ws); std::swap(cur_centroids_ptr, new_centroids_ptr); diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index bb96c9115b..881c0bd92a 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -43,6 +43,7 @@ #include #include +#include #include #include #include @@ -640,7 +641,7 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT average, MathT centroid_offset, IdxT seed, - IdxT* search_count, + IdxT stride, IdxT* update_count, MappingOpT mapping_op) { @@ -652,12 +653,17 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT i = n_rows; IdxT j = raft::laneId(); if (j == 0) { + // The candidate sequence of each receiver is derived from its index (not from a shared atomic + // counter), so the donor points do not depend on the order in which the warps run. + // `stride >= n_clusters`, so different receivers use different counters (and thus, before + // wrapping around, different candidates); `stride` is coprime with `n_rows`, so each receiver + // visits every row within `n_rows` attempts and finds a donor whenever one exists. IdxT attempt = 0; do { - auto old = atomicAdd(search_count, IdxT{1}); - auto candidate = - static_cast((static_cast(seed) * static_cast(old + 1)) % - static_cast(n_rows)); + auto counter = static_cast(attempt) * static_cast(stride) + + static_cast(receiver_cluster); + auto candidate = static_cast((static_cast(seed) * (counter + 1)) % + static_cast(n_rows)); if (static_cast(cluster_sizes[labels[candidate]]) >= average) { i = candidate; } ++attempt; } while (i >= n_rows && attempt < n_rows); @@ -768,6 +774,8 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL * max_cluster_size > average_size * balance_upper_tolerance * balance_upper_tolerance > 1 * @param[in] centroid_offset offset from the donor cluster centroid towards a donor point + * @param[in] donor_selection how the donor points are chosen (see above) + * @param[in] balancing_step index of the current EM iteration; selects the pseudo-random offset * @param[in] mapping_op Mapping operation from T to MathT * @param[inout] device_memory memory resource to use for temporary allocations * @@ -791,6 +799,7 @@ auto adjust_centers(const raft::resources& handle, MathT balance_upper_tolerance, MathT centroid_offset, cuvs::cluster::kmeans::balanced_donor_selection donor_selection, + uint32_t balancing_step, MappingOpT mapping_op, rmm::device_async_resource_ref device_memory) -> bool { @@ -802,7 +811,6 @@ auto adjust_centers(const raft::resources& handle, 601, 659, 733, 809, 863, 941, 1013, 1069, 1151, 1223, 1291, 1373, 1451, 1511, 1583, 1657, 1733, 1811, 1889, 1987, 2053, 2129, 2213, 2287, 2357, 2423, 2531, 2617, 2687, 2741}; - static IdxT i_primes = 0; auto average = static_cast(n_rows) / static_cast(n_clusters); auto lower_threshold = average * balance_lower_tolerance; @@ -837,10 +845,12 @@ auto adjust_centers(const raft::resources& handle, auto n_pairs = static_cast(host_receiver_clusters.size()); if (n_pairs == 0) { return false; } + // Rotate the pseudo-random offset with the balancing step rather than with a process-wide + // counter, so that the donor points do not depend on how many fits ran before in the process. IdxT ofst; + size_t i_primes = balancing_step; do { - i_primes = (i_primes + 1) % kPrimes.size(); - ofst = kPrimes[i_primes]; + ofst = kPrimes[i_primes++ % kPrimes.size()]; } while (n_rows % ofst == 0); rmm::device_uvector receiver_clusters(n_pairs, stream, device_memory); @@ -851,8 +861,11 @@ auto adjust_centers(const raft::resources& handle, update_count.set_value_to_zero_async(stream); if (donor_selection == cuvs::cluster::kmeans::balanced_donor_selection::Random) { - rmm::device_scalar search_count(stream, device_memory); - search_count.set_value_to_zero_async(stream); + // The smallest stride >= n_clusters that is coprime with n_rows (see the kernel). + auto stride = std::max(n_clusters, 1); + while (std::gcd(static_cast(stride), static_cast(n_rows)) != 1) { + ++stride; + } const dim3 grid_dim(raft::ceildiv(n_clusters, static_cast(kBlockDimY)), 1, 1); adjust_centers_random_donor_kernel <<>>(centers, @@ -866,7 +879,7 @@ auto adjust_centers(const raft::resources& handle, static_cast(n_rows / n_clusters), centroid_offset, ofst, - search_count.data(), + stride, update_count.data(), mapping_op); return update_count.value(stream) > 0; // NB: rmm scalar performs the sync @@ -979,6 +992,7 @@ void balancing_em_iters(const raft::resources& handle, balance_upper_tolerance, static_cast(params.centroid_offset), params.donor_selection, + iter, mapping_op, device_memory)) { if (balancing_counter++ >= balancing_pullback) { diff --git a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh new file mode 100644 index 0000000000..6248e00ae1 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh @@ -0,0 +1,416 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +/** + * Deterministic per-cluster (weighted) row sums for the k-means centroid updates. + * + * Accumulating the rows of each cluster with atomics (as raft::linalg::reduce_rows_by_key does) + * makes the floating-point summation order, and therefore the centroids, vary from run to run. + * Through the convergence check this can change the number of Lloyd iterations, so two identical + * fits can return noticeably different centroids. Here the summation order is a function of the + * input only: + * + * 1. Row indices are stably sorted by label (cub::DeviceRadixSort), so each cluster becomes a + * contiguous segment. + * 2. Each segment is cut into pieces of at most `kRowsPerPiece` rows. One thread block reduces a + * piece with a fixed assignment of rows to threads and a fixed-order combine. + * 3. One thread block per cluster adds up that cluster's piece sums in the same fixed way, and + * adds the result to the output. + * + * The order depends only on the labels and the sizes, not on scheduling or on the launch + * configuration. + * + * Why not a CCCL reduction: cub::DeviceSegmentedReduce (run_to_run) and + * cub::DeviceReduce::ReduceByKey (run_to_run) reduce one scalar per segment or run. A + * [n_rows x n_cols] matrix would need n_clusters * n_cols segments that read the gathered rows + * column by column (strided, uncoalesced), and DeviceSegmentedReduce gives each segment to one + * thread block, which serializes large clusters when n_clusters is small. The per-cluster weight + * sum alone could use them, but it comes for free here in the same pass over the rows. + */ +namespace cuvs::cluster::kmeans::detail::sums_by_label { + +constexpr int kBlockThreads = 256; +constexpr int kRowsPerPiece = 256; + +/** Number of threads along the columns: smallest power of two >= n_cols (at most kBlockThreads). */ +inline int column_threads(int64_t n_cols) +{ + int tc = 1; + while (tc < n_cols && tc < kBlockThreads) { + tc <<= 1; + } + return tc; +} + +/** + * Sums `values(i, j)` over i in [begin, end) for each column j in [0, n_cols), using a fixed + * assignment of items to threads and a fixed-order combine, so the result does not depend on + * scheduling. The result for column j is passed to `write(j, sum)` by a single thread. + */ +template +__device__ void fixed_order_column_sums( + IndexT begin, IndexT end, IndexT n_cols, int col_threads, ValueOp values, WriteOp write) +{ + __shared__ DataT smem[kBlockThreads]; + const int lanes = kBlockThreads / col_threads; + const int tc = threadIdx.x % col_threads; + const int lane = threadIdx.x / col_threads; + for (IndexT col0 = 0; col0 < n_cols; col0 += col_threads) { + const IndexT j = col0 + tc; + DataT acc = DataT{0}; + if (j < n_cols) { +#pragma unroll 4 + for (IndexT i = begin + lane; i < end; i += lanes) { + acc += values(i, j); + } + } + smem[threadIdx.x] = acc; + __syncthreads(); + if (lane == 0 && j < n_cols) { + DataT sum = smem[tc]; + for (int l = 1; l < lanes; ++l) { + sum += smem[l * col_threads + tc]; + } + write(j, sum); + } + __syncthreads(); + } +} + +template +RAFT_KERNEL init_sort_kernel(LabelsIterator labels, IndexT n_rows, uint32_t* keys, IndexT* rows) +{ + for (IndexT i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_rows; + i += static_cast(blockDim.x) * gridDim.x) { + keys[i] = static_cast(labels[i]); + rows[i] = i; + } +} + +/** Finds each cluster's segment of the sorted labels and the number of pieces it is cut into. */ +template +RAFT_KERNEL segment_kernel(const uint32_t* sorted_keys, + IndexT n_rows, + IndexT n_clusters, + IndexT* cluster_offsets, + IndexT* piece_counts) +{ + auto lower_bound = [=](IndexT c) { + IndexT lo = 0, hi = n_rows; + while (lo < hi) { + IndexT mid = lo + (hi - lo) / 2; + if (static_cast(sorted_keys[mid]) < c) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + for (IndexT c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; c <= n_clusters; + c += static_cast(blockDim.x) * gridDim.x) { + IndexT begin = lower_bound(c); + cluster_offsets[c] = begin; + if (c < n_clusters) { + IndexT end = lower_bound(c + 1); + piece_counts[c] = raft::ceildiv(end - begin, kRowsPerPiece); + } else { + piece_counts[c] = 0; + } + } +} + +template +RAFT_KERNEL piece_sums_kernel(DataIterator X, + IndexT n_cols, + const DataT* weights, + bool sum_weights, + const IndexT* sorted_rows, + const IndexT* cluster_offsets, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + DataT* piece_sums) +{ + const IndexT n_pieces = piece_offsets[n_clusters]; + for (IndexT piece = blockIdx.x; piece < n_pieces; piece += gridDim.x) { + // The cluster of this piece is the last c with piece_offsets[c] <= piece. + IndexT lo = 0, hi = n_clusters; + while (hi - lo > 1) { + IndexT mid = lo + (hi - lo) / 2; + if (piece_offsets[mid] <= piece) { + lo = mid; + } else { + hi = mid; + } + } + const IndexT c = lo; + const IndexT begin = cluster_offsets[c] + (piece - piece_offsets[c]) * kRowsPerPiece; + const IndexT end = std::min(begin + kRowsPerPiece, cluster_offsets[c + 1]); + DataT* out = piece_sums + static_cast(piece) * (n_cols + 1); + fixed_order_column_sums( + begin, + end, + n_cols, + col_threads, + [=](IndexT i, IndexT j) { + const IndexT row = sorted_rows[i]; + const auto x = static_cast(X[static_cast(row) * n_cols + j]); + return weights != nullptr ? x * weights[row] : x; + }, + [=](IndexT j, DataT sum) { out[j] = sum; }); + // The sum of weights is stored after the column sums. + if (sum_weights) { + fixed_order_column_sums( + begin, + end, + IndexT{1}, + 1, + [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, + [=](IndexT, DataT sum) { out[n_cols] = sum; }); + } + } +} + +template +RAFT_KERNEL cluster_sums_kernel(const DataT* piece_sums, + IndexT n_cols, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + bool reset_sums, + DataT* centroid_sums, + DataT* weight_per_cluster) +{ + for (IndexT c = blockIdx.x; c < n_clusters; c += gridDim.x) { + const IndexT p_begin = piece_offsets[c]; + const IndexT p_end = piece_offsets[c + 1]; + const IndexT stride = n_cols + 1; + DataT* out = centroid_sums + static_cast(c) * n_cols; + fixed_order_column_sums( + p_begin, + p_end, + n_cols, + col_threads, + [=](IndexT p, IndexT j) { return piece_sums[static_cast(p) * stride + j]; }, + [=](IndexT j, DataT sum) { out[j] = reset_sums ? sum : out[j] + sum; }); + if (weight_per_cluster != nullptr) { + fixed_order_column_sums( + p_begin, + p_end, + IndexT{1}, + 1, + [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, + [=](IndexT, DataT sum) { + weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; + }); + } + } +} + +inline size_t align_bytes(size_t bytes) { return raft::alignTo(bytes, 256); } + +/** Rows per pass of `weighted_sums_by_label`; bounds the workspace and fits 32-bit indices. */ +constexpr int64_t kMaxRowsPerPass = int64_t{1} << 22; + +/** One pass of `weighted_sums_by_label` over at most kMaxRowsPerPass rows. */ +template +void weighted_sums_by_label_pass(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + uint32_t n_rows, + uint32_t n_cols, + uint32_t n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + using IdxT = uint32_t; + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); + + int end_bit = 1; + while (end_bit < 32 && (uint64_t{1} << end_bit) < static_cast(n_clusters)) { + ++end_bit; + } + const IdxT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; + + // Workspace layout + size_t sort_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + sort_bytes, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + n_rows, + 0, + end_bit, + stream)); + size_t scan_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + scan_bytes, + static_cast(nullptr), + static_cast(nullptr), + n_clusters + 1, + stream)); + const size_t keys_bytes = align_bytes(sizeof(uint32_t) * n_rows); + const size_t rows_bytes = align_bytes(sizeof(IdxT) * n_rows); + const size_t offsets_bytes = align_bytes(sizeof(IdxT) * (n_clusters + 1)); + const size_t pieces_bytes = + align_bytes(sizeof(DataT) * static_cast(max_pieces) * (n_cols + 1)); + const size_t temp_bytes = align_bytes(std::max(sort_bytes, scan_bytes)); + const size_t total_bytes = + 2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes; + if (workspace.size() < total_bytes) { workspace.resize(total_bytes, stream); } + + char* ptr = workspace.data(); + auto take = [&ptr](size_t bytes) { + char* p = ptr; + ptr += bytes; + return p; + }; + auto* keys_in = reinterpret_cast(take(keys_bytes)); + auto* keys_out = reinterpret_cast(take(keys_bytes)); + auto* rows_in = reinterpret_cast(take(rows_bytes)); + auto* rows_out = reinterpret_cast(take(rows_bytes)); + auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_counts = reinterpret_cast(take(offsets_bytes)); + auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_sums = reinterpret_cast(take(pieces_bytes)); + void* temp = take(temp_bytes); + + if (n_rows > 0) { + const int grid = static_cast( + std::min(raft::ceildiv(n_rows, kBlockThreads), int64_t{65536})); + init_sort_kernel<<>>(labels, n_rows, keys_in, rows_in); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + temp, sort_bytes, keys_in, keys_out, rows_in, rows_out, n_rows, 0, end_bit, stream)); + } + { + const int grid = static_cast( + std::min(raft::ceildiv(n_clusters + 1, kBlockThreads), int64_t{65536})); + segment_kernel<<>>( + keys_out, n_rows, n_clusters, cluster_offsets, piece_counts); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + temp, scan_bytes, piece_counts, piece_offsets, n_clusters + 1, stream)); + + const int col_threads = column_threads(n_cols); + { + const int grid = static_cast(std::min(max_pieces, int64_t{65536})); + piece_sums_kernel<<>>(X, + n_cols, + weights, + weight_per_cluster != nullptr, + rows_out, + cluster_offsets, + piece_offsets, + n_clusters, + col_threads, + piece_sums); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + { + const int grid = static_cast(std::min(n_clusters, int64_t{65536})); + cluster_sums_kernel<<>>(piece_sums, + n_cols, + piece_offsets, + n_clusters, + col_threads, + reset_sums, + centroid_sums, + weight_per_cluster); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +/** + * @brief Deterministic (weighted) sum of the rows of X per label, and of the weights per label. + * + * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] + * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] + * + * The summation order depends only on the labels and the sizes, not on scheduling, + * so the result is bitwise reproducible. Rows are processed in passes of at most kMaxRowsPerPass + * rows, which bounds the workspace (16 bytes per row of a pass, plus the piece sums). + * + * @tparam DataT Accumulation and output type + * @tparam IndexT Index type + * @tparam DataIterator Random-access iterator over the row-major input; its values are converted + * to DataT (e.g. a transform iterator that maps int8 data to float) + * @tparam LabelsIterator Random-access iterator over the labels + * + * @param[in] handle RAFT resources handle + * @param[in] X Row-major input [n_rows x n_cols] + * @param[in] weights Per-row weights [n_rows], or nullptr for unit weights + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows of X + * @param[in] n_cols Number of columns of X + * @param[in] n_clusters Number of labels + * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] + * @param[inout] weight_per_cluster Output [n_clusters], or nullptr to skip the weight sums + * @param[inout] workspace Scratch buffer, resized as needed + * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise + */ +template +void weighted_sums_by_label(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + IndexT n_rows, + IndexT n_cols, + IndexT n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + if (n_clusters <= 0) { return; } + RAFT_EXPECTS(static_cast(n_clusters) < (uint64_t{1} << 31) && + static_cast(n_cols) < (uint64_t{1} << 31), + "weighted_sums_by_label: too many clusters or columns"); + // Passes are accumulated in a fixed order, so the result stays deterministic. There is always at + // least one pass, so that the outputs are reset even when there are no rows. + const auto total_rows = static_cast(n_rows); + int64_t start = 0; + do { + const int64_t pass_rows = std::min(kMaxRowsPerPass, total_rows - start); + weighted_sums_by_label_pass(handle, + X + start * static_cast(n_cols), + weights != nullptr ? weights + start : nullptr, + labels + start, + static_cast(pass_rows), + static_cast(n_cols), + static_cast(n_clusters), + centroid_sums, + weight_per_cluster, + workspace, + reset_sums && start == 0); + start += pass_rows; + } while (start < total_rows); +} + +} // namespace cuvs::cluster::kmeans::detail::sums_by_label diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 78edf8759e..efedd05b1b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -6,6 +6,7 @@ #include "../../distance/distance.cuh" #include "../../distance/top_1_nn.cuh" +#include "kmeans_centroid_sums.cuh" #include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +44,13 @@ #include #include #include +#include +#include #include +#include +#include +#include +#include #include #include @@ -121,6 +129,85 @@ struct WeightedDistanceOp { } }; +/** + * @brief A `cuda::mr::resource` that hands out temporary storage from one caller-owned buffer. + * + * The buffer only grows, and deallocation is a no-op, so repeated calls reuse the same memory + * instead of allocating (and possibly synchronizing) every time. Only one allocation may be live + * at a time. + */ +class workspace_buffer_resource { + public: + explicit workspace_buffer_resource(rmm::device_uvector& buffer) : buffer_{&buffer} {} + + void* allocate(cuda::stream_ref stream, size_t bytes, size_t alignment) + { + const size_t needed = bytes + alignment; + if (buffer_->size() < needed) { buffer_->resize(needed, stream); } + const auto ptr = reinterpret_cast(buffer_->data()); + return reinterpret_cast(raft::alignTo(ptr, alignment)); + } + void deallocate(cuda::stream_ref, void*, size_t, size_t) noexcept {} + void* allocate_sync(size_t bytes, size_t alignment) + { + return allocate(buffer_->stream(), bytes, alignment); + } + void deallocate_sync(void*, size_t, size_t) noexcept {} + bool operator==(const workspace_buffer_resource& other) const noexcept + { + return buffer_ == other.buffer_; + } + bool operator!=(const workspace_buffer_resource& other) const noexcept + { + return !(*this == other); + } + + private: + rmm::device_uvector* buffer_; +}; + +/** + * @brief Device-wide reduction that is bitwise reproducible from run to run. + * + * Requests `cuda::execution::determinism::run_to_run` explicitly (also cub's current default): + * cub combines the partial results in a fixed order on a given GPU, with no atomics. This is + * stronger than raft's mapThenSumReduce, which combines block results with atomics. + * + * `gpu_to_gpu` (an RFA-based sum that is identical across architectures) is not requested: the + * inputs of these reductions (distances from GEMM / cuTile TF32 kernels and the centroids derived + * from them) already differ between architectures, so a fit cannot be reproducible across GPUs. + * + * Temporary storage comes from `workspace`, which is resized as needed. + */ +template +void run_to_run_reduce(raft::resources const& handle, + InputIterator input, + IndexT size, + OutputT* output, + ReductionOpT reduction_op, + OutputT init, + rmm::device_uvector& workspace) +{ + auto env = + cuda::std::execution::env{cuda::stream_ref{raft::resource::get_cuda_stream(handle)}, + cuda::execution::require(cuda::execution::determinism::run_to_run), + cuda::std::execution::prop{cuda::mr::get_memory_resource, + workspace_buffer_resource{workspace}}}; + RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(input, output, size, reduction_op, init, env)); +} + +template +struct SquaredDiffOp { + const DataT* a; + const DataT* b; + + __device__ DataT operator()(IndexT offset) const + { + DataT diff = a[offset] - b[offset]; + return diff * diff; + } +}; + template void copyClusterLabels(raft::resources const& handle, const MinClusterAndDistanceResult& result, @@ -194,8 +281,14 @@ void weightSum( DataT wt_sum_h = DataT{0}; if constexpr (raft::is_device_mdspan_v) { - raft::linalg::mapThenSumReduce( - d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream.get(), weight.data_handle()); + rmm::device_uvector workspace(0, stream); + run_to_run_reduce(handle, + weight.data_handle(), + n_samples, + d_wt_sum.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -270,29 +363,9 @@ void computeClusterCost(raft::resources const& handle, MainOpT main_op, ReductionOpT reduction_op) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuda::transform_iterator itr(input, main_op); - - size_t temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); - - workspace.resize(temp_storage_bytes, stream); - - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); + run_to_run_reduce( + handle, itr, size, clusterCost.data_handle(), reduction_op, OutputT(), workspace); } template & workspace, bool reset_sums = true) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - auto n_samples = X.extent(0); - - workspace.resize(n_samples, stream); - - raft::linalg::reduce_rows_by_key(X.data_handle(), - X.extent(1), - cluster_labels, - sample_weights.data_handle(), - workspace.data(), - X.extent(0), - X.extent(1), - n_clusters, - centroid_sums.data_handle(), - stream, - reset_sums); - - raft::linalg::reduce_cols_by_key(sample_weights.data_handle(), - cluster_labels, - weight_per_cluster.data_handle(), - static_cast(1), - static_cast(n_samples), - n_clusters, - stream, - reset_sums); + sums_by_label::weighted_sums_by_label(handle, + X.data_handle(), + sample_weights.data_handle(), + cluster_labels, + X.extent(0), + X.extent(1), + n_clusters, + centroid_sums.data_handle(), + weight_per_cluster.data_handle(), + workspace, + reset_sums); } /** * @brief Finalize centroids by dividing accumulated sums by counts. @@ -707,21 +769,26 @@ void finalize_centroids(raft::resources const& handle, * @brief Compute the squared norm difference between two centroid sets. * * Writes sum((old_centroids - new_centroids)^2) into @p sqrd_norm_out. - * Used for convergence checking. Fully asynchronous — no stream sync. + * Used for convergence checking. Fully asynchronous — no stream sync. Temporary storage comes + * from @p workspace, which is resized as needed. */ template void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view old_centroids, raft::device_matrix_view new_centroids, - raft::device_scalar_view sqrd_norm_out) + raft::device_scalar_view sqrd_norm_out, + rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - raft::linalg::mapThenSumReduce(sqrd_norm_out.data_handle(), - old_centroids.size(), - raft::sqdiff_op{}, - stream, - old_centroids.data_handle(), - new_centroids.data_handle()); + run_to_run_reduce( + handle, + cuda::transform_iterator( + cuda::counting_iterator(0), + SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}), + static_cast(old_centroids.size()), + sqrd_norm_out.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); } /** diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index a86c30a087..4b25221f63 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -605,7 +605,8 @@ void mnmg_fit( dev_res, rank_centroids_const, raft::make_const_mdspan(new_centroids.view()), - sqrd_norm_error_dev.view()); + sqrd_norm_error_dev.view(), + workspace); raft::copy( rank_centroids.data_handle(), new_centroids.data_handle(), n_clusters * n_features, stream); diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index c0433c676a..d059836ed1 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,23 @@ __device__ constexpr AccT get_clamp_precision() } } +/** + * Whether an expanded L2 distance is round-off error of a self-distance and should be clamped to + * zero. Self-neighboring points have (aNorm == bNorm) == accVal, but the dot product (accVal) can + * have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. + * + * The round-off error scales with the norms, so the tolerance is relative to them for squared + * norms below one. Otherwise, distinct points with equal norms and a small magnitude (for example, + * a spectral embedding) would have their distances incorrectly clamped to zero. + */ +template +DI bool is_self_distance_round_off(AccT dist, AccT aNorm, AccT bNorm) +{ + const AccT norm_product = aNorm * bNorm; + const AccT norm_scale = norm_product < AccT(1) ? norm_product : AccT(1); + return (aNorm == bNorm) && (dist * dist < get_clamp_precision() * norm_scale); +} + // Epilogue operator for CUTLASS based kernel template struct l2_exp_cutlass_op { @@ -38,13 +55,7 @@ struct l2_exp_cutlass_op { inline __device__ AccT operator()(AccT aNorm, AccT bNorm, AccT accVal) const noexcept { AccT outVal = aNorm + bNorm - AccT(2.0) * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product (accVal) - * can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. - */ - outVal = - outVal * AccT(!((outVal * outVal < get_clamp_precision()) * (aNorm == bNorm))); + outVal = outVal * AccT(!is_self_distance_round_off(outVal, aNorm, bNorm)); return sqrt ? raft::sqrt(outVal * static_cast(outVal > AccT(0))) : outVal; } @@ -112,15 +123,9 @@ struct l2_exp_distance_op { for (int j = 0; j < Policy::AccColsPerTh; ++j) { AccT accVal = acc[i][j]; AccT val = regxn[i] + regyn[j] - (AccT)2.0 * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product - * (accVal) can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal - * instead. - */ - acc[i][j] = val * static_cast((val > AccT(0))) * - static_cast( - !((val * val < get_clamp_precision()) * (regxn[i] == regyn[j]))); + acc[i][j] = + val * static_cast((val > AccT(0))) * + static_cast(!is_self_distance_round_off(val, regxn[i], regyn[j])); } } if (sqrt) { diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 0e1454d987..c781d1e427 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -606,10 +606,38 @@ void search( cuvs::neighbors::filtering::none_sample_filter>( res, params, indices, queries, partition_ids, neighbors, distances, partition_bitsets); } else { + // As in the single-partition bitset path, derive the filtering rate from the bitsets when the + // caller did not set one, so that the search plan enlarges itopk_size for selective filters. + // Without this, a selective filter leaves too few valid candidates and the search can return + // fewer than k neighbors. One plan serves every partition, so size it for the most selective + // partition that still accepts at least one row; a partition that accepts nothing contributes + // no neighbors whatever the itopk_size, and an unfiltered partition has a rate of zero. + search_params params_copy = params; + if (params.filtering_rate < 0.0) { + float filtering_rate = 0.0f; + for (size_t i = 0; i < partition_bitsets.size() && i < indices.size(); i++) { + const auto& v = partition_bitsets[i]; + if (v.data() == nullptr || v.size() == 0) { continue; } + const int64_t n_rows = indices[i]->dataset().n_rows(); + if (n_rows == 0) { continue; } + // Count only this partition's rows; the bitset may be padded past n_rows. + const auto num_set_bits = + cuvs::core::bitset_view(const_cast(v.data()), + std::min(n_rows, v.size())) + .count(res); + if (num_set_bits == 0) { continue; } + filtering_rate = + std::max(filtering_rate, static_cast(n_rows - num_set_bits) / n_rows); + } + const float min_filtering_rate = 0.0; + const float max_filtering_rate = 0.999; + params_copy.filtering_rate = + std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate); + } using bitset_filter_t = cuvs::neighbors::filtering::bitset_filter; cagra::detail::search_multi_partition( res, - params, + params_copy, indices, queries, partition_ids, diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index deddc47dcb..47ea5278e8 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include #include @@ -60,7 +61,8 @@ struct pointer_residency_count { auto [on_device, on_host] = pointer_residency_count::run(ptrs...); cudaPointerAttributes attr; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); - if (attr.devicePointer || attr.type == cudaMemoryTypeDevice) { ++on_device; } + // Not `attr.devicePointer != nullptr`: with HMM it can be set for pageable host memory. + if (cuvs::neighbors::detail::is_device_accessible(attr)) { ++on_device; } if (attr.hostPointer || attr.type == cudaMemoryTypeUnregistered) { ++on_host; } return std::make_tuple(on_device, on_host); } diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 9a22676334..7d2cdb7517 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -164,10 +164,15 @@ class ace_disk_workspace { template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { + // Every row needs at least one neighbor other than itself. With a single row both degrees would + // drop to zero, and the kNN graph build kernels would then access memory out of bounds. + RAFT_EXPECTS(dataset_size >= 2, + "CAGRA graph build requires a dataset with at least 2 rows, got %zu", + dataset_size); if (intermediate_degree >= static_cast(dataset_size)) { RAFT_LOG_WARN( "Intermediate graph degree cannot be larger than dataset size, reducing it to %lu", - dataset_size); + dataset_size - 1); intermediate_degree = dataset_size - 1; } if (intermediate_degree < graph_degree) { diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index c1d0033a01..99bd6fc78d 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -253,10 +253,13 @@ struct dataset_descriptor_host { { if (!ready.load(std::memory_order_acquire)) { eval(stream); } // value is immutable at this point. - auto& [ptr, ready_stream] = std::get(value); - if (ready_stream != stream) { - RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); - } + auto* ptr = std::get<0>(std::get(value)); + // Always order `stream` after the initialization kernel. Comparing the stream handles is not + // enough: the per-thread default stream (`cudaStreamPerThread`, which is also raft's default + // stream) has the same handle in every host thread but refers to a different stream in each. + // When several threads share a raft::resources (and thus this cached descriptor), skipping + // the wait lets a search kernel read the descriptor before it is initialized. + RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); return ptr; } }; diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index ace28b7842..3107f841ae 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -16,6 +16,8 @@ #include +#include + namespace cuvs::neighbors::cagra::detail { template , dataset_descriptor_host> value{kDefaultSize}; + /** Serializes lookup-or-insert, so that concurrent searches share a single descriptor. */ + std::mutex mutex; }; } // namespace descriptor_cache @@ -170,9 +174,15 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, -> dataset_descriptor_host { auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms); - auto& cache = - raft::resource::get_custom_resource>(res) - ->value; + auto* store = + raft::resource::get_custom_resource>(res); + auto& cache = store->value; + // Hold the lock across the lookup and the insertion. Otherwise threads that miss concurrently + // each create a descriptor and the later insertion replaces the earlier one, which may already + // be in use by another thread; the last owner then frees it on its own stream while a kernel on + // a different stream can still be reading it. Only host-side work happens under the lock (the + // device-side initialization is deferred to the first `dev_ptr()` call). + std::lock_guard lock(store->mutex); dataset_descriptor_host desc; if (!cache.get(key, &desc)) { desc = diff --git a/cpp/src/neighbors/detail/pointer_attributes.cpp b/cpp/src/neighbors/detail/pointer_attributes.cpp new file mode 100644 index 0000000000..29e213b6c8 --- /dev/null +++ b/cpp/src/neighbors/detail/pointer_attributes.cpp @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include + +#include +#include + +namespace cuvs::neighbors::detail { + +namespace { + +/** + * Whether `device` reads pageable (unregistered) host memory through the host page tables (ATS, + * e.g. Grace Hopper / Grace Blackwell). The result is cached per device. + */ +auto pageable_memory_uses_host_page_tables(int device) -> bool +{ + static std::mutex mutex; + static std::unordered_map cache; + std::lock_guard lock(mutex); + if (auto it = cache.find(device); it != cache.end()) { return it->second; } + int uses_host_page_tables = 0; + RAFT_CUDA_TRY(cudaDeviceGetAttribute( + &uses_host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, device)); + cache.emplace(device, uses_host_page_tables != 0); + return uses_host_page_tables != 0; +} + +} // namespace + +auto is_device_accessible(const cudaPointerAttributes& attr) -> bool +{ + switch (attr.type) { + case cudaMemoryTypeDevice: + case cudaMemoryTypeManaged: return true; + case cudaMemoryTypeHost: return attr.devicePointer != nullptr; + default: { + if (attr.devicePointer == nullptr) { return false; } + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + return pageable_memory_uses_host_page_tables(device); + } + } +} + +} // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh index 6117935712..6bb3b2ffe8 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh @@ -171,7 +171,11 @@ void search_impl(raft::resources const& handle, [=] __device__(const uint32_t idx, const float dist) { const auto query = idx / n_lists; const auto cluster = idx % n_lists; - return dist / (q_norm_ptr[query] * index_center_norm_ptr[cluster]); + // A list whose k-means cluster ended up empty has a zero center. Define the cosine + // distance to a zero-norm vector as 1 (similarity 0), as IVF-SQ does: a NaN here is never + // selected, and with fewer finite lists than n_probes the probes get padded with list 0. + const float denom = q_norm_ptr[query] * index_center_norm_ptr[cluster]; + return denom > 0.0f ? dist / denom : 0.0f; }, raft::make_const_mdspan(distance_buffer_dev_view)); } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 85024d2de2..8f6cdccd61 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -58,6 +58,8 @@ #include #include #include +#include +#include #include #include @@ -205,17 +207,6 @@ void flat_compute_residuals( stream.get()); } -template -__launch_bounds__(BlockDim) static __global__ void fill_indices_kernel(IdxT n_rows, - IdxT* data_indices, - IdxT* data_offsets, - const uint32_t* labels) -{ - const auto i = IdxT(BlockDim) * IdxT(blockIdx.x) + IdxT(threadIdx.x); - if (i >= n_rows) { return; } - data_indices[atomicAdd(data_offsets + labels[i], 1)] = i; -} - /** * @brief Calculate cluster offsets and arrange data indices into clusters. * @@ -249,13 +240,12 @@ auto calculate_offsets_and_indices(IdxT n_rows, stream.sync(); RAFT_EXPECTS(cumsum == n_rows, "cluster sizes do not add up."); RAFT_LOG_DEBUG("Max cluster size %d", max_cluster_size); - rmm::device_uvector data_offsets_buf(n_lists, stream); - auto data_offsets = data_offsets_buf.data(); - raft::copy(data_offsets, cluster_offsets, n_lists, stream); - constexpr uint32_t n_threads = 128; // NOLINT - const IdxT n_blocks = raft::div_rounding_up_unsafe(n_rows, n_threads); - fill_indices_kernel - <<>>(n_rows, data_indices, data_offsets, labels); + // Group the row indices by cluster with a stable sort, so that the order of the rows within a + // cluster (which decides the per-cluster training subsets) is the same in every build. + rmm::device_uvector sorted_labels(n_rows, stream); + raft::copy(sorted_labels.data(), labels, n_rows, stream); + thrust::sequence(exec_policy, data_indices, data_indices + n_rows); + thrust::stable_sort_by_key(exec_policy, sorted_labels.begin(), sorted_labels.end(), data_indices); return max_cluster_size; } 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 e27d34557a..b3383ebeca 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -19,9 +19,11 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -719,38 +721,44 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, cluster_idx++; } - if (batch_vectors == 0) break; // No more clusters to process + // Every cluster fits into a batch (checked above), so each batch makes progress. + RAFT_EXPECTS(cluster_idx > batch_start_cluster, "IVF-RaBitQ streaming batch is empty"); // ------------------------- // 11. Gather and transfer batch data to GPU using reordered PIDs // ------------------------- - // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) - raft::copy(batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); - raft::resource::sync_stream(handle_); + // A batch that only contains empty clusters has no data to transfer, but the clusters still + // need their rotated centroids below. + if (batch_vectors > 0) { + // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) + raft::copy( + batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); + raft::resource::sync_stream(handle_); - // Gather data on host using the reordered PIDs (reuse preallocated buffer) - // Use OpenMP parallel region to leverage persistent thread pool across batches + // Gather data on host using the reordered PIDs (reuse preallocated buffer) + // Use OpenMP parallel region to leverage persistent thread pool across batches #pragma omp parallel num_threads(num_threads) - { - size_t tid = omp_get_thread_num(); - size_t chunk_size = batch_vectors / num_threads; - size_t start = tid * chunk_size; - size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; - - for (size_t i = start; i < end; ++i) { - std::memcpy(&batch_data.data_handle()[i * num_dimensions], - &host_data[batch_pids.data_handle()[i] * num_dimensions], - num_dimensions * sizeof(float)); + { + size_t tid = omp_get_thread_num(); + size_t chunk_size = batch_vectors / num_threads; + size_t start = tid * chunk_size; + size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; + + for (size_t i = start; i < end; ++i) { + std::memcpy(&batch_data.data_handle()[i * num_dimensions], + &host_data[batch_pids.data_handle()[i] * num_dimensions], + num_dimensions * sizeof(float)); + } } - } - // OpenMP implicit barrier ensures all threads complete before continuing + // OpenMP implicit barrier ensures all threads complete before continuing - // Transfer batch to GPU (reuse preallocated buffer) - raft::copy(d_batch_data.data_handle(), - batch_data.data_handle(), - batch_vectors * num_dimensions, - stream_); + // Transfer batch to GPU (reuse preallocated buffer) + raft::copy(d_batch_data.data_handle(), + batch_data.data_handle(), + batch_vectors * num_dimensions, + stream_); + } // ------------------------- // 12. Quantize each cluster in the batch @@ -758,8 +766,8 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, size_t batch_offset = 0; for (size_t c = batch_start_cluster; c < cluster_idx; ++c) { + // Empty clusters are quantized too, so that their rotated centroids are initialized. size_t cluster_size = h_cluster_meta[c].num; - if (cluster_size == 0) continue; const float* cur_centroid = device_centroids + c * num_dimensions; float* cur_rotated_c = d_rotated_centroids.data_handle() + c * num_padded_dim; @@ -1016,6 +1024,14 @@ void IVFGPU::PrepareClusterSearchInputs( &beta, searcher.get_centroid_distances(), num_centroids); + // Rounding in the expanded form can make ||q - c||^2 slightly negative when the query is on (or + // very near) a centroid. The search kernels take sqrt(||q - c||^2) for the error bound, which is + // NaN for a negative value, and a NaN lower bound drops every vector of that list, so clamp to 0. + raft::linalg::unaryOp(searcher.get_centroid_distances(), + searcher.get_centroid_distances(), + batch_size * num_centroids, + raft::plug_const_op{0.f}, + searcher_stream.get()); // Step 4: select top-nprobe clusters per query auto d_raft_vals = raft::make_device_matrix(searcher_handle, batch_size, nprobe); 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 b35ac9fab9..5eef2179d4 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -520,6 +520,9 @@ void data_transformation_batch_opt(const float* d_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -637,6 +640,9 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), @@ -714,6 +720,9 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -758,6 +767,9 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..b99acb50fd 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -5,6 +5,7 @@ #include "./detail/pq.cuh" +#include #include #include @@ -88,7 +89,9 @@ auto train_from_rows(raft::resources const& res, { cudaPointerAttributes ptr_attrs; RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src_ptr)); - auto const* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + auto const* device_ptr = cuvs::neighbors::detail::is_device_accessible(ptr_attrs) + ? reinterpret_cast(ptr_attrs.devicePointer) + : nullptr; if (device_ptr == nullptr) { // A host mdspan makes training subsample the rows and encoding stream them in bounded batches, // so the dense dataset is never staged on the device. diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index eee4150eed..66fff7c5c8 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -195,7 +195,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_IVF_RABITQ_TEST - PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu + PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu neighbors/ann_ivf_rabitq/test_exact_match.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 23fc248917..a6e5b4c0f6 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -23,6 +23,7 @@ #include +#include #include #include @@ -495,6 +496,82 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(handle, n_clusters, n_features); + raft::random::RngState rng(p.rng_state.seed); + raft::random::uniform(handle, rng, d_init.data_handle(), n_clusters * n_features, T(-1), T(1)); + + auto d_sw = d_sw_view(); + std::optional> h_sw = std::nullopt; + auto h_sample_weight = raft::make_host_vector(testparams.weighted ? n_samples : 0); + if (testparams.weighted) { + std::fill_n(h_sample_weight.data_handle(), n_samples, T(1)); + h_sw = std::make_optional(raft::make_const_mdspan(h_sample_weight.view())); + } + + constexpr int n_runs = 3; + for (int host : {0, 1}) { + std::vector> centroids(n_runs, std::vector(n_clusters * n_features)); + std::vector inertias(n_runs); + std::vector n_iters(n_runs); + for (int r = 0; r < n_runs; ++r) { + auto d_c = raft::make_device_matrix(handle, n_clusters, n_features); + raft::copy(d_c.data_handle(), d_init.data_handle(), n_clusters * n_features, stream); + T inertia = 0; + if (host) { + auto pb = p; + pb.device_buffer_samples = testparams.device_buffer_samples; + int64_t n_iter = 0; + cuvs::cluster::kmeans::fit( + handle, + pb, + raft::make_const_mdspan(h_X->view()), + h_sw, + raft::make_device_matrix_view(d_c.data_handle(), n_clusters, n_features), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } else { + int n_iter = 0; + cuvs::cluster::kmeans::fit(handle, + p, + raft::make_const_mdspan(d_X->view()), + d_sw, + d_c.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } + raft::copy(centroids[r].data(), d_c.data_handle(), n_clusters * n_features, stream); + raft::resource::sync_stream(handle, stream); + inertias[r] = inertia; + } + for (int r = 1; r < n_runs; ++r) { + EXPECT_EQ(n_iters[0], n_iters[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ(inertias[0], inertias[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ( + 0, std::memcmp(centroids[0].data(), centroids[r].data(), centroids[0].size() * sizeof(T))) + << (host ? "host" : "device") << " run " << r << ": centroids differ"; + } + } + } + T fitHostWithInitSize(int64_t init_size_value) { int n_features = testparams.n_col; @@ -709,6 +786,18 @@ TEST_P(KmeansFitBatchedTestD, Result) runZeroCost(); } +TEST_P(KmeansFitBatchedTestF, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + +TEST_P(KmeansFitBatchedTestD, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + INSTANTIATE_TEST_CASE_P(KmeansFitBatchedTests, KmeansFitBatchedTestF, ::testing::ValuesIn(batched_inputsf2)); diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 3b7988e76b..185f400dff 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/cluster/detail/kmeans_balanced.cuh" #include "../test_utils.h" #include @@ -12,16 +13,20 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include +#include #include #include @@ -301,4 +306,90 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); +/* + * The balancing step must pick the same donor points for identical inputs: they may depend neither + * on the order in which the GPU runs the warps nor on how many fits ran before in the process. + * Most rows are packed in a few clusters and many clusters are empty, so that many receivers need a + * donor. + */ +void check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection donor) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + const int64_t n_rows = 10007; + const int64_t dim = 8; + const int64_t n_lists = 200; + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, dataset.data_handle(), n_rows * dim, -1.0f, 1.0f); + std::vector h_centers(n_lists * dim); + for (size_t i = 0; i < h_centers.size(); i++) { + h_centers[i] = static_cast(i % 97) / 97.0f; + } + + // 95% of the rows in clusters [0, 10), the rest spread over [10, 60); clusters [60, 200) empty. + std::vector h_labels(n_rows); + std::vector h_sizes(n_lists, 0); + for (int64_t i = 0; i < n_rows; i++) { + h_labels[i] = i < n_rows * 95 / 100 ? i % 10 : 10 + i % 50; + h_sizes[h_labels[i]]++; + } + auto labels = raft::make_device_vector(handle, n_rows); + auto sizes = raft::make_device_vector(handle, n_lists); + raft::update_device(labels.data_handle(), h_labels.data(), n_rows, stream); + raft::update_device(sizes.data_handle(), h_sizes.data(), n_lists, stream); + + const bool random = donor == cuvs::cluster::kmeans::balanced_donor_selection::Random; + const float lower_tol = random ? 0.25f : 0.7f; + const float upper_tol = random ? 4.0f : 1.5f; + const float offset = random ? 0.05f : 0.2f; + + std::vector> results; + for (int run = 0; run < 3; run++) { + auto centers = raft::make_device_matrix(handle, n_lists, dim); + raft::update_device(centers.data_handle(), h_centers.data(), h_centers.size(), stream); + bool adjusted = + cuvs::cluster::kmeans::detail::adjust_centers(handle, + centers.data_handle(), + n_lists, + dim, + dataset.data_handle(), + n_rows, + labels.data_handle(), + sizes.data_handle(), + lower_tol, + upper_tol, + offset, + donor, + /*balancing_step=*/3u, + raft::identity_op{}, + rmm::mr::get_current_device_resource_ref()); + ASSERT_TRUE(adjusted); + std::vector h(centers.size()); + raft::update_host(h.data(), centers.data_handle(), h.size(), stream); + raft::resource::sync_stream(handle); + results.push_back(std::move(h)); + } + // Both strategies move (at least) the first 60 empty clusters. + for (int64_t c = 60; c < 120; c++) { + ASSERT_NE(0, std::memcmp(&results[0][c * dim], &h_centers[c * dim], dim * sizeof(float))) + << "cluster " << c << " was not moved"; + } + for (int run = 1; run < 3; run++) { + ASSERT_EQ( + 0, std::memcmp(results[0].data(), results[run].data(), results[0].size() * sizeof(float))) + << "adjust_centers moved the centers differently for identical inputs (call " << run << ")"; + } +} + +TEST(KmeansBalancedAdjustCenters, ReproducibleSizeSorted) +{ + check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted); +} + +TEST(KmeansBalancedAdjustCenters, ReproducibleRandomDonor) +{ + check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::Random); +} + } // namespace cuvs diff --git a/cpp/tests/distance/dist_l2_exp.cu b/cpp/tests/distance/dist_l2_exp.cu index a91daff1b3..1c9dbe4e28 100644 --- a/cpp/tests/distance/dist_l2_exp.cu +++ b/cpp/tests/distance/dist_l2_exp.cu @@ -1,11 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../test_utils.cuh" #include "distance_base.cuh" +#include +#include + +#include +#include + namespace cuvs { namespace distance { @@ -162,5 +168,46 @@ INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn( class BigMatrixEucExp : public BigMatrixDistanceTest {}; TEST_F(BigMatrixEucExp, Result) {} + +// Distinct points with equal norms must keep their (small) distances. The clamp of self-distance +// round-off to zero used to apply an absolute tolerance, which collapsed such points when the data +// has a small magnitude (e.g. a spectral embedding). +template +void test_small_scale_equal_norms(T scale) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n = 8; + std::vector h_x(n * n, T{0}); + for (int64_t i = 0; i < n; i++) { + h_x[i * n + i] = scale; + } + auto x = raft::make_device_matrix(handle, n, n); + auto d = raft::make_device_matrix(handle, n, n); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::L2SqrtExpanded}) { + cuvs::distance::pairwise_distance(handle, + raft::make_const_mdspan(x.view()), + raft::make_const_mdspan(x.view()), + d.view(), + metric); + std::vector h_d(n * n); + raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream); + raft::resource::sync_stream(handle, stream); + + const T expected = + metric == DistanceType::L2Expanded ? T(2) * scale * scale : std::sqrt(T(2)) * scale; + for (int64_t i = 0; i < n; i++) { + for (int64_t j = 0; j < n; j++) { + EXPECT_NEAR(h_d[i * n + j], i == j ? T(0) : expected, T(1e-3) * expected) + << "metric=" << static_cast(metric) << " i=" << i << " j=" << j; + } + } + } +} + +TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms(1e-2f); } +TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms(1e-5); } } // end namespace distance } // namespace cuvs diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu index 5dccce21da..2b445df7bc 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -7,6 +7,8 @@ #include "../ann_cagra.cuh" +#include +#include #include #include @@ -135,6 +137,152 @@ TEST(AnnCagraMultiPartition, MultiKernelRejected) search_params); } +namespace { +// Runs a filtered multi-partition MULTI_CTA search where every partition keeps only ~4% of its rows +// (still more than k each) and reports the number of queries that returned fewer than k neighbors +// and the recall against an exact search over the kept rows. +void search_with_selective_filter(float filtering_rate, int& short_queries, double& recall) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int num_partitions = 4, part_size = 500, dim = 128, n_queries = 500, k = 10; + // Keep one row in `keep_every` (~4% of each partition, i.e. ~21 rows per partition). + constexpr int keep_every = 24; + constexpr int n_rows = num_partitions * part_size; + + rmm::device_uvector database(static_cast(n_rows) * dim, stream); + rmm::device_uvector queries(static_cast(n_queries) * dim, stream); + raft::random::RngState r(1234ULL); + InitDataset(handle, database.data(), n_rows, dim, cuvs::distance::DistanceType::L2Expanded, r); + InitDataset(handle, queries.data(), n_queries, dim, cuvs::distance::DistanceType::L2Expanded, r); + raft::resource::sync_stream(handle); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + index_params.graph_build_params = graph_build_params::nn_descent_params( + index_params.intermediate_graph_degree, index_params.metric); + + std::vector> part_indices; + // An index only holds a view, so any padded copy has to outlive it. + std::vector> part_padded; + part_padded.reserve(num_partitions); + std::vector> part_bitsets; + part_bitsets.reserve(num_partitions); + std::vector> partition_bitsets; + std::vector removed_local; + for (int64_t i = 0; i < part_size; i++) { + if (i % keep_every != 0) { removed_local.push_back(i); } + } + auto removed = + raft::make_device_vector(handle, static_cast(removed_local.size())); + raft::update_device(removed.data_handle(), removed_local.data(), removed_local.size(), stream); + for (int p = 0; p < num_partitions; p++) { + auto view = raft::make_device_matrix_view( + database.data() + static_cast(p) * part_size * dim, part_size, dim); + part_padded.emplace_back(handle, view); + auto const& padded = part_padded.back().view; + part_indices.push_back(cagra::build(handle, index_params, padded)); + auto& part_index = part_indices.back(); + part_index = cagra::update_dataset(handle, std::move(part_index), padded); + part_bitsets.emplace_back(handle, removed.view(), part_size); + partition_bitsets.push_back(part_bitsets.back().view()); + } + std::vector*> index_ptrs; + for (auto& idx : part_indices) { + index_ptrs.push_back(&idx); + } + + const size_t out_size = static_cast(n_queries) * k; + rmm::device_uvector partition_ids_dev(out_size, stream); + rmm::device_uvector neighbors_dev(out_size, stream); + rmm::device_uvector distances_dev(out_size, stream); + + // A small itopk_size, as used for small k, on the MULTI_CTA algorithm that AUTO picks for a + // handful of queries over a handful of partitions. + cagra::search_params search_params; + search_params.algo = search_algo::MULTI_CTA; + search_params.itopk_size = k; + search_params.filtering_rate = filtering_rate; + + cagra::search( + handle, + search_params, + index_ptrs, + raft::make_device_matrix_view(queries.data(), n_queries, dim), + raft::make_device_matrix_view(partition_ids_dev.data(), n_queries, k), + raft::make_device_matrix_view(neighbors_dev.data(), n_queries, k), + raft::make_device_matrix_view(distances_dev.data(), n_queries, k), + partition_bitsets); + + std::vector partition_ids(out_size); + std::vector neighbors(out_size); + std::vector distances(out_size); + std::vector host_database(static_cast(n_rows) * dim); + std::vector host_queries(static_cast(n_queries) * dim); + raft::update_host(partition_ids.data(), partition_ids_dev.data(), out_size, stream); + raft::update_host(neighbors.data(), neighbors_dev.data(), out_size, stream); + raft::update_host(distances.data(), distances_dev.data(), out_size, stream); + raft::update_host(host_database.data(), database.data(), host_database.size(), stream); + raft::update_host(host_queries.data(), queries.data(), host_queries.size(), stream); + raft::resource::sync_stream(handle); + + // Exact top-k over the rows that pass the filter, for recall. + std::vector kept_rows; + for (int64_t g = 0; g < n_rows; g++) { + if ((g % part_size) % keep_every == 0) { kept_rows.push_back(g); } + } + ASSERT_GT(kept_rows.size(), static_cast(k)); + + short_queries = 0; + size_t matches = 0; + for (int q = 0; q < n_queries; q++) { + std::vector> exact; + for (auto g : kept_rows) { + float d = 0; + for (int j = 0; j < dim; j++) { + const float diff = host_queries[q * dim + j] - host_database[g * dim + j]; + d += diff * diff; + } + exact.emplace_back(d, g); + } + std::partial_sort(exact.begin(), exact.begin() + k, exact.end()); + std::vector expected; + for (int j = 0; j < k; j++) { + expected.push_back(exact[j].second); + } + + int valid = 0; + for (int j = 0; j < k; j++) { + const size_t i = static_cast(q) * k + j; + if (distances[i] == std::numeric_limits::max()) { continue; } + ASSERT_LT(partition_ids[i], static_cast(num_partitions)); + ASSERT_LT(neighbors[i], static_cast(part_size)); + ASSERT_EQ(neighbors[i] % keep_every, 0u) << "filtered-out row returned"; + valid++; + const int64_t g = static_cast(partition_ids[i]) * part_size + neighbors[i]; + if (std::find(expected.begin(), expected.end(), g) != expected.end()) { matches++; } + } + if (valid < k) { short_queries++; } + } + recall = static_cast(matches) / out_size; +} +} // namespace + +// A selective per-partition filter must not leave the search short of neighbors. Single-partition +// search derives the filtering rate from the bitset and enlarges itopk_size accordingly; the +// multi-partition search must do the same, or MULTI_CTA can return fewer than k valid neighbors +// even though every partition holds more than k rows that pass the filter. +TEST(AnnCagraMultiPartition, SelectiveFilterReturnsKNeighbors) +{ + int short_queries = 0; + double recall = 0; + search_with_selective_filter(-1.0f, short_queries, recall); + EXPECT_EQ(short_queries, 0) << "queries returning fewer than k neighbors"; + EXPECT_GE(recall, 0.95); +} + // The shared plan descriptor and the cross-partition select_k direction are derived from // indices[0], so all partitions must share one metric; a mismatch must be rejected. TEST(AnnCagraMultiPartition, MixedMetricRejected) diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index fee22db3c5..a70d41a09f 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -10,10 +10,13 @@ #include #include #include +#include +#include #include #include #include +#include namespace cuvs::neighbors::ivf_flat { @@ -48,6 +51,83 @@ TEST(IVFFlatSerialization, StreamRoundTrip) EXPECT_EQ(loaded.n_lists(), index.n_lists()); } +// k-means gives an empty cluster a zero center. With the cosine metric, the coarse search must +// still rank such lists instead of producing NaN distances: when fewer than n_probes lists have a +// finite distance, the NaN columns are never selected and the remaining probe slots are padded +// with list 0, which then is scanned several times and returns duplicated neighbors. +TEST(AnnIVFFlatTest, CosineZeroNormCenters) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t n_rows = 10000; + constexpr int64_t n_queries = 500; + constexpr int64_t dim = 8; + constexpr uint32_t n_lists = 1024; + constexpr uint32_t n_zero_lists = 1000; + constexpr uint32_t n_probes = 40; + constexpr uint32_t k = 16; + + auto data = raft::make_device_matrix(handle, n_rows, dim); + auto queries = raft::make_device_matrix(handle, n_queries, dim); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, data.data_handle(), data.size(), 0.1f, 2.0f); + raft::random::uniform(handle, rng, queries.data_handle(), queries.size(), 0.1f, 2.0f); + + index_params params; + params.n_lists = n_lists; + params.metric = cuvs::distance::DistanceType::CosineExpanded; + params.add_data_on_build = false; + auto idx = build(handle, params, raft::make_const_mdspan(data.view())); + + // Zero the centers of lists 1..n_zero_lists, as if k-means had left these clusters empty. Fewer + // lists than n_probes keep a non-zero center, and list 0 is one of them. + std::vector centers_h(idx.centers().size()); + raft::update_host(centers_h.data(), idx.centers().data_handle(), centers_h.size(), stream); + raft::resource::sync_stream(handle); + std::fill(centers_h.begin() + dim, centers_h.begin() + (n_zero_lists + 1) * dim, 0.0f); + raft::update_device(idx.centers().data_handle(), centers_h.data(), centers_h.size(), stream); + extend(handle, raft::make_const_mdspan(data.view()), std::nullopt, &idx); + + // All data rows are in the lists with a non-zero center, and all of those lists are probed. + search_params sparams; + sparams.n_probes = n_probes; + auto neighbors = raft::make_device_matrix(handle, n_queries, k); + auto distances = raft::make_device_matrix(handle, n_queries, k); + search(handle, + sparams, + idx, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + auto naive_neighbors = raft::make_device_matrix(handle, n_queries, k); + auto naive_distances = raft::make_device_matrix(handle, n_queries, k); + cuvs::neighbors::naive_knn(handle, + naive_distances.data_handle(), + naive_neighbors.data_handle(), + queries.data_handle(), + data.data_handle(), + n_queries, + n_rows, + dim, + k, + params.metric); + + std::vector neighbors_h(n_queries * k), naive_neighbors_h(n_queries * k); + std::vector distances_h(n_queries * k), naive_distances_h(n_queries * k); + raft::update_host(neighbors_h.data(), neighbors.data_handle(), neighbors_h.size(), stream); + raft::update_host(distances_h.data(), distances.data_handle(), distances_h.size(), stream); + raft::update_host( + naive_neighbors_h.data(), naive_neighbors.data_handle(), naive_neighbors_h.size(), stream); + raft::update_host( + naive_distances_h.data(), naive_distances.data_handle(), naive_distances_h.size(), stream); + raft::resource::sync_stream(handle); + + EXPECT_TRUE(eval_neighbours( + naive_neighbors_h, neighbors_h, naive_distances_h, distances_h, n_queries, k, 0.001, 0.99)); +} + TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) { raft::resources handle; diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 2755258600..7296e67ddf 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -396,4 +396,123 @@ inline auto var_search_mode_1_bit() -> test_cases_t #define INSTANTIATE(type, vals) \ INSTANTIATE_TEST_SUITE_P(IvfRabitq, type, ::testing::ValuesIn(vals)); /* NOLINT */ +/* Empty lists */ + +enum class dataset_input { device, host, host_streaming }; + +struct ivf_rabitq_empty_lists_inputs { + uint32_t dim; + dataset_input input; +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq_empty_lists_inputs& p) -> std::ostream& +{ + const char* input = p.input == dataset_input::device ? "device" + : p.input == dataset_input::host ? "host" + : "host_streaming"; + return os << "ivf_rabitq_empty_lists_inputs {.dim = " << p.dim << ", .input = " << input << "}"; +} + +/** + * The dataset has fewer distinct rows than there are lists. Identical rows are always assigned to + * the same list, so k-means is guaranteed to leave some lists empty. The index must still build + * from device and host input, and search must still find the nearest rows. + */ +class ivf_rabitq_empty_lists_test : public ::testing::TestWithParam { + protected: + void run() + { + constexpr int64_t kRows = 4096; + constexpr int64_t kDistinct = 8; + constexpr int64_t kQueries = 2 * kDistinct; + constexpr int64_t kK = 10; + constexpr uint32_t kLists = 32; + constexpr float kSpacing = 0.25f; + constexpr float kOffset = 0.1f; + + auto ps = GetParam(); + auto dim = static_cast(ps.dim); + auto stream = raft::resource::get_cuda_stream(handle_); + + // Row i is the point kSpacing * (i % kDistinct) on the diagonal. Query q is offset from the + // distinct row q % kDistinct by +-kOffset in every dimension, so that row is its unique nearest + // distinct row at squared L2 distance kOffset^2 * dim. + auto h_dataset = raft::make_host_matrix(kRows, dim); + for (int64_t i = 0; i < kRows; i++) { + for (int64_t j = 0; j < dim; j++) { + h_dataset(i, j) = kSpacing * static_cast(i % kDistinct); + } + } + auto h_queries = raft::make_host_matrix(kQueries, dim); + for (int64_t i = 0; i < kQueries; i++) { + for (int64_t j = 0; j < dim; j++) { + h_queries(i, j) = + kSpacing * static_cast(i % kDistinct) + (i < kDistinct ? kOffset : -kOffset); + } + } + auto d_dataset = raft::make_device_matrix(handle_, kRows, dim); + auto d_queries = raft::make_device_matrix(handle_, kQueries, dim); + raft::copy(d_dataset.data_handle(), h_dataset.data_handle(), h_dataset.size(), stream); + raft::copy(d_queries.data_handle(), h_queries.data_handle(), h_queries.size(), stream); + + index_params ipams; + ipams.n_lists = kLists; + if (ps.input == dataset_input::host_streaming) { + ipams.force_streaming = true; + ipams.streaming_batch_size = kRows; + } + auto idx = ps.input == dataset_input::device + ? build(handle_, ipams, raft::make_const_mdspan(d_dataset.view())) + : build(handle_, ipams, raft::make_const_mdspan(h_dataset.view())); + + // Probe every list, because an empty list's centroid can be closer to a query than the + // centroid of the list that holds its neighbors. + search_params sparams; + sparams.n_probes = kLists; + auto neighbors = raft::make_device_matrix(handle_, kQueries, kK); + auto distances = raft::make_device_matrix(handle_, kQueries, kK); + search(handle_, + sparams, + idx, + raft::make_const_mdspan(d_queries.view()), + neighbors.view(), + distances.view()); + + std::vector h_neighbors(kQueries * kK); + std::vector h_distances(kQueries * kK); + raft::copy(h_neighbors.data(), neighbors.data_handle(), h_neighbors.size(), stream); + raft::copy(h_distances.data(), distances.data_handle(), h_distances.size(), stream); + raft::resource::sync_stream(handle_); + + // Every row has kRows / kDistinct > kK copies, so all kK neighbors are copies of the nearest + // distinct row. The next nearest distinct row is (kSpacing - kOffset)^2 * dim away. + const float expected_dist = kOffset * kOffset * static_cast(dim); + for (int64_t q = 0; q < kQueries; q++) { + for (int64_t j = 0; j < kK; j++) { + auto neighbor = h_neighbors[q * kK + j]; + auto dist = h_distances[q * kK + j]; + ASSERT_TRUE(neighbor >= 0 && neighbor < kRows) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_EQ(neighbor % kDistinct, q % kDistinct) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_NEAR(dist, expected_dist, 0.1f * expected_dist) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + } + } + } + + raft::resources handle_; +}; + +inline auto empty_lists_inputs() -> std::vector +{ + std::vector xs; + for (uint32_t dim : {1, 2, 7, 64, 100}) { + for (auto input : {dataset_input::device, dataset_input::host, dataset_input::host_streaming}) { + xs.push_back({dim, input}); + } + } + return xs; +} + } // namespace cuvs::neighbors::ivf_rabitq diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu new file mode 100644 index 0000000000..d0fd00ef95 --- /dev/null +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 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 + +namespace cuvs::neighbors::ivf_rabitq { + +struct ivf_rabitq_exact_match_inputs { + uint32_t dim; + uint32_t bits_per_dim; + search_mode mode; +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq_exact_match_inputs& p) -> std::ostream& +{ + const char* mode = p.mode == search_mode::LUT16 ? "LUT16" + : p.mode == search_mode::LUT32 ? "LUT32" + : p.mode == search_mode::QUANT4 ? "QUANT4" + : "QUANT8"; + return os << "ivf_rabitq_exact_match_inputs {.dim = " << p.dim + << ", .bits_per_dim = " << p.bits_per_dim << ", .mode = " << mode << "}"; +} + +/** + * Queries that coincide with a list centroid. + * + * The dataset has as many distinct rows as there are lists, and each row is repeated many times, so + * each list holds the copies of one row and its centroid is that row. Each query is one of the + * distinct rows, so the probe distance ||q - c||^2 to its own list is zero up to rounding. Search + * must return copies of the query at distance ~0, for every search mode. + */ +class ivf_rabitq_exact_match_test : public ::testing::TestWithParam { + protected: + void run() + { + constexpr int64_t kDistinct = 32; + constexpr int64_t kCopies = 128; + constexpr int64_t kRows = kDistinct * kCopies; + constexpr int64_t kK = 10; + constexpr uint32_t kLists = kDistinct; + + auto ps = GetParam(); + auto dim = static_cast(ps.dim); + auto stream = raft::resource::get_cuda_stream(handle_); + + std::mt19937 gen(42); + std::uniform_real_distribution uniform(1.0f, 2.0f); + auto h_queries = raft::make_host_matrix(kDistinct, dim); + for (int64_t i = 0; i < kDistinct; i++) { + for (int64_t j = 0; j < dim; j++) { + h_queries(i, j) = uniform(gen); + } + } + auto h_dataset = raft::make_host_matrix(kRows, dim); + for (int64_t i = 0; i < kRows; i++) { + for (int64_t j = 0; j < dim; j++) { + h_dataset(i, j) = h_queries(i % kDistinct, j); + } + } + + // The squared distance from each distinct row to the nearest other distinct row. + std::vector gap(kDistinct, std::numeric_limits::max()); + for (int64_t a = 0; a < kDistinct; a++) { + for (int64_t b = 0; b < kDistinct; b++) { + if (a == b) { continue; } + float d = 0; + for (int64_t j = 0; j < dim; j++) { + float diff = h_queries(a, j) - h_queries(b, j); + d += diff * diff; + } + gap[a] = std::min(gap[a], d); + } + } + + auto d_dataset = raft::make_device_matrix(handle_, kRows, dim); + auto d_queries = raft::make_device_matrix(handle_, kDistinct, dim); + raft::copy(d_dataset.data_handle(), h_dataset.data_handle(), h_dataset.size(), stream); + raft::copy(d_queries.data_handle(), h_queries.data_handle(), h_queries.size(), stream); + + index_params ipams; + ipams.n_lists = kLists; + ipams.bits_per_dim = ps.bits_per_dim; + auto idx = build(handle_, ipams, raft::make_const_mdspan(d_dataset.view())); + + search_params sparams; + sparams.n_probes = kLists; + sparams.mode = ps.mode; + auto neighbors = raft::make_device_matrix(handle_, kDistinct, kK); + auto distances = raft::make_device_matrix(handle_, kDistinct, kK); + search(handle_, + sparams, + idx, + raft::make_const_mdspan(d_queries.view()), + neighbors.view(), + distances.view()); + + std::vector h_neighbors(kDistinct * kK); + std::vector h_distances(kDistinct * kK); + raft::copy(h_neighbors.data(), neighbors.data_handle(), h_neighbors.size(), stream); + raft::copy(h_distances.data(), distances.data_handle(), h_distances.size(), stream); + raft::resource::sync_stream(handle_); + + // Every row has kCopies > kK copies, so all kK neighbors must be exact copies of the query. + int64_t n_bad = 0; + for (int64_t q = 0; q < kDistinct; q++) { + bool bad = false; + for (int64_t j = 0; j < kK; j++) { + auto neighbor = h_neighbors[q * kK + j]; + auto dist = h_distances[q * kK + j]; + if (neighbor < 0 || neighbor >= kRows || neighbor % kDistinct != q || + !(std::abs(dist) <= 0.25f * gap[q])) { + bad = true; + } + } + if (bad) { + n_bad++; + auto msg = ::testing::Message() << ps << ": query " << q << " (gap " << gap[q] << "):"; + for (int64_t j = 0; j < kK; j++) { + msg << " " << h_neighbors[q * kK + j] << "(row " << h_neighbors[q * kK + j] % kDistinct + << ", " << h_distances[q * kK + j] << ")"; + } + ADD_FAILURE() << msg; + } + } + EXPECT_EQ(n_bad, 0) << ps << ": " << n_bad << " of " << kDistinct + << " queries did not return copies of themselves"; + } + + raft::resources handle_; +}; + +inline auto exact_match_inputs() -> std::vector +{ + std::vector xs; + for (uint32_t dim : {64, 100, 128}) { + for (uint32_t bits_per_dim : {1, 3}) { + for (auto mode : + {search_mode::LUT16, search_mode::LUT32, search_mode::QUANT4, search_mode::QUANT8}) { + xs.push_back({dim, bits_per_dim, mode}); + } + } + } + return xs; +} + +TEST_P(ivf_rabitq_exact_match_test, build_search) { this->run(); } // NOLINT +INSTANTIATE_TEST_SUITE_P(IvfRabitq, + ivf_rabitq_exact_match_test, + ::testing::ValuesIn(exact_match_inputs())); // NOLINT + +} // namespace cuvs::neighbors::ivf_rabitq 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 42bf84e6a3..28771bd785 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 @@ -18,4 +18,9 @@ INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims() + var_n_probes() + var_k() + var_bits_per_dim() + var_search_mode() + var_search_mode_1_bit()); +TEST_P(ivf_rabitq_empty_lists_test, build_search) { this->run(); } // NOLINT +INSTANTIATE_TEST_SUITE_P(IvfRabitq, + ivf_rabitq_empty_lists_test, + ::testing::ValuesIn(empty_lists_inputs())); // NOLINT + } // namespace cuvs::neighbors::ivf_rabitq diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5ac80fda73..f30d001f65 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -10,14 +10,17 @@ #include "../../src/distance/unfused_distance_nn.cuh" #include +#include #include #include #include #include #include #include +#include #include +#include namespace cuvs::neighbors { @@ -116,7 +119,17 @@ class NNTest : public ::testing::TestWithParam> { workspace_size = m * n * sizeof(AccT); } - if (backend_unavailable) { GTEST_SKIP() << "Requested top_1_nn path is unavailable"; } + if (backend_unavailable) { + // cuTile cubins are embedded only for specific architectures and are not forward compatible + // across minor versions, so cuTile (and FP16 Auto, which has no non-cuTile fallback) is + // expected to be unavailable on devices such as GB300 (sm_103) and GB10 (sm_121). + const auto& prop = raft::resource::get_device_properties(handle); + GTEST_SKIP() << "Requested top_1_nn path is unavailable on compute capability " << prop.major + << "." << prop.minor + << (cuvs::detail::jit_lto::cutile_launch_available_on_current_device() + ? "" + : " (this build has no cuTile cubin for this architecture)"); + } // Reset buffer if constexpr (std::is_same_v>) { @@ -446,6 +459,77 @@ TEST(Top1nnPlan, RejectsMismatchedLaunch) plan))); } +// Distinct points with equal norms and a small magnitude (as in a spectral embedding) must not be +// treated as self-neighbors when finding the nearest neighbor. +TEST(Top1nnEqualNorms, SmallScale) +{ + using Backend = cuvs::distance::detail::Top1nnBackend; + using Kvp = raft::KeyValuePair; + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int m = 1000; + constexpr int n = 5; + constexpr int k = 5; + constexpr float s = 1e-2f; + std::vector h_x(m * k, 0.0f); + std::vector h_y(n * k, 0.0f); + for (int i = 0; i < m; i++) { + h_x[i * k + i % n] = s; + } + for (int j = 0; j < n; j++) { + h_y[j * k + j] = s; + } + auto x = raft::make_device_matrix(handle, m, k); + auto y = raft::make_device_matrix(handle, n, k); + auto x_norm = raft::make_device_vector(handle, m); + auto y_norm = raft::make_device_vector(handle, n); + auto output = raft::make_device_vector(handle, m); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream.get()); + raft::update_device(y.data_handle(), h_y.data(), h_y.size(), stream.get()); + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream.get()); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, n, stream.get()); + + cuvs::distance::detail::Top1nnTuning tuning{}; + for (auto backend : {Backend::Cutlass, Backend::Unfused}) { + if (!cuvs::distance::detail::is_top_1_nn_backend_available( + backend, x.data_handle(), y.data_handle(), m, n, k, DistanceType::L2Expanded)) { + continue; + } + const auto plan = cuvs::distance::probe_top_1_nn( + handle, x.data_handle(), y.data_handle(), m, n, k, tuning, DistanceType::L2Expanded, backend); + ASSERT_TRUE(plan.available); + auto workspace = raft::make_device_vector(handle, plan.workspace_bytes); + cuvs::distance::top_1_nn(handle, + output.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + tuning, + workspace.data_handle(), + workspace.size(), + false, + true, + true, + DistanceType::L2Expanded, + 0.0f, + plan); + std::vector h_out(m); + raft::update_host(h_out.data(), output.data_handle(), m, stream.get()); + raft::resource::sync_stream(handle, stream); + int n_wrong = 0; + for (int i = 0; i < m; i++) { + n_wrong += h_out[i].key != i % n; + } + EXPECT_EQ(n_wrong, 0) << "backend=" << static_cast(backend); + } +} + TEST(Top1nnPlan, ExplicitUnavailableBackendIsStrict) { raft::resources handle; diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md index 5f3ff89fe9..a19381a5a5 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md @@ -35,7 +35,7 @@ Initializes the `CuVS2510GPUVectorsReader`, checks and loads the index. | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:87`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:90`_ ### CuVS2510GPUVectorsReader @@ -59,7 +59,7 @@ Initializes the reader with the cache owned by its vectors format. | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:100`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:103`_ ### readSimilarityFunction @@ -85,7 +85,7 @@ an instance of VectorSimilarityFunction | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:228`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:231`_ ### readVectorEncoding @@ -111,7 +111,7 @@ the vector encoding | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:243`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:246`_ ### getFieldEntry @@ -132,7 +132,7 @@ cuVS index for it. the meta information for the field, or `null` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:306`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:309`_ ### openCagraIndexForMerge @@ -165,7 +165,7 @@ a freshly loaded CAGRA index, or `null` if this segment has none for the field | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:325`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:328`_ ### close @@ -175,7 +175,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:394`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:430`_ ### checkIntegrity @@ -185,7 +185,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Checks consistency of this reader. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:453`_ ### getFloatVectorValues @@ -195,7 +195,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the FloatVectorValues for the given field. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:421`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:462`_ ### getByteVectorValues @@ -207,7 +207,7 @@ Returns the ByteVectorValues for the given field. This is not supported. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:431`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:472`_ ### search @@ -217,7 +217,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the k nearest neighbor documents using cuVS's CAGRA or brute force algorithm for this field, to the given vector. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:455`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:496`_ ### search @@ -229,7 +229,7 @@ Return the k nearest neighbor documents as determined by comparison of their vec This is not supported. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:615`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:656`_ ### readEntry @@ -257,16 +257,16 @@ an instance of FieldEntry | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:643`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:684`_ ### getCagraIndexForField ```java -public CagraIndex getCagraIndexForField(String field) +public CagraIndex getCagraIndexForField(String field) throws IOException ``` -Returns the `CagraIndex` for the given field, or `null` if unavailable -(e.g., during a merge or when the field is missing). +Returns the `CagraIndex` for the given field, or `null` if this segment has none +(e.g., when the field is missing or the segment only holds a brute force index). **Parameters** @@ -278,7 +278,13 @@ Returns the `CagraIndex` for the given field, or `null` if unavailable the CAGRA index, or `null` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:695`_ +**Throws** + +| Type | Description | +| --- | --- | +| `IOException` | I/O exception | + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:737`_ ### getFilterBitsetCache @@ -288,7 +294,7 @@ FilterBitsetCache getFilterBitsetCache() Returns the filter cache owned by the vectors format that created this reader. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:705`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:746`_ ### getFieldInfos @@ -302,7 +308,7 @@ Gets the instance of FieldInfos. the instance of FieldInfos -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:714`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:755`_ ### getCuvsIndexes @@ -314,9 +320,9 @@ Gets the map of `GPUIndex` objects. **Returns** -the map of GPU index objects +the map of GPU index objects, or `null` if this reader was opened for a merge and has not been searched yet -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:723`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:765`_ ### getFieldEntries @@ -330,6 +336,6 @@ Gets the map of FieldEntry objects that hold the meta information for the field. the map of FieldEntry objects -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:732`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:774`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:59`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:60`_ diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 491c5213fc..55ab018e18 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -14,7 +14,6 @@ import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; -import static com.nvidia.cuvs.internal.common.Util.concatenate; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; @@ -468,8 +467,6 @@ private static boolean isCagraPaddedLayout(CuVSMatrixInternal dataset) { == cagraRequiredRowWidth(logicalColumns, elementSizeBytes(dataset.dataType())); } - private static final BitSet[] EMPTY_PREFILTER_BITSET = new BitSet[0]; - /** * Invokes the native search_cagra_index via the Panama API for searching a * CAGRA index. @@ -495,9 +492,7 @@ public SearchResults search(CagraQuery query) throws Throwable { final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; - final BitSet[] prefilters = - hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; - final long prefilterDataLength = hasPreFilter ? query.getNumDocs() * prefilters.length : 0; + final long prefilterDataLength = hasPreFilter ? query.getNumDocs() : 0; final long prefilterLen = hasPreFilter ? (prefilterDataLength + 31) / 32 : 0; final long prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; @@ -531,8 +526,9 @@ public SearchResults search(CagraQuery query) throws Throwable { cuvsFilter.type(prefilter, 0); // NO_FILTER cuvsFilter.addr(prefilter, 0); } else { - BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); + // CAGRA takes a single filter for all queries. Use its bits as they are: concatenate() + // would read a BitSet with no bits set as "no filter" and return excluded rows. + long[] filters = query.getPrefilter().get(0, query.getNumDocs()).toLongArray(); var prefilterDataMemorySegment = buildMemorySegment(localArena, filters, (prefilterDataLength + 63) / 64); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index a4f08c47d4..4e276aeb11 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -692,6 +692,53 @@ public void testPrefilteringReducesResults() throws Throwable { } } + /** + * A prefilter that excludes every row leaves nothing to return. The search used to read a + * {@link BitSet} with no bits set as "no filter" and returned the excluded rows. + */ + @Test + public void testPrefilterExcludingEveryRowReturnsNothing() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = {{0.48216683f, 0.0428398f}, {0.74021935f, 0.9209938f}}; + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + try (CuVSResources resources = CheckedCuVSResources.create(); + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + var hostVectors = CuVSMatrix.ofArray(dataset); + var deviceVectors = hostVectors.toDevice(resources); + var indexDataset = index.makePaddedDataset(deviceVectors); + var queryVectors = CuVSMatrix.ofArray(queries)) { + index.updateDataset(indexDataset); + + for (var algo : + new CagraSearchParams.SearchAlgo[] { + CagraSearchParams.SearchAlgo.SINGLE_CTA, CagraSearchParams.SearchAlgo.MULTI_CTA + }) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(2) + .withSearchParams(new CagraSearchParams.Builder().withAlgo(algo).build()) + .withQueryVectors(queryVectors) + .withPrefilter(new BitSet(dataset.length), dataset.length) + .build(); + + List> results = index.search(query).getResults(); + assertEquals(algo.toString(), List.of(Map.of(), Map.of()), results); + } + } + } + /** * Regression test for a bug in {@code CagraIndexImpl.search}: the {@code distances} output tensor * was described (dtype and buffer size) using the query vectors' data type instead of the diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index 9110f6ddc2..4cb1bdc89e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -5,29 +5,55 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.util.Arrays; import java.util.BitSet; import java.util.List; +import java.util.Map; import org.junit.Before; -import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * Builds CAGRA indexes over random datasets and compares the search results with an exact search + * on the CPU. + * + *

CAGRA is an approximate index, and on uniformly random data in hundreds or thousands of + * dimensions the nearest neighbors are barely closer than the rest, so a single query can + * legitimately miss some or all of its exact neighbors. The test therefore checks what must hold + * for every query (the right number of distinct, valid rows that pass the prefilter, returned in + * order with their correct distances), and checks recall only as a mean over many queries. Each + * case also searches for a copy of one of its rows, which the search should return first. + * + *

This test used to require the first 5 results of every query to be among the exact {@code + * 2k + 10} nearest neighbors. That failed now and then for legitimate approximate results (#631), + * and it also caught the Java bindings leaking filtered-out rows (#1048, #1468, fixed in #2375), + * which the prefilter check here reports directly. + */ @RunWith(RandomizedRunner.class) public class CagraRandomizedIT extends CuVSTestCase { private static final Logger log = LoggerFactory.getLogger(CagraRandomizedIT.class); - @Before - public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); - initializeRandom(); - log.trace("Random context initialized for test."); - } + // Lower bounds for the mean recall@k over the queries of each dataset memory kind, without and + // with a prefilter. In CI (RTX PRO 6000, CUDA 13), about 300 runs of this test measured means of + // 0.83 (standard deviation 0.022, lowest 0.76) without a prefilter and 0.94 (standard deviation + // 0.011, lowest 0.90) with one. The bounds sit 8 and 12 standard deviations below the means, so + // they only fail when results are wrong rather than approximate. + private static final double MIN_MEAN_RECALL_UNFILTERED = 0.65; + private static final double MIN_MEAN_RECALL_FILTERED = 0.8; + // Minimum fraction of the searches for a copy of a row that return that row first. All of more + // than 100,000 such searches did in CI. + private static final double MIN_SELF_MATCH_RATE = 0.99; + // Returned distances are computed in float on the GPU; allow for the difference in rounding + // from the double-precision distances computed here (at most 2.4e-7 in CI). + private static final double DISTANCE_RELATIVE_TOLERANCE = 1e-4; enum TestDatasetMemoryKind { HEAP, @@ -35,17 +61,70 @@ enum TestDatasetMemoryKind { DEVICE } - @Ignore // https://github.com/nvidia/cuvs/issues/1468 + /** Recall accumulated per dataset memory kind, with and without a prefilter. */ + private final double[][] recallSum = new double[TestDatasetMemoryKind.values().length][2]; + + private final int[][] queryCount = new int[TestDatasetMemoryKind.values().length][2]; + + private int selfQueries; + private int selfMatches; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + } + @Test public void testResultsTopKWithRandomValues() throws Throwable { - TestDatasetMemoryKind[] testDatasetMemoryKinds = { - TestDatasetMemoryKind.HEAP, TestDatasetMemoryKind.NATIVE, TestDatasetMemoryKind.DEVICE - }; for (int i = 0; i < 100; i++) { - for (var datasetMemoryKind : testDatasetMemoryKinds) { + for (var datasetMemoryKind : TestDatasetMemoryKind.values()) { tmpResultsTopKWithRandomValues(datasetMemoryKind); } } + + assertTrue( + String.format( + "Only %d of %d searches for a row returned it first", selfMatches, selfQueries), + selfMatches >= MIN_SELF_MATCH_RATE * selfQueries); + for (var kind : TestDatasetMemoryKind.values()) { + for (int filtered = 0; filtered < 2; filtered++) { + int n = queryCount[kind.ordinal()][filtered]; + if (n == 0) { + continue; + } + double meanRecall = recallSum[kind.ordinal()][filtered] / n; + double minMeanRecall = + filtered == 1 ? MIN_MEAN_RECALL_FILTERED : MIN_MEAN_RECALL_UNFILTERED; + log.info( + "{} {} prefilter: mean recall@k {} over {} queries", + kind, + filtered == 1 ? "with" : "without", + meanRecall, + n); + assertTrue( + String.format( + "Mean recall@k %f over %d queries (%s, prefilter=%b) is below %f", + meanRecall, n, kind, filtered == 1, minMeanRecall), + meanRecall >= minMeanRecall); + } + } + } + + /** + * Gives the index the device-padded vectors CAGRA search needs. A device matrix whose rows are + * already padded can only be viewed, not copied. + */ + private static AutoCloseable attachPaddedDataset(CagraIndex index, CuVSMatrix deviceVectors) + throws Throwable { + if (CagraIndex.isPaddedDataset(deviceVectors)) { + var view = index.makePaddedDatasetView(deviceVectors); + index.updateDataset(view); + return view; + } + var dataset = index.makePaddedDataset(deviceVectors); + index.updateDataset(dataset); + return dataset; } private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryKind) @@ -53,66 +132,43 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK int DATASET_SIZE_LIMIT = 10_000; int DIMENSIONS_LIMIT = 2048; int NUM_QUERIES_LIMIT = 10; - int TOP_K_LIMIT = 64; // nocommit This fails beyond 64 + // The default CAGRA search parameters (itopk_size = 64) bound the number of results. + int TOP_K_LIMIT = 64; - int datasetSize = - random.nextInt(DATASET_SIZE_LIMIT) - + 2; // datasetSize of 1 fails/crashed due to a bug, hence adding 2 here. + // A CAGRA graph needs at least 2 rows. + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 2; int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, datasetSize); boolean usePrefilter = random.nextBoolean(); - if (datasetSize < topK) datasetSize = topK; - - BitSet sharedPrefilter = null; + BitSet prefilter = null; if (usePrefilter) { - sharedPrefilter = new BitSet(datasetSize); + prefilter = new BitSet(datasetSize); for (int j = 0; j < datasetSize; j++) { - sharedPrefilter.set(j, random.nextBoolean()); + prefilter.set(j, random.nextBoolean()); } } - BitSet[] prefilters = null; - if (sharedPrefilter != null) { - prefilters = new BitSet[numQueries]; - for (int i = 0; i < numQueries; i++) { - prefilters[i] = sharedPrefilter; - } - } - - // Generate a random dataset float[][] vectors = generateData(random, datasetSize, dimensions); - - // Generate random query vectors float[][] queries = generateData(random, numQueries, dimensions); - - log.debug("Dataset size: {}x{}", datasetSize, dimensions); - log.debug("Query size: {}x{}", numQueries, dimensions); - log.debug("TopK: {}", topK); - log.debug("Use memory dataset: " + datasetMemoryKind.name()); - - // Debugging: Log dataset and queries - if (log.isDebugEnabled()) { - log.debug("Dataset:"); - for (float[] row : vectors) { - log.debug(java.util.Arrays.toString(row)); - } - log.debug("Queries:"); - for (float[] query : queries) { - log.debug(java.util.Arrays.toString(query)); - } + // One more query is a copy of a row that passes the prefilter, if any does. + int selfRow = random.nextInt(datasetSize); + if (prefilter != null && !prefilter.get(selfRow)) { + int next = prefilter.nextSetBit(selfRow); + selfRow = next >= 0 ? next : prefilter.previousSetBit(selfRow); + } + if (selfRow >= 0) { + queries = Arrays.copyOf(queries, numQueries + 1); + queries[numQueries] = vectors[selfRow].clone(); } - // Sanity checks - assert vectors.length > 0 : "Dataset is empty."; - assert queries.length > 0 : "Queries are empty."; - assert dimensions > 0 : "Invalid dimensions."; - assert topK > 0 && topK <= datasetSize : "Invalid topK value."; - // Generate expected results using brute force - List> expected = generateExpectedResults(topK, vectors, queries, prefilters, log); + String caseDescription = + String.format( + "dataset %dx%d, %d queries, k=%d, prefilter=%b, %s", + datasetSize, dimensions, numQueries, topK, usePrefilter, datasetMemoryKind); + log.debug("Case: {}", caseDescription); - // Create CuVS index and query try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndexParams indexParams = new CagraIndexParams.Builder() @@ -144,35 +200,117 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK .withIndexParams(indexParams) .build(); } else { - assert datasetMemoryKind == TestDatasetMemoryKind.HEAP; index = CagraIndex.newBuilder(resources) .withDataset(vectors) .withIndexParams(indexParams) .build(); } - log.trace("Index built successfully."); - try (var queryVectors = CuVSMatrix.ofArray(queries)) { - // Execute search and retrieve results + // CAGRA search needs the device-padded vectors, whatever memory the index was built from. + try (var hostVectors = CuVSMatrix.ofArray(vectors); + var deviceVectors = hostVectors.toDevice(resources); + var indexDataset = attachPaddedDataset(index, deviceVectors); + var queryVectors = CuVSMatrix.ofArray(queries)) { CagraQuery.Builder queryBuilder = new CagraQuery.Builder(resources) .withQueryVectors(queryVectors) .withTopK(topK) .withSearchParams(new CagraSearchParams.Builder().build()); - - if (sharedPrefilter != null) { - queryBuilder.withPrefilter(sharedPrefilter, datasetSize); + if (prefilter != null) { + queryBuilder.withPrefilter(prefilter, datasetSize); + } + SearchResults results = index.search(queryBuilder.build()); + checkResults( + caseDescription, + datasetMemoryKind, + results, + vectors, + queries, + numQueries, + prefilter, + topK); + if (selfRow >= 0) { + var selfResult = results.getResults().get(numQueries); + selfQueries++; + if (!selfResult.isEmpty() && selfResult.keySet().iterator().next() == selfRow) { + selfMatches++; + } } - - CagraQuery query = queryBuilder.build(); - log.trace("Query built successfully. Executing search..."); - SearchResults results = index.search(query); - - compareResults(results, expected, topK, datasetSize, numQueries); } finally { index.close(); } } } + + /** Squared L2 distances from the query to every row, +inf for the rows the filter removes. */ + private static double[] exactDistances(float[] query, float[][] vectors, BitSet prefilter) { + double[] distances = new double[vectors.length]; + for (int j = 0; j < vectors.length; j++) { + if (prefilter != null && !prefilter.get(j)) { + distances[j] = Double.POSITIVE_INFINITY; + continue; + } + double sum = 0; + for (int d = 0; d < query.length; d++) { + double diff = (double) query[d] - (double) vectors[j][d]; + sum += diff * diff; + } + distances[j] = sum; + } + return distances; + } + + private void checkResults( + String caseDescription, + TestDatasetMemoryKind datasetMemoryKind, + SearchResults results, + float[][] vectors, + float[][] queries, + int numRandomQueries, + BitSet prefilter, + int topK) { + List> resultMaps = results.getResults(); + assertEquals(caseDescription, queries.length, resultMaps.size()); + int numPassing = prefilter == null ? vectors.length : prefilter.cardinality(); + int expectedCount = Math.min(topK, numPassing); + int filtered = prefilter == null ? 0 : 1; + + for (int q = 0; q < queries.length; q++) { + String where = caseDescription + ", query " + q; + double[] exact = exactDistances(queries[q], vectors, prefilter); + double[] sorted = exact.clone(); + Arrays.sort(sorted); + + // The results map holds one entry per distinct row, so the count also rules out duplicates. + Map result = resultMaps.get(q); + assertEquals(where + ": number of distinct results", expectedCount, result.size()); + + float previous = Float.NEGATIVE_INFINITY; + int found = 0; + for (var entry : result.entrySet()) { + int id = entry.getKey(); + float distance = entry.getValue(); + assertTrue(where + ": invalid row id " + id, id >= 0 && id < vectors.length); + assertTrue( + where + ": row " + id + " is excluded by the prefilter", + prefilter == null || prefilter.get(id)); + assertTrue(where + ": results are not sorted by distance", distance >= previous); + previous = distance; + double tolerance = DISTANCE_RELATIVE_TOLERANCE * Math.max(exact[id], 1.0); + assertTrue( + where + ": distance " + distance + " of row " + id + " should be " + exact[id], + Math.abs(distance - exact[id]) <= tolerance); + // Rows exactly as close as the k-th nearest neighbor count as found. + if (exact[id] <= sorted[expectedCount - 1]) { + found++; + } + } + // Recall is measured on the random queries only, not on the copy of a row. + if (expectedCount > 0 && q < numRandomQueries) { + recallSum[datasetMemoryKind.ordinal()][filtered] += (double) found / expectedCount; + queryCount[datasetMemoryKind.ordinal()][filtered]++; + } + } + } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java new file mode 100644 index 0000000000..6410dafc43 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +/** + * A CAGRA graph needs at least two rows: with a single row there is no neighbor to link to, so + * both graph degrees collapse to zero. Building one used to launch the kNN graph kernels anyway, + * which hit an illegal memory access and aborted the JVM (exit code 134). It has to fail with an + * exception instead. + */ +@RunWith(RandomizedRunner.class) +public class CagraSingleRowIT extends CuVSTestCase { + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + } + + private void assertSingleRowBuildFails(CagraGraphBuildAlgo algo) throws Throwable { + float[][] vectors = generateData(random, 1, 16); + try (CuVSResources resources = CheckedCuVSResources.create()) { + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(algo) + .withGraphDegree(64) + .withIntermediateGraphDegree(128) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + RuntimeException e = + assertThrows( + RuntimeException.class, + () -> + CagraIndex.newBuilder(resources) + .withDataset(vectors) + .withIndexParams(indexParams) + .build()); + assertTrue(e.getMessage(), e.getMessage().contains("at least 2 rows")); + } + } + + @Test + public void testSingleRowNnDescentBuildThrows() throws Throwable { + assertSingleRowBuildFails(CagraGraphBuildAlgo.NN_DESCENT); + } + + @Test + public void testSingleRowIvfPqBuildThrows() throws Throwable { + assertSingleRowBuildFails(CagraGraphBuildAlgo.IVF_PQ); + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index 4ace8fc92e..e35f3ce4a3 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -49,7 +49,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro int NUM_QUERIES_LIMIT = 10; int TOP_K_LIMIT = 64; // nocommit This fails beyond 64 - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; + // A CAGRA graph needs at least 2 rows, so the HNSW index converted from it does too. + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 2; int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, datasetSize); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 1efa2d1174..949e09ed22 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -41,6 +41,7 @@ import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.internal.hppc.IntObjectHashMap; import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.store.AlreadyClosedException; import org.apache.lucene.store.ChecksumIndexInput; import org.apache.lucene.store.DataInput; import org.apache.lucene.store.IOContext; @@ -64,7 +65,9 @@ public class CuVS2510GPUVectorsReader extends KnnVectorsReader { private final FlatVectorsReader flatVectorsReader; private final FieldInfos fieldInfos; private final IntObjectHashMap fields; - private final IntObjectHashMap cuvsIndices; + // Loaded eagerly unless this reader was opened for a merge; see getCuVSIndices(). + private volatile IntObjectHashMap cuvsIndices; + private boolean closed; private final IndexInput cuvsIndexInput; private final FilterBitsetCache filterBitsetCache; @@ -130,13 +133,13 @@ public CuVS2510GPUVectorsReader(SegmentReadState state, FlatVectorsReader flatRe var ioContext = state.context.withReadAdvice(ReadAdvice.SEQUENTIAL); cuvsIndexInput = openCuVSInput(state, versionMeta, ioContext); /* - * Only load indexes on the GPU when this reader is opening for searches. - * Do not load indexes on the GPU when this reader is opening during merge calls. - * With this approach we reduce device memory usage by approximately 50% during merges. + * Do not load indexes on the GPU when this reader is opened for a merge, which reads them + * through openCagraIndexForMerge instead. This reduces device memory usage by approximately + * 50% during merges. IndexWriter pools the reader it opens for a merge, though, and hands it + * (or a clone sharing this core) to near-real-time readers opened while the merge runs, so + * such a reader can still be searched; it then loads its indexes on first use. */ - if (state.context.context().equals(Context.MERGE)) { - cuvsIndices = null; - } else { + if (!state.context.context().equals(Context.MERGE)) { cuvsIndices = loadCuVSIndices(); } success = true; @@ -347,11 +350,44 @@ CagraIndex openCagraIndexForMerge(String field) throws IOException { */ private IntObjectHashMap loadCuVSIndices() throws IOException { var indices = new IntObjectHashMap(); - for (var field : fields) { - var fieldEntry = field.value; - int fieldNumber = field.key; - var cuvsIndex = loadCuVSIndex(fieldEntry); - indices.put(fieldNumber, cuvsIndex); + boolean success = false; + try { + for (var field : fields) { + var fieldEntry = field.value; + int fieldNumber = field.key; + var cuvsIndex = loadCuVSIndex(fieldEntry); + indices.put(fieldNumber, cuvsIndex); + } + success = true; + return indices; + } finally { + if (success == false) { + var loaded = stream(indices.values().iterator()).map(cursor -> cursor.value); + IOUtils.closeWhileHandlingException(loaded::iterator); + } + } + } + + /** + * Returns the GPU indexes of this segment, loading them first if this reader was opened for a + * merge and has not been searched yet. + * + * @return the map of {@link GPUIndex} objects, keyed by field number + * @throws IOException I/O exception + */ + private IntObjectHashMap getCuVSIndices() throws IOException { + var indices = cuvsIndices; + if (indices == null) { + synchronized (this) { + if (closed) { + throw new AlreadyClosedException("this reader is closed"); + } + indices = cuvsIndices; + if (indices == null) { + indices = loadCuVSIndices(); + cuvsIndices = indices; + } + } } return indices; } @@ -394,10 +430,15 @@ private GPUIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { */ @Override public void close() throws IOException { + final IntObjectHashMap indices; + synchronized (this) { + closed = true; + indices = cuvsIndices; + } var closeableStream = Stream.of(flatVectorsReader, cuvsIndexInput); IOUtils.close(closeableStream::iterator); - if (cuvsIndices != null) { - var indexClosableStream = stream(cuvsIndices.values().iterator()).map(cursor -> cursor.value); + if (indices != null) { + var indexClosableStream = stream(indices.values().iterator()).map(cursor -> cursor.value); IOUtils.close(indexClosableStream::iterator); } closeCuVSResourcesInstance(); @@ -462,7 +503,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits } var fieldNumber = fieldInfos.fieldInfo(field).number; - GPUIndex cuvsIndex = cuvsIndices != null ? cuvsIndices.get(fieldNumber) : null; + GPUIndex cuvsIndex = getCuVSIndices().get(fieldNumber); if (cuvsIndex == null) { throw new IllegalStateException("Index not found for field:" + field); } @@ -687,17 +728,17 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn } /** - * Returns the {@link CagraIndex} for the given field, or {@code null} if unavailable - * (e.g., during a merge or when the field is missing). + * Returns the {@link CagraIndex} for the given field, or {@code null} if this segment has none + * (e.g., when the field is missing or the segment only holds a brute force index). * * @param field the vector field name * @return the CAGRA index, or {@code null} + * @throws IOException I/O exception */ - public CagraIndex getCagraIndexForField(String field) { - if (cuvsIndices == null) return null; + public CagraIndex getCagraIndexForField(String field) throws IOException { FieldInfo info = fieldInfos.fieldInfo(field); if (info == null) return null; - GPUIndex gpuIndex = cuvsIndices.get(info.number); + GPUIndex gpuIndex = getCuVSIndices().get(info.number); if (gpuIndex == null) return null; return gpuIndex.getCagraIndex(); } @@ -719,7 +760,8 @@ public FieldInfos getFieldInfos() { /** * Gets the map of {@link GPUIndex} objects. * - * @return the map of GPU index objects + * @return the map of GPU index objects, or {@code null} if this reader was opened for a merge and + * has not been searched yet */ public IntObjectHashMap getCuvsIndexes() { return cuvsIndices; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java index b9f39e6c05..5ec3814cd9 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java @@ -111,13 +111,17 @@ public void testThreeLevelGraphCanBeReopenedAndSearched() throws Exception { int queryNode = graph.entryNode(); assertTrue(nodesByLevel.get(2).contains(queryNode)); + // Graph nodes are vector ordinals, which equal docIDs in this dense single-segment index. + // The randomized merge policy can reorder documents during forceMerge, so the entry node's + // vector must be looked up through its stored id rather than by indexing the input array + // with the ordinal. + int queryId = Integer.parseInt(leaf.storedFields().document(queryNode).get("id")); IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[queryNode], 10), 10); + var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[queryId], 10), 10); assertEquals(10, hits.scoreDocs.length); - String queryNodeId = Integer.toString(queryNode); boolean foundQueryNode = false; for (var hit : hits.scoreDocs) { - foundQueryNode |= queryNodeId.equals(searcher.storedFields().document(hit.doc).get("id")); + foundQueryNode |= hit.doc == queryNode; } assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 31958ef0db..e9bdfea9aa 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -114,7 +114,9 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + // Only even-numbered documents have vectors, so never ask for more results than exist. + int vectorCount = (datasetSize + 1) / 2; + int topK = Math.min(random.nextInt(5, TOP_K_LIMIT), vectorCount); GPUKnnFloatVectorQuery query = new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java index 2f77f19b8d..9f5f01d9c1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java @@ -4,17 +4,13 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertSearchFindsNearestNeighbors; import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -57,7 +53,9 @@ public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase { static int DATASET_SIZE_LIMIT = 1000; static int DIMENSIONS_LIMIT = 2048; static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 + // A wider search is more accurate, not less: in the measurements described in testVectorSearch, + // no result was outside the closest 3 * topK for any topK from 8 to 200. + static int TOP_K_LIMIT = 128; static float[][] dataset; // Whether each document of the dataset was indexed with a vector; only those can be hits. static boolean[] hasVector; @@ -121,66 +119,25 @@ public void testVectorSearch() throws IOException { int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - if (dataset.length < topK) topK = dataset.length; - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); log.log(Level.FINE, "TopK: " + topK); - Query query = new KnnFloatVectorQuery("vector", queries[0], topK); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topK*3: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - if (hasVector[j] == false) { - continue; - } - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); - } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, neighbors.size()))); - } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; + // The search is approximate. Lucene searches each segment's HNSW graph with a beam as wide as + // topK, and in this data (uniform random vectors with up to 2048 dimensions, so all distances + // are within about 10% of each other) a narrow beam sometimes stops just short of the true + // neighbors. Over 27,003 random indexes from this test's distribution, with 10 queries each, a + // result fell outside the closest 3 * topK in 1.7% of queries at topK = 1, 0.21% at topK = 2, + // 0.03% at topK = 3, 3 times in 270,000 queries at topK = 5, and never at topK >= 8. Those + // results were at most rank 17, and typically only 1.3% farther away than the nearest + // neighbor. On the same segments, Lucene's own HNSW format missed in 22% of queries at + // topK = 1. Allowing 2 misses across all the queries keeps the false failure rate below 1e-5 + // per run. + int allowedMisses = 2; + assertSearchFindsNearestNeighbors( + searcher, "vector", "id", dataset, hasVector, queries, topK, allowedMisses); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 6db380a3ce..a09b306152 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -18,6 +18,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; @@ -42,7 +43,9 @@ protected Codec getCodec() { public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { Document doc1 = new Document(); doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 707ee35709..8f56c7f45e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -18,6 +18,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; @@ -42,7 +43,9 @@ protected Codec getCodec() { public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { Document doc1 = new Document(); doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 3b05beaacf..67cae2b4e9 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -26,6 +26,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; @@ -75,7 +76,9 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { } try (Directory dir = newDirectory(new ByteBuffersDirectory()); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { for (int i = 0; i < R; i++) { Document doc = new Document(); doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java new file mode 100644 index 0000000000..679eb7c56e --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.ConcurrentMergeScheduler; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FilterMergePolicy; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.MergeTrigger; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * IndexWriter pools the {@link org.apache.lucene.index.SegmentReader} it opens to merge a segment + * and hands that same reader (or a clone sharing its core) to near-real-time readers opened while + * the merge is still running. Searching such a reader must work even though its vectors reader was + * opened with a merge {@link org.apache.lucene.store.IOContext}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestSearchSegmentsOpenedForMerge extends LuceneTestCase { + + private static final int NUM_SEGMENTS = 3; + private static final int DOCS_PER_SEGMENT = 64; + private static final int DIMENSION = 16; + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS not supported", isSupported()); + } + + /** Forced merge of all segments that blocks once its input readers are open. */ + private static final class BlockingMergePolicy extends FilterMergePolicy { + final CountDownLatch mergeReadersOpened = new CountDownLatch(1); + final CountDownLatch releaseMerge = new CountDownLatch(1); + + BlockingMergePolicy() { + super(NoMergePolicy.INSTANCE); + } + + @Override + public MergeSpecification findForcedMerges( + SegmentInfos segmentInfos, + int maxSegmentCount, + Map segmentsToMerge, + MergeContext mergeContext) { + List segments = new ArrayList<>(); + for (SegmentCommitInfo info : segmentInfos) { + if (segmentsToMerge.containsKey(info) + && !mergeContext.getMergingSegments().contains(info)) { + segments.add(info); + } + } + if (segments.size() < 2) { + return null; + } + MergeSpecification spec = new MergeSpecification(); + spec.add( + new OneMerge(segments) { + @Override + public CodecReader wrapForMerge(CodecReader reader) throws IOException { + // Called by IndexWriter#mergeMiddle once the input readers have been opened (with a + // merge IOContext) and pooled. + mergeReadersOpened.countDown(); + try { + if (!releaseMerge.await(5, TimeUnit.MINUTES)) { + throw new IOException("merge was never released"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + return reader; + } + }); + return spec; + } + + @Override + public MergeSpecification findMerges( + MergeTrigger mergeTrigger, SegmentInfos segmentInfos, MergeContext mergeContext) { + return null; + } + } + + @Test + public void testSearchNRTReaderDuringMerge() throws Exception { + BlockingMergePolicy mergePolicy = new BlockingMergePolicy(); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) + .setMergePolicy(mergePolicy) + .setMergeScheduler(new ConcurrentMergeScheduler()) + .setMaxBufferedDocs(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setRAMBufferSizeMB(64); + + float[][] vectors = new float[NUM_SEGMENTS * DOCS_PER_SEGMENT][]; + try (Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int seg = 0; seg < NUM_SEGMENTS; seg++) { + for (int i = 0; i < DOCS_PER_SEGMENT; i++) { + int id = seg * DOCS_PER_SEGMENT + i; + vectors[id] = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + vectors[id][d] = random().nextFloat(); + } + Document doc = new Document(); + doc.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", vectors[id], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + } + + try { + writer.forceMerge(1, false); + assertTrue( + "merge did not start", mergePolicy.mergeReadersOpened.await(5, TimeUnit.MINUTES)); + + // The segments being merged are still live, so this NRT reader shares the cores the merge + // opened. + try (DirectoryReader reader = DirectoryReader.open(writer)) { + assertEquals(NUM_SEGMENTS, reader.leaves().size()); + IndexSearcher searcher = new IndexSearcher(reader); + int topK = 10; + float[] target = vectors[random().nextInt(vectors.length)]; + List queries = + List.of( + new GPUKnnFloatVectorQuery("vector", target, topK, null, topK, 1), + new KnnFloatVectorQuery("vector", target, topK)); + for (Query query : queries) { + TopDocs hits = searcher.search(query, topK); + assertEquals(query.toString(), topK, hits.scoreDocs.length); + } + } + } finally { + mergePolicy.releaseMerge.countDown(); + } + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index ce58fa908c..42a26583e0 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -90,6 +90,38 @@ public static void assertSearchFindsNearestNeighbors( int topK, int allowedMisses) throws IOException { + assertSearchFindsNearestNeighbors( + searcher, vectorField, idField, dataset, null, queries, topK, allowedMisses); + } + + /** + * Like {@link #assertSearchFindsNearestNeighbors(IndexSearcher, String, String, float[][], + * float[][], int, int)}, for an index where only some rows of {@code dataset} were indexed with a + * vector: {@code indexed[i]} says whether row {@code i} was, and {@code null} means all of them. + * The closest vectors are ranked among the indexed rows only, since the others can never be + * returned. + * + *

Besides the approximate check, every query must return {@code min(topK, indexed rows)} + * distinct indexed documents in order of decreasing score, each scored by its own vector. + */ + public static void assertSearchFindsNearestNeighbors( + IndexSearcher searcher, + String vectorField, + String idField, + float[][] dataset, + boolean[] indexed, + float[][] queries, + int topK, + int allowedMisses) + throws IOException { + int indexedCount = dataset.length; + if (indexed != null) { + indexedCount = 0; + for (boolean isIndexed : indexed) { + if (isIndexed) indexedCount++; + } + } + int expectedResults = Math.min(topK, indexedCount); int acceptableRank = topK * 3; int misses = 0; StringBuilder found = new StringBuilder(); @@ -97,13 +129,23 @@ public static void assertSearchFindsNearestNeighbors( for (int q = 0; q < queries.length; q++) { KnnFloatVectorQuery query = new KnnFloatVectorQuery(vectorField, queries[q], topK); TopDocs results = searcher.search(query, topK); - assertEquals("results for query " + q, topK, results.scoreDocs.length); + assertEquals("results for query " + q, expectedResults, results.scoreDocs.length); - int[] ranks = distanceRanks(dataset, queries[q]); - int[] ids = new int[topK]; - int[] resultRanks = new int[topK]; - for (int i = 0; i < topK; i++) { + int[] ranks = distanceRanks(dataset, indexed, queries[q]); + int[] ids = new int[expectedResults]; + int[] resultRanks = new int[expectedResults]; + Set seen = new HashSet<>(); + for (int i = 0; i < expectedResults; i++) { ids[i] = Integer.parseInt(storedFields.document(results.scoreDocs[i].doc).get(idField)); + assertTrue( + "id " + ids[i] + " for query " + q + " has no vector", + indexed == null || indexed[ids[i]]); + assertTrue("id " + ids[i] + " returned twice for query " + q, seen.add(ids[i])); + if (i > 0) { + assertTrue( + "results for query " + q + " are not sorted by score at position " + i, + results.scoreDocs[i].score <= results.scoreDocs[i - 1].score); + } float expectedScore = VectorSimilarityFunction.EUCLIDEAN.compare(queries[q], dataset[ids[i]]); assertEquals( @@ -142,6 +184,14 @@ public static void assertSearchFindsNearestNeighbors( * @return the 1-based rank of each vector, indexed like {@code dataset} */ public static int[] distanceRanks(float[][] dataset, float[] query) { + return distanceRanks(dataset, null, query); + } + + /** + * Like {@link #distanceRanks(float[][], float[])}, ranking only the rows where {@code indexed} is + * true ({@code null} means all rows). The other rows get rank {@link Integer#MAX_VALUE}. + */ + public static int[] distanceRanks(float[][] dataset, boolean[] indexed, float[] query) { double[] distances = new double[dataset.length]; for (int i = 0; i < dataset.length; i++) { for (int j = 0; j < query.length; j++) { @@ -151,11 +201,13 @@ public static int[] distanceRanks(float[][] dataset, float[] query) { } int[] byDistance = IntStream.range(0, dataset.length) + .filter(i -> indexed == null || indexed[i]) .boxed() .sorted(Comparator.comparingDouble(i -> distances[i])) .mapToInt(Integer::intValue) .toArray(); int[] ranks = new int[dataset.length]; + Arrays.fill(ranks, Integer.MAX_VALUE); for (int rank = 0; rank < byDistance.length; rank++) { ranks[byDistance[rank]] = rank + 1; } diff --git a/python/cuvs/cuvs/tests/test_ivf_pq.py b/python/cuvs/cuvs/tests/test_ivf_pq.py index da43d22ea2..d0371c5525 100644 --- a/python/cuvs/cuvs/tests/test_ivf_pq.py +++ b/python/cuvs/cuvs/tests/test_ivf_pq.py @@ -1,4 +1,4 @@ -# 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 # @@ -149,7 +149,13 @@ def run_ivf_pq_build_search_test( "metric", ["sqeuclidean", "inner_product", "euclidean"] ) def test_ivf_pq(inplace, dtype, metric): + # One dimension per PQ subspace. With the default pq_dim (8 for 10 + # columns), the inner-product recall on this data is only about 0.74 and + # varies by about 0.01 from build to build, so `recall > 0.7` fails in + # roughly 1 of 1000 runs. run_ivf_pq_build_search_test( + n_cols=10, + pq_dim=10, dtype=dtype, inplace=inplace, metric=metric, @@ -304,15 +310,14 @@ def test_build_precomputed(codebook_kind, metric): queries = generate_data((n_queries, n_cols), dtype) queries_device = device_ndarray(queries) - # Search regular index + # Search both indexes for one extra neighbor, so that a distance tie + # across the k-th position can be detected. search_params = ivf_pq.SearchParams(n_probes=n_probes) regular_dist, regular_idx = ivf_pq.search( - search_params, regular_index, queries_device, k + search_params, regular_index, queries_device, k + 1 ) - - # Search precomputed index precomputed_dist, precomputed_idx = ivf_pq.search( - search_params, precomputed_index, queries_device, k + search_params, precomputed_index, queries_device, k + 1 ) # Copy results to host for comparison @@ -321,20 +326,35 @@ def test_build_precomputed(codebook_kind, metric): precomputed_idx_host = precomputed_idx.copy_to_host() precomputed_dist_host = precomputed_dist.copy_to_host() - # Compare results for exact match + # Both indexes hold the same centers and PQ codes, so the distances must + # match exactly. np.testing.assert_array_equal( - regular_idx_host, - precomputed_idx_host, - err_msg="Neighbor indices should match exactly", - ) - np.testing.assert_allclose( - regular_dist_host, - precomputed_dist_host, - rtol=1e-5, - atol=1e-5, - err_msg="Distances should match closely", + regular_dist_host[:, :k], + precomputed_dist_host[:, :k], + err_msg="Distances should match exactly", ) + # The order of the vectors within each list is not deterministic (extend + # assigns the slots atomically), and the search breaks distance ties by + # that order. Neighbors with bitwise-equal distances may therefore come + # back in a different order, and a tie across the k-th position may + # select different neighbors. Compare the indices of each group of equal + # distances as a set, except for a group that crosses the k-th position. + for row in range(n_queries): + dist = regular_dist_host[row] + start = 0 + while start < k: + end = start + 1 + while end <= k and dist[end] == dist[start]: + end += 1 + if end <= k: + np.testing.assert_array_equal( + np.sort(regular_idx_host[row, start:end]), + np.sort(precomputed_idx_host[row, start:end]), + err_msg=f"Neighbor indices should match (query {row})", + ) + start = end + @pytest.mark.parametrize("codebook_kind", ["subspace", "cluster"]) @pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8])