diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index 487ada503d..52bd3175a6 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -127,10 +127,10 @@ typedef struct cuvsIvfPqParams* cuvsIvfPqParams_t; /** * Parameters for ACE (Augmented Core Extraction) graph build. - * ACE enables building indices for datasets too large to fit in GPU memory by: + * ACE enables building indexes for datasets too large to fit in GPU memory by: * 1. Partitioning the dataset in core (closest) and augmented (second-closest) * partitions using balanced k-means. - * 2. Building sub-indices for each partition independently + * 2. Building sub-indexes for each partition independently * 3. Concatenating sub-graphs into a final unified index */ struct cuvsAceParams { @@ -251,22 +251,6 @@ cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params); */ cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params); -/** - * @brief Allocate ACE params, and populate with default values - * - * @param[in] params cuvsAceParams_t to allocate - * @return cuvsError_t - */ -cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params); - -/** - * @brief De-allocate ACE params - * - * @param[in] params - * @return cuvsError_t - */ -cuvsError_t cuvsAceParamsDestroy(cuvsAceParams_t params); - /** * @brief Create CAGRA index parameters similar to an HNSW index * diff --git a/c/include/cuvs/neighbors/hnsw.h b/c/include/cuvs/neighbors/hnsw.h index d526a4f877..af08a3058c 100644 --- a/c/include/cuvs/neighbors/hnsw.h +++ b/c/include/cuvs/neighbors/hnsw.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -36,6 +36,50 @@ enum cuvsHnswHierarchy { GPU = 2 }; +/** + * Parameters for ACE (Augmented Core Extraction) graph build for HNSW. + * ACE enables building indexes for datasets too large to fit in GPU memory by: + * 1. Partitioning the dataset in core and augmented partitions using balanced k-means + * 2. Building sub-indexes for each partition independently + * 3. Concatenating sub-graphs into a final unified index + */ +struct cuvsHnswAceParams { + /** + * Number of partitions for ACE partitioned build. + * Small values might improve recall but potentially degrade performance and + * increase memory usage. 100k - 5M vectors per partition is recommended. + */ + size_t npartitions; + /** + * Directory to store ACE build artifacts (e.g., KNN graph, optimized graph). + * Used when `use_disk` is true or when the graph does not fit in memory. + */ + const char* build_dir; + /** + * Whether to use disk-based storage for ACE build. + * When true, enables disk-based operations for memory-efficient graph construction. + */ + bool use_disk; +}; + +typedef struct cuvsHnswAceParams* cuvsHnswAceParams_t; + +/** + * @brief Allocate HNSW ACE params, and populate with default values + * + * @param[in] params cuvsHnswAceParams_t to allocate + * @return cuvsError_t + */ +cuvsError_t cuvsHnswAceParamsCreate(cuvsHnswAceParams_t* params); + +/** + * @brief De-allocate HNSW ACE params + * + * @param[in] params + * @return cuvsError_t + */ +cuvsError_t cuvsHnswAceParamsDestroy(cuvsHnswAceParams_t params); + struct cuvsHnswIndexParams { /* hierarchy of the hnsw index */ enum cuvsHnswHierarchy hierarchy; @@ -49,6 +93,17 @@ struct cuvsHnswIndexParams { is parallelized with the help of CPU threads. */ int num_threads; + /** HNSW M parameter: number of bi-directional links per node (used when building with ACE). + * graph_degree = m * 2, intermediate_graph_degree = m * 3. + */ + size_t M; + /** Distance type for the index. */ + cuvsDistanceType metric; + /** + * Optional: specify ACE parameters for building HNSW index using ACE algorithm. + * Set to nullptr for default behavior (from_cagra conversion). + */ + cuvsHnswAceParams_t ace_params; }; typedef struct cuvsHnswIndexParams* cuvsHnswIndexParams_t; @@ -203,6 +258,77 @@ cuvsError_t cuvsHnswFromCagraWithDataset(cuvsResources_t res, * @} */ +/** + * @defgroup hnsw_c_index_build Build HNSW index using ACE algorithm + * @{ + */ + +/** + * @brief Build an HNSW index using ACE (Augmented Core Extraction) algorithm. + * + * ACE enables building HNSW indexes for datasets too large to fit in GPU memory by: + * 1. Partitioning the dataset using balanced k-means into core and augmented partitions + * 2. Building sub-indexes for each partition independently + * 3. Concatenating sub-graphs into a final unified index + * + * NOTE: This function requires CUDA to be available at runtime. + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params cuvsHnswIndexParams_t with ACE parameters configured + * @param[in] dataset DLManagedTensor* host dataset to build index from + * @param[out] index cuvsHnswIndex_t to return the built HNSW index + * + * @return cuvsError_t + * + * @code{.c} + * #include + * #include + * + * // Create cuvsResources_t + * cuvsResources_t res; + * cuvsResourcesCreate(&res); + * + * // Create ACE parameters + * cuvsHnswAceParams_t ace_params; + * cuvsHnswAceParamsCreate(&ace_params); + * ace_params->npartitions = 4; + * ace_params->use_disk = true; + * ace_params->build_dir = "/tmp/hnsw_ace_build"; + * + * // Create index parameters + * cuvsHnswIndexParams_t params; + * cuvsHnswIndexParamsCreate(¶ms); + * params->hierarchy = GPU; + * params->ace_params = ace_params; + * params->M = 32; + * params->ef_construction = 120; + * + * // Create HNSW index + * cuvsHnswIndex_t hnsw_index; + * cuvsHnswIndexCreate(&hnsw_index); + * + * // Assume dataset is a populated DLManagedTensor with host data + * DLManagedTensor dataset; + * + * // Build the index + * cuvsHnswBuild(res, params, &dataset, hnsw_index); + * + * // Clean up + * cuvsHnswAceParamsDestroy(ace_params); + * cuvsHnswIndexParamsDestroy(params); + * cuvsHnswIndexDestroy(hnsw_index); + * cuvsResourcesDestroy(res); + * @endcode + */ +cuvsError_t cuvsHnswBuild(cuvsResources_t res, + cuvsHnswIndexParams_t params, + DLManagedTensor* dataset, + cuvsHnswIndex_t index); + +/** + * @} + */ + /** * @defgroup hnsw_c_index_extend Extend HNSW index with additional vectors * @{ diff --git a/c/src/neighbors/hnsw.cpp b/c/src/neighbors/hnsw.cpp index ac517689fd..547dc55146 100644 --- a/c/src/neighbors/hnsw.cpp +++ b/c/src/neighbors/hnsw.cpp @@ -1,6 +1,6 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -23,6 +23,36 @@ namespace { +template +void _build(cuvsResources_t res, + cuvsHnswIndexParams_t params, + DLManagedTensor* dataset_tensor, + cuvsHnswIndex_t hnsw_index) +{ + auto res_ptr = reinterpret_cast(res); + auto cpp_params = cuvs::neighbors::hnsw::index_params(); + cpp_params.hierarchy = static_cast(params->hierarchy); + cpp_params.ef_construction = params->ef_construction; + cpp_params.num_threads = params->num_threads; + cpp_params.M = params->M; + cpp_params.metric = static_cast(params->metric); + + // Configure ACE parameters + RAFT_EXPECTS(params->ace_params != nullptr, "ACE parameters must be set for hnsw::build"); + auto ace_params = cuvs::neighbors::hnsw::graph_build_params::ace_params(); + ace_params.npartitions = params->ace_params->npartitions; + ace_params.build_dir = params->ace_params->build_dir ? params->ace_params->build_dir : "/tmp/hnsw_ace_build"; + ace_params.use_disk = params->ace_params->use_disk; + cpp_params.graph_build_params = ace_params; + + using dataset_mdspan_type = raft::host_matrix_view; + auto dataset_mds = cuvs::core::from_dlpack(dataset_tensor); + + auto hnsw_index_unique_ptr = cuvs::neighbors::hnsw::build(*res_ptr, cpp_params, dataset_mds); + auto hnsw_index_ptr = hnsw_index_unique_ptr.release(); + hnsw_index->addr = reinterpret_cast(hnsw_index_ptr); +} + template void _from_cagra(cuvsResources_t res, cuvsHnswIndexParams_t params, @@ -118,11 +148,29 @@ void* _deserialize(cuvsResources_t res, } } // namespace +extern "C" cuvsError_t cuvsHnswAceParamsCreate(cuvsHnswAceParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + *params = new cuvsHnswAceParams{.npartitions = 1, + .build_dir = "/tmp/hnsw_ace_build", + .use_disk = false}; + }); +} + +extern "C" cuvsError_t cuvsHnswAceParamsDestroy(cuvsHnswAceParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} + extern "C" cuvsError_t cuvsHnswIndexParamsCreate(cuvsHnswIndexParams_t* params) { return cuvs::core::translate_exceptions([=] { - *params = new cuvsHnswIndexParams{ - .hierarchy = cuvsHnswHierarchy::NONE, .ef_construction = 200, .num_threads = 0}; + *params = new cuvsHnswIndexParams{.hierarchy = cuvsHnswHierarchy::NONE, + .ef_construction = 200, + .num_threads = 0, + .M = 32, + .metric = L2Expanded, + .ace_params = nullptr}; }); } @@ -213,6 +261,29 @@ extern "C" cuvsError_t cuvsHnswFromCagraWithDataset(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsHnswBuild(cuvsResources_t res, + cuvsHnswIndexParams_t params, + DLManagedTensor* dataset, + cuvsHnswIndex_t index) +{ + return cuvs::core::translate_exceptions([=] { + auto dataset_dl = dataset->dl_tensor; + index->dtype = dataset_dl.dtype; + + if (dataset_dl.dtype.code == kDLFloat && dataset_dl.dtype.bits == 32) { + _build(res, params, dataset, index); + } else if (dataset_dl.dtype.code == kDLFloat && dataset_dl.dtype.bits == 16) { + _build(res, params, dataset, index); + } else if (dataset_dl.dtype.code == kDLUInt && dataset_dl.dtype.bits == 8) { + _build(res, params, dataset, index); + } else if (dataset_dl.dtype.code == kDLInt && dataset_dl.dtype.bits == 8) { + _build(res, params, dataset, index); + } else { + RAFT_FAIL("Unsupported dtype: code=%d, bits=%d", dataset_dl.dtype.code, dataset_dl.dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsHnswExtend(cuvsResources_t res, cuvsHnswExtendParams_t params, DLManagedTensor* additional_dataset, diff --git a/c/tests/neighbors/ann_hnsw_c.cu b/c/tests/neighbors/ann_hnsw_c.cu index 8cbfb30417..3b8d3375f9 100644 --- a/c/tests/neighbors/ann_hnsw_c.cu +++ b/c/tests/neighbors/ann_hnsw_c.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -120,3 +120,216 @@ TEST(CagraHnswC, BuildSearch) cuvsHnswIndexDestroy(hnsw_index); cuvsResourcesDestroy(res); } + +TEST(HnswAceC, BuildSearch) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create ACE params + cuvsHnswAceParams_t ace_params; + cuvsHnswAceParamsCreate(&ace_params); + ace_params->npartitions = 2; + ace_params->build_dir = "/tmp/hnsw_ace_test"; + ace_params->use_disk = false; + + // create index params + cuvsHnswIndexParams_t hnsw_params; + cuvsHnswIndexParamsCreate(&hnsw_params); + hnsw_params->hierarchy = GPU; + hnsw_params->ef_construction = 100; + hnsw_params->ace_params = ace_params; + hnsw_params->metric = L2Expanded; + hnsw_params->M = 16; + + // create HNSW index + cuvsHnswIndex_t hnsw_index; + cuvsHnswIndexCreate(&hnsw_index); + + // build index using ACE + cuvsError_t build_status = cuvsHnswBuild(res, hnsw_params, &dataset_tensor, hnsw_index); + ASSERT_EQ(build_status, CUVS_SUCCESS) << "cuvsHnswBuild failed"; + + // create queries DLTensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries; + queries_tensor.dl_tensor.device.device_type = kDLCPU; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + std::vector neighbors(4); + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + std::vector distances(4); + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances.data(); + distances_tensor.dl_tensor.device.device_type = kDLCPU; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + // search index + cuvsHnswSearchParams_t search_params; + cuvsHnswSearchParamsCreate(&search_params); + search_params->ef = 100; + + cuvsError_t search_status = cuvsHnswSearch( + res, search_params, hnsw_index, &queries_tensor, &neighbors_tensor, &distances_tensor); + ASSERT_EQ(search_status, CUVS_SUCCESS) << "cuvsHnswSearch failed"; + + // verify output + ASSERT_TRUE(cuvs::hostVecMatch(neighbors_exp, neighbors, cuvs::Compare())); + ASSERT_TRUE(cuvs::hostVecMatch(distances_exp, distances, cuvs::CompareApprox(0.001f))); + + // cleanup + cuvsHnswAceParamsDestroy(ace_params); + cuvsHnswIndexParamsDestroy(hnsw_params); + cuvsHnswSearchParamsDestroy(search_params); + cuvsHnswIndexDestroy(hnsw_index); + cuvsResourcesDestroy(res); +} + +TEST(HnswAceDiskC, BuildSerializeDeserializeSearch) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create ACE params with use_disk = true + cuvsHnswAceParams_t ace_params; + cuvsHnswAceParamsCreate(&ace_params); + ace_params->npartitions = 2; + ace_params->build_dir = "/tmp/hnsw_ace_disk_test"; + ace_params->use_disk = true; + + // create index params + cuvsHnswIndexParams_t hnsw_params; + cuvsHnswIndexParamsCreate(&hnsw_params); + hnsw_params->hierarchy = GPU; + hnsw_params->ef_construction = 100; + hnsw_params->ace_params = ace_params; + hnsw_params->metric = L2Expanded; + hnsw_params->M = 16; + + // create HNSW index + cuvsHnswIndex_t hnsw_index; + cuvsHnswIndexCreate(&hnsw_index); + + // build index using ACE with disk mode (index is serialized to disk by the build function) + cuvsError_t build_status = cuvsHnswBuild(res, hnsw_params, &dataset_tensor, hnsw_index); + ASSERT_EQ(build_status, CUVS_SUCCESS) << "cuvsHnswBuild failed"; + + // create a new HNSW index for deserialization + cuvsHnswIndex_t deserialized_index; + cuvsHnswIndexCreate(&deserialized_index); + deserialized_index->dtype = hnsw_index->dtype; + + // deserialize from disk + cuvsError_t deserialize_status = + cuvsHnswDeserialize(res, hnsw_params, "/tmp/hnsw_ace_disk_test/hnsw_index.bin", 2, L2Expanded, deserialized_index); + ASSERT_EQ(deserialize_status, CUVS_SUCCESS) << "cuvsHnswDeserialize failed"; + + // create queries DLTensor + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries; + queries_tensor.dl_tensor.device.device_type = kDLCPU; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + std::vector neighbors(4); + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCPU; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 64; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + std::vector distances(4); + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances.data(); + distances_tensor.dl_tensor.device.device_type = kDLCPU; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + // search the deserialized index + cuvsHnswSearchParams_t search_params; + cuvsHnswSearchParamsCreate(&search_params); + search_params->ef = 100; + + cuvsError_t search_status = cuvsHnswSearch( + res, search_params, deserialized_index, &queries_tensor, &neighbors_tensor, &distances_tensor); + ASSERT_EQ(search_status, CUVS_SUCCESS) << "cuvsHnswSearch failed"; + + // verify output + ASSERT_TRUE(cuvs::hostVecMatch(neighbors_exp, neighbors, cuvs::Compare())); + ASSERT_TRUE(cuvs::hostVecMatch(distances_exp, distances, cuvs::CompareApprox(0.001f))); + + // cleanup + cuvsHnswAceParamsDestroy(ace_params); + cuvsHnswIndexParamsDestroy(hnsw_params); + cuvsHnswSearchParamsDestroy(search_params); + cuvsHnswIndexDestroy(hnsw_index); + cuvsHnswIndexDestroy(deserialized_index); + cuvsResourcesDestroy(res); +} diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index 00c70d72ad..a503ec21ac 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -10,6 +10,7 @@ #include "common.hpp" #include +#include #include "cagra.hpp" #include @@ -19,9 +20,13 @@ #include #include #include +#include namespace cuvs::neighbors::hnsw { +// Re-export graph_build_params into hnsw namespace for convenience +namespace graph_build_params = cuvs::neighbors::graph_build_params; + /** * @defgroup hnsw_cpp_index_params hnswlib index wrapper params * @{ @@ -51,6 +56,32 @@ struct index_params : cuvs::neighbors::index_params { is parallelized with the help of CPU threads. */ int num_threads = 0; + + /** HNSW M parameter: number of bi-directional links per node (used when building with ACE). + */ + size_t M = 32; + + /** Parameters to fine tune GPU graph building. By default we select the parameters based on + * dataset shape and HNSW build parameters. You can override these parameters to fine tune the + * graph building process as described in the CAGRA build docs. + * + * Set ace_params to configure ACE (Augmented Core Extraction) parameters for building + * a GPU-accelerated HNSW index. ACE enables building indexes for datasets too large + * to fit in GPU memory. + * + * @code{.cpp} + * hnsw::index_params params; + * params.M = 32; + * params.ef_construction = 120; + * // Configure ACE parameters + * params.graph_build_params = hnsw::graph_build_params::ace_params(); + * auto& ace = std::get(params.graph_build_params); + * ace.npartitions = 4; + * ace.use_disk = true; + * ace.build_dir = "/tmp/hnsw_ace_build"; + * @endcode + */ + std::variant graph_build_params; }; /** @@ -158,6 +189,247 @@ struct extend_params { int num_threads = 0; }; +/** + * @} + */ + +/** + * @defgroup hnsw_cpp_index_build Build HNSW index on the GPU + * @{ + */ + +/** + * @brief Build an HNSW index on the GPU + * + * The resulting graph is compatible for HNSW search, but is not an exact equivalent of the graph + * built by the HNSW. + * + * The HNSW index construction parameters `M` and `ef_construction` are the main parameters to + * control the graph degree and graph quality. We have additional options that can be used to fine + * tune graph building on the GPU (see `cuvs::neighbors::cagra::index_params`). In case the index + * does not fit the host or GPU memory, we would use disk as temporary storage. In such cases it is + * important to set `ace_params.build_dir` to a fast disk with sufficient storage size. + * + * NOTE: This function requires CUDA headers to be available at compile time. + * + * @param[in] res raft resources + * @param[in] params hnsw index parameters including ACE configuration + * @param[in] dataset a host matrix view to a row-major matrix [n_rows, dim] + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * + * // Create index parameters with ACE configuration + * hnsw::index_params params; + * params.metric = cuvs::distance::DistanceType::L2Expanded; + * params.hierarchy = hnsw::HnswHierarchy::GPU; + * params.M = 32; + * params.ef_construction = 120; + * + * // Configure GPU graph building parameters + * auto ace_params = hnsw::graph_build_params::ace_params(); + * ace_params.npartitions = 4; + * ace_params.use_disk = true; + * ace_params.build_dir = "/tmp/hnsw_ace_build"; + * params.graph_build_params = ace_params; + * + * // Build the index + * auto dataset = raft::make_host_matrix(res, N, D); + * // ... fill dataset ... + * auto hnsw_index = hnsw::build(res, params, raft::make_const_mdspan(dataset.view())); + * + * // Search the index + * hnsw::search_params search_params; + * search_params.ef = 200; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *hnsw_index, queries, neighbors.view(), distances.view()); + * + * // Serialize the index + * hnsw::serialize(res, "index.bin", *hnsw_index); + * @endcode + */ +std::unique_ptr> build( + raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset); + +/** + * @brief Build an HNSW index on the GPU + * + * The resulting graph is compatible for HNSW search, but is not an exact equivalent of the graph + * built by the HNSW. + * + * The HNSW index construction parameters `M` and `ef_construction` are the main parameters to + * control the graph degree and graph quality. We have additional options that can be used to fine + * tune graph building on the GPU (see `cuvs::neighbors::cagra::index_params`). In case the index + * does not fit the host or GPU memory, we would use disk as temporary storage. In such cases it is + * important to set `ace_params.build_dir` to a fast disk with sufficient storage size. + * + * NOTE: This function requires CUDA headers to be available at compile time. + * + * @param[in] res raft resources + * @param[in] params hnsw index parameters including ACE configuration + * @param[in] dataset a host matrix view to a row-major matrix [n_rows, dim] + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * + * // Create index parameters with ACE configuration + * hnsw::index_params params; + * params.metric = cuvs::distance::DistanceType::L2Expanded; + * params.hierarchy = hnsw::HnswHierarchy::GPU; + * params.M = 32; + * params.ef_construction = 120; + * + * // Configure GPU graph building parameters + * auto ace_params = hnsw::graph_build_params::ace_params(); + * ace_params.npartitions = 4; + * ace_params.use_disk = true; + * ace_params.build_dir = "/tmp/hnsw_ace_build"; + * params.graph_build_params = ace_params; + * + * // Build the index + * auto dataset = raft::make_host_matrix(res, N, D); + * // ... fill dataset ... + * auto hnsw_index = hnsw::build(res, params, raft::make_const_mdspan(dataset.view())); + * + * // Search the index + * hnsw::search_params search_params; + * search_params.ef = 200; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *hnsw_index, queries, neighbors.view(), distances.view()); + * + * // Serialize the index + * hnsw::serialize(res, "index.bin", *hnsw_index); + * @endcode + */ +std::unique_ptr> build( + raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset); + +/** + * @brief Build an HNSW index on the GPU + * + * The resulting graph is compatible for HNSW search, but is not an exact equivalent of the graph + * built by the HNSW. + * + * The HNSW index construction parameters `M` and `ef_construction` are the main parameters to + * control the graph degree and graph quality. We have additional options that can be used to fine + * tune graph building on the GPU (see `cuvs::neighbors::cagra::index_params`). In case the index + * does not fit the host or GPU memory, we would use disk as temporary storage. In such cases it is + * important to set `ace_params.build_dir` to a fast disk with sufficient storage size. + * + * NOTE: This function requires CUDA headers to be available at compile time. + * + * @param[in] res raft resources + * @param[in] params hnsw index parameters including ACE configuration + * @param[in] dataset a host matrix view to a row-major matrix [n_rows, dim] + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * + * // Create index parameters with ACE configuration + * hnsw::index_params params; + * params.metric = cuvs::distance::DistanceType::L2Expanded; + * params.hierarchy = hnsw::HnswHierarchy::GPU; + * params.M = 32; + * params.ef_construction = 120; + * + * // Configure GPU graph building parameters + * auto ace_params = hnsw::graph_build_params::ace_params(); + * ace_params.npartitions = 4; + * ace_params.use_disk = true; + * ace_params.build_dir = "/tmp/hnsw_ace_build"; + * params.graph_build_params = ace_params; + * + * // Build the index + * auto dataset = raft::make_host_matrix(res, N, D); + * // ... fill dataset ... + * auto hnsw_index = hnsw::build(res, params, raft::make_const_mdspan(dataset.view())); + * + * // Search the index + * hnsw::search_params search_params; + * search_params.ef = 200; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *hnsw_index, queries, neighbors.view(), distances.view()); + * + * // Serialize the index + * hnsw::serialize(res, "index.bin", *hnsw_index); + * @endcode + */ +std::unique_ptr> build( + raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset); + +/** + * @brief Build an HNSW index on the GPU + * + * The resulting graph is compatible for HNSW search, but is not an exact equivalent of the graph + * built by the HNSW. + * + * The HNSW index construction parameters `M` and `ef_construction` are the main parameters to + * control the graph degree and graph quality. We have additional options that can be used to fine + * tune graph building on the GPU (see `cuvs::neighbors::cagra::index_params`). In case the index + * does not fit the host or GPU memory, we would use disk as temporary storage. In such cases it is + * important to set `ace_params.build_dir` to a fast disk with sufficient storage size. + * + * NOTE: This function requires CUDA headers to be available at compile time. + * + * @param[in] res raft resources + * @param[in] params hnsw index parameters including ACE configuration + * @param[in] dataset a host matrix view to a row-major matrix [n_rows, dim] + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * + * // Create index parameters with ACE configuration + * hnsw::index_params params; + * params.metric = cuvs::distance::DistanceType::L2Expanded; + * params.hierarchy = hnsw::HnswHierarchy::GPU; + * params.M = 32; + * params.ef_construction = 120; + * + * // Configure GPU graph building parameters + * auto ace_params = hnsw::graph_build_params::ace_params(); + * ace_params.npartitions = 4; + * ace_params.use_disk = true; + * ace_params.build_dir = "/tmp/hnsw_ace_build"; + * params.graph_build_params = ace_params; + * + * // Build the index + * auto dataset = raft::make_host_matrix(res, N, D); + * // ... fill dataset ... + * auto hnsw_index = hnsw::build(res, params, raft::make_const_mdspan(dataset.view())); + * + * // Search the index + * hnsw::search_params search_params; + * search_params.ef = 200; + * auto neighbors = raft::make_host_matrix(res, n_queries, k); + * auto distances = raft::make_host_matrix(res, n_queries, k); + * hnsw::search(res, search_params, *hnsw_index, queries, neighbors.view(), distances.view()); + * + * // Serialize the index + * hnsw::serialize(res, "index.bin", *hnsw_index); + * @endcode + */ +std::unique_ptr> build( + raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset); + /** * @} */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 675ce34c57..8cc1ea9cee 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -793,10 +793,10 @@ void ace_load_partition_dataset_from_disk( } // Build CAGRA index using ACE (Augmented Core Extraction) partitioning -// ACE enables building indices for datasets too large to fit in GPU memory by: +// ACE enables building indexes for datasets too large to fit in GPU memory by: // 1. Partitioning the dataset using balanced k-means in core (non-overlapping) and augmented // (second-closest) partitions -// 2. Building sub-indices for each partition independently +// 2. Building sub-indexes for each partition independently // 3. Concatenating sub-graphs (of core partitions) into a final unified index // Supports both in-memory and disk-based modes depending on available host memory. // In disk mode, the graph is stored in build_dir and dataset is reordered on disk. diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 186216a4da..d9b8b2216a 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -96,7 +97,11 @@ struct index_impl : index { /** @brief Set ef for search */ - void set_ef(int ef) const override { appr_alg_->ef_ = ef; } + void set_ef(int ef) const override + { + ensure_loaded(); + appr_alg_->ef_ = ef; + } /** @brief Set index @@ -136,8 +141,42 @@ struct index_impl : index { return ""; } + /** + @brief Ensure the index is loaded into memory. + If the index is disk-backed and not yet loaded, this will load it from the file. + */ + void ensure_loaded() const + { + if (appr_alg_ != nullptr) { return; } // Already loaded + + // Check if we have a file descriptor to load from + if (!hnsw_fd_.has_value() || !hnsw_fd_->is_valid()) { + RAFT_FAIL("Cannot load HNSW index: no file descriptor available and index not in memory"); + } + + std::string filepath = hnsw_fd_->get_path(); + RAFT_EXPECTS(!filepath.empty(), "Cannot load HNSW index: file path is empty"); + RAFT_EXPECTS(std::filesystem::exists(filepath), + "Cannot load HNSW index: file does not exist: %s", + filepath.c_str()); + + RAFT_LOG_INFO("Loading HNSW index from disk: %s", filepath.c_str()); + + try { + appr_alg_ = std::make_unique::type>>( + space_.get(), filepath); + if (this->hierarchy() == HnswHierarchy::NONE) { appr_alg_->base_layer_only = true; } + } catch (const std::bad_alloc& e) { + RAFT_FAIL( + "Failed to load HNSW index from '%s': insufficient host memory. " + "The index is too large to fit in available RAM. " + "Consider using a machine with more memory or reducing the dataset size.", + filepath.c_str()); + } + } + private: - std::unique_ptr::type>> appr_alg_; + mutable std::unique_ptr::type>> appr_alg_; std::unique_ptr::type>> space_; std::optional hnsw_fd_; }; @@ -1090,10 +1129,10 @@ void extend(raft::resources const& res, raft::host_matrix_view additional_dataset, index& idx) { + // If the index is disk-backed, load it into memory first auto* idx_impl = dynamic_cast*>(&idx); - RAFT_EXPECTS(!idx_impl || !idx_impl->file_descriptor().has_value(), - "Cannot extend an HNSW index that is stored on disk. " - "The index must be deserialized into memory first using hnsw::deserialize()."); + if (idx_impl) { idx_impl->ensure_loaded(); } + auto* hnswlib_index = reinterpret_cast::type>*>( const_cast(idx.get_index())); auto current_element_count = hnswlib_index->getCurrentElementCount(); @@ -1135,10 +1174,9 @@ void search(raft::resources const& res, raft::host_matrix_view neighbors, raft::host_matrix_view distances) { + // If the index is disk-backed, load it into memory first auto* idx_impl = dynamic_cast*>(&idx); - RAFT_EXPECTS(!idx_impl || !idx_impl->file_descriptor().has_value(), - "Cannot search an HNSW index that is stored on disk. " - "The index must be deserialized into memory first using hnsw::deserialize()."); + if (idx_impl) { idx_impl->ensure_loaded(); } RAFT_EXPECTS(queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0), "Number of rows in output neighbors and distances matrices must equal the number of " @@ -1180,9 +1218,24 @@ template void serialize(raft::resources const& res, const std::string& filename, const index& idx) { auto* idx_impl = dynamic_cast*>(&idx); - RAFT_EXPECTS(!idx_impl || !idx_impl->file_descriptor().has_value(), - "Cannot serialize an HNSW index that is stored on disk. " - "The index must be deserialized into memory first using hnsw::deserialize()."); + + // Check if this is a disk-based index (created from disk-backed CAGRA) + if (idx_impl && idx_impl->file_descriptor().has_value()) { + // For disk-based indexes, copy the existing file to the new location + std::string source_path = idx_impl->file_path(); + RAFT_EXPECTS(!source_path.empty(), "Disk-based index has invalid file path"); + RAFT_EXPECTS(std::filesystem::exists(source_path), + "Disk-based index file does not exist: %s", + source_path.c_str()); + + // Copy the file to the new location + std::filesystem::copy_file( + source_path, filename, std::filesystem::copy_options::overwrite_existing); + RAFT_LOG_INFO( + "Copied disk-based HNSW index from %s to %s", source_path.c_str(), filename.c_str()); + return; + } + auto* hnswlib_index = reinterpret_cast::type>*>( const_cast(idx.get_index())); hnswlib_index->saveIndex(filename); @@ -1196,12 +1249,69 @@ void deserialize(raft::resources const& res, cuvs::distance::DistanceType metric, index** idx) { - auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); - auto appr_algo = std::make_unique::type>>( - hnsw_index->get_space(), filename); - if (params.hierarchy == HnswHierarchy::NONE) { appr_algo->base_layer_only = true; } - hnsw_index->set_index(std::move(appr_algo)); - *idx = hnsw_index.release(); + try { + auto hnsw_index = std::make_unique>(dim, metric, params.hierarchy); + auto appr_algo = std::make_unique::type>>( + hnsw_index->get_space(), filename); + if (params.hierarchy == HnswHierarchy::NONE) { appr_algo->base_layer_only = true; } + hnsw_index->set_index(std::move(appr_algo)); + *idx = hnsw_index.release(); + } catch (const std::bad_alloc& e) { + RAFT_FAIL( + "Failed to deserialize HNSW index from '%s': insufficient host memory. " + "The index is too large to fit in available RAM. " + "Consider using a machine with more memory or reducing the dataset size.", + filename.c_str()); + } +} + +/** + * @brief Build an HNSW index on the GPU using CAGRA graph building algorithm + * + * This function builds an HNSW index + * 1. Converting HNSW parameters to CAGRA parameters (ACE configuration by default) + * 2. Building a CAGRA index + * 3. Converting the CAGRA index to HNSW format + */ +template +std::unique_ptr> build(raft::resources const& res, + const index_params& params, + raft::host_matrix_view dataset) +{ + common::nvtx::range fun_scope("hnsw::build"); + + // Use provided ACE parameters or default ones if not specified + auto ace_params = + std::holds_alternative(params.graph_build_params) + ? std::get(params.graph_build_params) + : graph_build_params::ace_params{}; + + // Create CAGRA index parameters from HNSW parameters + cuvs::neighbors::cagra::index_params cagra_params; + cagra_params.metric = params.metric; + cagra_params.intermediate_graph_degree = params.M * 3; + cagra_params.graph_degree = params.M * 2; + + // Configure ACE parameters for CAGRA + cuvs::neighbors::cagra::graph_build_params::ace_params cagra_ace_params; + cagra_ace_params.npartitions = ace_params.npartitions; + cagra_ace_params.ef_construction = params.ef_construction; + cagra_ace_params.build_dir = ace_params.build_dir; + cagra_ace_params.use_disk = ace_params.use_disk; + cagra_params.graph_build_params = cagra_ace_params; + + RAFT_LOG_INFO( + "hnsw::build - Building HNSW index using ACE with %zu partitions, ef_construction=%zu", + ace_params.npartitions, + ace_params.ef_construction); + + // Build CAGRA index using ACE + auto cagra_index = cuvs::neighbors::cagra::build(res, cagra_params, dataset); + + RAFT_LOG_INFO("hnsw::build - Converting CAGRA index to HNSW format"); + + // Convert CAGRA index to HNSW index + return from_cagra(res, params, cagra_index, dataset); } } // namespace cuvs::neighbors::hnsw::detail diff --git a/cpp/src/neighbors/hnsw.cpp b/cpp/src/neighbors/hnsw.cpp index 5d26630358..54e9dcf12a 100644 --- a/cpp/src/neighbors/hnsw.cpp +++ b/cpp/src/neighbors/hnsw.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -26,6 +26,22 @@ auto to_cagra_params(raft::matrix_extent dataset, metric); } +#define CUVS_INST_HNSW_BUILD(T) \ + std::unique_ptr> build( \ + raft::resources const& res, \ + const index_params& params, \ + raft::host_matrix_view dataset) \ + { \ + return detail::build(res, params, dataset); \ + } + +CUVS_INST_HNSW_BUILD(float); +CUVS_INST_HNSW_BUILD(half); +CUVS_INST_HNSW_BUILD(uint8_t); +CUVS_INST_HNSW_BUILD(int8_t); + +#undef CUVS_INST_HNSW_BUILD + #define CUVS_INST_HNSW_FROM_CAGRA(T) \ std::unique_ptr> from_cagra( \ raft::resources const& res, \ diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 3c47db4a16..41894e92ce 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -195,34 +195,6 @@ ConfigureTest( PERCENT 100 ) -ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_ACE_FLOAT_UINT32_TEST - PATH neighbors/ann_cagra_ace/test_float_uint32_t.cu - GPUS 1 - PERCENT 100 -) - -ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_ACE_HALF_UINT32_TEST - PATH neighbors/ann_cagra_ace/test_half_uint32_t.cu - GPUS 1 - PERCENT 100 -) - -ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_ACE_INT8_UINT32_TEST - PATH neighbors/ann_cagra_ace/test_int8_t_uint32_t.cu - GPUS 1 - PERCENT 100 -) - -ConfigureTest( - NAME NEIGHBORS_ANN_CAGRA_ACE_UINT8_UINT32_TEST - PATH neighbors/ann_cagra_ace/test_uint8_t_uint32_t.cu - GPUS 1 - PERCENT 100 -) - ConfigureTest( NAME NEIGHBORS_ANN_NN_DESCENT_TEST PATH neighbors/ann_nn_descent/test_float_uint32_t.cu @@ -279,6 +251,50 @@ if(BUILD_CAGRA_HNSWLIB) ) target_link_libraries(NEIGHBORS_HNSW_TEST PRIVATE hnswlib::hnswlib) target_compile_definitions(NEIGHBORS_HNSW_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB) + + ConfigureTest( + NAME NEIGHBORS_ANN_HNSW_ACE_FLOAT_UINT32_TEST + PATH neighbors/ann_hnsw_ace/test_float_uint32_t.cu + GPUS 1 + PERCENT 100 + ) + target_link_libraries(NEIGHBORS_ANN_HNSW_ACE_FLOAT_UINT32_TEST PRIVATE hnswlib::hnswlib) + target_compile_definitions( + NEIGHBORS_ANN_HNSW_ACE_FLOAT_UINT32_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB + ) + + ConfigureTest( + NAME NEIGHBORS_ANN_HNSW_ACE_HALF_UINT32_TEST + PATH neighbors/ann_hnsw_ace/test_half_uint32_t.cu + GPUS 1 + PERCENT 100 + ) + target_link_libraries(NEIGHBORS_ANN_HNSW_ACE_HALF_UINT32_TEST PRIVATE hnswlib::hnswlib) + target_compile_definitions( + NEIGHBORS_ANN_HNSW_ACE_HALF_UINT32_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB + ) + + ConfigureTest( + NAME NEIGHBORS_ANN_HNSW_ACE_INT8_UINT32_TEST + PATH neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu + GPUS 1 + PERCENT 100 + ) + target_link_libraries(NEIGHBORS_ANN_HNSW_ACE_INT8_UINT32_TEST PRIVATE hnswlib::hnswlib) + target_compile_definitions( + NEIGHBORS_ANN_HNSW_ACE_INT8_UINT32_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB + ) + + ConfigureTest( + NAME NEIGHBORS_ANN_HNSW_ACE_UINT8_UINT32_TEST + PATH neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu + GPUS 1 + PERCENT 100 + ) + target_link_libraries(NEIGHBORS_ANN_HNSW_ACE_UINT8_UINT32_TEST PRIVATE hnswlib::hnswlib) + target_compile_definitions( + NEIGHBORS_ANN_HNSW_ACE_UINT8_UINT32_TEST PUBLIC CUVS_BUILD_CAGRA_HNSWLIB + ) endif() if(BUILD_MG_ALGOS) diff --git a/cpp/tests/neighbors/ann_cagra_ace.cuh b/cpp/tests/neighbors/ann_cagra_ace.cuh deleted file mode 100644 index 4c6d96050b..0000000000 --- a/cpp/tests/neighbors/ann_cagra_ace.cuh +++ /dev/null @@ -1,270 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include "ann_cagra.cuh" - -#include - -#include -#include - -namespace cuvs::neighbors::cagra { - -struct AnnCagraAceInputs { - int n_queries; - int n_rows; - int dim; - int k; - int npartitions; - int ef_construction; - bool use_disk; - cuvs::distance::DistanceType metric; - double min_recall; -}; - -inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraAceInputs& p) -{ - os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim - << ", k=" << p.k << ", npartitions=" << p.npartitions - << ", ef_construction=" << p.ef_construction - << ", use_disk=" << (p.use_disk ? "true" : "false") << ", metric="; - switch (p.metric) { - case cuvs::distance::DistanceType::L2Expanded: os << "L2"; break; - case cuvs::distance::DistanceType::InnerProduct: os << "InnerProduct"; break; - default: os << "Unknown"; break; - } - os << ", min_recall=" << p.min_recall << "}"; - return os; -} - -template -class AnnCagraAceTest : public ::testing::TestWithParam { - public: - AnnCagraAceTest() - : stream_(raft::resource::get_cuda_stream(handle_)), - ps(::testing::TestWithParam::GetParam()), - database_dev(0, stream_), - search_queries(0, stream_) - { - } - - protected: - void testAce() - { - size_t queries_size = ps.n_queries * ps.k; - std::vector indices_ace(queries_size); - std::vector indices_naive(queries_size); - std::vector distances_ace(queries_size); - std::vector distances_naive(queries_size); - - { - rmm::device_uvector distances_naive_dev(queries_size, stream_); - rmm::device_uvector indices_naive_dev(queries_size, stream_); - - cuvs::neighbors::naive_knn(handle_, - distances_naive_dev.data(), - indices_naive_dev.data(), - search_queries.data(), - database_dev.data(), - ps.n_queries, - ps.n_rows, - ps.dim, - ps.k, - ps.metric); - raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); - raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); - raft::resource::sync_stream(handle_); - } - - // Create temporary directory for ACE build - std::string temp_dir = std::string("/tmp/cuvs_ace_test_") + std::to_string(std::time(nullptr)) + - "_" + std::to_string(reinterpret_cast(this)); - std::filesystem::create_directories(temp_dir); - - { - auto database_host = raft::make_host_matrix(ps.n_rows, ps.dim); - raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); - raft::resource::sync_stream(handle_); - - cagra::index_params index_params; - index_params.metric = ps.metric; - index_params.intermediate_graph_degree = 128; - index_params.graph_degree = 64; - auto ace_params = graph_build_params::ace_params(); - ace_params.npartitions = ps.npartitions; - ace_params.ef_construction = ps.ef_construction; - ace_params.build_dir = temp_dir; - ace_params.use_disk = ps.use_disk; - index_params.graph_build_params = ace_params; - - auto index = - cagra::build(handle_, index_params, raft::make_const_mdspan(database_host.view())); - - ASSERT_EQ(index.size(), ps.n_rows); - - if (ps.use_disk) { - // Verify disk-based ACE index using HNSW index from disk - EXPECT_TRUE(index.dataset_fd().has_value() && index.graph_fd().has_value()); - - // Verify file directory from graph file descriptor - const auto& graph_fd = index.graph_fd(); - EXPECT_TRUE(graph_fd.has_value() && graph_fd->is_valid()); - std::string graph_path = graph_fd->get_path(); - std::string file_dir = std::filesystem::path(graph_path).parent_path().string(); - EXPECT_EQ(file_dir, temp_dir); - - EXPECT_TRUE(std::filesystem::exists(temp_dir + "/cagra_graph.npy")); - EXPECT_GE(std::filesystem::file_size(temp_dir + "/cagra_graph.npy"), - ps.n_rows * index_params.graph_degree * sizeof(IdxT)); - - EXPECT_TRUE(std::filesystem::exists(temp_dir + "/reordered_dataset.npy")); - EXPECT_GE(std::filesystem::file_size(temp_dir + "/reordered_dataset.npy"), - ps.n_rows * ps.dim * sizeof(DataT)); - - EXPECT_TRUE(std::filesystem::exists(temp_dir + "/dataset_mapping.npy")); - EXPECT_GE(std::filesystem::file_size(temp_dir + "/dataset_mapping.npy"), - ps.n_rows * sizeof(IdxT)); - - hnsw::index_params hnsw_params; - hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; - - auto hnsw_index = hnsw::from_cagra(handle_, hnsw_params, index); - ASSERT_NE(hnsw_index, nullptr); - - std::string hnsw_index_path = temp_dir + "/hnsw_index.bin"; - EXPECT_TRUE(std::filesystem::exists(hnsw_index_path)); - // For GPU hierarchy, HNSW index includes multi-layer structure - // The size should be at least the base layer size - auto hnsw_file_size = std::filesystem::file_size(hnsw_index_path); - EXPECT_GE(hnsw_file_size, ps.n_rows * index_params.graph_degree * sizeof(IdxT)); - - hnsw::index* hnsw_index_raw = nullptr; - hnsw::deserialize( - handle_, hnsw_params, hnsw_index_path, ps.dim, ps.metric, &hnsw_index_raw); - ASSERT_NE(hnsw_index_raw, nullptr); - - std::unique_ptr> hnsw_index_deserialized(hnsw_index_raw); - EXPECT_EQ(hnsw_index_deserialized->dim(), ps.dim); - EXPECT_EQ(hnsw_index_deserialized->metric(), ps.metric); - - auto queries_host = raft::make_host_matrix(ps.n_queries, ps.dim); - raft::copy( - queries_host.data_handle(), search_queries.data(), ps.n_queries * ps.dim, stream_); - raft::resource::sync_stream(handle_); - - auto indices_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); - auto distances_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); - - hnsw::search_params search_params; - search_params.ef = std::max(ps.ef_construction, ps.k * 2); - search_params.num_threads = 1; - - hnsw::search(handle_, - search_params, - *hnsw_index_deserialized, - queries_host.view(), - indices_hnsw_host.view(), - distances_hnsw_host.view()); - - for (size_t i = 0; i < queries_size; i++) { - indices_ace[i] = static_cast(indices_hnsw_host.data_handle()[i]); - distances_ace[i] = distances_hnsw_host.data_handle()[i]; - } - - EXPECT_TRUE(eval_neighbours(indices_naive, - indices_ace, - distances_naive, - distances_ace, - ps.n_queries, - ps.k, - 0.003, - ps.min_recall)) - << "Disk-based ACE index loaded via HNSW failed recall check"; - } else { - // For in-memory ACE, we can search directly - EXPECT_FALSE(index.dataset_fd().has_value() || index.graph_fd().has_value()); - ASSERT_GT(index.graph().size(), 0); - EXPECT_EQ(index.graph_degree(), 64); - - rmm::device_uvector distances_dev(queries_size, stream_); - rmm::device_uvector indices_dev(queries_size, stream_); - - auto queries_view = raft::make_device_matrix_view( - search_queries.data(), ps.n_queries, ps.dim); - auto indices_view = - raft::make_device_matrix_view(indices_dev.data(), ps.n_queries, ps.k); - auto distances_view = raft::make_device_matrix_view( - distances_dev.data(), ps.n_queries, ps.k); - - cagra::search_params search_params; - search_params.itopk_size = 64; - - cagra::search(handle_, search_params, index, queries_view, indices_view, distances_view); - - raft::update_host(distances_ace.data(), distances_dev.data(), queries_size, stream_); - raft::update_host(indices_ace.data(), indices_dev.data(), queries_size, stream_); - raft::resource::sync_stream(handle_); - - EXPECT_TRUE(eval_neighbours(indices_naive, - indices_ace, - distances_naive, - distances_ace, - ps.n_queries, - ps.k, - 0.003, - ps.min_recall)) - << "In-memory ACE index failed recall check"; - } - } - - // Clean up temporary directory - std::filesystem::remove_all(temp_dir); - } - - void SetUp() override - { - database_dev.resize(((size_t)ps.n_rows) * ps.dim, stream_); - search_queries.resize(ps.n_queries * ps.dim, stream_); - raft::random::RngState r(1234ULL); - InitDataset(handle_, database_dev.data(), ps.n_rows, ps.dim, ps.metric, r); - InitDataset(handle_, search_queries.data(), ps.n_queries, ps.dim, ps.metric, r); - raft::resource::sync_stream(handle_); - } - - void TearDown() override - { - raft::resource::sync_stream(handle_); - database_dev.resize(0, stream_); - search_queries.resize(0, stream_); - } - - private: - raft::resources handle_; - rmm::cuda_stream_view stream_; - AnnCagraAceInputs ps; - rmm::device_uvector database_dev; - rmm::device_uvector search_queries; -}; - -inline std::vector generate_ace_inputs() -{ - return raft::util::itertools::product( - {10}, // n_queries - {5000}, // n_rows - {64, 128}, // dim - {10}, // k - {2, 4}, // npartitions - {100}, // ef_construction - {false, true}, // use_disk (test both modes) - {cuvs::distance::DistanceType::L2Expanded, - cuvs::distance::DistanceType::InnerProduct}, // metric - {0.9} // min_recall - ); -} - -const std::vector ace_inputs = generate_ace_inputs(); - -} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra_ace/test_float_uint32_t.cu deleted file mode 100644 index de96a40339..0000000000 --- a/cpp/tests/neighbors/ann_cagra_ace/test_float_uint32_t.cu +++ /dev/null @@ -1,17 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include "../ann_cagra_ace.cuh" - -namespace cuvs::neighbors::cagra { - -typedef AnnCagraAceTest AnnCagraAceTestF_U32; -TEST_P(AnnCagraAceTestF_U32, AnnCagraAce) { this->testAce(); } - -INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, AnnCagraAceTestF_U32, ::testing::ValuesIn(ace_inputs)); - -} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_cagra_ace/test_half_uint32_t.cu deleted file mode 100644 index a1a6ec1397..0000000000 --- a/cpp/tests/neighbors/ann_cagra_ace/test_half_uint32_t.cu +++ /dev/null @@ -1,17 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include "../ann_cagra_ace.cuh" - -namespace cuvs::neighbors::cagra { - -typedef AnnCagraAceTest AnnCagraAceTestF16_U32; -TEST_P(AnnCagraAceTestF16_U32, AnnCagraAce) { this->testAce(); } - -INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, AnnCagraAceTestF16_U32, ::testing::ValuesIn(ace_inputs)); - -} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_cagra_ace/test_int8_t_uint32_t.cu deleted file mode 100644 index 3973b72cd6..0000000000 --- a/cpp/tests/neighbors/ann_cagra_ace/test_int8_t_uint32_t.cu +++ /dev/null @@ -1,17 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include "../ann_cagra_ace.cuh" - -namespace cuvs::neighbors::cagra { - -typedef AnnCagraAceTest AnnCagraAceTestI8_U32; -TEST_P(AnnCagraAceTestI8_U32, AnnCagraAce) { this->testAce(); } - -INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, AnnCagraAceTestI8_U32, ::testing::ValuesIn(ace_inputs)); - -} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_cagra_ace/test_uint8_t_uint32_t.cu deleted file mode 100644 index 5ca6f038df..0000000000 --- a/cpp/tests/neighbors/ann_cagra_ace/test_uint8_t_uint32_t.cu +++ /dev/null @@ -1,17 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include "../ann_cagra_ace.cuh" - -namespace cuvs::neighbors::cagra { - -typedef AnnCagraAceTest AnnCagraAceTestU8_U32; -TEST_P(AnnCagraAceTestU8_U32, AnnCagraAce) { this->testAce(); } - -INSTANTIATE_TEST_CASE_P(AnnCagraAceTest, AnnCagraAceTestU8_U32, ::testing::ValuesIn(ace_inputs)); - -} // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh new file mode 100644 index 0000000000..9885f77b14 --- /dev/null +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -0,0 +1,248 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "ann_cagra.cuh" + +#include + +#include +#include + +namespace cuvs::neighbors::hnsw { + +struct AnnHnswAceInputs { + int n_queries; + int n_rows; + int dim; + int k; + int npartitions; + int ef_construction; + bool use_disk; + cuvs::distance::DistanceType metric; + double min_recall; +}; + +inline ::std::ostream& operator<<(::std::ostream& os, const AnnHnswAceInputs& p) +{ + os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim + << ", k=" << p.k << ", npartitions=" << p.npartitions + << ", ef_construction=" << p.ef_construction + << ", use_disk=" << (p.use_disk ? "true" : "false") << ", metric="; + switch (p.metric) { + case cuvs::distance::DistanceType::L2Expanded: os << "L2"; break; + case cuvs::distance::DistanceType::InnerProduct: os << "InnerProduct"; break; + default: os << "Unknown"; break; + } + os << ", min_recall=" << p.min_recall << "}"; + return os; +} + +template +class AnnHnswAceTest : public ::testing::TestWithParam { + public: + AnnHnswAceTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database_dev(0, stream_), + search_queries(0, stream_) + { + } + + protected: + void testHnswAceBuild() + { + size_t queries_size = ps.n_queries * ps.k; + std::vector indexes_hnsw(queries_size); + std::vector indexes_naive(queries_size); + std::vector distances_hnsw(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indexes_naive_dev(queries_size, stream_); + + cuvs::neighbors::naive_knn(handle_, + distances_naive_dev.data(), + indexes_naive_dev.data(), + search_queries.data(), + database_dev.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::update_host(indexes_naive.data(), indexes_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + // Create temporary directory for ACE build + std::string temp_dir = std::string("/tmp/cuvs_hnsw_ace_test_") + + std::to_string(std::time(nullptr)) + "_" + + std::to_string(reinterpret_cast(this)); + std::filesystem::create_directories(temp_dir); + + { + // Copy dataset to host for hnsw::build + auto database_host = raft::make_host_matrix(ps.n_rows, ps.dim); + raft::copy(database_host.data_handle(), database_dev.data(), ps.n_rows * ps.dim, stream_); + raft::resource::sync_stream(handle_); + + // Configure HNSW index parameters with ACE + hnsw::index_params hnsw_params; + hnsw_params.metric = ps.metric; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; + hnsw_params.M = 32; + hnsw_params.ef_construction = ps.ef_construction; + + // Configure ACE parameters + auto ace_params = graph_build_params::ace_params(); + ace_params.npartitions = ps.npartitions; + ace_params.build_dir = temp_dir; + ace_params.use_disk = ps.use_disk; + hnsw_params.graph_build_params = ace_params; + + // Build HNSW index using ACE + auto hnsw_index = + hnsw::build(handle_, hnsw_params, raft::make_const_mdspan(database_host.view())); + + ASSERT_NE(hnsw_index, nullptr); + + // Prepare queries on host + auto queries_host = raft::make_host_matrix(ps.n_queries, ps.dim); + raft::copy(queries_host.data_handle(), search_queries.data(), ps.n_queries * ps.dim, stream_); + raft::resource::sync_stream(handle_); + + auto indexes_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + auto distances_hnsw_host = raft::make_host_matrix(ps.n_queries, ps.k); + + // Search the HNSW index + hnsw::search_params search_params; + search_params.ef = std::max(ps.ef_construction, ps.k * 2); + search_params.num_threads = 1; + + if (!ps.use_disk) { + hnsw::search(handle_, + search_params, + *hnsw_index, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw[i] = indexes_hnsw_host.data_handle()[i]; + distances_hnsw[i] = distances_hnsw_host.data_handle()[i]; + } + + // Convert indexes for comparison + std::vector indexes_hnsw_converted(queries_size); + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw_converted[i] = static_cast(indexes_hnsw[i]); + } + + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_naive, + indexes_hnsw_converted, + distances_naive, + distances_hnsw, + ps.n_queries, + ps.k, + 0.003, + ps.min_recall)) + << "HNSW ACE build and search failed recall check"; + } + + tmp_index_file index_file; + hnsw::serialize(handle_, index_file.filename, *hnsw_index); + + hnsw::index* deserialized_index = nullptr; + hnsw::deserialize( + handle_, hnsw_params, index_file.filename, ps.dim, ps.metric, &deserialized_index); + ASSERT_NE(deserialized_index, nullptr); + + // Reset search results + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw[i] = 0; + distances_hnsw[i] = 0; + } + hnsw::search(handle_, + search_params, + *deserialized_index, + queries_host.view(), + indexes_hnsw_host.view(), + distances_hnsw_host.view()); + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw[i] = indexes_hnsw_host.data_handle()[i]; + distances_hnsw[i] = distances_hnsw_host.data_handle()[i]; + } + + // Convert indexes for comparison + std::vector indexes_hnsw_converted(queries_size); + for (size_t i = 0; i < queries_size; i++) { + indexes_hnsw_converted[i] = static_cast(indexes_hnsw[i]); + } + + EXPECT_TRUE(cuvs::neighbors::eval_neighbours(indexes_naive, + indexes_hnsw_converted, + distances_naive, + distances_hnsw, + ps.n_queries, + ps.k, + 0.003, + ps.min_recall)) + << "HNSW ACE build and search failed recall check"; + // Clean up deserialized index + delete deserialized_index; + } + + // Clean up temporary directory + std::filesystem::remove_all(temp_dir); + } + + void SetUp() override + { + database_dev.resize(((size_t)ps.n_rows) * ps.dim, stream_); + search_queries.resize(ps.n_queries * ps.dim, stream_); + raft::random::RngState r(1234ULL); + cuvs::neighbors::cagra::InitDataset( + handle_, database_dev.data(), ps.n_rows, ps.dim, ps.metric, r); + cuvs::neighbors::cagra::InitDataset( + handle_, search_queries.data(), ps.n_queries, ps.dim, ps.metric, r); + raft::resource::sync_stream(handle_); + } + + void TearDown() override + { + raft::resource::sync_stream(handle_); + database_dev.resize(0, stream_); + search_queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnHnswAceInputs ps; + rmm::device_uvector database_dev; + rmm::device_uvector search_queries; +}; + +inline std::vector generate_hnsw_ace_inputs() +{ + return raft::util::itertools::product( + {10}, // n_queries + {5000}, // n_rows + {64, 128}, // dim + {10}, // k + {2, 4}, // npartitions + {100}, // ef_construction + {false, true}, // use_disk (test both modes) + {cuvs::distance::DistanceType::L2Expanded, + cuvs::distance::DistanceType::InnerProduct}, // metric + {0.9} // min_recall + ); +} + +const std::vector hnsw_ace_inputs = generate_hnsw_ace_inputs(); + +} // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu new file mode 100644 index 0000000000..827bb9f106 --- /dev/null +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_float_uint32_t.cu @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_hnsw_ace.cuh" + +namespace cuvs::neighbors::hnsw { + +typedef AnnHnswAceTest AnnHnswAceTest_float; +TEST_P(AnnHnswAceTest_float, AnnHnswAceBuild) { this->testHnswAceBuild(); } + +INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_float, ::testing::ValuesIn(hnsw_ace_inputs)); + +} // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu new file mode 100644 index 0000000000..38f75b2afa --- /dev/null +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_half_uint32_t.cu @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_hnsw_ace.cuh" + +namespace cuvs::neighbors::hnsw { + +typedef AnnHnswAceTest AnnHnswAceTest_half; +TEST_P(AnnHnswAceTest_half, AnnHnswAceBuild) { this->testHnswAceBuild(); } + +INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, AnnHnswAceTest_half, ::testing::ValuesIn(hnsw_ace_inputs)); + +} // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu new file mode 100644 index 0000000000..279df6555f --- /dev/null +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_int8_t_uint32_t.cu @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_hnsw_ace.cuh" + +namespace cuvs::neighbors::hnsw { + +typedef AnnHnswAceTest AnnHnswAceTest_int8_t; +TEST_P(AnnHnswAceTest_int8_t, AnnHnswAceBuild) { this->testHnswAceBuild(); } + +INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, + AnnHnswAceTest_int8_t, + ::testing::ValuesIn(hnsw_ace_inputs)); + +} // namespace cuvs::neighbors::hnsw diff --git a/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu new file mode 100644 index 0000000000..7e68dc4b17 --- /dev/null +++ b/cpp/tests/neighbors/ann_hnsw_ace/test_uint8_t_uint32_t.cu @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../ann_hnsw_ace.cuh" + +namespace cuvs::neighbors::hnsw { + +typedef AnnHnswAceTest AnnHnswAceTest_uint8_t; +TEST_P(AnnHnswAceTest_uint8_t, AnnHnswAceBuild) { this->testHnswAceBuild(); } + +INSTANTIATE_TEST_CASE_P(AnnHnswAceTest, + AnnHnswAceTest_uint8_t, + ::testing::ValuesIn(hnsw_ace_inputs)); + +} // namespace cuvs::neighbors::hnsw diff --git a/docs/source/cuvs_bench/param_tuning.rst b/docs/source/cuvs_bench/param_tuning.rst index 12e804dd56..01b9bc76f0 100644 --- a/docs/source/cuvs_bench/param_tuning.rst +++ b/docs/source/cuvs_bench/param_tuning.rst @@ -700,6 +700,90 @@ Use FAISS IVF-PQ index on CPU HNSW ==== +cuvs_hnsw +--------- + +cuVS HNSW builds an HNSW index using the ACE (Augmented Core Extraction) algorithm, which enables GPU-accelerated HNSW index construction for datasets too large to fit in GPU memory. + +.. list-table:: + + * - Parameter + - Type + - Required + - Data Type + - Default + - Description + + * - `hierarchy` + - `build` + - N + - [`NONE`, `CPU`, `GPU`] + - `NONE` + - Type of HNSW hierarchy to build. `NONE` creates a base-layer-only index, `CPU` builds full hierarchy on CPU, `GPU` builds full hierarchy on GPU. + + * - `efConstruction` + - `build` + - Y + - Positive integer >0 + - + - Controls index time and accuracy. Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. + + * - `M` + - `build` + - Y + - Positive integer. Often between 2-100 + - + - Number of bi-directional links create for every new element during construction. Higher values work for higher intrinsic dimensionality and/or high recall, low values can work for datasets with low intrinsic dimensionality and/or low recalls. Also affects the algorithm's memory consumption. + + * - `numThreads` + - `build` + - N + - Positive integer >0 + - 1 + - Number of threads to use to build the index. + + * - `npartitions` + - `build` + - N + - Positive integer >0 + - 1 + - Number of partitions to use for the ACE build. Small values might improve recall but potentially degrade performance and increase memory usage. 100k - 5M vectors per partition is recommended depending on the available host and GPU memory. + + * - `ef_construction` + - `build` + - N + - Positive integer >0 + - 120 + - Controls index time and accuracy when using ACE build. Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. + + * - `build_dir` + - `build` + - N + - String + - "/tmp/ace_build" + - The directory to use for the ACE build. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. + + * - `use_disk` + - `build` + - N + - Boolean + - `false` + - Whether to use disk-based storage for ACE build. When true, forces ACE to use disk-based storage even if the graph fits in host and GPU memory. When false, ACE will use in-memory storage if the graph fits in host and GPU memory and disk-based storage otherwise. + + * - `ef` + - `search` + - Y + - Positive integer >0 + - + - Size of the dynamic list for the nearest neighbors used for search. Higher value leads to more accurate but slower search. Cannot be lower than `k`. + + * - `numThreads` + - `search` + - N + - Positive integer >0 + - 1 + - Number of threads to use for queries. + hnswlib ------- @@ -724,7 +808,7 @@ hnswlib - Y - Positive integer. Often between 2-100 - - - umber of bi-directional links create for every new element during construction. Higher values work for higher intrinsic dimensionality and/or high recall, low values can work for datasets with low intrinsic dimensionality and/or low recalls. Also affects the algorithm's memory consumption. + - Number of bi-directional links create for every new element during construction. Higher values work for higher intrinsic dimensionality and/or high recall, low values can work for datasets with low intrinsic dimensionality and/or low recalls. Also affects the algorithm's memory consumption. * - `numThreads` - `build` diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 619583a83e..11c581f767 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on @@ -30,10 +30,11 @@ include(../cmake/thirdparty/get_cuvs.cmake) # -------------- compile tasks ----------------- # add_executable(BRUTE_FORCE_EXAMPLE src/brute_force_bitmap.cu) -add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_EXAMPLE src/cagra_example.cu) +add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) +add_executable(HNSW_ACE_EXAMPLE src/hnsw_ace_example.cu) add_executable(IVF_FLAT_EXAMPLE src/ivf_flat_example.cu) add_executable(IVF_PQ_EXAMPLE src/ivf_pq_example.cu) add_executable(VAMANA_EXAMPLE src/vamana_example.cu) @@ -42,14 +43,15 @@ add_executable(SCANN_EXAMPLE src/scann_example.cu) # `$` is a generator expression that ensures that targets are # installed in a conda environment, if one exists target_link_libraries(BRUTE_FORCE_EXAMPLE PRIVATE cuvs::cuvs $) -target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(CAGRA_EXAMPLE PRIVATE cuvs::cuvs $) +target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries( CAGRA_PERSISTENT_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads ) target_link_libraries( DYNAMIC_BATCHING_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads ) +target_link_libraries(HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(IVF_PQ_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(IVF_FLAT_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries(VAMANA_EXAMPLE PRIVATE cuvs::cuvs $) diff --git a/examples/cpp/src/hnsw_ace_example.cu b/examples/cpp/src/hnsw_ace_example.cu new file mode 100644 index 0000000000..7fd93bc95a --- /dev/null +++ b/examples/cpp/src/hnsw_ace_example.cu @@ -0,0 +1,165 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include "common.cuh" + +void hnsw_build_search_ace(raft::device_resources const& dev_resources, + raft::device_matrix_view dataset, + raft::device_matrix_view queries) +{ + using namespace cuvs::neighbors; + + int64_t topk = 12; + int64_t n_queries = queries.extent(0); + + // HNSW ACE build requires the dataset to be on the host + auto dataset_host = raft::make_host_matrix(dataset.extent(0), dataset.extent(1)); + raft::copy(dataset_host.data_handle(), + dataset.data_handle(), + dataset.extent(0) * dataset.extent(1), + raft::resource::get_cuda_stream(dev_resources)); + raft::resource::sync_stream(dev_resources); + + // HNSW index parameters + hnsw::index_params hnsw_params; + hnsw_params.metric = cuvs::distance::DistanceType::L2Expanded; + hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; + + // Parameters for GPU accelerated HNSW index building + auto ace_params = hnsw::graph_build_params::ace_params(); + // Set the number of partitions. Small values might improve recall but potentially degrade + // performance and increase memory usage. Partitions should not be too small to prevent issues in + // KNN graph construction. 100k - 5M vectors per partition is recommended depending on the + // available host and GPU memory. The partition size is on average 2 * (n_rows / npartitions) * + // dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance + // in the partition sizes (up to 3x in our tests). + ace_params.npartitions = 4; + // Set the directory to store the ACE build artifacts. This should be the fastest disk in the + // system and hold enough space for twice the dataset, final graph, and label mapping. + ace_params.build_dir = "/tmp/hnsw_ace_build"; + // Set whether to use disk-based storage for ACE build. When true, enables disk-based operations + // for memory-efficient graph construction. If not set, the index will be built in memory if the + // graph fits in host and GPU memory, and on disk otherwise. + ace_params.use_disk = true; + hnsw_params.graph_build_params = ace_params; + // Set M parameter to control the graph degree (graph_degree = m * 2, intermediate_graph_degree = + // m * 3). Higher values work for higher intrinsic dimensionality and/or high recall, low values + // can work for datasets with low intrinsic dimensionality and/or low recalls. Higher values lead + // to higher memory consumption. + hnsw_params.M = 32; + // Set the index quality for the ACE build. Bigger values increase the index quality. At some + // point, increasing this will no longer improve the quality. + hnsw_params.ef_construction = 120; + + // Build the HNSW index using ACE + std::cout << "Building HNSW index using ACE" << std::endl; + auto hnsw_index = + hnsw::build(dev_resources, hnsw_params, raft::make_const_mdspan(dataset_host.view())); + + // For disk-based indexes, the build function serializes the index to disk + // We need to deserialize it before searching + std::string hnsw_index_path = hnsw_index->file_path(); + std::cout << "Deserializing HNSW index from disk for search" << std::endl; + hnsw::index* hnsw_index_deserialized = nullptr; + hnsw::deserialize(dev_resources, + hnsw_params, + hnsw_index_path, + dataset.extent(1), + hnsw_params.metric, + &hnsw_index_deserialized); + + // Prepare queries on host for HNSW search + auto queries_host = raft::make_host_matrix(n_queries, queries.extent(1)); + raft::copy(queries_host.data_handle(), + queries.data_handle(), + n_queries * queries.extent(1), + raft::resource::get_cuda_stream(dev_resources)); + raft::resource::sync_stream(dev_resources); + + // Create output arrays for HNSW search (uses uint64_t indices) + auto indices_host = raft::make_host_matrix(n_queries, topk); + auto distances_host = raft::make_host_matrix(n_queries, topk); + + // Configure search parameters + hnsw::search_params search_params; + search_params.ef = std::max(200, static_cast(topk) * 2); + search_params.num_threads = 1; + + // Search the HNSW index + std::cout << "Searching HNSW index" << std::endl; + hnsw::search(dev_resources, + search_params, + *hnsw_index_deserialized, + queries_host.view(), + indices_host.view(), + distances_host.view()); + + // Convert results to device for printing + auto neighbors = raft::make_device_matrix(dev_resources, n_queries, topk); + auto distances = raft::make_device_matrix(dev_resources, n_queries, topk); + + // Convert HNSW uint64_t indices to uint32_t + auto neighbors_host = raft::make_host_matrix(n_queries, topk); + for (int64_t i = 0; i < n_queries; i++) { + for (int64_t j = 0; j < topk; j++) { + neighbors_host(i, j) = static_cast(indices_host(i, j)); + } + } + + // Copy results to device + raft::copy(neighbors.data_handle(), + neighbors_host.data_handle(), + n_queries * topk, + raft::resource::get_cuda_stream(dev_resources)); + raft::copy(distances.data_handle(), + distances_host.data_handle(), + n_queries * topk, + raft::resource::get_cuda_stream(dev_resources)); + 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()); + + // HNSW ACE build and search example. + hnsw_build_search_ace(dev_resources, + raft::make_const_mdspan(dataset.view()), + raft::make_const_mdspan(queries.view())); +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java index 1e26d49d48..f78d6957b2 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java @@ -1,15 +1,15 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; /** * Parameters for ACE (Augmented Core Extraction) graph build algorithm. - * ACE enables building indices for datasets too large to fit in GPU memory by: + * ACE enables building indexes for datasets too large to fit in GPU memory by: * 1. Partitioning the dataset in core (closest) and augmented (second-closest) * partitions using balanced k-means. - * 2. Building sub-indices for each partition independently + * 2. Building sub-indexes for each partition independently * 3. Concatenating sub-graphs into a final unified index * * @since 25.12 diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java new file mode 100644 index 0000000000..0e30801ef0 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +/** + * Parameters for ACE (Augmented Core Extraction) graph build for HNSW. + * ACE enables building indexes for datasets too large to fit in GPU memory by: + * 1. Partitioning the dataset in core and augmented partitions using balanced k-means + * 2. Building sub-indexes for each partition independently + * 3. Concatenating sub-graphs into a final unified index + * + * @since 25.02 + */ +public class HnswAceParams { + + private long npartitions; + private String buildDir; + private boolean useDisk; + + private HnswAceParams(long npartitions, String buildDir, boolean useDisk) { + this.npartitions = npartitions; + this.buildDir = buildDir; + this.useDisk = useDisk; + } + + /** + * Gets the number of partitions for ACE partitioned build. + * + * @return the number of partitions + */ + public long getNpartitions() { + return npartitions; + } + + /** + * Gets the directory to store ACE build artifacts. + * + * @return the build directory path + */ + public String getBuildDir() { + return buildDir; + } + + /** + * Gets whether disk-based storage is enabled for ACE build. + * + * @return true if disk mode is enabled + */ + public boolean isUseDisk() { + return useDisk; + } + + @Override + public String toString() { + return "HnswAceParams [npartitions=" + + npartitions + + ", buildDir=" + + buildDir + + ", useDisk=" + + useDisk + + "]"; + } + + /** + * Builder configures and creates an instance of {@link HnswAceParams}. + */ + public static class Builder { + + private long npartitions = 1; + private String buildDir = "/tmp/hnsw_ace_build"; + private boolean useDisk = false; + + /** + * Constructs this Builder. + */ + public Builder() {} + + /** + * Sets the number of partitions for ACE partitioned build. + * Small values might improve recall but potentially degrade performance. + * 100k - 5M vectors per partition is recommended. + * + * @param npartitions the number of partitions + * @return an instance of Builder + */ + public Builder withNpartitions(long npartitions) { + this.npartitions = npartitions; + return this; + } + + /** + * Sets the directory to store ACE build artifacts. + * Used when useDisk is true or when the graph does not fit in memory. + * + * @param buildDir the build directory path + * @return an instance of Builder + */ + public Builder withBuildDir(String buildDir) { + this.buildDir = buildDir; + return this; + } + + /** + * Sets whether to use disk-based storage for ACE build. + * When true, enables disk-based operations for memory-efficient graph construction. + * + * @param useDisk true to enable disk mode + * @return an instance of Builder + */ + public Builder withUseDisk(boolean useDisk) { + this.useDisk = useDisk; + return this; + } + + /** + * Builds an instance of {@link HnswAceParams}. + * + * @return an instance of {@link HnswAceParams} + */ + public HnswAceParams build() { + return new HnswAceParams(npartitions, buildDir, useDisk); + } + } +} diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java index 56d235a563..3eef491b62 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndex.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -57,6 +57,31 @@ static HnswIndex fromCagra(HnswIndexParams hnswParams, CagraIndex cagraIndex) th return CuVSProvider.provider().hnswIndexFromCagra(hnswParams, cagraIndex); } + /** + * Builds an HNSW index using the ACE (Augmented Core Extraction) algorithm. + * + * ACE enables building HNSW indexes for datasets too large to fit in GPU + * memory by partitioning the dataset and building sub-indexes for each + * partition independently. + * + * NOTE: This method requires `hnswParams.getAceParams()` to be set with + * an instance of HnswAceParams. + * + * @param resources The CuVS resources + * @param hnswParams Parameters for the HNSW index with ACE configuration + * @param dataset The dataset to build the index from + * @return A new HNSW index ready for search + * @throws Throwable if an error occurs during building + */ + static HnswIndex build(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) + throws Throwable { + Objects.requireNonNull(resources); + Objects.requireNonNull(hnswParams); + Objects.requireNonNull(dataset); + Objects.requireNonNull(hnswParams.getAceParams(), "ACE parameters must be set for build()"); + return CuVSProvider.provider().hnswIndexBuild(resources, hnswParams, dataset); + } + /** * Builder helps configure and create an instance of {@link HnswIndex}. */ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java index dd2cbb7577..6dc7d7ee6b 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswIndexParams.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -11,6 +11,20 @@ */ public class HnswIndexParams { + /** + * Distance metric types + */ + public enum CuvsDistanceType { + L2Expanded(0), + InnerProduct(2); + + public final int value; + + private CuvsDistanceType(int value) { + this.value = value; + } + } + /** * Hierarchy for HNSW index when converting from CAGRA index * @@ -27,7 +41,12 @@ public enum CuvsHnswHierarchy { /** * Full hierarchy is built using the CPU */ - CPU(1); + CPU(1), + + /** + * Full hierarchy is built using the GPU + */ + GPU(2); /** * The value for the enum choice. @@ -43,13 +62,25 @@ private CuvsHnswHierarchy(int value) { private int efConstruction = 200; private int numThreads = 2; private int vectorDimension; + private long m = 32; + private CuvsDistanceType metric = CuvsDistanceType.L2Expanded; + private HnswAceParams aceParams; private HnswIndexParams( - CuvsHnswHierarchy hierarchy, int efConstruction, int numThreads, int vectorDimension) { + CuvsHnswHierarchy hierarchy, + int efConstruction, + int numThreads, + int vectorDimension, + long m, + CuvsDistanceType metric, + HnswAceParams aceParams) { this.hierarchy = hierarchy; this.efConstruction = efConstruction; this.numThreads = numThreads; this.vectorDimension = vectorDimension; + this.m = m; + this.metric = metric; + this.aceParams = aceParams; } /** @@ -84,6 +115,35 @@ public int getVectorDimension() { return vectorDimension; } + /** + * Gets the HNSW M parameter: number of bi-directional links per node + * (used when building with ACE). graph_degree = m * 2, + * intermediate_graph_degree = m * 3. + * + * @return the M parameter + */ + public long getM() { + return m; + } + + /** + * Gets the distance metric type. + * + * @return the metric type + */ + public CuvsDistanceType getMetric() { + return metric; + } + + /** + * Gets the ACE parameters for building HNSW index using ACE algorithm. + * + * @return the ACE parameters, or null if not set + */ + public HnswAceParams getAceParams() { + return aceParams; + } + @Override public String toString() { return "HnswIndexParams [hierarchy=" @@ -94,6 +154,12 @@ public String toString() { + numThreads + ", vectorDimension=" + vectorDimension + + ", m=" + + m + + ", metric=" + + metric + + ", aceParams=" + + aceParams + "]"; } @@ -106,6 +172,9 @@ public static class Builder { private int efConstruction = 200; private int numThreads = 2; private int vectorDimension; + private long m = 32; + private CuvsDistanceType metric = CuvsDistanceType.L2Expanded; + private HnswAceParams aceParams; /** * Constructs this Builder with an instance of Arena. @@ -163,13 +232,55 @@ public Builder withVectorDimension(int vectorDimension) { return this; } + /** + * Sets the HNSW M parameter: number of bi-directional links per node + * (used when building with ACE). graph_degree = m * 2, + * intermediate_graph_degree = m * 3. + * + * @param m the M parameter + * @return an instance of Builder + */ + public Builder withM(long m) { + this.m = m; + return this; + } + + /** + * Sets the distance metric type. + * + * @param metric the metric type + * @return an instance of Builder + */ + public Builder withMetric(CuvsDistanceType metric) { + this.metric = metric; + return this; + } + + /** + * Sets the ACE parameters for building HNSW index using ACE algorithm. + * + * @param aceParams the ACE parameters + * @return an instance of Builder + */ + public Builder withAceParams(HnswAceParams aceParams) { + this.aceParams = aceParams; + return this; + } + /** * Builds an instance of {@link HnswIndexParams}. * * @return an instance of {@link HnswIndexParams} */ public HnswIndexParams build() { - return new HnswIndexParams(hierarchy, efConstruction, numThreads, vectorDimension); + return new HnswIndexParams( + hierarchy, + efConstruction, + numThreads, + vectorDimension, + m, + metric, + aceParams); } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 642815cb1e..424fe2238d 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.spi; @@ -123,6 +123,18 @@ HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) */ HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagraIndex) throws Throwable; + /** + * Builds an HNSW index using the ACE (Augmented Core Extraction) algorithm. + * + * @param resources The CuVS resources + * @param hnswParams Parameters for the HNSW index with ACE configuration + * @param dataset The dataset to build the index from + * @return A new HNSW index ready for search + * @throws Throwable if an error occurs during building + */ + HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) + throws Throwable; + /** Creates a new TieredIndex Builder. */ TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) throws UnsupportedOperationException; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java index 660dfb4cea..7cbeee4e75 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.spi; @@ -46,6 +46,12 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagra throw new UnsupportedOperationException(reasons); } + @Override + public HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) + throws Throwable { + throw new UnsupportedOperationException(reasons); + } + @Override public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { throw new UnsupportedOperationException(reasons); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java index 9cfc7e2f49..e142db8f04 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal; @@ -160,6 +160,25 @@ public void close() { } } + static CloseableHandle createHnswAceParamsNative() { + try (var localArena = Arena.ofConfined()) { + var paramsPtrPtr = localArena.allocate(cuvsHnswAceParams_t); + checkCuVSError(cuvsHnswAceParamsCreate(paramsPtrPtr), "cuvsHnswAceParamsCreate"); + var paramsPtr = paramsPtrPtr.get(cuvsHnswAceParams_t, 0L); + return new CloseableHandle() { + @Override + public MemorySegment handle() { + return paramsPtr; + } + + @Override + public void close() { + checkCuVSError(cuvsHnswAceParamsDestroy(paramsPtr), "cuvsHnswAceParamsDestroy"); + } + }; + } + } + static CloseableHandle createTieredIndexParams() { try (var localArena = Arena.ofConfined()) { var paramsPtrPtr = localArena.allocate(cuvsTieredIndexParams_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index 1426af407c..f09f7a61ac 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.internal; +import static com.nvidia.cuvs.internal.CuVSParamsHelper.createHnswAceParamsNative; import static com.nvidia.cuvs.internal.CuVSParamsHelper.createHnswIndexParams; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT; import static com.nvidia.cuvs.internal.common.LinkerHelper.C_LONG; @@ -15,6 +16,7 @@ import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.HnswAceParams; import com.nvidia.cuvs.HnswIndex; import com.nvidia.cuvs.HnswIndexParams; import com.nvidia.cuvs.HnswQuery; @@ -22,6 +24,7 @@ import com.nvidia.cuvs.SearchResults; import com.nvidia.cuvs.internal.common.CloseableHandle; import com.nvidia.cuvs.internal.panama.DLDataType; +import com.nvidia.cuvs.internal.panama.cuvsHnswAceParams; import com.nvidia.cuvs.internal.panama.cuvsHnswIndex; import com.nvidia.cuvs.internal.panama.cuvsHnswIndexParams; import com.nvidia.cuvs.internal.panama.cuvsHnswSearchParams; @@ -238,12 +241,100 @@ public static HnswIndex.Builder newBuilder(CuVSResources cuvsResources) { return new HnswIndexImpl.Builder(Objects.requireNonNull(cuvsResources)); } + /** + * Builds an HNSW index using the ACE algorithm. + * + * @param resources The CuVS resources + * @param hnswParams Parameters for the HNSW index with ACE configuration + * @param dataset The dataset to build the index from + * @return A new HNSW index ready for search + * @throws Throwable if an error occurs during building + */ + public static HnswIndex build(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) + throws Throwable { + Objects.requireNonNull(resources); + Objects.requireNonNull(hnswParams); + Objects.requireNonNull(dataset); + Objects.requireNonNull(hnswParams.getAceParams(), "ACE parameters must be set for build()"); + + // Create HNSW index + MemorySegment hnswIndex = createHnswIndexHandle(); + initializeIndexDType(hnswIndex, dataset); + + try (var localArena = Arena.ofConfined(); + var hnswParamsHandle = createHnswIndexParamsForBuild(localArena, hnswParams); + var aceParamsHandle = createHnswAceParams(localArena, hnswParams.getAceParams())) { + + MemorySegment hnswParamsMemorySegment = hnswParamsHandle.handle(); + + // Link ACE params to HNSW index params + cuvsHnswIndexParams.ace_params(hnswParamsMemorySegment, aceParamsHandle.handle()); + + // Prepare dataset tensor + MemorySegment datasetTensor = prepareTensorFromMatrix(localArena, dataset); + + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + + // Call cuvsHnswBuild + int returnValue = cuvsHnswBuild(cuvsRes, hnswParamsMemorySegment, datasetTensor, hnswIndex); + checkCuVSError(returnValue, "cuvsHnswBuild"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + } + } + return new HnswIndexImpl(new IndexReference(hnswIndex), resources, hnswParams); + } + + private static CloseableHandle createHnswIndexParamsForBuild(Arena arena, HnswIndexParams params) { + var hnswParams = createHnswIndexParams(); + MemorySegment seg = hnswParams.handle(); + + cuvsHnswIndexParams.hierarchy(seg, params.getHierarchy().value); + cuvsHnswIndexParams.ef_construction(seg, params.getEfConstruction()); + cuvsHnswIndexParams.num_threads(seg, params.getNumThreads()); + cuvsHnswIndexParams.M(seg, params.getM()); + cuvsHnswIndexParams.metric(seg, params.getMetric().value); + + return hnswParams; + } + + private static CloseableHandle createHnswAceParams(Arena arena, HnswAceParams aceParams) { + var params = createHnswAceParamsNative(); + MemorySegment seg = params.handle(); + + cuvsHnswAceParams.npartitions(seg, aceParams.getNpartitions()); + cuvsHnswAceParams.use_disk(seg, aceParams.isUseDisk()); + + String buildDir = aceParams.getBuildDir(); + if (buildDir != null) { + MemorySegment buildDirSeg = arena.allocateFrom(buildDir); + cuvsHnswAceParams.build_dir(seg, buildDirSeg); + } + + return params; + } + + private static MemorySegment prepareTensorFromMatrix(Arena arena, CuVSMatrix dataset) { + if (dataset instanceof CuVSMatrixInternal matrixInternal) { + return prepareTensor( + arena, + matrixInternal.memorySegment(), + new long[]{dataset.size(), dataset.columns()}, + matrixInternal.code(), + matrixInternal.bits(), + kDLCPU()); + } + throw new IllegalArgumentException("Unsupported matrix type for build"); + } + /** * Creates an HNSW index from an existing CAGRA index. * * @param hnswParams Parameters for the HNSW index * @param cagraIndex The CAGRA index to convert from - * @return A new HNSW index for in-memory indices, or null for disk-based indices + * @return A new HNSW index for in-memory indexes, or null for disk-based indexes * @throws Throwable if an error occurs during conversion */ public static HnswIndex fromCagra(HnswIndexParams hnswParams, CagraIndex cagraIndex) diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index b0ad22cd57..9aeb01d8b4 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs.spi; @@ -254,6 +254,12 @@ public HnswIndex hnswIndexFromCagra(HnswIndexParams hnswParams, CagraIndex cagra return HnswIndexImpl.fromCagra(hnswParams, cagraIndex); } + @Override + public HnswIndex hnswIndexBuild(CuVSResources resources, HnswIndexParams hnswParams, CuVSMatrix dataset) + throws Throwable { + return HnswIndexImpl.build(resources, hnswParams, dataset); + } + @Override public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) { return TieredIndexImpl.newBuilder(Objects.requireNonNull(cuVSResources)); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java index be90a747a5..7cc3769cf6 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraAceBuildAndSearchIT.java @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ package com.nvidia.cuvs; @@ -23,8 +23,8 @@ /** * Integration tests for CAGRA index using ACE (Augmented Core Extraction) build algorithm. - * ACE enables building indices for datasets too large to fit in GPU memory by partitioning - * the dataset and building sub-indices. + * ACE enables building indexes for datasets too large to fit in GPU memory by partitioning + * the dataset and building sub-indexes. * * @since 25.12 */ diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java new file mode 100644 index 0000000000..17d64838dc --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswAceBuildAndSearchIT.java @@ -0,0 +1,367 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.*; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import com.nvidia.cuvs.HnswIndexParams.CuvsDistanceType; +import com.nvidia.cuvs.HnswIndexParams.CuvsHnswHierarchy; +import java.io.InputStream; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Integration tests for HNSW index using ACE (Augmented Core Extraction) build algorithm. + * ACE enables building HNSW indexes for datasets too large to fit in GPU memory. + * + * Tests follow the same approach as C++ and C tests: + * - Build HNSW index using ACE via HnswIndex.build() + * - For disk mode: serialize -> deserialize -> search + * - For in-memory mode: search directly + * + * @since 25.12 + */ +@RunWith(RandomizedRunner.class) +public class HnswAceBuildAndSearchIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(HnswAceBuildAndSearchIT.class); + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + initializeRandom(); + log.trace("Random context initialized for test."); + } + + private static List> getExpectedResults() { + return Arrays.asList( + Map.of(3, 0.038782578f, 2, 0.3590463f, 0, 0.83774555f), + Map.of(0, 0.12472608f, 2, 0.21700792f, 1, 0.31918612f), + Map.of(3, 0.047766715f, 2, 0.20332818f, 0, 0.48305473f), + Map.of(1, 0.15224178f, 0, 0.59063464f, 3, 0.5986642f)); + } + + private static float[][] createSampleQueries() { + return new float[][] { + {0.48216683f, 0.0428398f}, + {0.5084142f, 0.6545497f}, + {0.51260436f, 0.2643005f}, + {0.05198065f, 0.5789965f} + }; + } + + private static float[][] createSampleData() { + return new float[][] { + {0.74021935f, 0.9209938f}, + {0.03902049f, 0.9689629f}, + {0.92514056f, 0.4463501f}, + {0.6673192f, 0.10993068f} + }; + } + + /** + * Test HNSW ACE build with in-memory mode (use_disk=false). + */ + @Test + public void testHnswAceInMemoryBuild() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = createSampleQueries(); + List> expectedResults = getExpectedResults(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + Path buildDir = Files.createTempDirectory("hnsw_ace_test"); + + try { + // Configure ACE parameters for in-memory mode + HnswAceParams aceParams = + new HnswAceParams.Builder() + .withNpartitions(2) + .withBuildDir(buildDir.toString()) + .withUseDisk(false) + .build(); + + // Configure HNSW index parameters with ACE + HnswIndexParams hnswParams = + new HnswIndexParams.Builder() + .withHierarchy(CuvsHnswHierarchy.GPU) + .withM(16) + .withEfConstruction(100) + .withMetric(CuvsDistanceType.L2Expanded) + .withVectorDimension(2) + .withAceParams(aceParams) + .build(); + + // Build the HNSW index using ACE + try (var datasetMatrix = CuVSMatrix.ofArray(dataset)) { + HnswIndex hnswIndex = HnswIndex.build(resources, hnswParams, datasetMatrix); + assertNotNull("HNSW index should not be null", hnswIndex); + log.debug("HNSW ACE index built successfully in memory"); + + // Search the index directly + HnswSearchParams searchParams = new HnswSearchParams.Builder().withEF(100).build(); + HnswQuery hnswQuery = + new HnswQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + SearchResults results = hnswIndex.search(hnswQuery); + log.debug("HNSW ACE search results: " + results.getResults().toString()); + + // Verify search results + checkResults(expectedResults, results.getResults()); + log.debug("HNSW ACE in-memory search verification passed"); + + hnswIndex.close(); + } + } finally { + deleteRecursively(buildDir); + } + } + } + + /** + * Test HNSW ACE build with disk-based mode (use_disk=true). + * This follows the same approach as C++ and C tests: + * build -> serialize -> deserialize -> search + */ + @Test + public void testHnswAceDiskBasedBuild() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = createSampleQueries(); + List> expectedResults = getExpectedResults(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + Path buildDir = Files.createTempDirectory("hnsw_ace_disk_test"); + + try { + // Configure ACE parameters for disk-based mode + HnswAceParams aceParams = + new HnswAceParams.Builder() + .withNpartitions(2) + .withBuildDir(buildDir.toString()) + .withUseDisk(true) + .build(); + + // Configure HNSW index parameters with ACE + HnswIndexParams hnswParams = + new HnswIndexParams.Builder() + .withHierarchy(CuvsHnswHierarchy.GPU) + .withM(16) + .withEfConstruction(100) + .withMetric(CuvsDistanceType.L2Expanded) + .withVectorDimension(2) + .withAceParams(aceParams) + .build(); + + // Build the HNSW index using ACE with disk mode + try (var datasetMatrix = CuVSMatrix.ofArray(dataset)) { + HnswIndex hnswIndex = HnswIndex.build(resources, hnswParams, datasetMatrix); + assertNotNull("HNSW index should not be null", hnswIndex); + log.debug("HNSW ACE index built with disk mode"); + + // For disk mode, the hnsw_index.bin file is created during build + Path hnswIndexPath = buildDir.resolve("hnsw_index.bin"); + + // Verify the index file was created by the disk-based build + assertTrue("HNSW index file should exist", Files.exists(hnswIndexPath)); + log.debug("HNSW index serialized to: " + hnswIndexPath); + + // Deserialize from disk for searching + try (InputStream inputStream = Files.newInputStream(hnswIndexPath)) { + HnswIndex deserializedIndex = + HnswIndex.newBuilder(resources) + .from(inputStream) + .withIndexParams(hnswParams) + .build(); + + assertNotNull("Deserialized index should not be null", deserializedIndex); + log.debug("HNSW index deserialized from disk"); + + // Search the deserialized index + HnswSearchParams searchParams = new HnswSearchParams.Builder().withEF(100).build(); + HnswQuery hnswQuery = + new HnswQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + SearchResults results = deserializedIndex.search(hnswQuery); + log.debug("HNSW ACE disk search results: " + results.getResults().toString()); + + // Verify search results + checkResults(expectedResults, results.getResults()); + log.debug("HNSW ACE disk-based search verification passed"); + + deserializedIndex.close(); + } + + hnswIndex.close(); + } + } finally { + deleteRecursively(buildDir); + } + } + } + + /** + * Test HNSW ACE build with different hierarchy options. + */ + @Test + public void testHnswAceWithDifferentHierarchy() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = createSampleQueries(); + List> expectedResults = getExpectedResults(); + + for (CuvsHnswHierarchy hierarchy : Arrays.asList(CuvsHnswHierarchy.NONE, CuvsHnswHierarchy.GPU)) { + try (CuVSResources resources = CheckedCuVSResources.create()) { + Path buildDir = Files.createTempDirectory("hnsw_ace_hierarchy_test"); + + try { + HnswAceParams aceParams = + new HnswAceParams.Builder() + .withNpartitions(2) + .withBuildDir(buildDir.toString()) + .withUseDisk(false) + .build(); + + HnswIndexParams hnswParams = + new HnswIndexParams.Builder() + .withHierarchy(hierarchy) + .withM(16) + .withEfConstruction(100) + .withMetric(CuvsDistanceType.L2Expanded) + .withVectorDimension(2) + .withAceParams(aceParams) + .build(); + + try (var datasetMatrix = CuVSMatrix.ofArray(dataset)) { + HnswIndex hnswIndex = HnswIndex.build(resources, hnswParams, datasetMatrix); + assertNotNull("HNSW index should not be null for hierarchy: " + hierarchy, hnswIndex); + + HnswSearchParams searchParams = new HnswSearchParams.Builder().withEF(100).build(); + HnswQuery hnswQuery = + new HnswQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + SearchResults results = hnswIndex.search(hnswQuery); + checkResults(expectedResults, results.getResults()); + log.debug("HNSW ACE with hierarchy {} verification passed", hierarchy); + + hnswIndex.close(); + } + } finally { + deleteRecursively(buildDir); + } + } + } + } + + /** + * Test the full disk-based ACE workflow explicitly. + * build -> serialize -> deserialize -> search + */ + @Test + public void testHnswAceDiskSerializeDeserialize() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = createSampleQueries(); + List> expectedResults = getExpectedResults(); + + try (CuVSResources resources = CheckedCuVSResources.create()) { + Path buildDir = Files.createTempDirectory("hnsw_ace_serialize_test"); + + try { + // Create ACE params with disk mode enabled + HnswAceParams aceParams = + new HnswAceParams.Builder() + .withNpartitions(2) + .withBuildDir(buildDir.toString()) + .withUseDisk(true) + .build(); + + // Create HNSW index params with ACE + HnswIndexParams hnswParams = + new HnswIndexParams.Builder() + .withHierarchy(CuvsHnswHierarchy.GPU) + .withM(16) + .withEfConstruction(100) + .withMetric(CuvsDistanceType.L2Expanded) + .withVectorDimension(2) + .withAceParams(aceParams) + .build(); + + // Build the index using ACE + try (var datasetMatrix = CuVSMatrix.ofArray(dataset)) { + HnswIndex hnswIndex = HnswIndex.build(resources, hnswParams, datasetMatrix); + assertNotNull("HNSW index should not be null", hnswIndex); + + // The disk-based build should create hnsw_index.bin + Path hnswFile = buildDir.resolve("hnsw_index.bin"); + assertTrue("HNSW index file should exist after disk build", Files.exists(hnswFile)); + + // Deserialize from disk + try (InputStream inputStream = Files.newInputStream(hnswFile)) { + HnswIndex loadedIndex = + HnswIndex.newBuilder(resources) + .from(inputStream) + .withIndexParams(hnswParams) + .build(); + + // Search the loaded index + HnswSearchParams searchParams = new HnswSearchParams.Builder().withEF(200).build(); + HnswQuery hnswQuery = + new HnswQuery.Builder(resources) + .withTopK(3) + .withSearchParams(searchParams) + .withQueryVectors(queries) + .build(); + + SearchResults results = loadedIndex.search(hnswQuery); + log.debug("Serialize/Deserialize test results: " + results.getResults().toString()); + + // Verify results + checkResults(expectedResults, results.getResults()); + log.debug("HNSW ACE serialize/deserialize test passed"); + + loadedIndex.close(); + } + + hnswIndex.close(); + } + } finally { + deleteRecursively(buildDir); + } + } + } + + /** + * Helper method to recursively delete a directory and its contents. + */ + private void deleteRecursively(Path path) { + try { + if (Files.isDirectory(path)) { + Files.list(path).forEach(this::deleteRecursively); + } + Files.deleteIfExists(path); + } catch (Exception e) { + log.warn("Failed to delete {}: {}", path, e.getMessage()); + } + } +} diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 1bd0cbaf17..1db657dab1 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -128,8 +128,8 @@ cdef class AceParams: """ Parameters for ACE (Augmented Core Extraction) graph building algorithm. - ACE enables building indices for datasets too large to fit in GPU memory by - partitioning the dataset using balanced k-means and building sub-indices + ACE enables building indexes for datasets too large to fit in GPU memory by + partitioning the dataset using balanced k-means and building sub-indexes for each partition independently. Parameters diff --git a/python/cuvs/cuvs/neighbors/hnsw/__init__.py b/python/cuvs/cuvs/neighbors/hnsw/__init__.py index 99d7d99ab2..f91835b7c5 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/__init__.py +++ b/python/cuvs/cuvs/neighbors/hnsw/__init__.py @@ -1,12 +1,14 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 from .hnsw import ( + AceParams, ExtendParams, Index, IndexParams, SearchParams, + build, extend, from_cagra, load, @@ -15,9 +17,11 @@ ) __all__ = [ + "AceParams", "IndexParams", "Index", "ExtendParams", + "build", "extend", "SearchParams", "load", diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd index 399fc06d47..cea142feca 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd @@ -1,9 +1,10 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 +from libc.stddef cimport size_t from libc.stdint cimport int32_t, uintptr_t from libcpp cimport bool @@ -20,10 +21,24 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil: CPU GPU + ctypedef struct cuvsHnswAceParams: + size_t npartitions + const char* build_dir + bool use_disk + + ctypedef cuvsHnswAceParams* cuvsHnswAceParams_t + + cuvsError_t cuvsHnswAceParamsCreate(cuvsHnswAceParams_t* params) + + cuvsError_t cuvsHnswAceParamsDestroy(cuvsHnswAceParams_t params) + ctypedef struct cuvsHnswIndexParams: cuvsHnswHierarchy hierarchy int32_t ef_construction int32_t num_threads + size_t M + cuvsDistanceType metric + cuvsHnswAceParams_t ace_params ctypedef cuvsHnswIndexParams* cuvsHnswIndexParams_t @@ -55,6 +70,11 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil: cuvsCagraIndex_t cagra_index, cuvsHnswIndex_t hnsw_index) except + + cuvsError_t cuvsHnswBuild(cuvsResources_t res, + cuvsHnswIndexParams_t params, + DLManagedTensor* dataset, + cuvsHnswIndex_t index) except + + cuvsError_t cuvsHnswExtend(cuvsResources_t res, cuvsHnswExtendParams_t params, DLManagedTensor* data, diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx index d43fd09651..e18c52a40a 100644 --- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx +++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # # cython: language_level=3 @@ -28,6 +28,63 @@ from pylibraft.common.cai_wrapper import wrap_array from pylibraft.common.interruptible import cuda_interruptible +cdef class AceParams: + """ + Parameters for ACE (Augmented Core Extraction) graph build for HNSW. + + ACE enables building HNSW indices for datasets too large to fit in GPU + memory by partitioning the dataset and building sub-indices for each + partition independently. + + Parameters + ---------- + npartitions : int, default = 1 (optional) + Number of partitions for ACE partitioned build. Small values might + improve recall but potentially degrade performance. 100k - 5M vectors + per partition is recommended depending on available memory. + build_dir : string, default = "/tmp/hnsw_ace_build" (optional) + Directory to store ACE build artifacts (KNN graph, optimized graph). + Used when `use_disk` is true or when the graph does not fit in memory. + use_disk : bool, default = False (optional) + Whether to use disk-based storage for ACE build. When true, enables + disk-based operations for memory-efficient graph construction. + """ + + cdef cuvsHnswAceParams* params + cdef object _build_dir_bytes + + def __cinit__(self): + check_cuvs(cuvsHnswAceParamsCreate(&self.params)) + self._build_dir_bytes = None + + def __dealloc__(self): + if self.params is not NULL: + check_cuvs(cuvsHnswAceParamsDestroy(self.params)) + + def __init__(self, *, + npartitions=1, + build_dir="/tmp/hnsw_ace_build", + use_disk=False): + self.params.npartitions = npartitions + self._build_dir_bytes = build_dir.encode('utf-8') + self.params.build_dir = self._build_dir_bytes + self.params.use_disk = use_disk + + @property + def npartitions(self): + return self.params.npartitions + + @property + def build_dir(self): + if self.params.build_dir is not NULL: + return self.params.build_dir.decode('utf-8') + return "/tmp/hnsw_ace_build" + + @property + def use_disk(self): + return self.params.use_disk + + cdef class IndexParams: """ Parameters to build index for HNSW nearest neighbor search @@ -50,12 +107,23 @@ cdef class IndexParams: NOTE: When hierarchy is `gpu`, while the majority of the work is done on the GPU, initialization of the HNSW index itself and some other work is parallelized with the help of CPU threads. + M : int, default = 32 (optional) + HNSW M parameter: number of bi-directional links per node + (used when building with ACE). graph_degree = m * 2, + intermediate_graph_degree = m * 3. + metric : string, default = "sqeuclidean" (optional) + Distance metric to use. Valid values: ["sqeuclidean", "inner_product"] + ace_params : AceParams, default = None (optional) + ACE parameters for building HNSW index using ACE algorithm. If set, + enables the build() function to use ACE for index construction. """ cdef cuvsHnswIndexParams* params + cdef AceParams _ace_params def __cinit__(self): check_cuvs(cuvsHnswIndexParamsCreate(&self.params)) + self._ace_params = None def __dealloc__(self): check_cuvs(cuvsHnswIndexParamsDestroy(self.params)) @@ -63,7 +131,10 @@ cdef class IndexParams: def __init__(self, *, hierarchy="none", ef_construction=200, - num_threads=0): + num_threads=0, + M=32, + metric="sqeuclidean", + ace_params=None): if hierarchy == "none": self.params.hierarchy = cuvsHnswHierarchy.NONE elif hierarchy == "cpu": @@ -72,9 +143,19 @@ cdef class IndexParams: self.params.hierarchy = cuvsHnswHierarchy.GPU else: raise ValueError("Invalid hierarchy type." - " Valid values are 'none' and 'cpu'.") + " Valid values are 'none', 'cpu', and 'gpu'.") self.params.ef_construction = ef_construction self.params.num_threads = num_threads + self.params.M = M + self.params.metric = DISTANCE_TYPES[metric] + + if ace_params is not None: + if not isinstance(ace_params, AceParams): + raise ValueError("ace_params must be an instance of AceParams") + self._ace_params = ace_params + self.params.ace_params = self._ace_params.params + else: + self.params.ace_params = NULL @property def hierarchy(self): @@ -93,6 +174,14 @@ cdef class IndexParams: def num_threads(self): return self.params.num_threads + @property + def m(self): + return self.params.M + + @property + def ace_params(self): + return self._ace_params + cdef class Index: """ @@ -347,6 +436,95 @@ def from_cagra(IndexParams index_params, cagra.Index cagra_index, return hnsw_index +@auto_sync_resources +def build(IndexParams index_params, dataset, resources=None): + """ + Build an HNSW index using the ACE (Augmented Core Extraction) algorithm. + + ACE enables building HNSW indices for datasets too large to fit in GPU + memory by partitioning the dataset and building sub-indices for each + partition independently. + + NOTE: This function requires `index_params.ace_params` to be set with + an instance of AceParams. + + Parameters + ---------- + index_params : IndexParams + Parameters for the HNSW index with ACE configuration. + Must have `ace_params` set. + dataset : Host array interface compliant matrix shape (n_samples, dim) + Supported dtype [float32, float16, int8, uint8] + {resources_docstring} + + Returns + ------- + index : Index + Trained HNSW index ready for search. + + Examples + -------- + >>> import numpy as np + >>> from cuvs.neighbors import hnsw + >>> + >>> n_samples = 50000 + >>> n_features = 50 + >>> dataset = np.random.random_sample((n_samples, n_features), + ... dtype=np.float32) + >>> + >>> # Create ACE parameters + >>> ace_params = hnsw.AceParams( + ... npartitions=4, + ... use_disk=True, + ... build_dir="/tmp/hnsw_ace_build" + ... ) + >>> + >>> # Create index parameters with ACE + >>> index_params = hnsw.IndexParams( + ... hierarchy="gpu", + ... ace_params=ace_params, + ... ef_construction=120, + ... M=32, + ... metric="sqeuclidean" + ... ) + >>> + >>> # Build the index + >>> index = hnsw.build(index_params, dataset) + >>> + >>> # Search the index + >>> queries = np.random.random_sample((10, n_features), dtype=np.float32) + >>> distances, neighbors = hnsw.search( + ... hnsw.SearchParams(ef=200), + ... index, + ... queries, + ... k=10 + ... ) + """ + if index_params.ace_params is None: + raise ValueError("index_params.ace_params must be set for hnsw.build(). " + "Use AceParams to configure ACE algorithm parameters.") + + dataset_ai = wrap_array(dataset) + _check_input_array(dataset_ai, [np.dtype('float32'), + np.dtype('float16'), + np.dtype('uint8'), + np.dtype('int8')]) + + cdef cydlpack.DLManagedTensor* dataset_dlpack = cydlpack.dlpack_c(dataset_ai) + cdef cuvsResources_t res = resources.get_c_obj() + + cdef Index hnsw_index = Index() + check_cuvs(cuvsHnswBuild( + res, + index_params.params, + dataset_dlpack, + hnsw_index.index + )) + + hnsw_index.trained = True + return hnsw_index + + @auto_sync_resources def extend(ExtendParams extend_params, Index index, data, resources=None): """ diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py new file mode 100644 index 0000000000..b99eb39fd1 --- /dev/null +++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# + +import os +import tempfile + +import numpy as np +import pytest +from sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from cuvs.neighbors import hnsw +from cuvs.tests.ann_utils import calc_recall, generate_data + + +def run_hnsw_ace_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + dtype=np.float32, + metric="sqeuclidean", + npartitions=2, + ef_construction=100, + use_disk=False, + hierarchy="gpu", + expected_recall=0.9, +): + """ + Test HNSW index build using ACE algorithm. + + - Build HNSW index using ACE via hnsw.build() + - For disk mode: serialize -> deserialize -> search + - For in-memory mode: search directly + """ + dataset = generate_data((n_rows, n_cols), dtype) + queries = generate_data((n_queries, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + queries = normalize(queries, norm="l2", axis=1) + if dtype in [np.int8, np.uint8]: + # Quantize the normalized data to the int8/uint8 range + dtype_max = np.iinfo(dtype).max + dataset = (dataset * dtype_max).astype(dtype) + queries = (queries * dtype_max).astype(dtype) + + # Create a temporary directory for ACE build + with tempfile.TemporaryDirectory() as temp_dir: + # Set up ACE parameters + ace_params = hnsw.AceParams( + npartitions=npartitions, + build_dir=temp_dir, + use_disk=use_disk, + ) + + # Build parameters with ACE configuration + index_params = hnsw.IndexParams( + hierarchy=hierarchy, + M=32, + ef_construction=ef_construction, + metric=metric, + ace_params=ace_params, + ) + + # Build the HNSW index using ACE + hnsw_index = hnsw.build(index_params, dataset) + + assert hnsw_index.trained + + if use_disk: + # For disk mode, the index is serialized to disk by the build function + # We need to deserialize it before searching + hnsw_file = os.path.join(temp_dir, "hnsw_index.bin") + assert os.path.exists(hnsw_file) + + # Deserialize from disk for searching + deserialized_index = hnsw.load( + index_params, + hnsw_file, + n_cols, + dtype, + metric=metric, + ) + + # Search the deserialized index + search_params = hnsw.SearchParams( + ef=max(ef_construction, k * 2), num_threads=1 + ) + out_dist, out_idx = hnsw.search( + search_params, deserialized_index, queries, k + ) + else: + # For in-memory mode, search directly + search_params = hnsw.SearchParams( + ef=max(ef_construction, k * 2), num_threads=1 + ) + out_dist, out_idx = hnsw.search( + search_params, hnsw_index, queries, k + ) + + # Calculate reference values with sklearn + skl_metric = { + "sqeuclidean": "sqeuclidean", + "inner_product": "cosine", + "euclidean": "euclidean", + }[metric] + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric=skl_metric + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(out_idx, skl_idx) + assert recall >= expected_recall, ( + f"Recall {recall:.3f} is below expected {expected_recall}" + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) +@pytest.mark.parametrize("metric", ["sqeuclidean", "inner_product"]) +@pytest.mark.parametrize("use_disk", [False, True]) +def test_hnsw_ace_build_search(dtype, metric, use_disk): + """Test HNSWACE with different data types and metrics.""" + run_hnsw_ace_build_search_test( + dtype=dtype, + metric=metric, + use_disk=use_disk, + ) + + +@pytest.mark.parametrize("npartitions", [2, 3, 8]) +def test_hnsw_ace_partitions(npartitions): + """Test HNSW ACE with different partition sizes (disk mode only).""" + run_hnsw_ace_build_search_test( + use_disk=True, + npartitions=npartitions, + ) + + +@pytest.mark.parametrize("ef_construction", [50, 100, 200]) +def test_hnsw_ace_ef_construction(ef_construction): + """Test HNSW ACE with different ef_construction values (disk mode only).""" + run_hnsw_ace_build_search_test( + use_disk=True, + ef_construction=ef_construction, + ) + + +@pytest.mark.parametrize("hierarchy", ["none", "gpu"]) +def test_hnsw_ace_hierarchy(hierarchy): + """Test HNSW ACE with different hierarchy modes (disk mode only).""" + run_hnsw_ace_build_search_test( + use_disk=True, + hierarchy=hierarchy, + ) + + +def test_hnsw_ace_disk_serialize_deserialize(): + """ + Test the full disk-based ACE workflow: + build -> serialize -> deserialize -> search + """ + n_rows = 10000 + n_cols = 10 + n_queries = 100 + k = 10 + dtype = np.float32 + metric = "sqeuclidean" + + dataset = generate_data((n_rows, n_cols), dtype) + queries = generate_data((n_queries, n_cols), dtype) + + with tempfile.TemporaryDirectory() as temp_dir: + # Create ACE params with disk mode enabled + ace_params = hnsw.AceParams( + npartitions=2, + build_dir=temp_dir, + use_disk=True, + ) + + # Create HNSW index params with ACE + index_params = hnsw.IndexParams( + hierarchy="gpu", + M=32, + ef_construction=120, + metric=metric, + ace_params=ace_params, + ) + + # Build the index using ACE + hnsw_index = hnsw.build(index_params, dataset) + assert hnsw_index.trained + + # Serialize to a specific file path + hnsw_file = os.path.join(temp_dir, "test_hnsw_index.bin") + hnsw.save(hnsw_file, hnsw_index) + assert os.path.exists(hnsw_file) + + # Deserialize from disk + loaded_index = hnsw.load( + index_params, + hnsw_file, + n_cols, + dtype, + metric=metric, + ) + + # Search the loaded index + search_params = hnsw.SearchParams(ef=200, num_threads=1) + out_dist, out_idx = hnsw.search( + search_params, loaded_index, queries, k + ) + + # Verify results against sklearn + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric="sqeuclidean" + ) + nn_skl.fit(dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(out_idx, skl_idx) + assert recall >= 0.7, f"Recall {recall:.3f} is below expected 0.7"