diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake index 9c6ce5eb3..ca3b64f46 100644 --- a/cmake/libs/libdiskann.cmake +++ b/cmake/libs/libdiskann.cmake @@ -1,6 +1,6 @@ add_definitions(-DKNOWHERE_WITH_DISKANN) -find_package(Boost REQUIRED COMPONENTS program_options) -include_directories(${Boost_INCLUDE_DIR}) +find_package(Boost REQUIRED COMPONENTS program_options CONFIG) +include_directories(${Boost_INCLUDE_DIRS}) find_package(aio REQUIRED) include_directories(${AIO_INCLUDE}) find_package(fmt REQUIRED) diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index bd49cedd7..e0c7e42ef 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -531,3 +531,26 @@ if(__PPC64) knowhere_utils) target_compile_definitions(faiss PRIVATE FINTEGER=int) endif() + +# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled +if(WITH_CUVS) + set(FAISS_GPU_HNSW_SRCS + thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu + thirdparty/faiss/faiss/gpu/GpuIndex.cu + thirdparty/faiss/faiss/gpu/GpuResources.cpp + thirdparty/faiss/faiss/gpu/StandardGpuResources.cpp + thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu + thirdparty/faiss/faiss/gpu/impl/IndexUtils.cu + thirdparty/faiss/faiss/gpu/utils/DeviceUtils.cu + thirdparty/faiss/faiss/gpu/utils/StackDeviceMemory.cpp + thirdparty/faiss/faiss/gpu/utils/Timer.cpp + ) + add_library(faiss_gpu_hnsw OBJECT ${FAISS_GPU_HNSW_SRCS}) + target_include_directories(faiss_gpu_hnsw PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/faiss + ${Boost_INCLUDE_DIRS} + ) + target_compile_definitions(faiss_gpu_hnsw PRIVATE FINTEGER=int) + target_link_libraries(faiss_gpu_hnsw PRIVATE CUDA::cudart) + target_link_libraries(faiss PUBLIC faiss_gpu_hnsw) +endif() diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9066b3641..ab26efe52 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -53,6 +53,7 @@ constexpr const char* INDEX_GPU_BRUTEFORCE = "GPU_BRUTE_FORCE"; constexpr const char* INDEX_GPU_IVFFLAT = "GPU_IVF_FLAT"; constexpr const char* INDEX_GPU_IVFPQ = "GPU_IVF_PQ"; constexpr const char* INDEX_GPU_CAGRA = "GPU_CAGRA"; +constexpr const char* INDEX_GPU_HNSW = "GPU_HNSW"; constexpr const char* INDEX_HNSW = "HNSW"; constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ"; diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index d807305e9..88e12ce73 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -82,6 +82,8 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_FLOAT16}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_INT8}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8}, // hnsw {IndexEnum::INDEX_HNSW, VecType::VECTOR_FLOAT}, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index d7f48ddb4..527c9404b 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3283,4 +3283,248 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH KNOWHERE_SIMPLE_REGISTER_DENSE_INT_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexHNSWPRQNodeTemplate, knowhere::feature::MMAP | knowhere::feature::MV | knowhere::feature::EMB_LIST) +#ifdef KNOWHERE_WITH_CUVS +} // namespace knowhere — temporarily close to include GPU headers at file scope +// ── GPU HNSW ───────────────────────────────────────────────────────────────── +#include +#include +namespace knowhere { // reopen namespace knowhere + +// Process-global StandardGpuResources shared by all GpuHnswIndexNode instances. +// Avoids per-segment 256 MiB pinned memory, cuBLAS handle, and CUDA stream +// allocations that accumulate to tens of GiB with many segments. +static std::shared_ptr& +GetSharedGpuResources() { + static std::once_flag flag; + static std::shared_ptr instance; + std::call_once(flag, [] { + instance = std::make_shared(); + instance->setTempMemory(0); + instance->setPinnedMemory(0); + }); + return instance; +} + +// Serialize GpuIndexHNSW construction across segments. +// StandardGpuResourcesImpl::initializeForDevice is not thread-safe; +// concurrent constructors race on the allocs_ map assertion. +static std::mutex& +GetGpuConstructionMutex() { + static std::mutex mtx; + return mtx; +} + +// Single GPU HNSW index node that handles all CPU storage formats (F32, SQ8, +// FP16, BF16) transparently. Uses faiss::gpu::GpuIndexHNSW for GPU search. +// Accepts CPU-serialized HNSW or HNSW_SQ binaries at load time. +class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { + public: + GpuHnswIndexNode(const int32_t& version, const Object& object) + : BaseFaissRegularIndexHNSWNode(version, object, DataFormatEnum::fp32) { + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + return true; + } + + static expected + StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, + const knowhere::BaseConfig& config, const IndexVersion& version) { + // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed + // after upload in Deserialize(). Report zero CPU memory cost so the + // Milvus segment loader does not over-commit host RAM reservations. + return Resource{.memoryCost = 0, .diskCost = 0}; + } + + std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + std::string + Type() const override { + return IndexEnum::INDEX_GPU_HNSW; + } + + protected: + Status + TrainInternal(const DataSetPtr /*dataset*/, const Config& /*cfg*/) override { + return Status::not_implemented; + } + + public: + Status + Deserialize(const BinarySet& binset, std::shared_ptr cfg) override { + std::unique_lock lock(gpu_mutex_); + gpu_index_.reset(); + + // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. + Status status; + if (!binset.Contains(IndexEnum::INDEX_GPU_HNSW)) { + BinarySet aliased = binset; + for (const char* key : {IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + if (binset.Contains(key)) { + aliased.Append(IndexEnum::INDEX_GPU_HNSW, binset.GetByName(key)); + break; + } + } + status = BaseFaissRegularIndexHNSWNode::Deserialize(aliased, cfg); + } else { + status = BaseFaissRegularIndexHNSWNode::Deserialize(binset, cfg); + } + if (status != Status::success) { + return status; + } + + // Eager GPU upload via faiss::gpu::GpuIndexHNSW. + const auto* faiss_idx = GetFaissHnswIndex(); + if (faiss_idx) { + try { + // Detect metric from the FAISS index type rather than config, + // because Deserialize may be called without metric_type in the config + // (e.g. empty json defaults metric_type to L2). + bool is_cosine = + dynamic_cast(faiss_idx) != nullptr; + bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); + + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + // Release CPU copy — vectors and graph are now on GPU. + indexes[0].reset(); + } catch (const std::exception& e) { + fprintf(stderr, "[gpu_hnsw] eager GPU upload failed: %s\n", e.what()); + gpu_index_.reset(); + } + } + return Status::success; + } + + expected + Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + if (!bitset.empty() && bitset.count() > 0) { + return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); + } + + // Fast path: gpu_index_ is set during Deserialize and never cleared. + if (!gpu_index_) { + std::unique_lock lock(gpu_mutex_); + if (!gpu_index_) { + const auto* faiss_idx = GetFaissHnswIndex(); + if (!faiss_idx) { + return expected::Err(Status::empty_index, "index not loaded"); + } + try { + const auto& hnsw_cfg = static_cast(*cfg); + bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); + bool use_ip = IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || is_cosine; + + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + const_cast&>(indexes[0]).reset(); + } catch (const std::exception& e) { + return expected::Err(Status::cuvs_inner_error, + std::string("failed to build GPU HNSW index: ") + e.what()); + } + } + } + + const auto& hnsw_cfg = static_cast(*cfg); + auto k = hnsw_cfg.k.value(); + auto nq = dataset->GetRows(); + auto dim = dataset->GetDim(); + auto ef = hnsw_cfg.ef.value_or(200); + const auto* h_queries_raw = reinterpret_cast(dataset->GetTensor()); + + // For COSINE metric, normalize queries to unit length. + const float* h_queries = h_queries_raw; + std::unique_ptr normalized_queries; + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + normalized_queries = std::make_unique(nq * dim); + for (int64_t i = 0; i < nq; i++) { + const float* src = h_queries_raw + i * dim; + float* dst = normalized_queries.get() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) sq_norm += src[d] * src[d]; + float inv = (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 1.0f; + for (int64_t d = 0; d < dim; d++) dst[d] = src[d] * inv; + } + h_queries = normalized_queries.get(); + } + + auto h_ids = std::make_unique(nq * k); + auto h_dist = std::make_unique(nq * k); + + try { + faiss::gpu::GpuHnswSearchParams gsp; + gsp.ef = ef; + gpu_index_->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); + } catch (const std::exception& e) { + LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); + return expected::Err(Status::cuvs_inner_error, + std::string("GPU HNSW search failed: ") + e.what()); + } + + // Negate back to positive for IP and COSINE. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || + IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq * k); i++) { + h_dist[i] = -h_dist[i]; + } + } + + return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); + } + + ~GpuHnswIndexNode() override = default; + + private: + const ::faiss::cppcontrib::knowhere::IndexHNSW* + GetFaissHnswIndex() const { + if (indexes.empty() || !indexes[0]) + return nullptr; + return dynamic_cast(indexes[0].get()); + } + + mutable std::mutex gpu_mutex_; + mutable std::shared_ptr gpu_resources_; + mutable std::unique_ptr gpu_index_; +}; + +// Register GPU_HNSW in the static config map at process startup. +__attribute__((constructor)) static void +register_gpu_hnsw_static_config() { + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); +} + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, + true, feature::GPU_ANN_FLOAT_INDEX); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + int8, true, (feature::INT8 | feature::GPU)); +#endif // KNOWHERE_WITH_CUVS + } // namespace knowhere diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 4f7f21ed7..136131912 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -443,5 +443,386 @@ TEST_CASE("Test All GPU Index", "[search]") { CHECK(GetRelativeLoss(gt_dist[i], dist[i]) < 0.1f); } } + + // GPU_HNSW tests: build on CPU as HNSW, serialize, deserialize as GPU_HNSW, search on GPU. + SECTION("Test GPU HNSW Search (CPU build -> GPU search)") { + // Build a CPU HNSW index + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize the CPU index + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // Deserialize as GPU_HNSW + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Search on GPU + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + // Compare against brute force + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW Search Cosine Metric") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 1); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.65f); + } + + SECTION("Test GPU HNSW Search TopK") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + const auto topk_values = { + std::make_tuple(5, 0.85f), + std::make_tuple(25, 0.85f), + std::make_tuple(100, 0.85f), + }; + + for (const auto& [topk, threshold] : topk_values) { + hnsw_json[knowhere::meta::TOPK] = topk; + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= threshold); + } + } + + SECTION("Test GPU HNSW Serialize/Deserialize Round Trip") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Self-search: each vector should find itself as nearest neighbor + auto results = gpu_idx.Search(train_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto ids = results.value()->GetIds(); + int correct = 0; + for (int i = 0; i < nq; ++i) { + if (ids[i] == i) + correct++; + } + float self_recall = static_cast(correct) / nq; + REQUIRE(self_recall >= 0.95f); + } + + SECTION("Test GPU HNSW SQ8 Deserialization") { + // Build a CPU HNSW_SQ index, then load as GPU_HNSW + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // GPU_HNSW should accept HNSW_SQ binaries + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // SQ8 has lower recall than flat due to quantization + REQUIRE(recall >= 0.7f); + } +} + +TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { + using Catch::Approx; + + int64_t nb = 10000, nq = 100; + int64_t dim = 128; + int64_t seed = 42; + + auto version = GenTestVersionList(); + + auto run_comparison = [&](const std::string& metric, float min_recall, float max_dist_drift) { + CAPTURE(metric, nb, nq, dim); + + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = metric; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + // --- Build CPU HNSW --- + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize for GPU deserialization + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // --- CPU Search with timing --- + auto cpu_start = std::chrono::high_resolution_clock::now(); + auto cpu_results = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto cpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(cpu_results.has_value()); + double cpu_ms = std::chrono::duration(cpu_end - cpu_start).count(); + + // --- Build GPU HNSW from serialized CPU index --- + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Warm-up search (first GPU call has kernel launch overhead) + auto warmup = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(warmup.has_value()); + + // --- GPU Search with timing --- + auto gpu_start = std::chrono::high_resolution_clock::now(); + auto gpu_results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto gpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(gpu_results.has_value()); + double gpu_ms = std::chrono::duration(gpu_end - gpu_start).count(); + + // --- Brute force ground truth --- + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + // --- Recall comparison --- + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_results.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_results.value()); + + // --- Distance accuracy: compare GPU distances against CPU distances --- + auto cpu_dist = cpu_results.value()->GetDistance(); + auto gpu_dist = gpu_results.value()->GetDistance(); + auto k = hnsw_json[knowhere::meta::TOPK].get(); + + double total_rel_error = 0.0; + int valid_pairs = 0; + for (int64_t i = 0; i < nq * k; ++i) { + if (std::abs(cpu_dist[i]) > 1e-6f) { + total_rel_error += std::abs((gpu_dist[i] - cpu_dist[i]) / cpu_dist[i]); + valid_pairs++; + } + } + double avg_rel_dist_error = (valid_pairs > 0) ? (total_rel_error / valid_pairs) : 0.0; + + // --- ID overlap: how many of the same results are returned --- + auto cpu_ids = cpu_results.value()->GetIds(); + auto gpu_ids = gpu_results.value()->GetIds(); + int id_overlap = 0; + for (int64_t q = 0; q < nq; ++q) { + std::set cpu_set(cpu_ids + q * k, cpu_ids + (q + 1) * k); + for (int64_t j = 0; j < k; ++j) { + if (cpu_set.count(gpu_ids[q * k + j])) { + id_overlap++; + } + } + } + float id_overlap_ratio = static_cast(id_overlap) / (nq * k); + + // --- Print comparison report --- + fprintf(stderr, "\n=== CPU vs GPU HNSW Comparison (%s, nb=%ld, nq=%ld, dim=%ld, k=%ld) ===\n", metric.c_str(), + (long)nb, (long)nq, (long)dim, (long)k); + fprintf(stderr, " CPU recall@%ld: %.4f\n", (long)k, cpu_recall); + fprintf(stderr, " GPU recall@%ld: %.4f\n", (long)k, gpu_recall); + fprintf(stderr, " Recall delta (GPU - CPU): %+.4f\n", gpu_recall - cpu_recall); + fprintf(stderr, " CPU search time: %.2f ms\n", cpu_ms); + fprintf(stderr, " GPU search time: %.2f ms (includes H2D/D2H transfers)\n", gpu_ms); + fprintf(stderr, " Speedup: %.2fx\n", cpu_ms / gpu_ms); + fprintf(stderr, " Avg relative distance error: %.6f\n", avg_rel_dist_error); + fprintf(stderr, " ID overlap (CPU vs GPU): %.4f (%d/%ld)\n", id_overlap_ratio, id_overlap, (long)(nq * k)); + fprintf(stderr, "===\n\n"); + + // --- Assertions --- + // GPU recall should be close to CPU recall (within a small tolerance) + REQUIRE(gpu_recall >= min_recall); + // GPU should return at least 80% of the same IDs as CPU + REQUIRE(id_overlap_ratio >= 0.8f); + // Average relative distance error should be small + REQUIRE(avg_rel_dist_error <= max_dist_drift); + }; + + SECTION("L2 metric comparison") { + run_comparison(knowhere::metric::L2, 0.85f, 0.05); + } + + SECTION("IP metric comparison") { + run_comparison(knowhere::metric::IP, 0.85f, 0.05); + } + + SECTION("COSINE metric comparison") { + run_comparison(knowhere::metric::COSINE, 0.60f, 0.10); + } + + SECTION("Varying ef comparison") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + fprintf(stderr, "\n=== ef sweep (L2, nb=%ld, nq=%ld, k=10) ===\n", (long)nb, (long)nq); + fprintf(stderr, " %6s %10s %10s %10s %10s\n", "ef", "cpu_recall", "gpu_recall", "cpu_ms", "gpu_ms"); + + for (int ef : {16, 32, 64, 128, 256, 512}) { + hnsw_json[knowhere::indexparam::EF] = ef; + + auto t0 = std::chrono::high_resolution_clock::now(); + auto cpu_res = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t1 = std::chrono::high_resolution_clock::now(); + auto gpu_res = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t2 = std::chrono::high_resolution_clock::now(); + + REQUIRE(cpu_res.has_value()); + REQUIRE(gpu_res.has_value()); + + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_res.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_res.value()); + double cpu_ms = std::chrono::duration(t1 - t0).count(); + double gpu_ms = std::chrono::duration(t2 - t1).count(); + + fprintf(stderr, " %6d %10.4f %10.4f %10.2f %10.2f\n", ef, cpu_recall, gpu_recall, cpu_ms, gpu_ms); + + // GPU recall should be reasonable at all ef values + REQUIRE(gpu_recall >= 0.5f); + } + fprintf(stderr, "===\n\n"); + } } #endif diff --git a/thirdparty/faiss/faiss/gpu/CMakeLists.txt b/thirdparty/faiss/faiss/gpu/CMakeLists.txt index 50740450a..59c7e7eab 100644 --- a/thirdparty/faiss/faiss/gpu/CMakeLists.txt +++ b/thirdparty/faiss/faiss/gpu/CMakeLists.txt @@ -29,9 +29,11 @@ set(FAISS_GPU_SRC GpuIndexIVF.cu GpuIndexIVFFlat.cu GpuIndexIVFPQ.cu + GpuIndexHNSW.cu GpuIndexIVFScalarQuantizer.cu GpuResources.cpp StandardGpuResources.cpp + impl/GpuHnswTypes.cu impl/BinaryDistance.cu impl/BinaryFlatIndex.cu impl/BroadcastSum.cu @@ -96,10 +98,15 @@ set(FAISS_GPU_HEADERS GpuIndexIVF.h GpuIndexIVFFlat.h GpuIndexIVFPQ.h + GpuIndexHNSW.h GpuIndexIVFScalarQuantizer.h GpuIndicesOptions.h GpuResources.h StandardGpuResources.h + impl/GpuHnswBuild.cuh + impl/GpuHnswSearch.cuh + impl/GpuHnswSearchKernel.cuh + impl/GpuHnswTypes.h impl/BinaryDistance.cuh impl/BinaryFlatIndex.cuh impl/BroadcastSum.cuh diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu new file mode 100644 index 000000000..e92db6d49 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -0,0 +1,276 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace faiss { +namespace gpu { + +GpuIndexHNSW::GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric, + GpuIndexHNSWConfig config) + : GpuIndex(provider->getResources(), dims, metric, 0.0f, config), + hnswConfig_(config) { + this->is_trained = false; +} + +GpuIndexHNSW::~GpuIndexHNSW() = default; + +void GpuIndexHNSW::copyFrom( + const faiss::cppcontrib::knowhere::IndexHNSW* index) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + // Detect cosine from index type (IndexHNSWFlatCosine / IndexHNSWSQCosine + // implement HasInverseL2Norms). + bool is_cosine = + dynamic_cast( + index) != nullptr; + bool use_ip = + is_cosine || (index->metric_type == faiss::METRIC_INNER_PRODUCT); + + if (dynamic_cast(index->storage)) { + deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + } else { + deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + if (dynamic_cast(index->storage)) { + deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + } else { + deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::reset() { + deviceIndex_.reset(); + this->ntotal = 0; + this->is_trained = false; +} + +void GpuIndexHNSW::setSearchParams(const GpuHnswSearchParams& params) const { + std::lock_guard lock(searchParamsMutex_); + directSearchParams_ = params; + hasDirectSearchParams_ = true; +} + +bool GpuIndexHNSW::addImplRequiresIDs_() const { + return false; +} + +void GpuIndexHNSW::addImpl_(idx_t, const float*, const idx_t*) { + FAISS_THROW_MSG( + "GpuIndexHNSW does not support add(). " + "Build on CPU with IndexHNSW, then call copyFrom()."); +} + +void GpuIndexHNSW::searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + auto& idx = *deviceIndex_; + + GpuHnswSearchParams sp; + bool got_params = false; + + // Prefer direct params set via setSearchParams() — avoids dynamic_cast. + { + std::lock_guard lock(searchParamsMutex_); + if (hasDirectSearchParams_) { + sp = directSearchParams_; + hasDirectSearchParams_ = false; + got_params = true; + } + } + + // Fallback: try dynamic_cast from SearchParameters. + if (!got_params && search_params) { + auto* params = + dynamic_cast(search_params); + if (params) { + sp.ef = params->ef; + sp.search_width = params->search_width; + sp.max_iterations = params->max_iterations; + sp.thread_block_size = params->thread_block_size; + sp.overflow_factor = params->overflow_factor; + got_params = true; + } + } + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int overflow_ef = sp.overflow_factor * sp.ef; + sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + + // D2D: query vectors (GpuIndex::search passes device pointers) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + // D2D: distances (output is a device pointer from GpuIndex::search) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + // Labels: D2H stage (uint64_t→idx_t conversion), then H2D back + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + auto h_labels = std::make_unique(nq * k); + for (int i = 0; i < nq * k; i++) { + h_labels[i] = (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + labels, + h_labels.get(), + static_cast(nq) * k * sizeof(idx_t), + cudaMemcpyHostToDevice, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +void GpuIndexHNSW::searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& sp) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); + DeviceScope scope(config_.device); + auto& idx = *deviceIndex_; + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int overflow_ef = sp.overflow_factor * sp.ef; + sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x_host, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDefault, + stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances_host, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToHost, + stream)); + + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int i = 0; i < nq * k; i++) { + labels_host[i] = + (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h new file mode 100644 index 000000000..a6fde1701 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -0,0 +1,135 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { +namespace cppcontrib { +namespace knowhere { +struct IndexHNSW; +} // namespace knowhere +} // namespace cppcontrib +namespace gpu { + +struct GpuHnswDeviceIndex; + +struct GpuIndexHNSWConfig : public GpuIndexConfig {}; + +struct SearchParametersGpuHNSW : SearchParameters { + /// Search ef — number of candidates maintained during search. + /// Higher values improve recall at the cost of speed. + int ef = 200; + + /// Number of candidates to expand per iteration. + int search_width = 4; + + /// Maximum search iterations (0 = auto). + int max_iterations = 0; + + /// Thread block size (0 = auto, default 128). + int thread_block_size = 0; + + /// Overflow queue factor: overflow_ef = overflow_factor * ef. + int overflow_factor = 2; +}; + +/// GPU implementation of HNSW search. +/// +/// This index type does NOT build an HNSW graph on GPU — it takes a +/// CPU-built faiss::IndexHNSW (Flat or SQ storage), converts the graph +/// to a GPU-friendly dense format, and runs the search on GPU using an +/// Overflow Candidate Queue (OCQ) beam search kernel. +/// +/// Supports L2, inner product, and cosine metrics. +/// Supports float32 and int8 (QT_8bit_direct_signed) data. +struct GpuIndexHNSW : public GpuIndex { + public: + GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric = faiss::METRIC_L2, + GpuIndexHNSWConfig config = GpuIndexHNSWConfig()); + + ~GpuIndexHNSW() override; + + /// Load an HNSW index from CPU to GPU. + /// The CPU index must have been built and trained already. + /// Supports IndexHNSWFlat (float32) and IndexHNSWSQ + /// (QT_8bit_direct_signed for INT8, or dequantized for other SQ types). + void copyFrom(const faiss::cppcontrib::knowhere::IndexHNSW* index); + + /// Load with explicit metric specification. + void copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine); + + void reset() override; + + /// Set search parameters directly, bypassing SearchParameters. + /// Thread-safe: uses atomic/mutex internally. + void setSearchParams(const GpuHnswSearchParams& params) const; + + /// Search with host pointers directly, bypassing GpuIndex::search. + /// All input/output pointers must be host memory. + /// This avoids the GpuIndex::search_ex temp allocation chain + /// which can cause SIGSEGV from pointer lifetime issues. + void searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& params) const; + + protected: + bool addImplRequiresIDs_() const override; + + void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + + void searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const override; + + private: + GpuIndexHNSWConfig hnswConfig_; + + std::unique_ptr deviceIndex_; + + mutable std::mutex searchParamsMutex_; + mutable GpuHnswSearchParams directSearchParams_; + mutable bool hasDirectSearchParams_ = false; +}; + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh new file mode 100644 index 000000000..8c4ea63fc --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -0,0 +1,341 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Converts a FAISS HNSW index (CSR graph) + vectors into a +// GpuHnswDeviceIndex on device memory. +// +// This version is adapted for Knowhere's FAISS fork which uses +// faiss::cppcontrib::knowhere::IndexHNSW / HNSW types. + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#define GPU_HNSW_BUILD_CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +/// Extract HNSW graph layers from a Knowhere HNSW struct. +/// Template parameter HnswT can be faiss::cppcontrib::knowhere::HNSW +/// or faiss::HNSW — any type with neighbor_range(), nb_neighbors(), +/// neighbors, levels, entry_point, max_level. +template +inline void extract_hnsw_layers( + const HnswT& hnsw, + int64_t n_rows, + std::vector& h_upper_layers, + std::vector& h_layer0_flat, + uint32_t& entry_point, + int& M, + int& max_degree0, + int& num_layers) { + const int maxM0 = hnsw.nb_neighbors(0); + const int maxM = hnsw.nb_neighbors(1); + const int max_lv = hnsw.max_level; + + entry_point = static_cast(hnsw.entry_point); + M = maxM; + max_degree0 = maxM0; + num_layers = max_lv + 1; + + // Layer 0: dense [n_rows x maxM0] + h_layer0_flat.assign(n_rows * maxM0, UINT32_MAX); + for (int64_t i = 0; i < n_rows; i++) { + size_t begin, end; + hnsw.neighbor_range(i, 0, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_layer0_flat[i * maxM0 + j] = static_cast(nb); + } + } + + // Upper layers (1..max_level): sparse [num_nodes_at_L x maxM] + h_upper_layers.resize(max_lv); + for (int layer = 1; layer <= max_lv; layer++) { + auto& ul = h_upper_layers[layer - 1]; + ul.max_degree = static_cast(maxM); + + std::vector node_ids; + for (int64_t i = 0; i < n_rows; i++) { + if (hnsw.levels[i] > layer) + node_ids.push_back(static_cast(i)); + } + ul.num_nodes = static_cast(node_ids.size()); + + std::vector h_neighbors(ul.num_nodes * maxM, UINT32_MAX); + std::vector h_node_ids = node_ids; + + for (uint32_t idx = 0; idx < ul.num_nodes; idx++) { + int64_t i = node_ids[idx]; + size_t begin, end; + hnsw.neighbor_range(i, layer, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_neighbors[idx * maxM + j] = static_cast(nb); + } + } + + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&ul.d_node_ids, ul.num_nodes * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_node_ids, + h_node_ids.data(), + ul.num_nodes * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc( + &ul.d_neighbors, + ul.num_nodes * maxM * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_neighbors, + h_neighbors.data(), + ul.num_nodes * maxM * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + } +} + +inline void normalize_vectors( + std::vector& h_vectors, + int64_t n_rows, + int64_t dim) { + for (int64_t i = 0; i < n_rows; i++) { + float* v = h_vectors.data() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) + sq_norm += v[d] * v[d]; + if (sq_norm > 0.0f) { + float inv = 1.0f / std::sqrt(sq_norm); + for (int64_t d = 0; d < dim; d++) + v[d] *= inv; + } + } +} + +template +inline void upload_graph_to_gpu( + GpuHnswDeviceIndex& idx, + const HnswT& hnsw, + int64_t n_rows) { + std::vector h_layer0_flat; + extract_hnsw_layers( + hnsw, + n_rows, + idx.upper_layers, + h_layer0_flat, + idx.entry_point, + idx.M, + idx.max_degree0, + idx.num_layers); + + size_t graph0_bytes = + static_cast(n_rows) * idx.max_degree0 * sizeof(uint32_t); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_layer0_graph, graph0_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_layer0_graph, + h_layer0_flat.data(), + graph0_bytes, + cudaMemcpyHostToDevice)); + + int num_upper = static_cast(idx.upper_layers.size()); + idx.num_upper_layers_built = num_upper; + if (num_upper > 0) { + using kernel_ptrs = hnsw_kernel::upper_layer_ptrs; + std::vector h_ptrs(num_upper); + for (int i = 0; i < num_upper; i++) { + const auto& ul = idx.upper_layers[i]; + h_ptrs[i] = { + ul.d_node_ids, ul.d_neighbors, ul.num_nodes, ul.max_degree}; + } + size_t ptrs_bytes = num_upper * sizeof(kernel_ptrs); + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&idx.d_upper_layer_ptrs, ptrs_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_upper_layer_ptrs, + h_ptrs.data(), + ptrs_bytes, + cudaMemcpyHostToDevice)); + } + +} + +inline void upload_fp32_dataset( + GpuHnswDeviceIndex& idx, + std::vector& h_vectors, + int64_t n_rows, + bool is_cosine) { + int64_t dim = idx.dim; + if (is_cosine) + normalize_vectors(h_vectors, n_rows, dim); + + size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + h_vectors.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_int8 = false; +} + +inline void upload_int8_dataset( + GpuHnswDeviceIndex& idx, + const uint8_t* codes, + int64_t n_rows, + bool is_cosine) { + int64_t dim = idx.dim; + size_t dataset_bytes = static_cast(n_rows) * dim; + + std::vector signed_codes(dataset_bytes); + for (size_t i = 0; i < dataset_bytes; i++) { + signed_codes[i] = + static_cast(static_cast(codes[i]) - 128); + } + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + signed_codes.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_int8 = true; + + if (is_cosine) { + std::vector h_inv_norms(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + const int8_t* row = signed_codes.data() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) { + float v = static_cast(row[d]); + sq_norm += v * v; + } + h_inv_norms[i] = + (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 0.0f; + } + size_t norms_bytes = static_cast(n_rows) * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&idx.d_inv_norms, norms_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_inv_norms, + h_inv_norms.data(), + norms_bytes, + cudaMemcpyHostToDevice)); + } +} + +/// Build from Knowhere's HNSW index with SQ storage. +inline std::unique_ptr from_faiss_hnsw_sq( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false) { + const auto* sq_storage = + dynamic_cast( + hnsw_index.storage); + if (!sq_storage) + throw std::runtime_error( + "gpu_hnsw: storage is not IndexScalarQuantizer"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->scratch_pool = std::make_unique(4, 0); + + bool is_direct_signed = (sq_storage->sq.qtype == + faiss::ScalarQuantizer::QT_8bit_direct_signed); + + if (is_direct_signed) { + upload_int8_dataset(*idx, sq_storage->codes.data(), n_rows, is_cosine); + } else { + std::vector h_vectors(n_rows * dim); + sq_storage->sa_decode( + n_rows, sq_storage->codes.data(), h_vectors.data()); + upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + } + + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +/// Build from Knowhere's HNSW index with Flat storage. +inline std::unique_ptr from_faiss_hnsw_flat( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false) { + const auto* flat_storage = + dynamic_cast(hnsw_index.storage); + if (!flat_storage) + throw std::runtime_error("gpu_hnsw: storage is not IndexFlat"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + // Use reconstruct_n instead of get_xb() — get_xb() can return a + // device pointer in GPU querynode context, causing SIGSEGV when + // accessed from CPU. + std::vector h_vectors(n_rows * dim); + flat_storage->reconstruct_n(0, n_rows, h_vectors.data()); + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->scratch_pool = std::make_unique(4, 0); + + upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh new file mode 100644 index 000000000..126c88674 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -0,0 +1,167 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include +#include + +#define GPU_HNSW_CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +inline void gpu_hnsw_search( + cudaStream_t stream, + const GpuHnswSearchParams& params, + const GpuHnswDeviceIndex& idx, + GpuHnswSearchScratch& sc, + int num_queries, + int k) { + + int ef = params.ef; + int sw = params.search_width; + int overflow_ef = params.overflow_factor * ef; + int max_iter = params.max_iterations > 0 + ? params.max_iterations + : 2 * ef / sw + 10; + int dim = static_cast(idx.dim); + int num_upper_layers = idx.num_upper_layers_built; + + auto launch_kernels = [&]( + const DataT* d_data, + const float* d_inv_norms) { + if (num_upper_layers > 0) { + auto* d_layer_ptrs = static_cast( + idx.d_upper_layer_ptrs); + + int warps_per_block = 4; + int threads_per_block = warps_per_block * 32; + int num_blocks = + (num_queries + warps_per_block - 1) / warps_per_block; + + hnsw_kernel::upper_layer_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + d_layer_ptrs, + sc.d_entry_points, + idx.entry_point, + num_queries, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(num_queries, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_entry_points, + h_eps.data(), + num_queries * sizeof(uint32_t), + cudaMemcpyHostToDevice, + stream)); + } + + int block_size = + params.thread_block_size > 0 ? params.thread_block_size : 128; + + { + int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; + int max_ef = (49152 - smem_overhead) / 12; + if (ef > max_ef) { + ef = max_ef; + } + } + + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0); + int N_int = static_cast(idx.n_rows); + size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( + num_queries, N_int); + + GPU_HNSW_CUDA_CHECK( + cudaMemsetAsync(sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + if (overflow_ef > 0) { + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_overflow_count, + 0, + static_cast(num_queries) * sizeof(int), + stream)); + } + + hnsw_kernel::layer0_beam_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + idx.d_layer0_graph, + sc.d_entry_points, + sc.d_visited_bitmaps, + sc.d_neighbors, + sc.d_distances, + num_queries, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip, + overflow_ef, + sc.d_overflow_ids, + sc.d_overflow_dists, + sc.d_overflow_expanded, + sc.d_overflow_count); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + }; + + if (idx.dataset_int8) { + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + } else { + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh new file mode 100644 index 000000000..5d0856777 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -0,0 +1,618 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { +namespace gpu { +namespace hnsw_kernel { + +// ============================================================================ +// Distance computation helpers (templated on dataset element type) +// ============================================================================ + +__device__ __forceinline__ float load_elem(const float* ptr, int idx) { + return __ldg(&ptr[idx]); +} +__device__ __forceinline__ float load_elem(const half* ptr, int idx) { + return __half2float(__ldg(&ptr[idx])); +} +__device__ __forceinline__ float load_elem(const int8_t* ptr, int idx) { + return static_cast(__ldg(&ptr[idx])); +} + +template +__device__ __forceinline__ float thread_l2_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + float diff = query[d] - load_elem(vec, d); + sum += diff * diff; + } + return sum; +} + +template +__device__ __forceinline__ float thread_ip_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + sum += query[d] * load_elem(vec, d); + } + return -sum; +} + +// ============================================================================ +// Warp-cooperative distance computation +// ============================================================================ + +__device__ __forceinline__ int select_threads_per_dist(int dim) { + // Always use 1 thread per distance: 128 concurrent distances per block + // vs 32 with 4-thread cooperative. For int8 dim=384, each vector is only + // 384 bytes — fits in L1 cache. The 4x concurrency gain outweighs the + // cooperative memory coalescing benefit at this data size. + (void)dim; + return 1; +} + +template +__device__ __forceinline__ float coop_l2_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim, + int lane_in_group, + int threads_per_dist, + uint32_t group_mask) { + int chunk = dim / threads_per_dist; + int start = lane_in_group * chunk; + int end = (lane_in_group == threads_per_dist - 1) ? dim : start + chunk; + + float partial = 0.0f; + for (int d = start; d < end; d++) { + float diff = query[d] - load_elem(vec, d); + partial += diff * diff; + } + + for (int offset = threads_per_dist / 2; offset > 0; offset >>= 1) { + partial += __shfl_down_sync(group_mask, partial, offset); + } + return partial; +} + +template +__device__ __forceinline__ float coop_ip_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim, + int lane_in_group, + int threads_per_dist, + uint32_t group_mask) { + int chunk = dim / threads_per_dist; + int start = lane_in_group * chunk; + int end = (lane_in_group == threads_per_dist - 1) ? dim : start + chunk; + + float partial = 0.0f; + for (int d = start; d < end; d++) { + partial += query[d] * load_elem(vec, d); + } + + for (int offset = threads_per_dist / 2; offset > 0; offset >>= 1) { + partial += __shfl_down_sync(group_mask, partial, offset); + } + return -partial; +} + +// ============================================================================ +// Phase 1: Upper-layer greedy search +// ============================================================================ + +struct upper_layer_ptrs { + const uint32_t* d_node_ids; + const uint32_t* d_neighbors; + uint32_t num_nodes; + uint32_t max_degree; +}; + +__device__ __forceinline__ uint32_t +binary_search_node(const uint32_t* d_node_ids, uint32_t n, uint32_t global_id) { + uint32_t lo = 0, hi = n; + while (lo < hi) { + uint32_t mid = (lo + hi) / 2; + if (__ldg(&d_node_ids[mid]) < global_id) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < n && __ldg(&d_node_ids[lo]) == global_id) + return lo; + return UINT32_MAX; +} + +template +__global__ void upper_layer_search_kernel( + const float* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const upper_layer_ptrs* __restrict__ d_layer_ptrs, + uint32_t* __restrict__ d_entry_points, + uint32_t global_entry_point, + int num_queries, + int dim, + int num_upper_layers, + bool use_inner_product) { + int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int lane = threadIdx.x % 32; + if (warp_id >= num_queries) + return; + + const float* query = d_queries + static_cast(warp_id) * dim; + uint32_t current = global_entry_point; + + float best_dist; + if (lane == 0) { + if (use_inner_product) { + best_dist = thread_ip_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + if (d_inv_norms) + best_dist *= __ldg(&d_inv_norms[current]); + } else { + best_dist = thread_l2_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + } + } + best_dist = __shfl_sync(0xffffffff, best_dist, 0); + + for (int li = num_upper_layers - 1; li >= 0; li--) { + const upper_layer_ptrs& lp = d_layer_ptrs[li]; + bool improved = true; + while (improved) { + improved = false; + uint32_t local_idx = + binary_search_node(lp.d_node_ids, lp.num_nodes, current); + if (local_idx == UINT32_MAX) + break; + + uint32_t best_nbr = UINT32_MAX; + float best_nbr_dist = best_dist; + + for (uint32_t j = lane; j < lp.max_degree; j += 32) { + uint32_t nbr = lp.d_neighbors + [static_cast(local_idx) * + lp.max_degree + + j]; + float dist = FLT_MAX; + if (nbr != UINT32_MAX) { + const DataT* nbr_vec = + d_dataset + static_cast(nbr) * dim; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + } + if (dist < best_nbr_dist) { + best_nbr_dist = dist; + best_nbr = nbr; + } + } + + for (int offset = 16; offset > 0; offset >>= 1) { + float other_dist = + __shfl_down_sync(0xffffffff, best_nbr_dist, offset); + uint32_t other_id = + __shfl_down_sync(0xffffffff, best_nbr, offset); + if (other_dist < best_nbr_dist) { + best_nbr_dist = other_dist; + best_nbr = other_id; + } + } + best_nbr_dist = __shfl_sync(0xffffffff, best_nbr_dist, 0); + best_nbr = __shfl_sync(0xffffffff, best_nbr, 0); + + if (best_nbr != UINT32_MAX && best_nbr_dist < best_dist) { + best_dist = best_nbr_dist; + current = best_nbr; + improved = true; + } + } + } + + if (lane == 0) { + d_entry_points[warp_id] = current; + } +} + +// ============================================================================ +// Phase 2: Layer-0 beam search kernel with Overflow Candidate Queue (OCQ) +// ============================================================================ + +__device__ __forceinline__ bool bitmap_visit( + uint32_t* bitmap, + uint32_t node_id) { + uint32_t word = node_id >> 5; + uint32_t bit = 1u << (node_id & 31); + uint32_t old = atomicOr(&bitmap[word], bit); + return (old & bit) == 0; +} + +__device__ __forceinline__ void overflow_insert( + uint32_t* ovf_ids, + float* ovf_dists, + uint32_t* ovf_exp, + int& ovf_rc, + int overflow_ef, + uint32_t id, + float dist, + uint32_t expanded) { + if (ovf_rc >= overflow_ef && dist >= ovf_dists[ovf_rc - 1]) + return; + + int lo = 0, hi = ovf_rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (ovf_dists[mid] < dist) + lo = mid + 1; + else + hi = mid; + } + + int insert_end = ovf_rc < overflow_ef ? ovf_rc : overflow_ef - 1; + for (int i = insert_end; i > lo; i--) { + ovf_ids[i] = ovf_ids[i - 1]; + ovf_dists[i] = ovf_dists[i - 1]; + ovf_exp[i] = ovf_exp[i - 1]; + } + ovf_ids[lo] = id; + ovf_dists[lo] = dist; + ovf_exp[lo] = expanded; + if (ovf_rc < overflow_ef) + ovf_rc++; +} + +template +__global__ void layer0_beam_search_kernel( + const float* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint32_t* __restrict__ d_layer0_graph, + const uint32_t* __restrict__ d_entry_points, + uint32_t* __restrict__ d_visited_bitmaps, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances, + int num_queries, + int N, + int dim, + int max_degree0, + int k, + int ef, + int search_width, + int max_iterations, + bool use_inner_product, + int overflow_ef, + uint32_t* __restrict__ d_overflow_ids, + float* __restrict__ d_overflow_dists, + uint32_t* __restrict__ d_overflow_expanded, + int* __restrict__ d_overflow_count) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + const float* query = d_queries + static_cast(query_idx) * dim; + + extern __shared__ char smem[]; + + int max_staging = search_width * max_degree0; + + uint32_t* result_ids = reinterpret_cast(smem); + float* result_dists = reinterpret_cast(result_ids + ef); + uint32_t* is_expanded = reinterpret_cast(result_dists + ef); + uint32_t* staging_ids = is_expanded + ef; + float* staging_dists = reinterpret_cast(staging_ids + max_staging); + uint32_t* parent_ids = + reinterpret_cast(staging_dists + max_staging); + int* meta = reinterpret_cast(parent_ids + search_width); + + int bitmap_words = (N + 31) / 32; + uint32_t* visited_bmap = + d_visited_bitmaps + static_cast(query_idx) * bitmap_words; + + uint32_t* ovf_ids = overflow_ef > 0 + ? d_overflow_ids + static_cast(query_idx) * overflow_ef + : nullptr; + float* ovf_dists = overflow_ef > 0 + ? d_overflow_dists + static_cast(query_idx) * overflow_ef + : nullptr; + uint32_t* ovf_exp = overflow_ef > 0 + ? d_overflow_expanded + + static_cast(query_idx) * overflow_ef + : nullptr; + + for (int i = threadIdx.x; i < ef; i += blockDim.x) { + result_ids[i] = UINT32_MAX; + result_dists[i] = FLT_MAX; + is_expanded[i] = 0; + } + if (threadIdx.x == 0) { + meta[0] = 0; + meta[1] = 0; + meta[2] = 0; + if (overflow_ef > 0) + d_overflow_count[query_idx] = 0; + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + if (threadIdx.x == 0) { + float ep_dist; + if (use_inner_product) { + ep_dist = thread_ip_distance( + query, d_dataset + static_cast(ep) * dim, dim); + if (d_inv_norms) + ep_dist *= __ldg(&d_inv_norms[ep]); + } else { + ep_dist = thread_l2_distance( + query, d_dataset + static_cast(ep) * dim, dim); + } + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + uint32_t nbr = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + break; + } + } + meta[0] = rc; + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + for (int i = 0; i < rc && num_parents < search_width; i++) { + if (!is_expanded[i]) { + parent_ids[num_parents++] = result_ids[i]; + is_expanded[i] = 1; + } + } + + if (num_parents == 0 && overflow_ef > 0) { + int ovf_rc = d_overflow_count[query_idx]; + for (int i = 0; i < ovf_rc && num_parents < search_width; + i++) { + if (!ovf_exp[i]) { + parent_ids[num_parents++] = ovf_ids[i]; + ovf_exp[i] = 1; + } + } + } + meta[2] = num_parents; + } + __syncthreads(); + + int num_parents = meta[2]; + if (num_parents == 0) + break; + + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + int total_work = num_parents * max_degree0; + for (int wi = threadIdx.x; wi < total_work; wi += blockDim.x) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + uint32_t nbr = + d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + meta[0] = rc; + } + __syncthreads(); + + // Stagnation detected when num_parents==0 (all expanded) → break above + } + + // --- Copy top-k results to global memory --- + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = FLT_MAX; + } + } +} + +inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { + int max_staging = search_width * max_degree0; + + size_t size = 0; + size += ef * sizeof(uint32_t); + size += ef * sizeof(float); + size += ef * sizeof(uint32_t); + size += max_staging * sizeof(uint32_t); + size += max_staging * sizeof(float); + size += search_width * sizeof(uint32_t); + size += 3 * sizeof(int); + return size; +} + +inline size_t calc_visited_bitmap_size(int num_queries, int N) { + int bitmap_words = (N + 31) / 32; + return static_cast(num_queries) * bitmap_words * sizeof(uint32_t); +} + +} // namespace hnsw_kernel +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu new file mode 100644 index 000000000..1cd98bbdb --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -0,0 +1,200 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +namespace faiss { +namespace gpu { + +namespace { + +inline void check_cuda(cudaError_t err, const char* file, int line) { + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("CUDA error in scratch alloc: ") + + cudaGetErrorString(err) + " at " + file + ":" + + std::to_string(line)); + } +} + +#define SCRATCH_CUDA_CHECK(expr) check_cuda((expr), __FILE__, __LINE__) + +} // namespace + +void GpuHnswSearchScratch::ensure( + int nq, + int k, + int dim, + int N, + int overflow_ef) { + size_t need_q = static_cast(nq) * dim * sizeof(float); + if (need_q > queries_bytes) { + if (d_queries) + cudaFree(d_queries); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_queries, need_q)); + queries_bytes = need_q; + } + size_t need_n = static_cast(nq) * k * sizeof(uint64_t); + if (need_n > neighbors_bytes) { + if (d_neighbors) + cudaFree(d_neighbors); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_neighbors, need_n)); + neighbors_bytes = need_n; + } + size_t need_d = static_cast(nq) * k * sizeof(float); + if (need_d > distances_bytes) { + if (d_distances) + cudaFree(d_distances); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_distances, need_d)); + distances_bytes = need_d; + } + if (nq > entry_cap) { + if (d_entry_points) + cudaFree(d_entry_points); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_entry_points, + static_cast(nq) * sizeof(uint32_t))); + entry_cap = nq; + } + int bitmap_words = (N + 31) / 32; + size_t need_bm = + static_cast(nq) * bitmap_words * sizeof(uint32_t); + if (need_bm > bitmap_bytes) { + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_visited_bitmaps, need_bm)); + bitmap_bytes = need_bm; + } + if (overflow_ef > 0) { + size_t ovf_entries = static_cast(nq) * overflow_ef; + size_t need_ovf = + ovf_entries * + (sizeof(uint32_t) + sizeof(float) + sizeof(uint32_t)) + + static_cast(nq) * sizeof(int); + if (need_ovf > overflow_bytes) { + if (d_overflow_ids) + cudaFree(d_overflow_ids); + if (d_overflow_dists) + cudaFree(d_overflow_dists); + if (d_overflow_expanded) + cudaFree(d_overflow_expanded); + if (d_overflow_count) + cudaFree(d_overflow_count); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_ids, ovf_entries * sizeof(uint32_t))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_dists, ovf_entries * sizeof(float))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_expanded, ovf_entries * sizeof(uint32_t))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_count, + static_cast(nq) * sizeof(int))); + overflow_bytes = need_ovf; + } + } +} + +GpuHnswSearchScratch::~GpuHnswSearchScratch() { + if (d_queries) + cudaFree(d_queries); + if (d_neighbors) + cudaFree(d_neighbors); + if (d_distances) + cudaFree(d_distances); + if (d_entry_points) + cudaFree(d_entry_points); + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + if (d_overflow_ids) + cudaFree(d_overflow_ids); + if (d_overflow_dists) + cudaFree(d_overflow_dists); + if (d_overflow_expanded) + cudaFree(d_overflow_expanded); + if (d_overflow_count) + cudaFree(d_overflow_count); +} + +GpuHnswScratchSlot::~GpuHnswScratchSlot() { + if (stream) + cudaStreamDestroy(stream); +} + +GpuHnswScratchPool::GpuHnswScratchPool(int pool_size, int device) + : pool_size_(pool_size), device_(device) {} + +void GpuHnswScratchPool::init_once() { + if (initialized_) + return; + slots_.reserve(pool_size_); + available_.reserve(pool_size_); + for (int i = 0; i < pool_size_; i++) { + auto slot = std::make_unique(); + SCRATCH_CUDA_CHECK(cudaSetDevice(device_)); + SCRATCH_CUDA_CHECK( + cudaStreamCreateWithFlags(&slot->stream, cudaStreamNonBlocking)); + available_.push_back(slot.get()); + slots_.push_back(std::move(slot)); + } + initialized_ = true; +} + +GpuHnswScratchSlot* GpuHnswScratchPool::acquire() { + std::unique_lock lock(mutex_); + init_once(); + cv_.wait(lock, [this] { return !available_.empty(); }); + auto* slot = available_.back(); + available_.pop_back(); + return slot; +} + +void GpuHnswScratchPool::release(GpuHnswScratchSlot* slot) { + { + std::lock_guard lock(mutex_); + available_.push_back(slot); + } + cv_.notify_one(); +} + +GpuHnswDeviceIndex::~GpuHnswDeviceIndex() { + if (d_dataset) + cudaFree(d_dataset); + if (d_inv_norms) + cudaFree(d_inv_norms); + if (d_layer0_graph) + cudaFree(d_layer0_graph); + for (auto& ul : upper_layers) { + if (ul.d_node_ids) + cudaFree(ul.d_node_ids); + if (ul.d_neighbors) + cudaFree(ul.d_neighbors); + } + if (d_upper_layer_ptrs) + cudaFree(d_upper_layer_ptrs); +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h new file mode 100644 index 000000000..86b781282 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -0,0 +1,162 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace faiss { +namespace gpu { + +struct GpuHnswSearchParams { + int ef = 200; + int search_width = 4; + int max_iterations = 0; + int thread_block_size = 0; + int overflow_factor = 0; +}; + +struct GpuHnswDeviceUpperLayer { + uint32_t* d_node_ids = nullptr; + uint32_t* d_neighbors = nullptr; + uint32_t num_nodes = 0; + uint32_t max_degree = 0; +}; + +struct GpuHnswSearchScratch { + float* d_queries = nullptr; + uint64_t* d_neighbors = nullptr; + float* d_distances = nullptr; + uint32_t* d_entry_points = nullptr; + uint32_t* d_visited_bitmaps = nullptr; + + uint32_t* d_overflow_ids = nullptr; + float* d_overflow_dists = nullptr; + uint32_t* d_overflow_expanded = nullptr; + int* d_overflow_count = nullptr; + + size_t queries_bytes = 0; + size_t neighbors_bytes = 0; + size_t distances_bytes = 0; + int entry_cap = 0; + size_t bitmap_bytes = 0; + size_t overflow_bytes = 0; + + void ensure(int nq, int k, int dim, int N, int overflow_ef); + + ~GpuHnswSearchScratch(); + + GpuHnswSearchScratch() = default; + GpuHnswSearchScratch(const GpuHnswSearchScratch&) = delete; + GpuHnswSearchScratch& operator=(const GpuHnswSearchScratch&) = delete; +}; + +/// One slot in the scratch pool: a scratch buffer + its own CUDA stream. +struct GpuHnswScratchSlot { + GpuHnswSearchScratch scratch; + cudaStream_t stream = nullptr; + + ~GpuHnswScratchSlot(); + GpuHnswScratchSlot() = default; + GpuHnswScratchSlot(const GpuHnswScratchSlot&) = delete; + GpuHnswScratchSlot& operator=(const GpuHnswScratchSlot&) = delete; +}; + +/// Pool of scratch buffers allowing concurrent GPU searches on the same segment. +/// Each acquire() returns a slot with its own scratch + CUDA stream. +/// Callers block if all slots are in use. +class GpuHnswScratchPool { + public: + /// Create a pool. CUDA streams are allocated lazily on first acquire(). + explicit GpuHnswScratchPool(int pool_size = 4, int device = 0); + ~GpuHnswScratchPool() = default; + + GpuHnswScratchPool(const GpuHnswScratchPool&) = delete; + GpuHnswScratchPool& operator=(const GpuHnswScratchPool&) = delete; + + /// Acquire a scratch slot (blocks until one is available). + GpuHnswScratchSlot* acquire(); + /// Release a previously acquired scratch slot back to the pool. + void release(GpuHnswScratchSlot* slot); + + int pool_size() const { return pool_size_; } + + private: + void init_once(); + + std::mutex mutex_; + std::condition_variable cv_; + int pool_size_; + int device_; + bool initialized_ = false; + std::vector> slots_; + std::vector available_; +}; + +/// RAII guard: acquires a scratch slot on construction, releases on destruction. +class ScratchPoolGuard { + public: + ScratchPoolGuard(GpuHnswScratchPool& pool) + : pool_(pool), slot_(pool.acquire()) {} + ~ScratchPoolGuard() { pool_.release(slot_); } + GpuHnswScratchSlot* get() const { return slot_; } + + ScratchPoolGuard(const ScratchPoolGuard&) = delete; + ScratchPoolGuard& operator=(const ScratchPoolGuard&) = delete; + + private: + GpuHnswScratchPool& pool_; + GpuHnswScratchSlot* slot_; +}; + +struct GpuHnswDeviceIndex { + void* d_dataset = nullptr; + bool dataset_int8 = false; + float* d_inv_norms = nullptr; + uint32_t* d_layer0_graph = nullptr; + std::vector upper_layers; + + int64_t n_rows = 0; + int64_t dim = 0; + uint32_t entry_point = 0; + int num_layers = 0; + int M = 0; + int max_degree0 = 0; + bool use_ip = false; + + void* d_upper_layer_ptrs = nullptr; + int num_upper_layers_built = 0; + + mutable std::unique_ptr scratch_pool; + + ~GpuHnswDeviceIndex(); +}; + +} // namespace gpu +} // namespace faiss