diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 049d2c9d5..13b483e73 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -42,11 +42,13 @@ namespace cuvs::neighbors::vamana { * */ struct index_params : cuvs::neighbors::index_params { - /** Degree of output graph. */ + /** Maximum degree of output graph corresponds to the R parameter in the original Vamana + * literature. */ uint32_t graph_degree = 32; - /** Maximum number of visited nodes per search **/ + /** Maximum number of visited nodes per search corresponds to the L parameter in the Vamana + * literature **/ uint32_t visited_size = 64; - /** Number of Vamana iterations. */ + /** Number of Vamana vector insertion iterations (each iteration inserts all vectors). */ uint32_t vamana_iters = 1; /** Alpha for pruning parameter */ float alpha = 1.2; @@ -119,6 +121,9 @@ struct index : cuvs::neighbors::index { return graph_view_; } + /** Return the id of the vector selected as the medoid. */ + [[nodiscard]] inline auto medoid() const noexcept -> IdxT { return medoid_id_; } + // Don't allow copying the index for performance reasons (try avoiding copying data) index(const index&) = delete; index(index&&) = default; @@ -144,11 +149,13 @@ struct index : cuvs::neighbors::index { cuvs::distance::DistanceType metric, raft::mdspan, raft::row_major, data_accessor> dataset, raft::mdspan, raft::row_major, graph_accessor> - vamana_graph) + vamana_graph, + IdxT medoid_id) : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), - dataset_(make_aligned_dataset(res, dataset, 16)) + dataset_(make_aligned_dataset(res, dataset, 16)), + medoid_id_(medoid_id) { RAFT_EXPECTS(dataset.extent(0) == vamana_graph.extent(0), "Dataset and vamana_graph must have equal number of rows"); @@ -197,6 +204,7 @@ struct index : cuvs::neighbors::index { raft::device_matrix graph_; raft::device_matrix_view graph_view_; std::unique_ptr> dataset_; + IdxT medoid_id_; }; /** * @} diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index d1f0d7cb3..605eef2de 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -89,6 +89,7 @@ void batched_insert_vamana( const index_params& params, raft::mdspan, raft::row_major, Accessor> dataset, raft::host_matrix_view graph, + IdxT* medoid_id, cuvs::distance::DistanceType metric, int dim) { @@ -186,7 +187,9 @@ void batched_insert_vamana( } // Random medoid has minor impact on recall - int medoid_id = rand() % N; + // TODO: use heuristic for better medoid selection, issue: + // https://github.com/rapidsai/cuvs/issues/355 + *medoid_id = rand() % N; // size of current batch of inserts, increases logarithmically until max_batchsize int step_size = 1; @@ -213,7 +216,7 @@ void batched_insert_vamana( dataset, query_list_ptr.data_handle(), step_size, - medoid_id, + *medoid_id, degree, dataset.extent(0), visited_size, @@ -381,12 +384,13 @@ index build( cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded; + IdxT medoid_id; batched_insert_vamana( - res, params, dataset, vamana_graph.view(), metric, dim); + res, params, dataset, vamana_graph.view(), &medoid_id, metric, dim); try { return index( - res, params.metric, dataset, raft::make_const_mdspan(vamana_graph.view())); + res, params.metric, dataset, raft::make_const_mdspan(vamana_graph.view()), medoid_id); } catch (std::bad_alloc& e) { RAFT_LOG_DEBUG("Insufficient GPU memory to construct VAMANA index with dataset on GPU"); // We just add the graph. User is expected to update dataset separately (e.g allocating in diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index 484939db6..877680d0c 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -61,7 +61,7 @@ void serialize(raft::resources const& res, index_of.seekp(file_offset, index_of.beg); uint32_t max_degree = 0; size_t index_size = 24; // Starting metadata - uint32_t start = 0; + uint32_t start = static_cast(index_.medoid()); size_t num_frozen_points = 0; uint32_t max_observed_degree = 0; diff --git a/cpp/src/neighbors/detail/vamana/vamana_structs.cuh b/cpp/src/neighbors/detail/vamana/vamana_structs.cuh index 9284ac622..0142208af 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_structs.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_structs.cuh @@ -42,6 +42,7 @@ namespace cuvs::neighbors::vamana::detail { #define FULL_BITMASK 0xFFFFFFFF +// Currently supported values for graph_degree. static const int DEGREE_SIZES[4] = {32, 64, 128, 256}; // Object used to store id,distance combination graph construction operations @@ -114,27 +115,6 @@ class Point { id = other.id; return *this; } - - /* - // Computes Cosine dist. Uses 2 registers to increase pipeline efficiency and ILP - // Assumes coordinates are normalized so each vector is of unit length. This lets us - // perform a dot-product instead of the full cosine distance computation. -// __device__ SUMTYPE cosine(Point* other, bool test) {return NULL;} - __device__ SUMTYPE cosine(Point* other) { - SUMTYPE total[2]={0,0}; - - for(int i=0; icoords[i])); - total[1] += ((SUMTYPE)((SUMTYPE)coords[i+1] * (SUMTYPE)other->coords[i+1])); - } - return (SUMTYPE)1.0 - (total[0]+total[1]); - } - - __forceinline__ __device__ SUMTYPE dist(Point* other, int metric) { - if(metric == 0) return l2(other); - else return cosine(other); - } - */ }; /* L2 fallback for low dimension when ILP is not possible */