Skip to content

Navigation Menu

Sign in
Sign up

Update IVF parameters to match cuVS #8484

New issue

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

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

Already on GitHub? Sign in to your account

Open
NIne-WIngEd wants to merge 5 commits into NVIDIA:main
base: main
Choose a base branch
Loading
from NIne-WIngEd:enh-ivf-params-cuvs
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion cpp/include/cuml/neighbors/knn.hpp
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,14 +132,25 @@ struct knnIndexParam {
struct IVFParam : knnIndexParam {
int nlist;
int nprobe;
uint32_t kmeans_n_iters = 20; ///< Number of k-means iterations.
double kmeans_trainset_fraction = 0.5; ///< Fraction used for k-means training.
bool conservative_memory_allocation = false; ///< Allocate only required memory.
};

struct IVFFlatParam : IVFParam {};

struct IVFPQParam : IVFParam {
int M;
int n_bits;
bool usePrecomputedTables;
bool usePrecomputedTables = false; ///< Deprecated and ignored.
int codebook_kind = 0; ///< 0=PER_SUBSPACE, 1=PER_CLUSTER.
int codes_layout = 1; ///< 0=FLAT, 1=INTERLEAVED.
bool force_random_rotation = false; ///< Always apply a random rotation.
uint32_t max_train_points_per_pq_code = 256; ///< Training points per PQ code.
int lut_dtype = 0; ///< 0=float32, 1=float16, 2=uint8.
int internal_distance_dtype = 0; ///< 0=float32, 1=float16.
int coarse_search_dtype = 0; ///< 0=float32, 1=float16, 3=int8.
uint32_t max_internal_batch_size = 4096; ///< Maximum internal search batch.
};

/**
Expand Down
65 changes: 54 additions & 11 deletions cpp/src/knn/knn.cu
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include <cuml/neighbors/knn.hpp>

#include <raft/core/device_resources.hpp>
#include <raft/core/error.hpp>
#include <raft/core/handle.hpp>
#include <raft/core/operators.hpp>
#include <raft/label/classlabels.cuh>
Expand Down Expand Up @@ -40,10 +41,26 @@ struct knnIndexImpl {
std::unique_ptr<cuvs::neighbors::ivf_flat::index<float, int64_t>> ivf_flat;
std::unique_ptr<cuvs::neighbors::ivf_pq::index<int64_t>> ivf_pq;

int pq_lut_dtype = 0;
int pq_internal_distance_dtype = 0;
int pq_coarse_search_dtype = 0;
uint32_t pq_max_internal_batch_size = 4096;

std::unique_ptr<rmm::device_uvector<float>> corr_norms;
std::unique_ptr<rmm::device_uvector<float>> corr_means;
};

auto ivfpq_dtype_from_code(int code) -> cudaDataType_t
{
switch (code) {
case 0: return CUDA_R_32F;
case 1: return CUDA_R_16F;
case 2: return CUDA_R_8U;
case 3: return CUDA_R_8I;
default: RAFT_FAIL("Invalid IVF-PQ dtype code.");
}
}

knnIndex::knnIndex() : pimpl{std::make_unique<knnIndexImpl>()} {}
knnIndex::~knnIndex() = default;

Expand Down Expand Up @@ -280,21 +297,42 @@ void approx_knn_build_index(raft::handle_t& handle,
if (ivf_ft_pams) {
index->nprobe = ivf_ft_pams->nprobe;
cuvs::neighbors::ivf_flat::index_params params;
params.metric = static_cast<cuvs::distance::DistanceType>(metric);
params.metric_arg = metricArg;
params.n_lists = ivf_ft_pams->nlist;
params.metric = static_cast<cuvs::distance::DistanceType>(metric);
params.metric_arg = metricArg;
params.n_lists = ML::narrow_cast<uint32_t>(ivf_ft_pams->nlist);
params.kmeans_n_iters = ivf_ft_pams->kmeans_n_iters;
params.kmeans_trainset_fraction = ivf_ft_pams->kmeans_trainset_fraction;
params.conservative_memory_allocation = ivf_ft_pams->conservative_memory_allocation;

index->pimpl->ivf_flat = std::make_unique<cuvs::neighbors::ivf_flat::index<float, int64_t>>(
cuvs::neighbors::ivf_flat::build(handle, params, index_view));
} else if (ivf_pq_pams) {
index->nprobe = ivf_pq_pams->nprobe;
cuvs::neighbors::ivf_pq::index_params params;
params.metric = static_cast<cuvs::distance::DistanceType>(metric);
params.metric_arg = metricArg;
params.n_lists = ivf_pq_pams->nlist;
params.pq_bits = ivf_pq_pams->n_bits;
params.pq_dim = ivf_pq_pams->M;
// TODO: handle ivf_pq_pams.usePrecomputedTables ?
params.metric = static_cast<cuvs::distance::DistanceType>(metric);
params.metric_arg = metricArg;
params.n_lists = ML::narrow_cast<uint32_t>(ivf_pq_pams->nlist);
params.kmeans_n_iters = ivf_pq_pams->kmeans_n_iters;
params.kmeans_trainset_fraction = ivf_pq_pams->kmeans_trainset_fraction;
params.pq_bits = ML::narrow_cast<uint32_t>(ivf_pq_pams->n_bits);
params.pq_dim = ML::narrow_cast<uint32_t>(ivf_pq_pams->M);
RAFT_EXPECTS(ivf_pq_pams->codebook_kind == 0 || ivf_pq_pams->codebook_kind == 1,
"Invalid IVF-PQ codebook_kind.");
RAFT_EXPECTS(ivf_pq_pams->codes_layout == 0 || ivf_pq_pams->codes_layout == 1,
"Invalid IVF-PQ codes_layout.");

params.codebook_kind =
static_cast<cuvs::neighbors::ivf_pq::codebook_gen>(ivf_pq_pams->codebook_kind);
params.codes_layout =
static_cast<cuvs::neighbors::ivf_pq::list_layout>(ivf_pq_pams->codes_layout);
params.force_random_rotation = ivf_pq_pams->force_random_rotation;
params.conservative_memory_allocation = ivf_pq_pams->conservative_memory_allocation;
params.max_train_points_per_pq_code = ivf_pq_pams->max_train_points_per_pq_code;

index->pimpl->pq_lut_dtype = ivf_pq_pams->lut_dtype;
index->pimpl->pq_internal_distance_dtype = ivf_pq_pams->internal_distance_dtype;
index->pimpl->pq_coarse_search_dtype = ivf_pq_pams->coarse_search_dtype;
index->pimpl->pq_max_internal_batch_size = ivf_pq_pams->max_internal_batch_size;

index->pimpl->ivf_pq = std::make_unique<cuvs::neighbors::ivf_pq::index<int64_t>>(
cuvs::neighbors::ivf_pq::build(handle, params, index_view));
Expand Down Expand Up @@ -353,15 +391,20 @@ void approx_knn_search(raft::handle_t& handle,
auto query_view = raft::make_device_matrix_view<const float, int64_t>(
query_array, n, index->pimpl->ivf_flat->dim());
cuvs::neighbors::ivf_flat::search_params params;
params.n_probes = index->nprobe;
params.n_probes = ML::narrow_cast<uint32_t>(index->nprobe);

cuvs::neighbors::ivf_flat::search(
handle, params, *index->pimpl->ivf_flat, query_view, indices_view, distances_view);
} else if (index->pimpl->ivf_pq) {
auto query_view = raft::make_device_matrix_view<const float, int64_t>(
query_array, n, index->pimpl->ivf_pq->dim());
cuvs::neighbors::ivf_pq::search_params params;
params.n_probes = index->nprobe;
params.n_probes = ML::narrow_cast<uint32_t>(index->nprobe);
params.lut_dtype = ivfpq_dtype_from_code(index->pimpl->pq_lut_dtype);
params.internal_distance_dtype =
ivfpq_dtype_from_code(index->pimpl->pq_internal_distance_dtype);
params.coarse_search_dtype = ivfpq_dtype_from_code(index->pimpl->pq_coarse_search_dtype);
params.max_internal_batch_size = index->pimpl->pq_max_internal_batch_size;

cuvs::neighbors::ivf_pq::search(
handle, params, *(index->pimpl->ivf_pq), query_view, indices_view, distances_view);
Expand Down
Loading
Loading

AltStyle によって変換されたページ (->オリジナル) /