From 074353ea73e0f80b7cc918a99b39022ad3cc8468 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Mon, 18 Sep 2023 09:20:05 +0200 Subject: [PATCH 1/5] Add IVF-Flat C++ example --- cpp/template/CMakeLists.txt | 5 + cpp/template/src/ivf_flat_example.cu | 173 +++++++++++++++++++++++++ cpp/template/src/test_vector_search.cu | 16 +-- 3 files changed, 186 insertions(+), 8 deletions(-) create mode 100644 cpp/template/src/ivf_flat_example.cu diff --git a/cpp/template/CMakeLists.txt b/cpp/template/CMakeLists.txt index 44b06e1b5f..41f23fe4a1 100644 --- a/cpp/template/CMakeLists.txt +++ b/cpp/template/CMakeLists.txt @@ -36,3 +36,8 @@ 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) +target_compile_definitions(TEST_RAFT PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") + +add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) +target_link_libraries(IVF_FLAT_EXAMPLE PRIVATE raft::raft raft::compiled) +target_compile_definitions(IVF_FLAT_EXAMPLE PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu new file mode 100644 index 0000000000..b891c3cf4a --- /dev/null +++ b/cpp/template/src/ivf_flat_example.cu @@ -0,0 +1,173 @@ +/* + * 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 + +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 = 12; + 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()); +} + +/** Subsample the dataset to create a training set*/ +raft::device_matrix subsample( + raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + 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 data_indices = raft::make_device_vector(dev_resources, n_samples); + auto train_indices = raft::make_device_vector(dev_resources, n_train); + + thrust::counting_iterator first(0); + thrust::device_ptr ptr(data_indices.data_handle()); + thrust::copy(raft::resource::get_thrust_policy(dev_resources), first, first + n_samples, ptr); + + raft::random::sample_without_replacement(dev_resources, + rng, + raft::make_const_mdspan(data_indices.view()), + 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; +} + +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; + + // Sub-sample the dataset to create a training set. + auto trainset = subsample(dev_resources, dataset, 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; + + auto data_indices = raft::make_device_vector(dev_resources, dataset.extent(1)); + 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()); +} + +int main() +{ + raft::device_resources dev_resources; + + // Set pool allocator with 2 GiB limit for temporary arrays. + 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 labels = raft::make_device_vector(dev_resources, n_samples); + auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); + raft::random::make_blobs(dev_resources, dataset.view(), 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); + + // 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 index f54cfc03e7..cf9dcc2361 100644 --- a/cpp/template/src/test_vector_search.cu +++ b/cpp/template/src/test_vector_search.cu @@ -37,7 +37,7 @@ int main() 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 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()); @@ -45,15 +45,15 @@ int main() // 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( + 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()); + cagra::search(dev_resources, + search_params, + index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); } From c44d751feb6e9f6f1ce997add67f3ec836a4edce Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Mon, 18 Sep 2023 13:08:09 +0200 Subject: [PATCH 2/5] Fix dataset indices, Sync and print results --- cpp/template/src/ivf_flat_example.cu | 67 +++++++++++++++++++++------- 1 file changed, 51 insertions(+), 16 deletions(-) diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu index b891c3cf4a..1e4b790c0c 100644 --- a/cpp/template/src/ivf_flat_example.cu +++ b/cpp/template/src/ivf_flat_example.cu @@ -23,11 +23,38 @@ #include #include #include +#include #include #include #include +// Copy the results to host and print a few samples +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 ivf_flat::search and raft::copy is asyncronous. + // 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); + } +} + void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, raft::device_matrix_view dataset, raft::device_matrix_view queries) @@ -46,7 +73,7 @@ void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, << index.size() << std::endl; // Create output arrays. - int64_t topk = 12; + 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); @@ -58,12 +85,18 @@ void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, // 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 asyncronous. Before accessing the data, sync by calling + // raft::resource::sync_stream(dev_resources); + + print_results(dev_resources, neighbors.view(), distances.view()); } /** 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); @@ -73,19 +106,10 @@ raft::device_matrix subsample( int seed = 137; raft::random::RngState rng(seed); - auto data_indices = raft::make_device_vector(dev_resources, n_samples); auto train_indices = raft::make_device_vector(dev_resources, n_train); - thrust::counting_iterator first(0); - thrust::device_ptr ptr(data_indices.data_handle()); - thrust::copy(raft::resource::get_thrust_policy(dev_resources), first, first + n_samples, ptr); - - raft::random::sample_without_replacement(dev_resources, - rng, - raft::make_const_mdspan(data_indices.view()), - std::nullopt, - train_indices.view(), - std::nullopt); + 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())); @@ -99,8 +123,16 @@ void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, { 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, 0.1); + 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; @@ -115,9 +147,7 @@ void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, << index.size() << std::endl; std::cout << "Filling index with the dataset vectors" << std::endl; - - auto data_indices = raft::make_device_vector(dev_resources, dataset.extent(1)); - index = ivf_flat::extend(dev_resources, + index = ivf_flat::extend(dev_resources, dataset, std::make_optional(raft::make_const_mdspan(data_indices.view())), index); @@ -137,6 +167,11 @@ void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, // 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 asyncronous. Before accessing the data, sync using: + // raft::resource::sync_stream(dev_resources); + + print_results(dev_resources, neighbors.view(), distances.view()); } int main() From f126e51636bf013e1a06fdcb1fb9b73ae8d528a4 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Mon, 18 Sep 2023 16:35:52 +0200 Subject: [PATCH 3/5] Fix typos --- cpp/template/src/ivf_flat_example.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu index 1e4b790c0c..011eb0985d 100644 --- a/cpp/template/src/ivf_flat_example.cu +++ b/cpp/template/src/ivf_flat_example.cu @@ -43,7 +43,7 @@ void print_results(raft::device_resources const& 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 ivf_flat::search and raft::copy is asyncronous. + // The calls to ivf_flat::search and raft::copy is asynchronous. // We need to sync the stream before accessing the data. raft::resource::sync_stream(dev_resources, stream); @@ -86,7 +86,7 @@ void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, ivf_flat::search( dev_resources, search_params, index, queries, neighbors.view(), distances.view()); - // The call to ivf_flat::search is asyncronous. Before accessing the data, sync by calling + // 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()); @@ -168,7 +168,7 @@ void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, ivf_flat::search( dev_resources, search_params, index, queries, neighbors.view(), distances.view()); - // The call to ivf_flat::search is asyncronous. Before accessing the data, sync using: + // 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()); From 1f3e827b8486bc43685b6b57ab0f2491197d0796 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Thu, 21 Sep 2023 22:54:12 +0200 Subject: [PATCH 4/5] Rename CAGRA example, and refactor it --- cpp/template/CMakeLists.txt | 11 ++- cpp/template/src/cagra_example.cu | 90 ++++++++++++++++++++++++ cpp/template/src/common.cuh | 95 ++++++++++++++++++++++++++ cpp/template/src/ivf_flat_example.cu | 78 ++++----------------- cpp/template/src/test_vector_search.cu | 59 ---------------- 5 files changed, 208 insertions(+), 125 deletions(-) create mode 100644 cpp/template/src/cagra_example.cu create mode 100644 cpp/template/src/common.cuh delete mode 100644 cpp/template/src/test_vector_search.cu diff --git a/cpp/template/CMakeLists.txt b/cpp/template/CMakeLists.txt index 41f23fe4a1..24af8ae5dd 100644 --- a/cpp/template/CMakeLists.txt +++ b/cpp/template/CMakeLists.txt @@ -34,10 +34,15 @@ 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) -target_compile_definitions(TEST_RAFT PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") +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) + +# We link the example programs to raft::compiled, to reduce compile time. To ensure that we do not +# implicitly instantiate kernels that are expensive to compile, we add the +# RAF_EXPLICIT_INSTANTIATE_ONLY flag. This is optional, for details see +# https://github.com/rapidsai/raft/blob/branch-23.10/docs/source/developer_guide.md#header-organization-of-expensive-function-templates +target_compile_definitions(CAGRA_EXAMPLE PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") target_compile_definitions(IVF_FLAT_EXAMPLE PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") 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 index 011eb0985d..5d91f8fe8b 100644 --- a/cpp/template/src/ivf_flat_example.cu +++ b/cpp/template/src/ivf_flat_example.cu @@ -19,41 +19,17 @@ #include #include #include -#include #include -#include -#include #include +#include +#include + #include #include #include -// Copy the results to host and print a few samples -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 ivf_flat::search 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); - } -} +#include "common.cuh" void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, raft::device_matrix_view dataset, @@ -92,31 +68,6 @@ void ivf_flat_build_search_simple(raft::device_resources const& dev_resources, print_results(dev_resources, neighbors.view(), distances.view()); } -/** 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; -} - void ivf_flat_build_extend_search(raft::device_resources const& dev_resources, raft::device_matrix_view dataset, raft::device_matrix_view queries) @@ -178,23 +129,24 @@ int main() { raft::device_resources dev_resources; - // Set pool allocator with 2 GiB limit for temporary arrays. - raft::resource::set_workspace_to_pool_resource(dev_resources, 2 * 1024 * 1024 * 1024ull); + // 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 labels = raft::make_device_vector(dev_resources, n_samples); auto queries = raft::make_device_matrix(dev_resources, n_queries, n_dim); - raft::random::make_blobs(dev_resources, dataset.view(), 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); + generate_dataset(dev_resources, dataset.view(), queries.view()); // Simple build and search example. ivf_flat_build_search_simple(dev_resources, diff --git a/cpp/template/src/test_vector_search.cu b/cpp/template/src/test_vector_search.cu deleted file mode 100644 index cf9dcc2361..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()); -} From e9c5b341110541434b2ef9b123a1f180787c5fa6 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Fri, 22 Sep 2023 17:12:23 +0200 Subject: [PATCH 5/5] Remove RAFT_ENABLE_EXPLICIT_INSTANTIATE_ONLY flag --- cpp/template/CMakeLists.txt | 6 ------ 1 file changed, 6 deletions(-) diff --git a/cpp/template/CMakeLists.txt b/cpp/template/CMakeLists.txt index 24af8ae5dd..a1341f3609 100644 --- a/cpp/template/CMakeLists.txt +++ b/cpp/template/CMakeLists.txt @@ -40,9 +40,3 @@ 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) -# We link the example programs to raft::compiled, to reduce compile time. To ensure that we do not -# implicitly instantiate kernels that are expensive to compile, we add the -# RAF_EXPLICIT_INSTANTIATE_ONLY flag. This is optional, for details see -# https://github.com/rapidsai/raft/blob/branch-23.10/docs/source/developer_guide.md#header-organization-of-expensive-function-templates -target_compile_definitions(CAGRA_EXAMPLE PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY") -target_compile_definitions(IVF_FLAT_EXAMPLE PRIVATE "RAFT_EXPLICIT_INSTANTIATE_ONLY")