diff --git a/CMakeLists.txt b/CMakeLists.txt index 5b090f70d8..0a85cce027 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -125,6 +125,10 @@ if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs) find_package(cuvs) endif() +if(FAISS_ENABLE_CUVS AND NOT TARGET rmm::rmm) + find_package(rmm REQUIRED) +endif() + add_subdirectory(faiss) if(FAISS_ENABLE_GPU) diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index 42f8548a65..c67d0e941d 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -17,7 +17,6 @@ #include #include #include -#include #include #include "faiss/Index.h" @@ -1433,13 +1432,11 @@ void IndexHNSWCagra::search( // first real candidate will always be strictly better. nearest_d[i] = C::neutral(); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution distrib( - 0, this->ntotal - 1); + // Seeded per query so entrypoints are reproducible. + SplitMix64RandomGenerator gen(i); for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) { - auto idx = distrib(gen); + idx_t idx = gen.rand_int64() % this->ntotal; auto distance = (*dis)(idx); if (C::cmp(nearest_d[i], distance)) { nearest[i] = static_cast(idx); @@ -1498,12 +1495,11 @@ void IndexHNSWCagra::range_search( // real candidate will always be strictly better. float nearest_d = C::neutral(); - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_int_distribution distrib(0, ntotal - 1); + // For reproducible entrypoint. + SplitMix64RandomGenerator gen(i); for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) { - auto idx = distrib(gen); + idx_t idx = gen.rand_int64() % ntotal; auto distance = (*dis)(idx); // C::cmp(nearest_d, distance) is true iff distance is // strictly better than the current nearest_d. diff --git a/faiss/IndexHNSW.h b/faiss/IndexHNSW.h index 8e3d7f2906..405e14d68a 100644 --- a/faiss/IndexHNSW.h +++ b/faiss/IndexHNSW.h @@ -253,7 +253,7 @@ struct IndexHNSWCagra : IndexHNSW { /// searches only the base level knn graph of the HNSW index. /// This parameter selects the entry point by randomly selecting /// some points and using the best one. - int num_base_level_search_entrypoints = 32; + int num_base_level_search_entrypoints = 256; void add(idx_t n, const float* x) override; diff --git a/faiss/gpu/CMakeLists.txt b/faiss/gpu/CMakeLists.txt index c896a85a29..2387eada10 100644 --- a/faiss/gpu/CMakeLists.txt +++ b/faiss/gpu/CMakeLists.txt @@ -363,7 +363,7 @@ else() find_package(CUDAToolkit REQUIRED) - target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$:cuvs::cuvs> $<$:OpenMP::OpenMP_CXX>) + target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$:cuvs::cuvs> $<$:rmm::rmm> $<$:OpenMP::OpenMP_CXX>) target_compile_options(faiss_gpu_objs PRIVATE $<$:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr diff --git a/faiss/gpu/GpuIndexBinaryCagra.cu b/faiss/gpu/GpuIndexBinaryCagra.cu index 46b6efb636..d2d560f430 100644 --- a/faiss/gpu/GpuIndexBinaryCagra.cu +++ b/faiss/gpu/GpuIndexBinaryCagra.cu @@ -78,6 +78,13 @@ std::shared_ptr GpuIndexBinaryCagra::getResources() { } void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) { + // The config is shared with the float index; there is no binary + // multi-GPU build, so reject rather than silently ignoring the request. + FAISS_THROW_IF_MSG( + cagraConfig_.devices.size() > 1, + "binary CAGRA has no multi-GPU build; " + "GpuIndexCagraConfig::devices must name at most one device"); + DeviceScope scope(cagraConfig_.device); if (this->is_trained) { FAISS_ASSERT(index_); diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu index acd356c5aa..0eac9f8752 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -23,15 +23,12 @@ #include #include -#include #include #include #include #include #include -#include #include -#include #include #include @@ -40,56 +37,31 @@ #include #include #include -#include -#include +// clang-format off +#if __has_include() +// Some cuVS versions put helpers::optimize in its own header. +#include +#endif +// TEMPORARY. cuVS has no public API to prune a graph that is already on the +// device, so we reach into its internals when they happen to be available. +// Delete this and the branch it guards once that API ships upstream. +#if defined(FAISS_CAGRA_DEVICE_OPTIMIZE) +#include +#endif +// clang-format on namespace { -template -__global__ void kern_detour_count( - const uint32_t* __restrict__ knn_graph, - const uint32_t graph_size, - const uint32_t graph_degree, - const uint32_t output_degree, - const uint32_t batch_size, - const uint32_t offset, - uint8_t* __restrict__ detour_count) { - __shared__ uint32_t smem[MAX_DEGREE]; - - const uint64_t iA = blockIdx.x + (uint64_t)offset; - if (iA >= graph_size) - return; - - for (uint32_t k = threadIdx.x; k < graph_degree; k += blockDim.x) { - smem[k] = 0; - if (knn_graph[k + (uint64_t)graph_degree * iA] == (uint32_t)iA) - smem[k] = graph_degree; - } - __syncthreads(); - - for (uint32_t kAD = 0; kAD < graph_degree - 1; kAD++) { - const uint64_t iD = knn_graph[kAD + (uint64_t)graph_degree * iA]; - if (iD >= graph_size) - continue; - for (uint32_t kDB = threadIdx.x; kDB < graph_degree; - kDB += blockDim.x) { - const uint64_t iB_cand = - knn_graph[kDB + (uint64_t)graph_degree * iD]; - for (uint32_t kAB = kAD + 1; kAB < graph_degree; kAB++) { - const uint64_t iB = - knn_graph[kAB + (uint64_t)graph_degree * iA]; - if (iB == iB_cand) { - atomicAdd(smem + kAB, 1); - break; - } - } - } - __syncthreads(); - } - - for (uint32_t k = threadIdx.x; k < graph_degree; k += blockDim.x) { - detour_count[k + (uint64_t)graph_degree * iA] = - static_cast(min(smem[k], 255u)); +// all_neighbors emits int64 indices, optimize consumes uint32. Narrowing on +// device lets the graph stay there. +__global__ void kern_narrow_indices( + const int64_t* __restrict__ src, + uint32_t* __restrict__ dst, + size_t count) { + size_t stride = (size_t)gridDim.x * blockDim.x; + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < count; + i += stride) { + dst[i] = static_cast(src[i]); } } @@ -109,6 +81,19 @@ GpuIndexCagra::GpuIndexCagra( } void GpuIndexCagra::train_ex(idx_t n, const void* x, NumericType numeric_type) { + if (cagraConfig_.devices.size() > 1) { + FAISS_THROW_IF_NOT_MSG( + numeric_type == NumericType::Float32, + "multi-GPU CAGRA build only supports Float32"); + trainAllNeighbors_(n, static_cast(x)); + return; + } + + FAISS_THROW_IF_MSG( + cagraConfig_.build_algo == graph_build_algo::BRUTE_FORCE, + "graph_build_algo::BRUTE_FORCE is only available on the multi-GPU " + "build path (set GpuIndexCagraConfig::devices)"); + numeric_type_ = numeric_type; bool index_is_initialized = !std::holds_alternative(index_); @@ -370,461 +355,99 @@ void GpuIndexCagra::searchImpl_( search_params); } -void GpuIndexCagra::trainMultiGpu( - idx_t n, - const float* x, - std::vector& providers, - std::vector& devices, - idx_t stitch_per_shard, - int stitch_k, - int stitch_mode) { +void GpuIndexCagra::trainAllNeighbors_(idx_t n, const float* x) { FAISS_THROW_IF_MSG(is_trained, "index is already trained"); - FAISS_THROW_IF_MSG(devices.empty(), "must provide at least one GPU"); - FAISS_THROW_IF_NOT_MSG(stitch_k >= 1, "stitch_k must be >= 1"); - - numeric_type_ = NumericType::Float32; - int num_shards = static_cast(devices.size()); - idx_t shard_size = (n + num_shards - 1) / num_shards; - - int total_cross = (num_shards > 1) ? stitch_k * (num_shards - 1) : 0; - - auto t_phase = std::chrono::high_resolution_clock::now(); - auto t_start = t_phase; - - // Use cuVS native multi-GPU CAGRA build via SNMG - raft::device_resources_snmg clique(devices); - - cuvs::neighbors::mg_index_params - mg_params; - mg_params.mode = cuvs::neighbors::SHARDED; - mg_params.intermediate_graph_degree = - cagraConfig_.intermediate_graph_degree; - mg_params.graph_degree = cagraConfig_.graph_degree; - mg_params.guarantee_connectivity = cagraConfig_.guarantee_connectivity; - mg_params.metric = metricFaissToCuvs(this->metric_type, false); - - if (cagraConfig_.build_algo == graph_build_algo::IVF_PQ) { - cuvs::neighbors::cagra::graph_build_params::ivf_pq_params - graph_build_params; - if (cagraConfig_.ivf_pq_params) { - graph_build_params.build_params.n_lists = - cagraConfig_.ivf_pq_params->n_lists; - graph_build_params.build_params.kmeans_n_iters = - cagraConfig_.ivf_pq_params->kmeans_n_iters; - graph_build_params.build_params.kmeans_trainset_fraction = - cagraConfig_.ivf_pq_params->kmeans_trainset_fraction; - graph_build_params.build_params.pq_bits = - cagraConfig_.ivf_pq_params->pq_bits; - graph_build_params.build_params.pq_dim = - cagraConfig_.ivf_pq_params->pq_dim; - graph_build_params.build_params.codebook_kind = - static_cast( - cagraConfig_.ivf_pq_params->codebook_kind); - graph_build_params.build_params.force_random_rotation = - cagraConfig_.ivf_pq_params->force_random_rotation; - graph_build_params.build_params.conservative_memory_allocation = - cagraConfig_.ivf_pq_params->conservative_memory_allocation; - } - if (cagraConfig_.ivf_pq_search_params) { - graph_build_params.search_params.n_probes = - cagraConfig_.ivf_pq_search_params->n_probes; - graph_build_params.search_params.lut_dtype = - cagraConfig_.ivf_pq_search_params->lut_dtype; - graph_build_params.search_params.preferred_shmem_carveout = - cagraConfig_.ivf_pq_search_params->preferred_shmem_carveout; - } - graph_build_params.build_params.metric = - metricFaissToCuvs(this->metric_type, false); - graph_build_params.refinement_rate = cagraConfig_.refine_rate; - mg_params.graph_build_params = graph_build_params; - if (mg_params.graph_degree == mg_params.intermediate_graph_degree) { - mg_params.intermediate_graph_degree = 1.5 * mg_params.graph_degree; - } - } else { - cuvs::neighbors::cagra::graph_build_params::nn_descent_params - graph_build_params(mg_params.intermediate_graph_degree); - graph_build_params.max_iterations = cagraConfig_.nn_descent_niter; - graph_build_params.metric = metricFaissToCuvs(this->metric_type, false); - mg_params.graph_build_params = graph_build_params; - } - - auto dataset = raft::make_host_matrix_view( - x, n, static_cast(this->d)); - - auto mg_idx = cuvs::neighbors::cagra::build(clique, mg_params, dataset); - - auto t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] SNMG CAGRA build: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - t_phase = t_now; - - // Extract per-shard graphs into expanded layout: [CAGRA neighbors | - // cross-shard slots] - idx_t graph_degree = static_cast(cagraConfig_.graph_degree); - idx_t expanded_degree = graph_degree + total_cross; - merged_knngraph_.assign(n * expanded_degree, -1); - - for (int s = 0; s < num_shards; s++) { - idx_t offset = static_cast(s) * shard_size; - idx_t actual_size = std::min(shard_size, n - offset); - - auto& shard_idx = mg_idx.ann_interfaces_[s].index_.value(); - auto device_graph = shard_idx.graph(); - - const auto& dev_res = - raft::resource::set_current_device_to_rank(clique, s); - FAISS_THROW_IF_NOT_FMT( - device_graph.extent(1) == graph_degree, - "Shard %d has graph_degree %ld, expected %ld", - s, - (long)device_graph.extent(1), - (long)graph_degree); - - std::vector host_graph(actual_size * graph_degree); - raft::resource::sync_stream(dev_res); - thrust::copy( - thrust::device_ptr(device_graph.data_handle()), - thrust::device_ptr( - device_graph.data_handle() + - actual_size * graph_degree), - host_graph.data()); - -#pragma omp parallel for - for (idx_t i = 0; i < actual_size; i++) { - for (idx_t j = 0; j < graph_degree; j++) { - merged_knngraph_[(offset + i) * expanded_degree + j] = - static_cast(host_graph[i * graph_degree + j]) + - offset; - } - } - } - - merged_knngraph_degree_ = expanded_degree; - - t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] Graph extract + merge: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - t_phase = t_now; - - // Cross-shard stitching: appends cross-shard edges after CAGRA neighbors. - // Mode 0: CPU HNSW search (approximate). Mode 1: GPU brute-force (exact). - if (num_shards > 1) { - if (stitch_mode == 1 && stitch_per_shard == 0) { - stitch_per_shard = 100000; - fprintf(stderr, - " [trainMultiGpu] GPU brute-force mode: defaulting " - "stitch_per_shard to %ld (all-vectors is O(N^2))\n", - (long)stitch_per_shard); - } - - fprintf(stderr, - " [trainMultiGpu] Stitching: mode=%s, stitch_per_shard=%ld, " - "stitch_k=%d, total_cross=%d, expanded_degree=%ld\n", - stitch_mode == 1 ? "GPU-brute-force" : "CPU-HNSW", - (long)stitch_per_shard, - stitch_k, - total_cross, - (long)expanded_degree); - - if (stitch_mode == 1) { - // GPU brute-force stitching: exact cross-shard neighbors - for (int j = 0; j < num_shards; j++) { - idx_t j_offset = static_cast(j) * shard_size; - idx_t j_size = std::min(shard_size, n - j_offset); - if (j_size == 0) - continue; - - GpuIndexFlatConfig flat_config; - flat_config.device = devices[j]; - GpuIndexFlatL2 flat_idx(providers[j], this->d, flat_config); - flat_idx.add(j_size, x + j_offset * this->d); - - for (int i = 0; i < num_shards; i++) { - if (i == j) - continue; - idx_t i_offset = static_cast(i) * shard_size; - idx_t i_size = std::min(shard_size, n - i_offset); - if (i_size == 0) - continue; - - int target_pos = (j < i) ? j : j - 1; - idx_t slot_base = graph_degree + - static_cast(target_pos) * stitch_k; - - idx_t num_queries = i_size; - const float* query_data = x + i_offset * this->d; - std::vector sampled_indices; - std::vector sampled_vectors; - - if (stitch_per_shard > 0 && stitch_per_shard < i_size) { - num_queries = stitch_per_shard; - sampled_indices.resize(i_size); - std::iota( - sampled_indices.begin(), - sampled_indices.end(), - idx_t(0)); - std::mt19937 rng(42 + i * num_shards + j); - for (idx_t s = 0; s < stitch_per_shard; s++) { - std::uniform_int_distribution dist( - s, i_size - 1); - std::swap( - sampled_indices[s], - sampled_indices[dist(rng)]); - } - sampled_indices.resize(stitch_per_shard); - - sampled_vectors.resize(stitch_per_shard * this->d); -#pragma omp parallel for - for (idx_t s = 0; s < stitch_per_shard; s++) { - memcpy(sampled_vectors.data() + s * this->d, - x + - (i_offset + sampled_indices[s]) * - this->d, - this->d * sizeof(float)); - } - query_data = sampled_vectors.data(); - } - - std::vector distances(num_queries * stitch_k); - std::vector labels(num_queries * stitch_k); - flat_idx.search( - num_queries, - query_data, - stitch_k, - distances.data(), - labels.data()); - -#pragma omp parallel for - for (idx_t s = 0; s < num_queries; s++) { - idx_t orig_idx = (!sampled_indices.empty()) - ? sampled_indices[s] - : s; - idx_t global_id = i_offset + orig_idx; - for (int k = 0; k < stitch_k; k++) { - idx_t local_nb = labels[s * stitch_k + k]; - merged_knngraph_ - [global_id * expanded_degree + slot_base + - k] = (local_nb >= 0) - ? (local_nb + j_offset) - : local_nb; - } - } - } - fprintf(stderr, - " Shard %d/%d GPU-stitched\n", - j + 1, - num_shards); - } - } else { - // CPU HNSW stitching (Approach C) - auto M_temp = graph_degree / 2; - for (int j = 0; j < num_shards; j++) { - idx_t j_offset = static_cast(j) * shard_size; - idx_t j_size = std::min(shard_size, n - j_offset); - if (j_size == 0) - continue; - - IndexHNSWCagra temp_idx; - temp_idx.d = this->d; - temp_idx.metric_type = this->metric_type; - temp_idx.base_level_only = true; - temp_idx.num_base_level_search_entrypoints = 32; - if (this->metric_type == METRIC_L2) { - temp_idx.storage = new IndexFlatL2(this->d); - } else { - temp_idx.storage = new IndexFlatIP(this->d); - } - temp_idx.own_fields = true; - temp_idx.keep_max_size_level0 = true; - temp_idx.hnsw.reset(); - temp_idx.hnsw.assign_probas.clear(); - temp_idx.hnsw.cum_nneighbor_per_level.clear(); - temp_idx.hnsw.set_default_probas(M_temp, 1.0 / log(M_temp)); - temp_idx.init_level0 = false; - temp_idx.hnsw.prepare_level_tab(j_size, false); - temp_idx.storage->add(j_size, x + j_offset * this->d); - temp_idx.ntotal = j_size; - -#pragma omp parallel for - for (idx_t i = 0; i < j_size; i++) { - size_t begin, end; - temp_idx.hnsw.neighbor_range(i, 0, &begin, &end); - for (size_t k = begin; k < end; k++) { - idx_t global_nb = merged_knngraph_ - [(j_offset + i) * expanded_degree + - (k - begin)]; - temp_idx.hnsw.neighbors[k] = global_nb - j_offset; - } - } - - temp_idx.hnsw.efSearch = 64; - - for (int i = 0; i < num_shards; i++) { - if (i == j) - continue; - idx_t i_offset = static_cast(i) * shard_size; - idx_t i_size = std::min(shard_size, n - i_offset); - if (i_size == 0) - continue; - - int target_pos = (j < i) ? j : j - 1; - idx_t slot_base = graph_degree + - static_cast(target_pos) * stitch_k; - - idx_t num_queries = i_size; - const float* query_data = x + i_offset * this->d; - std::vector sampled_indices; - std::vector sampled_vectors; - - if (stitch_per_shard > 0 && stitch_per_shard < i_size) { - num_queries = stitch_per_shard; - sampled_indices.resize(i_size); - std::iota( - sampled_indices.begin(), - sampled_indices.end(), - idx_t(0)); - std::mt19937 rng(42 + i * num_shards + j); - for (idx_t s = 0; s < stitch_per_shard; s++) { - std::uniform_int_distribution dist( - s, i_size - 1); - std::swap( - sampled_indices[s], - sampled_indices[dist(rng)]); - } - sampled_indices.resize(stitch_per_shard); - - sampled_vectors.resize(stitch_per_shard * this->d); -#pragma omp parallel for - for (idx_t s = 0; s < stitch_per_shard; s++) { - memcpy(sampled_vectors.data() + s * this->d, - x + - (i_offset + sampled_indices[s]) * - this->d, - this->d * sizeof(float)); - } - query_data = sampled_vectors.data(); - } - - std::vector distances(num_queries * stitch_k); - std::vector labels(num_queries * stitch_k); - temp_idx.search( - num_queries, - query_data, - stitch_k, - distances.data(), - labels.data()); - -#pragma omp parallel for - for (idx_t s = 0; s < num_queries; s++) { - idx_t orig_idx = (!sampled_indices.empty()) - ? sampled_indices[s] - : s; - idx_t global_id = i_offset + orig_idx; - for (int k = 0; k < stitch_k; k++) { - idx_t local_nb = labels[s * stitch_k + k]; - merged_knngraph_ - [global_id * expanded_degree + slot_base + - k] = (local_nb >= 0) - ? (local_nb + j_offset) - : local_nb; - } - } - } - fprintf(stderr, - " Shard %d/%d stitched\n", - j + 1, - num_shards); - } - } // end else (CPU HNSW) - } - - t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] Stitching: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - fprintf(stderr, - " [trainMultiGpu] Total: %.2f seconds\n", - std::chrono::duration(t_now - t_start).count()); - - multi_gpu_dataset_ = x; - this->is_trained = true; - this->ntotal = n; -} -void GpuIndexCagra::trainAllNeighbors( - idx_t n, - const float* x, - std::vector& devices, - int n_clusters_override, - int overlap_factor_override, - bool multi_gpu_optimize, - int build_algo, - float refinement_rate, - int ivfpq_search_batch) { - FAISS_THROW_IF_MSG(is_trained, "index is already trained"); - FAISS_THROW_IF_MSG(devices.empty(), "must provide at least one GPU"); + const auto& devices = cagraConfig_.devices; + const auto& an_config = cagraConfig_.all_neighbors_params; numeric_type_ = NumericType::Float32; idx_t graph_degree = static_cast(cagraConfig_.graph_degree); idx_t intermediate_degree = static_cast(cagraConfig_.intermediate_graph_degree); + const size_t knn_count = (size_t)n * intermediate_degree; auto t_start = std::chrono::high_resolution_clock::now(); auto t_phase = t_start; - // Phase 1+2: Build kNN graph (multi-GPU) then D2H + cast. - // clique and d_indices are scoped so SNMG resources are fully destroyed - // before Phase 3, avoiding CUDA context interference with single-GPU - // optimize. - auto h_knn = - raft::make_host_matrix(n, intermediate_degree); + // Allocated off the clique so it survives the clique's destruction below. + CUDA_VERIFY(cudaSetDevice(devices[0])); + raft::device_resources single_gpu_res; + auto d_knn = raft::make_device_matrix( + single_gpu_res, n, intermediate_degree); + + // Build the knn graph across all devices. The clique is scoped: optimize + // hangs intermittently if multi-GPU state is still alive when it runs. { raft::device_resources_snmg clique(devices); cuvs::neighbors::all_neighbors::all_neighbors_params an_params; int num_gpus = static_cast(devices.size()); - an_params.n_clusters = n_clusters_override > 0 - ? static_cast(n_clusters_override) - : std::max(num_gpus * 2, 4); - an_params.overlap_factor = overlap_factor_override > 0 - ? static_cast(overlap_factor_override) - : 2; + an_params.n_clusters = an_config.n_clusters > 0 + ? an_config.n_clusters + : static_cast(std::max(num_gpus * 2, 4)); + an_params.overlap_factor = an_config.overlap_factor; an_params.metric = metricFaissToCuvs(this->metric_type, false); const char* algo_name = "nn_descent"; - if (build_algo == 1) { - // Brute-force: exact kNN via tiled GEMM, O(N²D) per cluster - cuvs::neighbors::all_neighbors::graph_build_params:: - brute_force_params bf_params; - an_params.graph_build_params = bf_params; - algo_name = "brute_force"; - } else if (build_algo == 2) { - // IVF-PQ: approximate kNN with refinement - auto dataset_ext = raft::make_extents( - n, static_cast(this->d)); - cuvs::neighbors::all_neighbors::graph_build_params::ivf_pq_params - ivfpq_params(dataset_ext); - ivfpq_params.refinement_rate = refinement_rate; - // Bound the IVF-PQ search workspace to avoid GPU OOM at scale. The - // cuVS default max_internal_batch_size is 128*1024, whose search - // buffers can exceed H100 memory for dense 100M clusters. Capping - // it batches the search into smaller query chunks (results - // unchanged). - if (ivfpq_search_batch > 0) { - ivfpq_params.search_params.max_internal_batch_size = - static_cast(ivfpq_search_batch); + switch (cagraConfig_.build_algo) { + case graph_build_algo::BRUTE_FORCE: { + // Exact kNN via tiled GEMM, O(N^2 D) per cluster + cuvs::neighbors::all_neighbors::graph_build_params:: + brute_force_params bf_params; + an_params.graph_build_params = bf_params; + algo_name = "brute_force"; + break; } - an_params.graph_build_params = ivfpq_params; - algo_name = "ivf_pq"; - } else { - // Default: NN-descent - cuvs::neighbors::all_neighbors::graph_build_params:: - nn_descent_params nn_params; - nn_params.graph_degree = intermediate_degree; - nn_params.max_iterations = cagraConfig_.nn_descent_niter; - an_params.graph_build_params = nn_params; + case graph_build_algo::IVF_PQ: { + // cuVS derives good IVF-PQ params from the dataset shape, + // so only override what it cannot infer. See + // AllNeighborsCagraConfig::ivf_pq_size_from_cluster. + idx_t sizing_rows = n; + if (an_config.ivf_pq_size_from_cluster) { + sizing_rows = std::max( + 1, + (idx_t)((double)n * an_params.overlap_factor / + (double)an_params.n_clusters)); + } + auto dataset_ext = raft::make_extents( + sizing_rows, static_cast(this->d)); + cuvs::neighbors::all_neighbors::graph_build_params:: + ivf_pq_params ivfpq_params(dataset_ext); + ivfpq_params.refinement_rate = an_config.refinement_rate; + // Bounds search memory; batches the search rather than + // changing its result. + if (an_config.ivf_pq_search_batch_size > 0) { + ivfpq_params.search_params.max_internal_batch_size = + an_config.ivf_pq_search_batch_size; + } + fprintf(stderr, + " [trainAllNeighbors] IVF-PQ sized from %s " + "(%ld rows): n_lists=%u n_probes=%u " + "kmeans_trainset_fraction=%.4f refine=%.2f\n", + an_config.ivf_pq_size_from_cluster ? "cluster" + : "full dataset", + (long)sizing_rows, + ivfpq_params.build_params.n_lists, + ivfpq_params.search_params.n_probes, + ivfpq_params.build_params.kmeans_trainset_fraction, + ivfpq_params.refinement_rate); + an_params.graph_build_params = ivfpq_params; + algo_name = "ivf_pq"; + break; + } + case graph_build_algo::NN_DESCENT: { + cuvs::neighbors::all_neighbors::graph_build_params:: + nn_descent_params nn_params; + nn_params.graph_degree = intermediate_degree; + nn_params.max_iterations = cagraConfig_.nn_descent_niter; + an_params.graph_build_params = nn_params; + break; + } + default: + FAISS_THROW_MSG( + "multi-GPU CAGRA build supports build_algo IVF_PQ, " + "NN_DESCENT or BRUTE_FORCE"); } auto dataset = raft::make_host_matrix_view( @@ -842,7 +465,9 @@ void GpuIndexCagra::trainAllNeighbors( an_params.n_clusters, an_params.overlap_factor, algo_name, - build_algo == 2 ? refinement_rate : 0.0f); + cagraConfig_.build_algo == graph_build_algo::IVF_PQ + ? an_config.refinement_rate + : 0.0f); cuvs::neighbors::all_neighbors::build( clique, an_params, dataset, d_indices.view()); @@ -853,258 +478,76 @@ void GpuIndexCagra::trainAllNeighbors( std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // D2H + cast int64 -> uint32 - std::vector h_indices_i64(n * intermediate_degree); - raft::copy( - h_indices_i64.data(), - d_indices.data_handle(), - n * intermediate_degree, - raft::resource::get_cuda_stream(clique)); - raft::resource::sync_stream(clique); - -#pragma omp parallel for - for (idx_t i = 0; i < n * intermediate_degree; i++) { - h_knn.data_handle()[i] = static_cast(h_indices_i64[i]); - } + CUDA_VERIFY(cudaSetDevice(devices[0])); + constexpr int kThreads = 256; + const int blocks = (int)std::min( + (knn_count + kThreads - 1) / kThreads, 65535); + kern_narrow_indices<<>>( + d_indices.data_handle(), d_knn.data_handle(), knn_count); + CUDA_VERIFY(cudaGetLastError()); + CUDA_VERIFY(cudaDeviceSynchronize()); } // clique + d_indices destroyed here, freeing all GPU resources auto t_now = std::chrono::high_resolution_clock::now(); fprintf(stderr, - " [trainAllNeighbors] D2H + cast: %.2f seconds\n", + " [trainAllNeighbors] narrow to uint32 (device): %.2f seconds\n", std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // Phase 3: Optimize kNN graph into CAGRA graph - auto h_cagra = raft::make_host_matrix(n, graph_degree); - memset(h_cagra.data_handle(), 0xff, n * graph_degree * sizeof(uint32_t)); - - if (multi_gpu_optimize && devices.size() > 1) { - // S2: Multi-GPU detour counting + CPU pruning/reverse/merge - int num_gpus = static_cast(devices.size()); - fprintf(stderr, - " [trainAllNeighbors] Multi-GPU optimize: %d GPUs, " - "n=%ld, %ld->%ld\n", - num_gpus, - (long)n, - (long)intermediate_degree, - (long)graph_degree); - - // Phase 3a: Multi-GPU detour counting - FAISS_THROW_IF_NOT_MSG( - intermediate_degree <= 1024, - "intermediate_graph_degree must be <= 1024 for " - "multi-GPU optimize"); - - auto h_detour = raft::make_host_matrix( - n, intermediate_degree); - memset(h_detour.data_handle(), - 0xff, - n * intermediate_degree * sizeof(uint8_t)); - - idx_t chunk = (n + num_gpus - 1) / num_gpus; - -#pragma omp parallel for num_threads(num_gpus) - for (int g = 0; g < num_gpus; g++) { - idx_t my_start = g * chunk; - idx_t my_end = std::min(my_start + chunk, n); - idx_t my_n = my_end - my_start; - if (my_n <= 0) - continue; - - CUDA_VERIFY(cudaSetDevice(devices[g])); - cudaStream_t stream; - CUDA_VERIFY(cudaStreamCreate(&stream)); - - uint32_t* d_knn_graph; - CUDA_VERIFY(cudaMalloc( - &d_knn_graph, - (size_t)n * intermediate_degree * sizeof(uint32_t))); - CUDA_VERIFY(cudaMemcpyAsync( - d_knn_graph, - h_knn.data_handle(), - (size_t)n * intermediate_degree * sizeof(uint32_t), - cudaMemcpyHostToDevice, - stream)); - - uint8_t* d_detour; - CUDA_VERIFY(cudaMalloc( - &d_detour, - (size_t)n * intermediate_degree * sizeof(uint8_t))); - CUDA_VERIFY(cudaMemsetAsync( - d_detour, - 0xff, - (size_t)n * intermediate_degree * sizeof(uint8_t), - stream)); - - // Launch detour counting for this GPU's node range - constexpr uint32_t BATCH = 262144; - dim3 threads(32, 1, 1); - uint32_t total_batches = - (static_cast(my_n) + BATCH - 1) / BATCH; - - for (uint32_t ib = 0; ib < total_batches; ib++) { - uint32_t batch_start = - static_cast(my_start) + ib * BATCH; - uint32_t batch_n = std::min( - BATCH, static_cast(my_end) - batch_start); - dim3 blocks(batch_n, 1, 1); - kern_detour_count<1024><<>>( - d_knn_graph, - static_cast(n), - static_cast(intermediate_degree), - static_cast(graph_degree), - BATCH, - batch_start, - d_detour); - } - - CUDA_VERIFY(cudaGetLastError()); - CUDA_VERIFY(cudaMemcpyAsync( - h_detour.data_handle() + my_start * intermediate_degree, - d_detour + my_start * intermediate_degree, - (size_t)my_n * intermediate_degree * sizeof(uint8_t), - cudaMemcpyDeviceToHost, - stream)); - - CUDA_VERIFY(cudaStreamSynchronize(stream)); - CUDA_VERIFY(cudaFree(d_knn_graph)); - CUDA_VERIFY(cudaFree(d_detour)); - CUDA_VERIFY(cudaStreamDestroy(stream)); - } - - auto t_opt1 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] multi-GPU detour counting: " - "%.2f seconds\n", - std::chrono::duration(t_opt1 - t_phase).count()); - - // Phase 3b: CPU pruning (select top-graph_degree by lowest detour - // count) - auto* output_ptr = h_cagra.data_handle(); -#pragma omp parallel for - for (idx_t i = 0; i < n; i++) { - idx_t pk = 0; - uint32_t num_detour = 0; - for (uint32_t l = 0; - l < (uint32_t)intermediate_degree && pk < graph_degree; - l++) { - uint32_t next_num_detour = std::numeric_limits::max(); - for (idx_t k = 0; k < intermediate_degree; k++) { - uint32_t dc = - h_detour.data_handle()[i * intermediate_degree + k]; - if (dc > num_detour) - next_num_detour = std::min(dc, next_num_detour); - if (dc != num_detour) - continue; - - uint32_t candidate = - h_knn.data_handle()[i * intermediate_degree + k]; - bool dup = false; - for (idx_t dk = 0; dk < pk; dk++) { - if (candidate == output_ptr[i * graph_degree + dk]) { - dup = true; - break; - } - } - if (!dup && candidate < (uint32_t)n) { - output_ptr[i * graph_degree + pk] = candidate; - pk++; - } - if (pk >= graph_degree) - break; - } - if (pk >= graph_degree) - break; - if (next_num_detour == std::numeric_limits::max()) - break; - num_detour = next_num_detour; - } - } - - auto t_opt2 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] CPU pruning: %.2f seconds\n", - std::chrono::duration(t_opt2 - t_opt1).count()); - - // Phase 3c: Build reverse graph and merge into output - // (matches graph_core.cuh Phase 4+5 logic) - std::vector rev_graph(n * graph_degree, UINT32_MAX); - std::vector rev_count(n, 0); - - // Build reverse adjacency: for each edge A→B, record B←A - for (idx_t i = 0; i < n; i++) { - for (idx_t k = 0; k < graph_degree; k++) { - uint32_t dest = output_ptr[i * graph_degree + k]; - if (dest >= (uint32_t)n) - continue; - uint32_t slot = rev_count[dest]; - if (slot < (uint32_t)graph_degree) { - rev_graph[dest * graph_degree + slot] = - static_cast(i); - rev_count[dest]++; - } - } - } - - // Merge reverse edges into output graph (replace tail edges) -#pragma omp parallel for - for (idx_t i = 0; i < n; i++) { - uint32_t num_protected = std::max(graph_degree / 2, 1); - uint32_t kr = - std::min(rev_count[i], static_cast(graph_degree)); - while (kr > 0) { - kr--; - uint32_t rev_node = rev_graph[i * graph_degree + kr]; - if (rev_node >= (uint32_t)n) - continue; - - // Check if already in neighbor list - bool found = false; - for (idx_t j = 0; j < graph_degree; j++) { - if (output_ptr[i * graph_degree + j] == rev_node) { - found = true; - break; - } - } - if (found) - continue; - - // Shift tail edges and insert at protected boundary - for (idx_t j = graph_degree - 1; j > (idx_t)num_protected; - j--) { - output_ptr[i * graph_degree + j] = - output_ptr[i * graph_degree + j - 1]; - } - output_ptr[i * graph_degree + num_protected] = rev_node; - } - } + // Phase 3: prune the knn graph into a CAGRA graph. + const size_t out_count = (size_t)n * graph_degree; + std::vector h_cagra(out_count); - auto t_opt3 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] reverse graph: %.2f seconds\n", - std::chrono::duration(t_opt3 - t_opt2).count()); - } else { - // Single-GPU optimize (original path) - cudaSetDevice(devices[0]); - cudaDeviceSynchronize(); - raft::device_resources single_gpu_res; - cuvs::neighbors::cagra::helpers::optimize( - single_gpu_res, h_knn.view(), h_cagra.view()); - } +#if !defined(FAISS_CAGRA_DEVICE_OPTIMIZE) + // Public API: host matrices only, so the graph makes a round trip it does + // not need. Some cuVS versions also read the input as host memory whatever + // accessor you hand them, so passing device pointers here is not an option. + auto h_knn = + raft::make_host_matrix(n, intermediate_degree); + raft::copy( + h_knn.data_handle(), + d_knn.data_handle(), + knn_count, + raft::resource::get_cuda_stream(single_gpu_res)); + raft::resource::sync_stream(single_gpu_res); + + auto h_cagra_view = raft::make_host_matrix_view( + h_cagra.data(), n, graph_degree); + cuvs::neighbors::cagra::helpers::optimize( + single_gpu_res, h_knn.view(), h_cagra_view); + const char* optimize_where = "host (public API)"; +#else + auto d_cagra = raft::make_device_matrix( + single_gpu_res, n, graph_degree); + + cuvs::neighbors::cagra::detail::graph::optimize( + single_gpu_res, + d_knn.view(), + d_cagra.view(), + cagraConfig_.guarantee_connectivity); + raft::resource::sync_stream(single_gpu_res); + raft::copy( + h_cagra.data(), + d_cagra.data_handle(), + out_count, + raft::resource::get_cuda_stream(single_gpu_res)); + raft::resource::sync_stream(single_gpu_res); + const char* optimize_where = "device-resident"; +#endif t_now = std::chrono::high_resolution_clock::now(); fprintf(stderr, - " [trainAllNeighbors] optimize: %.2f seconds\n", + " [trainAllNeighbors] optimize (%s): %.2f seconds\n", + optimize_where, std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // Phase 4: Store as merged_knngraph_ for copyTo() - merged_knngraph_.resize(n * graph_degree); + merged_knngraph_.resize(out_count); merged_knngraph_degree_ = graph_degree; #pragma omp parallel for - for (idx_t i = 0; i < n * graph_degree; i++) { - merged_knngraph_[i] = static_cast(h_cagra.data_handle()[i]); + for (idx_t i = 0; i < (idx_t)out_count; i++) { + merged_knngraph_[i] = static_cast(h_cagra[i]); } multi_gpu_dataset_ = x; @@ -1112,6 +555,9 @@ void GpuIndexCagra::trainAllNeighbors( this->ntotal = n; t_now = std::chrono::high_resolution_clock::now(); + fprintf(stderr, + " [trainAllNeighbors] D2H + widen: %.2f seconds\n", + std::chrono::duration(t_now - t_phase).count()); fprintf(stderr, " [trainAllNeighbors] Total: %.2f seconds\n", std::chrono::duration(t_now - t_start).count()); @@ -1411,6 +857,121 @@ void GpuIndexCagra::copyTo(faiss::IndexHNSWCagra* index) const { index->init_level0 = true; } +void GpuIndexCagra::buildHnswUpperLevelsGpu_( + faiss::IndexHNSWCagra* index, + int max_lvl) const { + auto& hnsw = index->hnsw; + const idx_t n_train = this->ntotal; + // Levels >= 1 all carry M neighbours; level 0 carries 2*M. + const int upper_degree = (int)hnsw.nb_neighbors(1); + + // BRUTE_FORCE only exists on the multi-GPU path. + auto sub_algo = cagraConfig_.build_algo == graph_build_algo::IVF_PQ + ? faiss::cagra_build_algo::IVF_PQ + : faiss::cagra_build_algo::NN_DESCENT; + + const int sub_intermediate = cagraConfig_.gpu_hnsw_intermediate_degree > 0 + ? (int)cagraConfig_.gpu_hnsw_intermediate_degree + : 2 * upper_degree; + + // Search starts at entry_point and immediately reads its neighbours at + // max_level, so the two must agree. Track the highest populated level + // rather than assuming max_lvl has any nodes. + int top_lvl = 0; + HNSW::storage_idx_t top_entry = 0; + + std::vector ids; + std::vector sub_x; + for (int lvl = 1; lvl <= max_lvl; lvl++) { + ids.clear(); + for (idx_t i = 0; i < n_train; i++) { + if (hnsw.levels[i] - 1 >= lvl) { + ids.push_back(i); + } + } + const idx_t n_lvl = (idx_t)ids.size(); + if (n_lvl < 1) { + continue; + } + top_lvl = lvl; + top_entry = static_cast(ids[0]); + if (n_lvl == 1) { + // Single node: no edges to add, but it is still a valid entry. + continue; + } + + if (n_lvl <= upper_degree + 1) { + // Too small for a meaningful proximity graph: connect all-to-all. + for (idx_t a = 0; a < n_lvl; a++) { + size_t begin, end; + hnsw.neighbor_range(ids[a], lvl, &begin, &end); + size_t slot = begin; + for (idx_t b = 0; b < n_lvl && slot < end; b++) { + if (b != a) { + hnsw.neighbors[slot++] = ids[b]; + } + } + } + continue; + } + + sub_x.resize((size_t)n_lvl * this->d); +#pragma omp parallel for + for (idx_t a = 0; a < n_lvl; a++) { + memcpy(sub_x.data() + (size_t)a * this->d, + multi_gpu_dataset_ + (size_t)ids[a] * this->d, + this->d * sizeof(float)); + } + + auto t0 = std::chrono::high_resolution_clock::now(); + std::vector sub_graph; + { + CuvsCagra sub( + this->resources_.get(), + this->d, + sub_intermediate, + upper_degree, + sub_algo, + cagraConfig_.nn_descent_niter, + /*store_dataset=*/false, + this->metric_type, + this->metric_arg, + INDICES_64_BIT, + std::nullopt, + std::nullopt, + cagraConfig_.refine_rate, + cagraConfig_.gpu_hnsw_guarantee_connectivity); + sub.train(n_lvl, sub_x.data()); + sub_graph = sub.get_knngraph(); + } + + // Map subgraph-local ids back to global ids. +#pragma omp parallel for + for (idx_t a = 0; a < n_lvl; a++) { + size_t begin, end; + hnsw.neighbor_range(ids[a], lvl, &begin, &end); + for (size_t j = begin; j < end; j++) { + idx_t local = sub_graph[(size_t)a * upper_degree + (j - begin)]; + hnsw.neighbors[j] = + (local >= 0 && local < n_lvl) ? ids[local] : -1; + } + } + + fprintf(stderr, + " [copyTo] level %d: %ld nodes, igd=%d gc=%d, " + "CAGRA subgraph %.2f s\n", + lvl, + (long)n_lvl, + sub_intermediate, + (int)cagraConfig_.gpu_hnsw_guarantee_connectivity, + std::chrono::duration( + std::chrono::high_resolution_clock::now() - t0) + .count()); + } + hnsw.entry_point = top_entry; + hnsw.max_level = top_lvl; +} + void GpuIndexCagra::copyToMultiGpu_(faiss::IndexHNSWCagra* index) const { FAISS_ASSERT(this->is_trained && index && !merged_knngraph_.empty()); FAISS_THROW_IF_NOT_MSG( @@ -1446,12 +1007,23 @@ void GpuIndexCagra::copyToMultiGpu_(faiss::IndexHNSWCagra* index) const { auto n_train = this->ntotal; index->init_level0 = false; - if (!index->base_level_only) { - index->add(n_train, multi_gpu_dataset_); - } else { + if (index->base_level_only) { index->hnsw.prepare_level_tab(n_train, false); index->storage->add(n_train, multi_gpu_dataset_); index->ntotal = n_train; + } else if (cagraConfig_.gpu_hnsw_upper_levels) { + auto t0 = std::chrono::high_resolution_clock::now(); + int max_lvl = index->hnsw.prepare_level_tab(n_train, false); + index->storage->add(n_train, multi_gpu_dataset_); + index->ntotal = n_train; + buildHnswUpperLevelsGpu_(index, max_lvl); + fprintf(stderr, + " [copyTo] GPU upper levels: %.2f seconds\n", + std::chrono::duration( + std::chrono::high_resolution_clock::now() - t0) + .count()); + } else { + index->add(n_train, multi_gpu_dataset_); } #pragma omp parallel for diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index 6f0ae93be2..0ec56b4e14 100644 --- a/faiss/gpu/GpuIndexCagra.h +++ b/faiss/gpu/GpuIndexCagra.h @@ -47,7 +47,9 @@ enum class graph_build_algo { /// Use NN-Descent to build all-neighbors knn graph NN_DESCENT, /// Use iterative search to build knn graph - ITERATIVE_SEARCH + ITERATIVE_SEARCH, + /// Exact knn graph via tiled brute force + BRUTE_FORCE }; /// A type for specifying how PQ codebooks are created. @@ -173,6 +175,34 @@ struct IVFPQSearchCagraConfig { uint32_t max_internal_batch_size = 4096; }; +/// Knobs for the multi-GPU build path, selected by listing more than one +/// device in GpuIndexCagraConfig::devices. Degrees, build_algo and metric +/// still come from GpuIndexCagraConfig. +struct AllNeighborsCagraConfig { + /// Number of overlapping clusters the dataset is partitioned into. + /// 0 selects max(2 * number of devices, 4). + size_t n_clusters = 0; + + /// Clusters each vector is assigned to. Must be >= 2: with 1 there are no + /// cross-cluster edges and recall collapses. 0 selects 2. + size_t overlap_factor = 2; + + /// Bounds IVF-PQ search memory during the knn build. cuVS's default can + /// exhaust device memory around 100M vectors; 8192 avoids that and does + /// not change recall. 0 keeps cuVS's default. + uint32_t ivf_pq_search_batch_size = 0; + + /// IVF-PQ refinement multiplier. The refine pass runs on the CPU, so cost + /// scales with this. Raising it above 1.0 measured both slower and less + /// accurate, so it is separate from the shared refine_rate. + float refinement_rate = 1.0f; + + /// Size the IVF-PQ index for one cluster rather than the whole dataset. + /// cuVS reuses these params for every cluster without rescaling, so + /// sizing from the full dataset badly over-partitions each one. + bool ivf_pq_size_from_cluster = true; +}; + struct GpuIndexCagraConfig : public GpuIndexConfig { /// Degree of input graph for pruning. size_t intermediate_graph_degree = 128; @@ -190,6 +220,27 @@ struct GpuIndexCagraConfig : public GpuIndexConfig { /// Whether to use MST optimization to guarantee graph connectivity. bool guarantee_connectivity = false; + + /// Devices to build on. More than one selects the multi-GPU build in + /// train(); see AllNeighborsCagraConfig and train() for its restrictions. + /// Empty or one device uses GpuIndexConfig::device as usual. + std::vector devices; + + AllNeighborsCagraConfig all_neighbors_params; + + /// Build the HNSW upper levels on the GPU during copyTo() instead of by + /// CPU insertion. Roughly 10x faster, but measured worse recall than + /// base_level_only, which is cheaper still. Off by default. + bool gpu_hnsw_upper_levels = false; + + /// intermediate_graph_degree for the per-level subgraph builds. + /// 0 = twice the upper-level degree. + size_t gpu_hnsw_intermediate_degree = 0; + + /// MST connectivity pass on the per-level subgraphs. On by default: + /// upper levels are walked greedily with no backtracking, so a + /// disconnected component is a trap the descent cannot escape. + bool gpu_hnsw_guarantee_connectivity = true; }; enum class search_algo { @@ -284,46 +335,6 @@ struct GpuIndexCagra : public GpuIndex { /// in the index instance void copyTo(faiss::IndexHNSWCagra* index) const; - /// Train CAGRA using multiple GPUs by sharding the dataset. - /// Uses cuVS native SNMG (single-node multi-GPU) CAGRA build. - /// Each device builds one shard in parallel via OpenMP. - /// Float32 only. After training, call copyTo() to produce a CPU - /// IndexHNSWCagra with full HNSW upper levels. - /// The training data pointer must remain valid until copyTo() completes. - /// stitch_mode: 0=CPU HNSW (Approach C), 1=GPU brute-force (Approach B) - void trainMultiGpu( - idx_t n, - const float* x, - std::vector& providers, - std::vector& devices, - idx_t stitch_per_shard = 0, - int stitch_k = 2, - int stitch_mode = 0); - - /// Build a unified CAGRA graph using cuVS all_neighbors - /// (multi-GPU kNN graph construction with overlapping clusters) followed - /// by cagra::optimize (graph pruning). Produces a single unified graph - /// without stitching. The training data pointer must remain valid until - /// copyTo() completes. - /// build_algo: 0=NN-descent (default), 1=brute-force, 2=IVF-PQ - /// refinement_rate: IVF-PQ refinement multiplier (only used when - /// build_algo==2). Higher values trade build time for recall; the cuVS - /// default is 2.0. - /// ivfpq_search_batch: cap the IVF-PQ search `max_internal_batch_size` used - /// during the all_neighbors kNN build (build_algo==2). 0 = cuVS default - /// (128*1024), which can OOM at 100M; a smaller value (e.g. 8192) bounds - /// the GPU search workspace with no effect on results (recall-neutral). - void trainAllNeighbors( - idx_t n, - const float* x, - std::vector& devices, - int n_clusters = 0, - int overlap_factor = 0, - bool multi_gpu_optimize = false, - int build_algo = 0, - float refinement_rate = 2.0f, - int ivfpq_search_batch = 0); - void reset() override; std::vector get_knngraph() const; @@ -357,6 +368,18 @@ struct GpuIndexCagra : public GpuIndex { idx_t* labels, const SearchParameters* search_params) const override; + /// Multi-GPU build path taken by train() when cagraConfig_.devices lists + /// more than one device: cuVS `all_neighbors` knn graph construction over + /// overlapping clusters, followed by graph pruning. Leaves the result in + /// merged_knngraph_ for copyTo(); index_ stays empty. + void trainAllNeighbors_(idx_t n, const float* x); + + /// Populate HNSW levels >= 1 of `index` by building a CAGRA graph over + /// each level's node subset on the GPU. Requires the level table to be + /// prepared and the storage populated. Returns the max level. + void buildHnswUpperLevelsGpu_(faiss::IndexHNSWCagra* index, int max_lvl) + const; + void copyToMultiGpu_(faiss::IndexHNSWCagra* index) const; /// Our configuration options @@ -372,7 +395,7 @@ struct GpuIndexCagra : public GpuIndex { std::shared_ptr>> index_; - /// Multi-GPU state: populated by trainMultiGpu(), used by copyTo() + /// Multi-GPU state: populated by trainAllNeighbors_(), used by copyTo() std::vector merged_knngraph_; idx_t merged_knngraph_degree_ = 0; const float* multi_gpu_dataset_ = nullptr; diff --git a/faiss/gpu/test/bench_approaches.py b/faiss/gpu/test/bench_approaches.py index b62ebd7252..5070f9f698 100644 --- a/faiss/gpu/test/bench_approaches.py +++ b/faiss/gpu/test/bench_approaches.py @@ -5,15 +5,16 @@ # LICENSE file in the root directory of this source tree. """ -Benchmark 4 approaches for multi-GPU CAGRA -> HNSW: - A. IndexShards (independent HNSW per shard, no stitching) - B. Unified + GPU brute-force sampled stitching (exact cross-shard NNs) - C. Unified + CPU HNSW stitching (approximate cross-shard NNs) - D. all_neighbors + optimize (multi-GPU overlapping clusters, no stitching) +Benchmark the multi-GPU CAGRA -> HNSW build: cuVS all_neighbors builds a knn +graph over overlapping clusters across all GPUs, cagra::optimize prunes it into +a single unified graph, and copyTo produces a CPU IndexHNSWCagra. + +Measures build/copyTo/serialize wall-clock, index size, and recall@10 vs +brute-force ground truth at several efSearch values. Usage: buck run @//mode/opt fbcode//faiss/gpu/test:bench_approaches -- \\ - --data /path/to/vectors.npy --approaches D --multi-gpu-optimize + --data /path/to/vectors.npy """ import argparse @@ -42,6 +43,9 @@ _t0 = time.time() +EF_VALUES = [16, 32, 64, 128, 256, 512] + + def compute_recall(I_test, I_gt, k): nq = I_test.shape[0] return np.mean([len(set(I_test[i]) & set(I_gt[i])) / k for i in range(nq)]) @@ -98,164 +102,18 @@ def load_from_hive( return xb[:i] -def approach_a_indexshards(xb, d, num_gpus, graph_degree=32, save_dir=None): - """Build N independent CAGRA->HNSW, wrap in IndexShards.""" - n = xb.shape[0] - shard_size = (n + num_gpus - 1) // num_gpus - timings = {} - - t0 = time.time() - shards = [] - for g in range(num_gpus): - start = g * shard_size - end = min(start + shard_size, n) - shard_data = xb[start:end] - if len(shard_data) == 0: - continue - - res = faiss.StandardGpuResources() - config = faiss.GpuIndexCagraConfig() - config.graph_degree = graph_degree - config.intermediate_graph_degree = graph_degree * 2 - config.build_algo = faiss.graph_build_algo_NN_DESCENT - idx = faiss.GpuIndexCagra(res, d, faiss.METRIC_L2, config) - idx.train(shard_data) - - cpu_idx = faiss.IndexHNSWCagra() - idx.copyTo(cpu_idx) - shards.append(cpu_idx) - timings["build"] = time.time() - t0 - - t0 = time.time() - index = faiss.IndexShards(d, True, True) - for s in shards: - index.add_shard(s) - timings["assemble"] = time.time() - t0 - - out_dir = save_dir or tempfile.mkdtemp() - t0 = time.time() - total_bytes = 0 - for i, s in enumerate(shards): - path = os.path.join(out_dir, f"A_shard_{i}_{n // 1_000_000}M.faiss") - faiss.write_index(s, path) - total_bytes += os.path.getsize(path) - timings["serialize"] = time.time() - t0 - timings["file_size_gb"] = total_bytes / 1e9 - - return index, shards, timings - - -def approach_b_unified_gpu_stitch( - xb, - d, - num_gpus, - graph_degree=32, - stitch_per_shard=100000, - stitch_k=16, - save_dir=None, -): - """Unified SNMG build + GPU brute-force sampled stitching (exact NNs).""" - n = xb.shape[0] - timings = {} - - resources = faiss.GpuResourcesVector() - devices = faiss.Int32Vector() - res_list = [] - for i in range(num_gpus): - res = faiss.StandardGpuResources() - res_list.append(res) - resources.push_back(res) - devices.push_back(i) - - config = faiss.GpuIndexCagraConfig() - config.graph_degree = graph_degree - config.intermediate_graph_degree = graph_degree * 2 - config.build_algo = faiss.graph_build_algo_NN_DESCENT - index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config) - - t0 = time.time() - index.trainMultiGpu( - n, - faiss.swig_ptr(xb), - resources, - devices, - stitch_per_shard, - stitch_k, - 1, - ) - timings["build"] = time.time() - t0 - - t0 = time.time() - cpu_index = faiss.IndexHNSWCagra() - index.copyTo(cpu_index) - timings["copyTo"] = time.time() - t0 - - out_dir = save_dir or tempfile.mkdtemp() - path = os.path.join(out_dir, f"B_unified_{n // 1_000_000}M.faiss") - t0 = time.time() - faiss.write_index(cpu_index, path) - timings["file_size_gb"] = os.path.getsize(path) / 1e9 - timings["serialize"] = time.time() - t0 - - return cpu_index, timings +BUILD_ALGOS = ("ivf_pq", "nn_descent", "brute_force") -def approach_c_unified_cpu_stitch( - xb, - d, - num_gpus, - graph_degree=32, - stitch_per_shard=0, - stitch_k=16, - save_dir=None, -): - """Unified SNMG build + CPU HNSW stitching (approximate NNs).""" - n = xb.shape[0] - timings = {} +def cagra_build_algo(name): + return getattr(faiss, f"graph_build_algo_{name.upper()}") - resources = faiss.GpuResourcesVector() - devices = faiss.Int32Vector() - res_list = [] - for i in range(num_gpus): - res = faiss.StandardGpuResources() - res_list.append(res) - resources.push_back(res) - devices.push_back(i) - config = faiss.GpuIndexCagraConfig() - config.graph_degree = graph_degree - config.intermediate_graph_degree = graph_degree * 2 - config.build_algo = faiss.graph_build_algo_NN_DESCENT - index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config) +def index_filename(n): + return f"cagra_hnsw_{n // 1_000_000}M.faiss" - t0 = time.time() - index.trainMultiGpu( - n, - faiss.swig_ptr(xb), - resources, - devices, - stitch_per_shard, - stitch_k, - 0, - ) - timings["build"] = time.time() - t0 - t0 = time.time() - cpu_index = faiss.IndexHNSWCagra() - index.copyTo(cpu_index) - timings["copyTo"] = time.time() - t0 - - out_dir = save_dir or tempfile.mkdtemp() - path = os.path.join(out_dir, f"C_unified_{n // 1_000_000}M.faiss") - t0 = time.time() - faiss.write_index(cpu_index, path) - timings["file_size_gb"] = os.path.getsize(path) / 1e9 - timings["serialize"] = time.time() - t0 - - return cpu_index, timings - - -def approach_d_all_neighbors( +def build_cagra_hnsw( xb, d, num_gpus, @@ -263,52 +121,68 @@ def approach_d_all_neighbors( save_dir=None, n_clusters=0, overlap_factor=0, - multi_gpu_optimize=False, - build_algo=0, + build_algo="ivf_pq", base_level_only=False, intermediate_graph_degree=48, - refinement_rate=2.0, + refinement_rate=1.0, ivfpq_search_batch=0, + guarantee_connectivity=False, + ivfpq_size_from_cluster=True, + gpu_hnsw_upper_levels=False, + gpu_hnsw_igd=0, + gpu_hnsw_guarantee_connectivity=True, + ef_construction=0, + entrypoints=0, ): - """all_neighbors + cagra::optimize → unified CAGRA graph. No stitching.""" + """Multi-GPU all_neighbors build -> optimize -> CPU IndexHNSWCagra.""" + # The multi-GPU build does not copy the dataset, so xb must stay alive and + # contiguous until copyTo() below has run. + xb = np.ascontiguousarray(xb, dtype=np.float32) n = xb.shape[0] timings = {} + res_list = [faiss.StandardGpuResources() for _ in range(num_gpus)] + devices = faiss.Int32Vector() - res_list = [] for i in range(num_gpus): - res = faiss.StandardGpuResources() - res_list.append(res) devices.push_back(i) + all_neighbors = faiss.AllNeighborsCagraConfig() + all_neighbors.n_clusters = n_clusters + all_neighbors.overlap_factor = overlap_factor + all_neighbors.ivf_pq_search_batch_size = ivfpq_search_batch + all_neighbors.refinement_rate = refinement_rate + all_neighbors.ivf_pq_size_from_cluster = ivfpq_size_from_cluster + config = faiss.GpuIndexCagraConfig() config.graph_degree = graph_degree config.intermediate_graph_degree = intermediate_graph_degree - config.build_algo = faiss.graph_build_algo_NN_DESCENT + config.build_algo = cagra_build_algo(build_algo) + config.guarantee_connectivity = guarantee_connectivity + config.gpu_hnsw_upper_levels = gpu_hnsw_upper_levels + config.gpu_hnsw_intermediate_degree = gpu_hnsw_igd + config.gpu_hnsw_guarantee_connectivity = gpu_hnsw_guarantee_connectivity + config.devices = devices + config.all_neighbors_params = all_neighbors + index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config) t0 = time.time() - index.trainAllNeighbors( - n, - faiss.swig_ptr(xb), - devices, - n_clusters, - overlap_factor, - multi_gpu_optimize, - build_algo, - refinement_rate, - ivfpq_search_batch, - ) + index.train(xb) timings["build"] = time.time() - t0 t0 = time.time() cpu_index = faiss.IndexHNSWCagra() cpu_index.base_level_only = base_level_only + if ef_construction > 0: + cpu_index.hnsw.efConstruction = ef_construction + if entrypoints > 0: + cpu_index.num_base_level_search_entrypoints = entrypoints index.copyTo(cpu_index) timings["copyTo"] = time.time() - t0 out_dir = save_dir or tempfile.mkdtemp() - path = os.path.join(out_dir, f"D_allneighbors_{n // 1_000_000}M.faiss") + path = os.path.join(out_dir, index_filename(n)) t0 = time.time() faiss.write_index(cpu_index, path) timings["file_size_gb"] = os.path.getsize(path) / 1e9 @@ -318,18 +192,12 @@ def approach_d_all_neighbors( def _set_ef(index, ef): - if isinstance(index, faiss.IndexShards): - for i in range(index.count()): - shard = faiss.downcast_index(index.at(i)) - if hasattr(shard, "hnsw"): - shard.hnsw.efSearch = ef - elif hasattr(index, "hnsw"): - index.hnsw.efSearch = ef + index.hnsw.efSearch = ef def eval_recall(index, xq, Igt, k=10, ef_values=None): if ef_values is None: - ef_values = [64, 128, 256] + ef_values = EF_VALUES results = {} for ef in ef_values: @@ -340,10 +208,30 @@ def eval_recall(index, xq, Igt, k=10, ef_values=None): return results +def eval_qps(index, xq, k=10, ef_values=None, repeat=5): + """Full-batch throughput at each efSearch, best of `repeat` runs.""" + if ef_values is None: + ef_values = EF_VALUES + nq = xq.shape[0] + results = {} + + for ef in ef_values: + _set_ef(index, ef) + index.search(xq, k) # warmup + best = float("inf") + for _ in range(repeat): + t0 = time.time() + index.search(xq, k) + best = min(best, time.time() - t0) + results[ef] = nq / best + + return results + + def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5): """Measure kcycles/query using vench CycleCounter.""" if ef_values is None: - ef_values = [64, 128, 256] + ef_values = EF_VALUES try: from vector_search.vench.perf_cycles import CycleCounter @@ -361,22 +249,15 @@ def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5): prev_threads = faiss.omp_get_max_threads() faiss.omp_set_num_threads(1) - measure_index = index - if isinstance(index, faiss.IndexShards) and index.count() > 0: - measure_index = faiss.IndexShards(index.d, False, True) - for i in range(index.count()): - measure_index.add_shard(index.at(i)) - for ef in ef_values: _set_ef(index, ef) - _set_ef(measure_index, ef) for _ in range(warmup): - measure_index.search(xq, k) + index.search(xq, k) best_kcycles = float("inf") for _ in range(repeat): c0 = cc.read() - measure_index.search(xq, k) + index.search(xq, k) c1 = cc.read() kc = (c1 - c0) / 1000.0 / nq best_kcycles = min(best_kcycles, kc) @@ -388,7 +269,7 @@ def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5): def main(): parser = argparse.ArgumentParser( - description="Benchmark multi-GPU CAGRA->HNSW approaches" + description="Benchmark the multi-GPU CAGRA->HNSW build" ) parser.add_argument( "--data", @@ -405,8 +286,6 @@ def main(): ) parser.add_argument("--num-gpus", type=int, default=0) parser.add_argument("--graph-degree", type=int, default=32) - parser.add_argument("--stitch-k", type=int, default=16) - parser.add_argument("--stitch-per-shard", type=int, default=100000) parser.add_argument( "--n-clusters", type=int, @@ -420,15 +299,66 @@ def main(): help="all_neighbors overlap_factor (0=default 2)", ) parser.add_argument( - "--multi-gpu-optimize", + "--build-algo", + choices=BUILD_ALGOS, + default="ivf_pq", + help="graph_build_algo used for the kNN graph (brute_force is " + "multi-GPU only and O(N^2 D) per cluster)", + ) + parser.add_argument( + "--ivfpq-size-from-cluster", + action=argparse.BooleanOptionalAction, + default=True, + help="Derive IVF-PQ params from the per-cluster subproblem. " + "--no-ivfpq-size-from-cluster sizes them from the full dataset, " + "which over-partitions every cluster", + ) + parser.add_argument( + "--guarantee-connectivity", action="store_true", - help="Partition detour counting across GPUs", + help="Run the MST pass in cagra::optimize so the pruned graph is " + "guaranteed connected (cuVS default is on; costs build time)", ) parser.add_argument( - "--build-algo", + "--entrypoints-sweep", + type=str, + default="", + help="Comma-separated num_base_level_search_entrypoints values to " + "sweep at search time (base-level-only only), e.g. 32,256,1024", + ) + parser.add_argument( + "--gpu-hnsw-upper-levels", + action="store_true", + help="Build HNSW levels >=1 as GPU CAGRA subgraphs in copyTo instead " + "of incremental CPU insertion (base-level-only off only)", + ) + parser.add_argument( + "--gpu-hnsw-igd", type=int, default=0, - help="0=nn_descent, 1=brute_force, 2=ivf_pq", + help="intermediate_graph_degree for the per-level GPU subgraphs " + "(0 = 2x the upper-level degree)", + ) + parser.add_argument( + "--gpu-hnsw-guarantee-connectivity", + action=argparse.BooleanOptionalAction, + default=True, + help="MST connectivity pass on the per-level GPU subgraphs " + "(on by default; greedy descent needs connectivity)", + ) + parser.add_argument( + "--ef-construction", + type=int, + default=0, + help="Override hnsw.efConstruction used when building upper levels " + "(0 = faiss default of 40)", + ) + parser.add_argument( + "--entrypoints", + type=int, + default=0, + help="Override num_base_level_search_entrypoints (0 = faiss default " + "of 256); only affects base-level-only search", ) parser.add_argument( "--base-level-only", @@ -444,9 +374,10 @@ def main(): parser.add_argument( "--refinement-rate", type=float, - default=2.0, - help="IVF-PQ refinement multiplier (build-algo=2 only); " - "higher trades build time for recall (cuVS default 2.0)", + default=1.0, + help="IVF-PQ refinement multiplier (build-algo=2 only). Sets " + "candidate_k = k * rate; the refine pass runs on the host so cost is " + "linear in this. Measured best at 1.0 (default)", ) parser.add_argument( "--ivfpq-search-batch", @@ -456,15 +387,6 @@ def main(): "build (build-algo=2). 0=cuVS default (131072); smaller (e.g. 8192) " "bounds GPU search workspace to avoid OOM at 100M (recall-neutral)", ) - parser.add_argument( - "--approaches", - type=str, - default="A,B,C,D", - help=( - "A (IndexShards), B (GPU stitch), " - "C (CPU stitch), D (all_neighbors)" - ), - ) parser.add_argument( "--index-dir", type=str, @@ -585,274 +507,101 @@ def main(): f"graph_degree={args.graph_degree}" ) print( - f" stitch_k={args.stitch_k}, " - f"stitch_per_shard={args.stitch_per_shard}" - ) - print( - f" [D] n_clusters={args.n_clusters or 'auto'}, " + f" n_clusters={args.n_clusters or 'auto'}, " f"overlap_factor={args.overlap_factor or 'auto'}, " f"intermediate_graph_degree={args.intermediate_graph_degree}, " f"build_algo={args.build_algo}, refinement_rate={args.refinement_rate}" ) print() - approaches_to_run = [a.strip().upper() for a in args.approaches.split(",")] - results_table = [] save_dir = args.index_dir os.makedirs(save_dir, exist_ok=True) if args.kcycles_only: - print(f"Loading persisted indices from {save_dir} (kcycles-only mode)") - n_tag = f"{n // 1_000_000}M" - for approach in approaches_to_run: - if approach == "A": - pre = "A_shard_" - suf = f"_{n_tag}.faiss" - shard_files = sorted( - f - for f in os.listdir(save_dir) - if f.startswith(pre) and f.endswith(suf) - ) - if not shard_files: - print(f" No A files for {n_tag}") - continue - shards = [ - faiss.read_index(os.path.join(save_dir, f)) - for f in shard_files - ] - idx = faiss.IndexShards(d, True, True) - for s in shards: - idx.add_shard(s) - label = f"A: IndexShards ({len(shard_files)})" - else: - prefixes = { - "B": "B_unified", - "C": "C_unified", - "D": "D_allneighbors", - } - prefix = prefixes.get(approach) - if not prefix: - prefix = f"{approach}_unified" - path = os.path.join(save_dir, f"{prefix}_{n_tag}.faiss") - if not os.path.exists(path): - print(f" No {prefix} for {n_tag}") - continue - idx = faiss.read_index(path) - label = f"{approach}: unified" - - recall = eval_recall(idx, xq, Igt, k) - kcycles = eval_kcycles(idx, xq, k) - for ef in sorted(recall.keys()): - kc = kcycles.get(ef, 0) - print( - f" {label:25s} ef={ef:>3d} recall={recall[ef]:.4f} " - f"kcyc/q={kc:.0f}" - ) - del idx - return - - print(f"Indices will be saved to: {save_dir}") - - if "A" in approaches_to_run: - print("=" * 60) - print("APPROACH A: IndexShards (independent HNSW per shard)") - print("=" * 60) - index_a, shards_a, timings_a = approach_a_indexshards( - xb, - d, - num_gpus, - args.graph_degree, - save_dir=save_dir, - ) - recall_a = eval_recall(index_a, xq, Igt, k) - kcycles_a = eval_kcycles(index_a, xq, k) - total_build_a = timings_a["build"] + timings_a.get("assemble", 0) - print(f" Build: {timings_a['build']:.1f}s") - print( - f" Serialize: {timings_a['serialize']:.1f}s " - f"({timings_a['file_size_gb']:.2f} GB, {num_gpus} files)" - ) - for ef in sorted(recall_a.keys()): - kc = kcycles_a.get(ef, 0) + print(f"Loading persisted index from {save_dir} (kcycles-only mode)") + path = os.path.join(save_dir, index_filename(n)) + if not os.path.exists(path): + print(f"ERROR: no index at {path}", file=sys.stderr) + sys.exit(1) + idx = faiss.read_index(path) + recall = eval_recall(idx, xq, Igt, k) + qps = eval_qps(idx, xq, k) + kcycles = eval_kcycles(idx, xq, k) + for ef in sorted(recall.keys()): print( - f" efSearch={ef}: recall@{k}={recall_a[ef]:.4f} " - f"kcycles/q={kc:.0f}" - ) - results_table.append( - ( - "A: IndexShards", - total_build_a, - timings_a["serialize"], - timings_a["file_size_gb"], - recall_a, - kcycles_a, + f" ef={ef:>4d} recall@{k}={recall[ef]:.4f} " + f"qps={qps[ef]:,.1f} kcyc/q={kcycles.get(ef, 0):.0f}" ) - ) - del index_a, shards_a - print() + return - if "B" in approaches_to_run: - print("=" * 60) - print( - f"APPROACH B: Unified + GPU brute-force stitch " - f"(sps={args.stitch_per_shard}, k={args.stitch_k})" - ) - print("=" * 60) - index_b, timings_b = approach_b_unified_gpu_stitch( - xb, - d, - num_gpus, - args.graph_degree, - args.stitch_per_shard, - args.stitch_k, - save_dir=save_dir, - ) - recall_b = eval_recall(index_b, xq, Igt, k) - kcycles_b = eval_kcycles(index_b, xq, k) - total_build_b = timings_b["build"] + timings_b["copyTo"] - print(f" Build (SNMG+GPU-stitch): {timings_b['build']:.1f}s") - print(f" copyTo: {timings_b['copyTo']:.1f}s") - print( - f" Serialize: {timings_b['serialize']:.1f}s " - f"({timings_b['file_size_gb']:.2f} GB)" - ) - for ef in sorted(recall_b.keys()): - kc = kcycles_b.get(ef, 0) - print( - f" efSearch={ef}: recall@{k}={recall_b[ef]:.4f} " - f"kcycles/q={kc:.0f}" - ) - results_table.append( - ( - "B: GPU-brute", - total_build_b, - timings_b["serialize"], - timings_b["file_size_gb"], - recall_b, - kcycles_b, - ) - ) - del index_b - print() - - if "C" in approaches_to_run: - print("=" * 60) - print(f"APPROACH C: Unified + CPU stitch (sps=0, k={args.stitch_k})") - print("=" * 60) - index_c, timings_c = approach_c_unified_cpu_stitch( - xb, - d, - num_gpus, - args.graph_degree, - 0, - args.stitch_k, - save_dir=save_dir, - ) - recall_c = eval_recall(index_c, xq, Igt, k) - kcycles_c = eval_kcycles(index_c, xq, k) - total_build_c = timings_c["build"] + timings_c["copyTo"] - print(f" Build (SNMG+CPU-stitch): {timings_c['build']:.1f}s") - print(f" copyTo: {timings_c['copyTo']:.1f}s") - print( - f" Serialize: {timings_c['serialize']:.1f}s " - f"({timings_c['file_size_gb']:.2f} GB)" - ) - for ef in sorted(recall_c.keys()): - kc = kcycles_c.get(ef, 0) - print( - f" efSearch={ef}: recall@{k}={recall_c[ef]:.4f} " - f"kcycles/q={kc:.0f}" - ) - results_table.append( - ( - "C: CPU-stitch", - total_build_c, - timings_c["serialize"], - timings_c["file_size_gb"], - recall_c, - kcycles_c, - ) - ) - del index_c - print() - - if "D" in approaches_to_run: - print("=" * 60) - print("APPROACH D: all_neighbors + cagra::optimize (no stitching)") - print("=" * 60) - index_d, timings_d = approach_d_all_neighbors( - xb, - d, - num_gpus, - args.graph_degree, - save_dir=save_dir, - n_clusters=args.n_clusters, - overlap_factor=args.overlap_factor, - multi_gpu_optimize=args.multi_gpu_optimize, - build_algo=args.build_algo, - base_level_only=args.base_level_only, - intermediate_graph_degree=args.intermediate_graph_degree, - refinement_rate=args.refinement_rate, - ivfpq_search_batch=args.ivfpq_search_batch, - ) - recall_d = eval_recall(index_d, xq, Igt, k) - kcycles_d = eval_kcycles(index_d, xq, k) - total_build_d = timings_d["build"] + timings_d["copyTo"] - print(f" Build (allneighbors+optimize): {timings_d['build']:.1f}s") - print(f" copyTo: {timings_d['copyTo']:.1f}s") - ser = timings_d["serialize"] - gb = timings_d["file_size_gb"] - print(f" Serialize: {ser:.1f}s " f"({gb:.2f} GB)") - # Index-build wall-clock, isolated from data load (Koski/Manifold) and - # ground-truth: this is the headline build->serialize number. - index_total_d = ( - timings_d["build"] + timings_d["copyTo"] + timings_d["serialize"] - ) - print( - f" >>> INDEX build->serialize total " - f"(excl. data load + ground truth): " - f"{index_total_d:.1f}s ({index_total_d / 60:.1f} min)" - ) - for ef in sorted(recall_d.keys()): - kc = kcycles_d.get(ef, 0) - print( - f" efSearch={ef}: recall@{k}=" - f"{recall_d[ef]:.4f} kcyc/q={kc:.0f}" - ) - results_table.append( - ( - "D: all_neighbors", - total_build_d, - timings_d["serialize"], - timings_d["file_size_gb"], - recall_d, - kcycles_d, - ) - ) - del index_d - print() + print(f"Index will be saved to: {save_dir}") + + index, timings = build_cagra_hnsw( + xb, + d, + num_gpus, + args.graph_degree, + save_dir=save_dir, + n_clusters=args.n_clusters, + overlap_factor=args.overlap_factor, + build_algo=args.build_algo, + base_level_only=args.base_level_only, + intermediate_graph_degree=args.intermediate_graph_degree, + refinement_rate=args.refinement_rate, + ivfpq_search_batch=args.ivfpq_search_batch, + guarantee_connectivity=args.guarantee_connectivity, + ivfpq_size_from_cluster=args.ivfpq_size_from_cluster, + gpu_hnsw_upper_levels=args.gpu_hnsw_upper_levels, + gpu_hnsw_igd=args.gpu_hnsw_igd, + gpu_hnsw_guarantee_connectivity=args.gpu_hnsw_guarantee_connectivity, + ef_construction=args.ef_construction, + entrypoints=args.entrypoints, + ) + sweep_eps = [ + int(v) for v in args.entrypoints_sweep.split(",") if v.strip() + ] + built_ep = index.num_base_level_search_entrypoints + for ep in sweep_eps: + index.num_base_level_search_entrypoints = ep + r = eval_recall(index, xq, Igt, k) + q = eval_qps(index, xq, k) + print(f"\n --- num_base_level_search_entrypoints={ep} ---") + print(f" {'efSearch':>8s} {'recall@' + str(k):>10s} {'QPS':>12s}") + for ef in sorted(r.keys()): + print(f" {ef:>8d} {r[ef]:>10.4f} {q[ef]:>12,.1f}") + index.num_base_level_search_entrypoints = built_ep + + recall = eval_recall(index, xq, Igt, k) + qps = eval_qps(index, xq, k) + kcycles = eval_kcycles(index, xq, k) print("=" * 60) - print("SUMMARY") + print("RESULTS") print("=" * 60) - header = ( - f"{'Approach':<20s} {'Build':>7s} {'Ser':>5s} {'GB':>5s}" - f" {'ef':>4s} {'recall':>7s} {'kcyc/q':>7s}" - ) - print(header) - print("-" * len(header)) - for name, build, ser, size, recalls, kcycles in results_table: - for ef in [64, 128, 256]: - r = recalls.get(ef, 0) - kc = kcycles.get(ef, 0) - bld = f"{build:.0f}" if ef == 64 else "" - sr = f"{ser:.0f}" if ef == 64 else "" - sz = f"{size:.1f}" if ef == 64 else "" - nm = name if ef == 64 else "" - print( - f"{nm:<20s} {bld:>7s} {sr:>5s} {sz:>5s}" - f" {ef:>4d} {r:>7.4f} {kc:>7.0f}" - ) + print(f" Build (all_neighbors+optimize): {timings['build']:.1f}s") + print(f" copyTo: {timings['copyTo']:.1f}s") + print( + f" Serialize: {timings['serialize']:.1f}s " + f"({timings['file_size_gb']:.2f} GB)" + ) + # Index-build wall-clock, isolated from data load (Koski/Manifold) and + # ground-truth: this is the headline build->serialize number. + index_total = ( + timings["build"] + timings["copyTo"] + timings["serialize"] + ) + print( + f" >>> INDEX build->serialize total " + f"(excl. data load + ground truth): " + f"{index_total:.1f}s ({index_total / 60:.1f} min)" + ) + print(f" {'efSearch':>8s} {'recall@' + str(k):>10s} {'QPS':>12s} " + f"{'us/query':>10s} {'kcyc/q':>9s}") + for ef in sorted(recall.keys()): + q = qps[ef] + print( + f" {ef:>8d} {recall[ef]:>10.4f} {q:>12,.1f} " + f"{1e6 / q:>10.1f} {kcycles.get(ef, 0):>9.0f}" + ) if __name__ == "__main__": diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 7cbd07447f..e18a5b6a65 100644 --- a/faiss/gpu/test/test_cagra.py +++ b/faiss/gpu/test/test_cagra.py @@ -248,58 +248,56 @@ def test_IDMapCagra_IP_Int8(self): @unittest.skipIf( "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in" ) -@unittest.skipIf( - faiss.get_num_gpus() < 2, "need at least 2 GPUs for multi-GPU test" -) -class TestMultiGpuCagra(unittest.TestCase): - - def test_multi_gpu_build_and_search(self): - ds = datasets.SyntheticDataset(128, 0, 100_000, 1000) - xb = ds.get_database() - xq = ds.get_queries() - k = 10 - - gt_index = faiss.IndexFlatL2(ds.d) - gt_index.add(xb) - Dref, Iref = gt_index.search(xq, k) +class TestCagraConfig(unittest.TestCase): + """Config surface for the multi-GPU build, checkable without 2 GPUs.""" - num_gpus = min(faiss.get_num_gpus(), 4) - resources = faiss.GpuResourcesVector() + def test_multi_gpu_knobs_survive_swig(self): devices = faiss.Int32Vector() - res_list = [] - for i in range(num_gpus): - res = faiss.StandardGpuResources() - res_list.append(res) - resources.push_back(res) - devices.push_back(i) + devices.push_back(0) + devices.push_back(1) + + all_neighbors = faiss.AllNeighborsCagraConfig() + all_neighbors.n_clusters = 16 + all_neighbors.overlap_factor = 3 + all_neighbors.ivf_pq_search_batch_size = 8192 config = faiss.GpuIndexCagraConfig() - config.graph_degree = 32 - index = faiss.GpuIndexCagra(res_list[0], ds.d, faiss.METRIC_L2, config) - index.trainMultiGpu(ds.nb, faiss.swig_ptr(xb), resources, devices, 0, 2) + config.devices = devices + config.all_neighbors_params = all_neighbors + + self.assertEqual(config.devices.size(), 2) + self.assertEqual(config.devices.at(1), 1) + stored = config.all_neighbors_params + self.assertEqual( + ( + stored.n_clusters, + stored.overlap_factor, + stored.ivf_pq_search_batch_size, + ), + (16, 3, 8192), + ) - cpu_index = faiss.IndexHNSWCagra() - index.copyTo(cpu_index) - self.assertEqual(cpu_index.ntotal, ds.nb) + def test_brute_force_rejected_without_multiple_devices(self): + ds = datasets.SyntheticDataset(32, 0, 1000, 1) + config = faiss.GpuIndexCagraConfig() + config.build_algo = faiss.graph_build_algo_BRUTE_FORCE + res = faiss.StandardGpuResources() + index = faiss.GpuIndexCagra(res, ds.d, faiss.METRIC_L2, config) - cpu_index.hnsw.efSearch = 128 - Dnew, Inew = cpu_index.search(xq, k) + with self.assertRaises(RuntimeError): + index.train(ds.get_database()) - recall = np.mean( - [len(set(Inew[i]) & set(Iref[i])) / k for i in range(ds.nq)] - ) - self.assertGreater( - recall, 0.80, f"Multi-GPU recall@{k} too low: {recall:.4f}" - ) - # Serialization roundtrip - data = faiss.serialize_index(cpu_index) - loaded = faiss.deserialize_index(data) - loaded.hnsw.efSearch = 128 - Dnew2, Inew2 = loaded.search(xq, k) - np.testing.assert_array_equal(Inew, Inew2) +@unittest.skipIf( + "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in" +) +@unittest.skipIf( + faiss.get_num_gpus() < 2, "need at least 2 GPUs for multi-GPU test" +) +class TestMultiGpuCagra(unittest.TestCase): def test_all_neighbors_build(self): + """train() on a multi-device config builds a unified CAGRA graph.""" ds = datasets.SyntheticDataset(32, 0, 50_000, 100) xb = ds.get_database() xq = ds.get_queries() @@ -309,19 +307,22 @@ def test_all_neighbors_build(self): gt_index.add(xb) Dref, Iref = gt_index.search(xq, k) - num_gpus = min(faiss.get_num_gpus(), 4) devices = faiss.Int32Vector() - for i in range(num_gpus): + for i in range(min(faiss.get_num_gpus(), 4)): devices.push_back(i) - res = faiss.StandardGpuResources() + all_neighbors = faiss.AllNeighborsCagraConfig() + config = faiss.GpuIndexCagraConfig() config.graph_degree = 32 config.intermediate_graph_degree = 48 + config.build_algo = faiss.graph_build_algo_NN_DESCENT + config.devices = devices + config.all_neighbors_params = all_neighbors + + res = faiss.StandardGpuResources() index = faiss.GpuIndexCagra(res, ds.d, faiss.METRIC_L2, config) - index.trainAllNeighbors( - ds.nb, faiss.swig_ptr(xb), devices, 0, 0, True, 0 - ) + index.train(xb) cpu_index = faiss.IndexHNSWCagra() cpu_index.base_level_only = True @@ -337,3 +338,9 @@ def test_all_neighbors_build(self): self.assertGreater( recall, 0.70, f"all_neighbors recall@{k} too low: {recall:.4f}" ) + + # Serialization roundtrip + loaded = faiss.deserialize_index(faiss.serialize_index(cpu_index)) + loaded.hnsw.efSearch = 128 + _, Inew2 = loaded.search(xq, k) + np.testing.assert_array_equal(Inew, Inew2) diff --git a/pyproject-gpu-cuvs.toml b/pyproject-gpu-cuvs.toml index 9e1b1d5d60..21899ede08 100644 --- a/pyproject-gpu-cuvs.toml +++ b/pyproject-gpu-cuvs.toml @@ -45,6 +45,7 @@ dependencies = [ "nvidia-curand>=10.3.7,<11", "nvidia-nvjitlink>=13.2,<14", "libcuvs-cu13>=26.06,<27", + "librmm-cu13>=26.06,<27", ] [project.urls] @@ -114,6 +115,7 @@ test-requires = [ "nvidia-curand>=10.3.7,<11", "nvidia-nvjitlink>=13.2,<14", "libcuvs-cu13>=26.06,<27", + "librmm-cu13>=26.06,<27", ] test-command = "python -m pytest {project}/tests/test_wheel_smoke_gpu.py -v"