From 25630a05581b003b7c0c393748af56fdaeb82f10 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Thu, 27 Jun 2024 16:22:35 -0400 Subject: [PATCH] Adding IVF examples (#203) Authors: - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Divye Gala (https://github.com/divyegala) URL: https://github.com/rapidsai/cuvs/pull/203 --- examples/build.sh | 2 +- examples/cpp/CMakeLists.txt | 4 + examples/cpp/src/ivf_flat_example.cu | 155 +++++++++++++++++++++++++++ examples/cpp/src/ivf_pq_example.cu | 110 +++++++++++++++++++ 4 files changed, 270 insertions(+), 1 deletion(-) create mode 100644 examples/cpp/src/ivf_flat_example.cu create mode 100644 examples/cpp/src/ivf_pq_example.cu diff --git a/examples/build.sh b/examples/build.sh index cd021ca21..59f589af2 100755 --- a/examples/build.sh +++ b/examples/build.sh @@ -33,7 +33,7 @@ if [ "$1" == "clean" ]; then fi ################################################################################ -# Add individual libcudf examples build scripts down below +# Add individual libcuvs examples build scripts down below build_example() { example_dir=${1} diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 15b12b118..84536597e 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -36,7 +36,11 @@ include(../cmake/thirdparty/get_cuvs.cmake) # -------------- compile tasks ----------------- # add_executable(CAGRA_EXAMPLE src/cagra_example.cu) +add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) +add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) # `$` is a generator expression that ensures that targets are # installed in a conda environment, if one exists target_link_libraries(CAGRA_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(IVF_PQ_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(IVF_FLAT_EXAMPLE PRIVATE cuvs::cuvs $) \ No newline at end of file diff --git a/examples/cpp/src/ivf_flat_example.cu b/examples/cpp/src/ivf_flat_example.cu new file mode 100644 index 000000000..0480fa1b9 --- /dev/null +++ b/examples/cpp/src/ivf_flat_example.cu @@ -0,0 +1,155 @@ +/* + * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "common.cuh" + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include + +void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace cuvs::neighbors; + + ivf_flat::index_params index_params; + index_params.n_lists = 1024; + index_params.kmeans_trainset_fraction = 0.1; + index_params.metric = cuvs::distance::DistanceType::L2Expanded; + + std::cout << "Building IVF-Flat index" << std::endl; + auto index = ivf_flat::build(dev_resources, index_params, dataset); + + std::cout << "Number of clusters " << index.n_lists() << ", number of vectors added to index " + << index.size() << std::endl; + + // Create output arrays. + int64_t topk = 10; + int64_t n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + + // Set search parameters. + ivf_flat::search_params search_params; + search_params.n_probes = 50; + + // Search K nearest neighbors for each of the queries. + ivf_flat::search( + dev_resources, search_params, index, queries, neighbors.view(), distances.view()); + + // The call to ivf_flat::search is asynchronous. Before accessing the data, sync by calling + raft::resource::sync_stream(dev_resources); + + print_results(dev_resources, neighbors.view(), distances.view()); +} + +void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace cuvs::neighbors; + + // Define dataset indices. + auto data_indices = raft::make_device_vector(dev_resources, dataset.extent(0)); + thrust::counting_iterator first(0); + thrust::device_ptr ptr(data_indices.data_handle()); + thrust::copy( + raft::resource::get_thrust_policy(dev_resources), first, first + dataset.extent(0), ptr); + + // Sub-sample the dataset to create a training set. + auto trainset = + subsample(dev_resources, dataset, raft::make_const_mdspan(data_indices.view()), 0.1); + + ivf_flat::index_params index_params; + index_params.n_lists = 100; + index_params.metric = cuvs::distance::DistanceType::L2Expanded; + index_params.add_data_on_build = false; + + std::cout << "\nRun k-means clustering using the training set" << std::endl; + auto index = + ivf_flat::build(dev_resources, index_params, raft::make_const_mdspan(trainset.view())); + + std::cout << "Number of clusters " << index.n_lists() << ", number of vectors added to index " + << index.size() << std::endl; + + std::cout << "Filling index with the dataset vectors" << std::endl; + index = ivf_flat::extend(dev_resources, + dataset, + std::make_optional(raft::make_const_mdspan(data_indices.view())), + index); + + std::cout << "Index size after addin dataset vectors " << index.size() << std::endl; + + // Set search parameters. + ivf_flat::search_params search_params; + search_params.n_probes = 10; + + // Create output arrays. + int64_t topk = 10; + int64_t n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + + // Search K nearest neighbors for each queries. + ivf_flat::search( + dev_resources, search_params, index, queries, neighbors.view(), distances.view()); + + // The call to ivf_flat::search is asynchronous. Before accessing the data, sync using: + raft::resource::sync_stream(dev_resources); + + print_results(dev_resources, neighbors.view(), distances.view()); +} + +int main() +{ + raft::device_resources dev_resources; + + // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool. + rmm::mr::pool_memory_resource pool_mr( + rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(&pool_mr); + + // Create input arrays. + int64_t n_samples = 10000; + int64_t n_dim = 3; + int64_t n_queries = 10; + auto dataset = raft::make_device_matrix(dev_resources, n_samples, n_dim); + auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); + generate_dataset(dev_resources, dataset.view(), queries.view()); + + // Simple build and search example. + ivf_flat_build_search_simple(dev_resources, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(queries.view())); + + // Build and extend example. + ivf_flat_build_extend_search(dev_resources, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(queries.view())); +} diff --git a/examples/cpp/src/ivf_pq_example.cu b/examples/cpp/src/ivf_pq_example.cu new file mode 100644 index 000000000..bef978714 --- /dev/null +++ b/examples/cpp/src/ivf_pq_example.cu @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "common.cuh" + +#include +#include +#include +#include + +#include +#include + +#include + +void ivf_pq_build_search(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace cuvs::neighbors; // NOLINT + + ivf_pq::index_params index_params; + index_params.n_lists = 1024; + index_params.kmeans_trainset_fraction = 0.1; + index_params.metric = cuvs::distance::DistanceType::L2Expanded; + index_params.pq_bits = 8; + index_params.pq_dim = 2; + + std::cout << "Building IVF-PQ index" << std::endl; + auto index = ivf_pq::build(dev_resources, index_params, dataset); + + std::cout << "Number of clusters " << index.n_lists() << ", number of vectors added to index " + << index.size() << std::endl; + + // Set search parameters. + ivf_pq::search_params search_params; + search_params.n_probes = 50; + // Set the internal search precision to 16-bit floats; + // usually, this improves the performance at a slight cost to the recall. + search_params.internal_distance_dtype = CUDA_R_16F; + search_params.lut_dtype = CUDA_R_16F; + + // Create output arrays. + int64_t topk = 10; + int64_t n_queries = queries.extent(0); + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + + // Search K nearest neighbors for each of the queries. + ivf_pq::search( + dev_resources, search_params, index, queries, neighbors.view(), distances.view()); + + // Re-ranking operation: refine the initial search results by computing exact distances + int64_t topk_refined = 7; + auto neighbors_refined = + raft::make_device_matrix(dev_resources, n_queries, topk_refined); + auto distances_refined = raft::make_device_matrix(dev_resources, n_queries, topk_refined); + + // Note, refinement requires the original dataset and the queries. + // Don't forget to specify the same distance metric as used by the index. + cuvs::neighbors::refine(dev_resources, + dataset, + queries, + raft::make_const_mdspan(neighbors.view()), + neighbors_refined.view(), + distances_refined.view(), + index.metric()); + + // Show both the original and the refined results + std::cout << std::endl << "Original results:" << std::endl; + print_results(dev_resources, neighbors.view(), distances.view()); + std::cout << std::endl << "Refined results:" << std::endl; + print_results(dev_resources, neighbors_refined.view(), distances_refined.view()); +} + +int main() +{ + raft::device_resources dev_resources; + + // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool. + rmm::mr::pool_memory_resource pool_mr( + rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(&pool_mr); + + // Create input arrays. + int64_t n_samples = 10000; + int64_t n_dim = 3; + int64_t n_queries = 10; + auto dataset = raft::make_device_matrix(dev_resources, n_samples, n_dim); + auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); + generate_dataset(dev_resources, dataset.view(), queries.view()); + + // Simple build and search example. + ivf_pq_build_search(dev_resources, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(queries.view())); +}