Skip to content

Commit 22f107c

Browse files
authored
Migrate IVF-PQ from RAFT to cuVS (#86)
List of changes I made during the migration: - The unit tests are testing IVFPQ codepacker and build functions. So these functions had to be exposed (through `cuvs/neighbors/details`) - The bitset filter is now located under `cuvs/neighbors/bitset_filter.cuh" - search_with_filter is added to the public API, with bitset_filter - Bitset is exposed in public API in a `.hpp` header, for inclusion in CUDA-free code. `bitset::test()` function has been temporarily disabled due to the switch from `.cuh` to `.hpp`. Authors: - Micka (https://github.com/lowener) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #86
1 parent 07791df commit 22f107c

285 files changed

Lines changed: 28539 additions & 1385 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

cpp/CMakeLists.txt

Lines changed: 189 additions & 20 deletions
Large diffs are not rendered by default.

cpp/include/cuvs/core/bitset.hpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
/*
2+
* Copyright (c) 2024, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#pragma once
18+
19+
#include <raft/core/bitset.hpp>
20+
21+
namespace cuvs::core {
22+
/* To use bitset functions containing CUDA code, include <raft/core/bitset.cuh> */
23+
24+
template <typename bitset_t, typename index_t>
25+
using bitset_view = raft::core::bitset_view<bitset_t, index_t>;
26+
27+
template <typename bitset_t, typename index_t>
28+
using bitset = raft::core::bitset<bitset_t, index_t>;
29+
30+
} // end namespace cuvs::core

cpp/include/cuvs/neighbors/ann_types.hpp

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,48 @@
1616

1717
#pragma once
1818

19+
#include <cstdint>
1920
#include <cuvs/distance/distance_types.hpp>
2021

2122
namespace cuvs::neighbors::ann {
2223

24+
/** Parameters for VPQ compression. */
25+
struct vpq_params {
26+
/**
27+
* The bit length of the vector element after compression by PQ.
28+
*
29+
* Possible values: [4, 5, 6, 7, 8].
30+
*
31+
* Hint: the smaller the 'pq_bits', the smaller the index size and the better the search
32+
* performance, but the lower the recall.
33+
*/
34+
uint32_t pq_bits = 8;
35+
/**
36+
* The dimensionality of the vector after compression by PQ.
37+
* When zero, an optimal value is selected using a heuristic.
38+
*
39+
* TODO: at the moment `dim` must be a multiple `pq_dim`.
40+
*/
41+
uint32_t pq_dim = 0;
42+
/**
43+
* Vector Quantization (VQ) codebook size - number of "coarse cluster centers".
44+
* When zero, an optimal value is selected using a heuristic.
45+
*/
46+
uint32_t vq_n_centers = 0;
47+
/** The number of iterations searching for kmeans centers (both VQ & PQ phases). */
48+
uint32_t kmeans_n_iters = 25;
49+
/**
50+
* The fraction of data to use during iterative kmeans building (VQ phase).
51+
* When zero, an optimal value is selected using a heuristic.
52+
*/
53+
double vq_kmeans_trainset_fraction = 0;
54+
/**
55+
* The fraction of data to use during iterative kmeans building (PQ phase).
56+
* When zero, an optimal value is selected using a heuristic.
57+
*/
58+
double pq_kmeans_trainset_fraction = 0;
59+
};
60+
2361
/**
2462
* @defgroup ann_types Approximate Nearest Neighbors Types
2563
* @{
@@ -31,7 +69,7 @@ struct index {};
3169
/** The base for KNN index parameters. */
3270
struct index_params {
3371
/** Distance type. */
34-
cuvs::distance::DistanceType metric = distance::DistanceType::L2Expanded;
72+
cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded;
3573
/** The argument used by some distance metrics. */
3674
float metric_arg = 2.0f;
3775
/**

0 commit comments

Comments
 (0)