diff --git a/cpp/template/CMakeLists.txt b/cpp/template/CMakeLists.txt index 44b06e1b5f..a1341f3609 100644 --- a/cpp/template/CMakeLists.txt +++ b/cpp/template/CMakeLists.txt @@ -34,5 +34,9 @@ rapids_cpm_init() include(cmake/thirdparty/get_raft.cmake) # -------------- compile tasks ----------------- # -add_executable(TEST_RAFT src/test_vector_search.cu) -target_link_libraries(TEST_RAFT PRIVATE raft::raft raft::compiled) +add_executable(CAGRA_EXAMPLE src/cagra_example.cu) +target_link_libraries(CAGRA_EXAMPLE PRIVATE raft::raft raft::compiled) + +add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) +target_link_libraries(IVF_FLAT_EXAMPLE PRIVATE raft::raft raft::compiled) + diff --git a/cpp/template/src/cagra_example.cu b/cpp/template/src/cagra_example.cu new file mode 100644 index 0000000000..7f3a7d6676 --- /dev/null +++ b/cpp/template/src/cagra_example.cu @@ -0,0 +1,90 @@ +/* + * Copyright (c) 2022-2023, 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 +#include +#include +#include +#include + +#include +#include + +#include "common.cuh" + +void cagra_build_search_simple(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace raft::neighbors; + + int64_t topk = 12; + int64_t n_queries = queries.extent(0); + + // create output arrays + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + + // use default index parameters + cagra::index_params index_params; + + std::cout << "Building CAGRA index (search graph)" << std::endl; + auto index = cagra::build(dev_resources, index_params, dataset); + + std::cout << "CAGRA index has " << index.size() << " vectors" << std::endl; + std::cout << "CAGRA graph has degree " << index.graph_degree() << ", graph size [" + << index.graph().extent(0) << ", " << index.graph().extent(1) << "]" << std::endl; + + // use default search parameters + cagra::search_params search_params; + // search K nearest neighbors + cagra::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()); +} + +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); + + // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT + // algorithms). In that case only the internal arrays would use the pool, any other allocation + // uses the default RMM memory resource. Here is how to change the workspace memory resource to + // a pool with 2 GiB upper limit. + // raft::resource::set_workspace_to_pool_resource(dev_resources, 2 * 1024 * 1024 * 1024ull); + + // Create input arrays. + int64_t n_samples = 10000; + int64_t n_dim = 90; + 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. + cagra_build_search_simple(dev_resources, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(queries.view())); +} diff --git a/cpp/template/src/common.cuh b/cpp/template/src/common.cuh new file mode 100644 index 0000000000..0b72d3bf3b --- /dev/null +++ b/cpp/template/src/common.cuh @@ -0,0 +1,95 @@ +/* + * Copyright (c) 2023, 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 +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// Fill dataset and queries with synthetic data. +void generate_dataset(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + auto labels = raft::make_device_vector(dev_resources, dataset.extent(0)); + raft::random::make_blobs(dev_resources, dataset, labels.view()); + raft::random::RngState r(1234ULL); + raft::random::uniform(dev_resources, + r, + raft::make_device_vector_view(queries.data_handle(), queries.size()), + -1.0f, + 1.0f); +} + +// Copy the results to host and print a few samples +template +void print_results(raft::device_resources const& dev_resources, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances) +{ + int64_t topk = neighbors.extent(1); + auto neighbors_host = raft::make_host_matrix(neighbors.extent(0), topk); + auto distances_host = raft::make_host_matrix(distances.extent(0), topk); + + cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); + + raft::copy(neighbors_host.data_handle(), neighbors.data_handle(), neighbors.size(), stream); + raft::copy(distances_host.data_handle(), distances.data_handle(), distances.size(), stream); + + // The calls to RAFT algorithms and raft::copy is asynchronous. + // We need to sync the stream before accessing the data. + raft::resource::sync_stream(dev_resources, stream); + + for (int query_id = 0; query_id < 2; query_id++) { + std::cout << "Query " << query_id << " neighbor indices: "; + raft::print_host_vector("", &neighbors_host(query_id, 0), topk, std::cout); + std::cout << "Query " << query_id << " neighbor distances: "; + raft::print_host_vector("", &distances_host(query_id, 0), topk, std::cout); + } +} + +/** Subsample the dataset to create a training set*/ +raft::device_matrix subsample( + raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_vector_view data_indices, + float fraction) +{ + int64_t n_samples = dataset.extent(0); + int64_t n_dim = dataset.extent(1); + int64_t n_train = n_samples * fraction; + auto trainset = raft::make_device_matrix(dev_resources, n_train, n_dim); + + int seed = 137; + raft::random::RngState rng(seed); + auto train_indices = raft::make_device_vector(dev_resources, n_train); + + raft::random::sample_without_replacement( + dev_resources, rng, data_indices, std::nullopt, train_indices.view(), std::nullopt); + + raft::matrix::copy_rows( + dev_resources, dataset, trainset.view(), raft::make_const_mdspan(train_indices.view())); + + return trainset; +} diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu new file mode 100644 index 0000000000..5d91f8fe8b --- /dev/null +++ b/cpp/template/src/ivf_flat_example.cu @@ -0,0 +1,160 @@ +/* + * Copyright (c) 2023, 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 +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include "common.cuh" + +void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace raft::neighbors; + + ivf_flat::index_params index_params; + index_params.n_lists = 1024; + index_params.kmeans_trainset_fraction = 0.1; + index_params.metric = raft::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 raft::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 = raft::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); + + // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT + // algorithms). In that case only the internal arrays would use the pool, any other allocation + // uses the default RMM memory resource. Here is how to change the workspace memory resource to + // a pool with 2 GiB upper limit. + // raft::resource::set_workspace_to_pool_resource(dev_resources, 2 * 1024 * 1024 * 1024ull); + + // 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/cpp/template/src/test_vector_search.cu b/cpp/template/src/test_vector_search.cu deleted file mode 100644 index f54cfc03e7..0000000000 --- a/cpp/template/src/test_vector_search.cu +++ /dev/null @@ -1,59 +0,0 @@ -/* - * Copyright (c) 2022-2023, 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 -#include -#include -#include -#include - -int main() -{ - using namespace raft::neighbors; - raft::device_resources dev_resources; - // Use 5 GB of pool memory - raft::resource::set_workspace_to_pool_resource( - dev_resources, std::make_optional(5 * 1024 * 1024 * 1024ull)); - - int64_t n_samples = 50000; - int64_t n_dim = 90; - int64_t topk = 12; - int64_t n_queries = 1; - - // create input and output arrays - auto input = raft::make_device_matrix(dev_resources, n_samples, n_dim); - auto labels = raft::make_device_vector(dev_resources, n_samples); - auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); - auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); - auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); - - raft::random::make_blobs(dev_resources, input.view(), labels.view()); - - // use default index parameters - cagra::index_params index_params; - // create and fill the index from a [n_samples, n_dim] input - auto index = cagra::build( - dev_resources, index_params, raft::make_const_mdspan(input.view())); - // use default search parameters - cagra::search_params search_params; - // search K nearest neighbors - cagra::search(dev_resources, - search_params, - index, - raft::make_const_mdspan(queries.view()), - neighbors.view(), - distances.view()); -}